[mlir][tblgen] Add custom parsing and printing within struct

Author: jtuyls

mlir/test/IR/custom-struct-attr-roundtrip.mlir

@@ -0,0 +1,66 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics| FileCheck %s
+
+// CHECK-LABEL: @test_struct_attr_roundtrip
+func.func @test_struct_attr_roundtrip() -> () {
+  // CHECK: attr = #test.custom_struct<type_str = "struct", value = 2, opt_value = [3, 3]>
+  "test.op"() {attr = #test.custom_struct<type_str = "struct", value = 2, opt_value = [3, 3]>} : () -> ()
+  // CHECK: attr = #test.custom_struct<type_str = "struct", value = 2, opt_value = [3, 3]>
+  "test.op"() {attr = #test.custom_struct<value = 2, type_str = "struct", opt_value = [3, 3]>} : () -> ()
+  // CHECK: attr = #test.custom_struct<type_str = "struct", value = 2>
+  "test.op"() {attr = #test.custom_struct<type_str = "struct", value = 2>} : () -> ()
+  // CHECK: attr = #test.custom_struct<type_str = "struct", value = 2>
+  "test.op"() {attr = #test.custom_struct<value = 2, type_str = "struct">} : () -> ()
+  return
+}
+
+// -----
+
+// Verify all required parameters must be provided. `value` is missing.
+
+// expected-error @below {{struct is missing required parameter: value}}
+"test.op"() {attr = #test.custom_struct<type_str = "struct">} : () -> ()
+
+// -----
+
+// Verify all keywords must be provided. All missing.
+
+// expected-error @below {{expected valid keyword}}
+// expected-error @below {{expected a parameter name in struct}}
+"test.op"() {attr = #test.custom_struct<"struct", 2>} : () -> ()
+
+// -----
+
+// Verify all keywords must be provided. `type_str` missing.
+
+// expected-error @below {{expected valid keyword}}
+// expected-error @below {{expected a parameter name in struct}}
+"test.op"() {attr = #test.custom_struct<"struct", value = 2, opt_value = [3, 3]>} : () -> ()
+
+// -----
+
+// Verify all keywords must be provided. `value` missing.
+
+// expected-error @below {{expected valid keyword}}
+// expected-error @below {{expected a parameter name in struct}}
+"test.op"() {attr = #test.custom_struct<type_str = "struct", 2>} : () -> ()
+
+// -----
+
+// Verify invalid keyword provided.
+
+// expected-error @below {{duplicate or unknown struct parameter name: type_str2}}
+"test.op"() {attr = #test.custom_struct<type_str2 = "struct", value = 2>} : () -> ()
+
+// -----
+
+// Verify duplicated keyword provided.
+
+// expected-error @below {{duplicate or unknown struct parameter name: type_str}}
+"test.op"() {attr = #test.custom_struct<type_str = "struct", type_str = "struct2", value = 2>} : () -> ()
+
+// -----
+
+// Verify equals missing.
+
+// expected-error @below {{expected '='}}
+"test.op"() {attr = #test.custom_struct<type_str "struct", value = 2>} : () -> ()

mlir/test/lib/Dialect/Test/TestAttrDefs.td

@@ -369,6 +369,16 @@ def TestCustomFloatAttr : Test_Attr<"TestCustomFloat"> {
   }];
 }
 
+// Test `struct` with nested `custom` assembly format.
+def TestCustomStructAttr : Test_Attr<"TestCustomStruct"> {
+  let mnemonic = "custom_struct";
+  let parameters = (ins "mlir::StringAttr":$type_str, "int64_t":$value,
+                        OptionalParameter<"mlir::ArrayAttr">:$opt_value);
+  let assemblyFormat = [{
+    `<` struct($type_str, custom<CustomStructAttr>($value), custom<CustomOptStructFieldAttr>($opt_value)) `>`
+  }];
+}
+
 def NestedPolynomialAttr : Test_Attr<"NestedPolynomialAttr"> {
   let mnemonic = "nested_polynomial";
   let parameters = (ins Polynomial_IntPolynomialAttr:$poly);

mlir/test/lib/Dialect/Test/TestAttributes.cpp

@@ -316,6 +316,33 @@ static ParseResult parseCustomFloatAttr(AsmParser &p, StringAttr &typeStrAttr,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// TestCustomStructAttr
+//===----------------------------------------------------------------------===//
+
+static void printCustomStructAttr(AsmPrinter &p, int64_t value) {
+  p.printStrippedAttrOrType(value);
+}
+
+static ParseResult parseCustomStructAttr(AsmParser &p, int64_t &value) {
+  if (failed(p.parseInteger(value))) {
+    return failure();
+  }
+  return success();
+}
+
+static void printCustomOptStructFieldAttr(AsmPrinter &p, ArrayAttr attr) {
+  p.printStrippedAttrOrType(attr);
+}
+
+static ParseResult parseCustomOptStructFieldAttr(AsmParser &p,
+                                                 ArrayAttr &attr) {
+  if (failed(p.parseAttribute(attr))) {
+    return failure();
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TestOpAsmAttrInterfaceAttr
 //===----------------------------------------------------------------------===//

mlir/test/mlir-tblgen/attr-or-type-format-invalid.td

@@ -37,7 +37,7 @@ def InvalidTypeC : InvalidType<"InvalidTypeC", "invalid_c"> {
 def InvalidTypeD : InvalidType<"InvalidTypeD", "invalid_d"> {
   let parameters = (ins "int":$v0);
   // CHECK: literals may only be used in the top-level section of the format
-  // CHECK: expected a variable in `struct` argument list
+  // CHECK: expected a variable or `custom` directive in `struct` argument list
   let assemblyFormat = "`<` struct($v0, `,`) `>`";
 }
 
@@ -144,3 +144,17 @@ def InvalidTypeT : InvalidType<"InvalidTypeT", "invalid_t"> {
   // CHECK: `custom` directive with no bound parameters cannot be used as optional group anchor
   let assemblyFormat = "$a (`(` custom<Foo>(ref($a))^ `)`)?";
 }
+
+// Test `struct` with nested `custom` directive with multiple fields.
+def InvalidTypeU : InvalidType<"InvalidTypeU", "invalid_u"> {
+  let parameters = (ins OptionalParameter<"int">:$a, "int":$b);
+  // CHECK: `struct` can only contain `custom` directives with a single argument
+  let assemblyFormat = "struct(custom<Foo>($a, $b))";
+}
+
+// Test `struct` with nested `custom` directive invalid parameter.
+def InvalidTypeV : InvalidType<"InvalidTypeV", "invalid_v"> {
+  let parameters = (ins OptionalParameter<"int">:$a);
+  // CHECK: a `custom` directive nested within a `struct` must be passed a parameter
+  let assemblyFormat = "struct($a, custom<Foo>(ref($a)))";
+}

mlir/test/mlir-tblgen/attr-or-type-format.td

@@ -736,6 +736,78 @@ def TypeS : TestType<"TestS"> {
   let assemblyFormat = "$a";
 }
 
+/// Test that a `struct` with nested `custom` parser and printer are generated correctly.
+
+// ATTR:    ::mlir::Attribute TestTAttr::parse(::mlir::AsmParser &odsParser,
+// ATTR:                                       ::mlir::Type odsType) {
+// ATTR:      bool _seen_v0 = false;
+// ATTR:      bool _seen_v1 = false;
+// ATTR:      bool _seen_v2 = false;
+// ATTR:      const auto _loop_body = [&](::llvm::StringRef _paramKey) -> bool {
+// ATTR:        if (odsParser.parseEqual())
+// ATTR:          return {};
+// ATTR:        if (!_seen_v0 && _paramKey == "v0") {
+// ATTR:          _seen_v0 = true;
+// ATTR:          _result_v0 = ::parseAttrParamA(odsParser, odsType);
+// ATTR:          if (::mlir::failed(_result_v0))
+// ATTR:            return {};
+// ATTR:        } else if (!_seen_v1 && _paramKey == "v1") {
+// ATTR:          _seen_v1 = true;
+// ATTR:          {
+// ATTR:            auto odsCustomResult = parseNestedCustom(odsParser,
+// ATTR-NEXT:         ::mlir::detail::unwrapForCustomParse(_result_v1));
+// ATTR:            if (::mlir::failed(odsCustomResult)) return {};
+// ATTR:            if (::mlir::failed(_result_v1)) {
+// ATTR:              odsParser.emitError(odsCustomLoc, "custom parser failed to parse parameter 'v1'");
+// ATTR:              return {};
+// ATTR:            }
+// ATTR:          }
+// ATTR:        } else if (!_seen_v2 && _paramKey == "v2") {
+// ATTR:        } else {
+// ATTR:          return {};
+// ATTR:        }
+// ATTR:        return true;
+// ATTR:      }
+// ATTR:      do {
+// ATTR:        ::llvm::StringRef _paramKey;
+// ATTR:        if (odsParser.parseKeyword(&_paramKey)) {
+// ATTR:          odsParser.emitError(odsParser.getCurrentLocation(),
+// ATTR-NEXT:                         "expected a parameter name in struct");
+// ATTR:          return {};
+// ATTR:        }
+// ATTR:        if (!_loop_body(_paramKey)) return {};
+// ATTR:      } while(!odsParser.parseOptionalComma());
+// ATTR:      if (!_seen_v0)
+// ATTR:      if (!_seen_v1)
+// ATTR:      return TestTAttr::get(odsParser.getContext(),
+// ATTR:          TestParamA((*_result_v0)),
+// ATTR:          TestParamB((*_result_v1)),
+// ATTR:           AttrParamB((_result_v2.value_or(AttrParamB()))));
+// ATTR:    }
+
+// ATTR:      void TestTAttr::print(::mlir::AsmPrinter &odsPrinter) const {
+// ATTR:        odsPrinter << "v0 = ";
+// ATTR:        ::printAttrParamA(odsPrinter, getV0());
+// ATTR:        odsPrinter << ", ";
+// ATTR:        odsPrinter << "v1 = ";
+// ATTR:        printNestedCustom(odsPrinter,
+// ATTR-NEXT:     getV1());
+// ATTR:        if (!(getV2() == AttrParamB())) {
+// ATTR:          odsPrinter << "v2 = ";
+// ATTR:          odsPrinter.printStrippedAttrOrType(getV2());
+// ATTR:      }
+
+def AttrT : TestAttr<"TestT"> {
+  let parameters = (ins
+      AttrParamA:$v0,
+      AttrParamB:$v1,
+      OptionalParameter<"AttrParamB">:$v2
+  );
+
+  let mnemonic = "attr_t";
+  let assemblyFormat = "`{` struct($v0, custom<NestedCustom>($v1), $v2) `}`";
+}
+
 // DEFAULT_TYPE_PARSER: TestDialect::parseType(::mlir::DialectAsmParser &parser)
 // DEFAULT_TYPE_PARSER: auto parseResult = parseOptionalDynamicType(mnemonic, parser, genType);
 // DEFAULT_TYPE_PARSER: if (parseResult.has_value()) {