[Mlir-commits] [mlir] 607f14d - [mlir][ods] Add ArrayOfAttr for creating custom array attributes

[Jeff Niu]

Author: Jeff Niu
Date: 2022-09-02T11:12:58-07:00
New Revision: 607f14d9605da801034e7119c297c3f58ebce603

URL: https://github.com/llvm/llvm-project/commit/607f14d9605da801034e7119c297c3f58ebce603
DIFF: https://github.com/llvm/llvm-project/commit/607f14d9605da801034e7119c297c3f58ebce603.diff

LOG: [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

Added: 
    mlir/test/IR/array-of-attr.mlir

Modified: 
    mlir/include/mlir/IR/AttrTypeBase.td
    mlir/test/lib/Dialect/Test/TestAttrDefs.td
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AttrTypeBase.td b/mlir/include/mlir/IR/AttrTypeBase.td
index 5bf1dfe3239b0..8f168064a362f 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 0000000000000..8b85aa79ee992
--- /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 1a7ccdc0694c3..dba035779c0df 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 f0dd4e29be0f8..80d3c5abdc84d 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2474,6 +2474,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
 //===----------------------------------------------------------------------===//