From 607f14d9605da801034e7119c297c3f58ebce603 Mon Sep 17 00:00:00 2001
From: Jeff Niu <jeff@modular.com>
Date: Thu, 1 Sep 2022 09:30:54 -0700
Subject: [PATCH] [mlir][ods] Add ArrayOfAttr for creating custom array
 attributes

`ArrayOfAttr` can be used to easily create an attribute that just
contains an array of something. The elements can be other attributes,
in which case the custom parsers and printers are invoked directly for
nice syntax, or any C++ type that supports parsing and printing, either
though custom `printer` and `parser` methods or `FieldParser`.

An array of integers:

```
def ArrayOfInts : ArrayOfAttr<Test_Dialect, "ArrayOfInts", "array_of_ints",
                              "int32_t">;
```

When embedded in an op's assembly format, it will look like

```
foo.ints value = [1, 2, 3]
```

An array of enums, when embedded in an op's assembly format, will look
like:

```
foo.enums value = [first, second, last]
```

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D133131
---
 mlir/include/mlir/IR/AttrTypeBase.td       | 29 +++++++++++++++++++++++++++++
 mlir/test/IR/array-of-attr.mlir            | 10 ++++++++++
 mlir/test/lib/Dialect/Test/TestAttrDefs.td | 25 +++++++++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td      | 11 +++++++++++
 4 files changed, 75 insertions(+)
 create mode 100644 mlir/test/IR/array-of-attr.mlir

diff --git a/mlir/include/mlir/IR/AttrTypeBase.td b/mlir/include/mlir/IR/AttrTypeBase.td
index 5bf1dfe..8f16806 100644
--- a/mlir/include/mlir/IR/AttrTypeBase.td
+++ b/mlir/include/mlir/IR/AttrTypeBase.td
@@ -401,4 +401,33 @@ class AttributeSelfTypeParameter<string desc,
                          "::mlir::NoneType::get($_ctxt)", typeBuilder);
 }
 
+//===----------------------------------------------------------------------===//
+// ArrayOfAttr
+//===----------------------------------------------------------------------===//
+
+/// This class defines an attribute that contains an array of elements. The
+/// elements can be any type, but if they are attributes, the nested elements
+/// are parsed and printed using the custom attribute syntax.
+class ArrayOfAttr<Dialect dialect, string attrName, string attrMnemonic,
+                  string eltName, list<Trait> traits = []>
+    : AttrDef<dialect, attrName, traits> {
+  let parameters = (ins ArrayRefParameter<eltName>:$value);
+  let mnemonic = attrMnemonic;
+  let assemblyFormat = "`[` $value `]`";
+
+  let returnType = "::llvm::ArrayRef<" # eltName # ">";
+  let constBuilderCall = "$_builder.getAttr<" # attrName # "Attr>($0)";
+  let convertFromStorage = "$_self.getValue()";
+
+  let extraClassDeclaration = [{
+    auto begin() const { return getValue().begin(); }
+    auto end() const { return getValue().end(); }
+    bool empty() const { return getValue().empty(); }
+    size_t size() const { return getValue().size(); }
+    auto &front() const { return getValue().front(); }
+    auto &back() const { return getValue().back(); }
+    auto &operator[](size_t index) { return getValue()[index]; }
+  }];
+}
+
 #endif // ATTRTYPEBASE_TD
diff --git a/mlir/test/IR/array-of-attr.mlir b/mlir/test/IR/array-of-attr.mlir
new file mode 100644
index 0000000..8b85aa7
--- /dev/null
+++ b/mlir/test/IR/array-of-attr.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK: test.array_of_attr_op
+test.array_of_attr_op
+    // CHECK-SAME: a = [ begin 0 : index end, begin 2 : index end ]
+    a = [begin 0 : index end, begin 2 : index end],
+    // CHECK-SAME: [0, 1, -42, 42]
+    b = [0, 1, -42, 42],
+    // CHECK-SAME: [a, b, b, a]
+    c = [a, b, b, a]
diff --git a/mlir/test/lib/Dialect/Test/TestAttrDefs.td b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
index 1a7ccdc..dba0357 100644
--- a/mlir/test/lib/Dialect/Test/TestAttrDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
@@ -17,6 +17,7 @@
 include "TestDialect.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SubElementInterfaces.td"
 
@@ -266,4 +267,28 @@ def TestExtern1DI64ElementsAttr : Test_Attr<"TestExtern1DI64Elements", [
   let assemblyFormat = "`<` $handle `>`";
 }
 
+// An array of nested attributes.
+def TestArrayOfUglyAttrs : ArrayOfAttr<Test_Dialect, "ArrayOfUglyAttrs",
+    "array_of_ugly", "TestAttrUglyAttr"> {
+  let assemblyFormat = "`[` $value ` ` `]`";
+}
+
+// An array of integers.
+def TestArrayOfInts : ArrayOfAttr<Test_Dialect, "ArrayOfInts",
+    "array_of_ints", "int32_t">;
+
+// An array of enum attributes.
+def TestSimpleEnum : I32EnumAttr<"SimpleEnum", "", [
+    I32EnumAttrCase<"a", 0>,
+    I32EnumAttrCase<"b", 1>
+  ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::test";
+}
+def TestSimpleEnumAttr : EnumAttr<Test_Dialect, TestSimpleEnum, "simple_enum"> {
+  let assemblyFormat = "`` $value";
+}
+def TestArrayOfEnums : ArrayOfAttr<Test_Dialect, "ArrayOfEnums",
+    "array_of_enums", "SimpleEnumAttr">;
+
 #endif // TEST_ATTRDEFS
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index f0dd4e2..80d3c5a 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2475,6 +2475,17 @@ def FormatInferTypeVariadicOperandsOp
 }
 
 //===----------------------------------------------------------------------===//
+// Test ArrayOfAttr
+//===----------------------------------------------------------------------===//
+
+// Embed the array attributes directly in the assembly format for a nice syntax.
+def ArrayOfAttrOp : TEST_Op<"array_of_attr_op"> {
+  let arguments = (ins TestArrayOfUglyAttrs:$a, TestArrayOfInts:$b,
+                       TestArrayOfEnums:$c);
+  let assemblyFormat = "`a` `=` $a `,` `b` `=` $b `,` `c` `=` $c attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//
 
-- 
2.7.4