[mlir][ods] Ignore AttributeSelfTypeParameter in assembly formats

[Mogball]

Summary

The attribute self type parameter is currently treated like any other attribute parameter in the assembly format. The self type parameter should be handled by the operation parser and printer and play no role in the generated parsers and printers of attributes.

[rriddle]

Thanks!

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

237

nit: Please drop the const here.

[Mogball]

There's a further wrinkle here: attribute parsers embedded in operation assembly formats won't be able to have types.

def MyAttr : ... {
  let parameters = (ins "int":$a, AttributeSelfTypeParameter<"">:$type);
  let assemblyFormat = "$a";
}

def MyOp : ... {
  let arguments = (ins MyAttr:$a);
  let assemblyFormat = "$a attr-dict"
}

The pretty format won't be able to have a type because the parser directly invokes MyAttr::parse and won't be able to look ahead for the type.

[Mogball]

drop const

[rriddle]

In D125724#3517058, @Mogball wrote:

There's a further wrinkle here: attribute parsers embedded in operation assembly formats won't be able to have types.

def MyAttr : ... {
  let parameters = (ins "int":$a, AttributeSelfTypeParameter<"">:$type);
  let assemblyFormat = "$a";
}

def MyOp : ... {
  let arguments = (ins MyAttr:$a);
  let assemblyFormat = "$a attr-dict"
}

The pretty format won't be able to have a type because the parser directly invokes MyAttr::parse and won't be able to look ahead for the type.

Can you just file a general bug for me to fix Attribute type handling in assembly formats (declarative and not)? There are also things surrounding type elision, providing types, etc. that should "just work" but we haven't invested in.

[Mogball]

Sure. I was thinking we needed a $type-elision directive or something in operation assembly formats.

[Mogball]

filed https://github.com/llvm/llvm-project/issues/55513

mlir/lib/IR/AsmPrinter.cpp

if (!attr) {
     os << "<<NULL ATTRIBUTE>>";
     return;
   }
 
   // Try to print an alias for this attribute.
   if (succeeded(printAlias(attr)))
     return;
 
-  if (!isa<BuiltinDialect>(attr.getDialect()))
-    return printDialectAttribute(attr);
-
   auto attrType = attr.getType();
-  if (auto opaqueAttr = attr.dyn_cast<OpaqueAttr>()) {
+  if (!isa<BuiltinDialect>(attr.getDialect())) {
+    printDialectAttribute(attr);
+  } else if (auto opaqueAttr = attr.dyn_cast<OpaqueAttr>()) {
     printDialectSymbol(os, "#", opaqueAttr.getDialectNamespace(),
                        opaqueAttr.getAttrData());
   } else if (attr.isa<UnitAttr>()) {
     os << "unit";
     return;
   } else if (auto dictAttr = attr.dyn_cast<DictionaryAttr>()) {
     os << '{';
     interleaveComma(dictAttr.getValue(),

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

def CompoundAttrNested : Test_Attr<"CompoundAttrNested"> {
   let parameters = (ins CompoundAttrA : $nested );
   let assemblyFormat = "`<` `nested` `=` $nested `>`";
 }
 
 // An attribute testing AttributeSelfTypeParameter.
 def AttrWithSelfTypeParam : Test_Attr<"AttrWithSelfTypeParam"> {
   let mnemonic = "attr_with_self_type_param";
   let parameters = (ins AttributeSelfTypeParameter<"">:$type);
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = "";
 }
 
 // An attribute testing AttributeSelfTypeParameter.
 def AttrWithTypeBuilder : Test_Attr<"AttrWithTypeBuilder"> {
   let mnemonic = "attr_with_type_builder";
   let parameters = (ins "::mlir::IntegerAttr":$attr);
   let typeBuilder = "$_attr.getType()";
   let hasCustomAssemblyFormat = 1;

 // Test types can be parsed/printed.
 def TestAttrWithTypeParam : Test_Attr<"TestAttrWithTypeParam"> {
   let parameters = (ins "::mlir::IntegerType":$int_type,
                         "::mlir::Type":$any_type);
   let mnemonic = "attr_with_type";
   let assemblyFormat = "`<` $int_type `,` $any_type `>`";
 }
 
+// Test self type parameter with assembly format.
+def TestAttrSelfTypeParameterFormat
+    : Test_Attr<"TestAttrSelfTypeParameterFormat"> {
+  let parameters = (ins "int":$a, AttributeSelfTypeParameter<"">:$type);
+
+  let mnemonic = "attr_self_type_format";
+  let assemblyFormat = "`<` $a `>`";
+}
+
 #endif // TEST_ATTRDEFS

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

 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/ADT/bit.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir;
 using namespace test;
 
 //===----------------------------------------------------------------------===//
-// AttrWithSelfTypeParamAttr
-//===----------------------------------------------------------------------===//
-
-Attribute AttrWithSelfTypeParamAttr::parse(AsmParser &parser, Type type) {
-  Type selfType;
-  if (parser.parseType(selfType))
-    return Attribute();
-  return get(parser.getContext(), selfType);
-}
-
-void AttrWithSelfTypeParamAttr::print(AsmPrinter &printer) const {
-  printer << " " << getType();
-}
-
-//===----------------------------------------------------------------------===//
 // AttrWithTypeBuilderAttr
 //===----------------------------------------------------------------------===//
 
 Attribute AttrWithTypeBuilderAttr::parse(AsmParser &parser, Type type) {
   IntegerAttr element;
   if (parser.parseAttribute(element))
     return Attribute();
   return get(parser.getContext(), element);

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

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: @test_roundtrip_parameter_parsers
 // CHECK: !test.type_with_format<111, three = #test<"attr_ugly begin 5 : index end">, two = "foo">
 // CHECK: !test.type_with_format<2147, three = "hi", two = "hi">
 func.func private @test_roundtrip_parameter_parsers(!test.type_with_format<111, three = #test<"attr_ugly begin 5 : index end">, two = "foo">) -> !test.type_with_format<2147, two = "hi", three = "hi">
 attributes {
   // CHECK: #test.attr_with_format<3 : two = "hello", four = [1, 2, 3] : 42 : i64>
   attr0 = #test.attr_with_format<3 : two = "hello", four = [1, 2, 3] : 42 : i64>,
   // CHECK: #test.attr_with_format<5 : two = "a_string", four = [4, 5, 6, 7, 8] : 8 : i8>
   attr1 = #test.attr_with_format<5 : two = "a_string", four = [4, 5, 6, 7, 8] : 8 : i8>,
   // CHECK: #test<"attr_ugly begin 5 : index end">
   attr2 = #test<"attr_ugly begin 5 : index end">,
   // CHECK: #test.attr_params<42, 24>
   attr3 = #test.attr_params<42, 24>,
   // CHECK: #test.attr_with_type<i32, vector<4xi32>>
-  attr4 = #test.attr_with_type<i32, vector<4xi32>>
+  attr4 = #test.attr_with_type<i32, vector<4xi32>>,
+  // CHECK: #test.attr_self_type_format<5> : i32
+  attr5 = #test.attr_self_type_format<5> : i32
 }
 
 // CHECK-LABEL: @test_roundtrip_default_parsers_struct
 // CHECK: !test.no_parser<255, [1, 2, 3, 4, 5], "foobar", 4>
 // CHECK: !test.struct_capture_all<v0 = 0, v1 = 1, v2 = 2, v3 = 3>
 // CHECK: !test.optional_param<, 6>
 // CHECK: !test.optional_param<5, 6>
 // CHECK: !test.optional_params<"a">

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

 
 def AttrC : TestAttr<"TestF"> {
   let parameters = (ins "int":$v0, "int":$v1);
 
   let mnemonic = "attr_c";
   let assemblyFormat = "params";
 }
 
+/// Test attribute with self type parameter
+
+// ATTR: TestGAttr::parse
+// ATTR:   return TestGAttr::get
+// ATTR:     odsType
+def AttrD : TestAttr<"TestG"> {
+  let parameters = (ins "int":$a, AttributeSelfTypeParameter<"">:$type);
+  let mnemonic = "attr_d";
+  let assemblyFormat = "$a";
+}
+
 /// Test type parser and printer that mix variables and struct are generated
 /// correctly.
 
 // TYPE: ::mlir::Type TestCType::parse(::mlir::AsmParser &odsParser) {
 // TYPE:  FailureOr<IntegerAttr> _result_value;
 // TYPE:  FailureOr<TestParamC> _result_complex;
 // TYPE:  if (odsParser.parseKeyword("foo"))
 // TYPE:    return {};

mlir/test/mlir-tblgen/testdialect-attrdefs.mlir

 // RUN: mlir-opt %s | mlir-opt -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func private @compoundA()
 // CHECK-SAME: #test.cmpnd_a<1, !test.smpla, [5, 6]>
 func.func private @compoundA() attributes {foo = #test.cmpnd_a<1, !test.smpla, [5, 6]>}
 
-// CHECK: test.result_has_same_type_as_attr #test<"attr_with_self_type_param i32"> -> i32
-%a = test.result_has_same_type_as_attr #test<"attr_with_self_type_param i32"> -> i32
+// CHECK: test.result_has_same_type_as_attr #test.attr_with_self_type_param : i32 -> i32
+%a = test.result_has_same_type_as_attr #test.attr_with_self_type_param : i32 -> i32
 
-// CHECK: test.result_has_same_type_as_attr #test<"attr_with_type_builder 10 : i16"> -> i16
+// CHECK: test.result_has_same_type_as_attr #test<"attr_with_type_builder 10 : i16"> : i16 -> i16
 %b = test.result_has_same_type_as_attr #test<"attr_with_type_builder 10 : i16"> -> i16
 
 // CHECK-LABEL: @qualifiedAttr()
 // CHECK-SAME: #test.cmpnd_nested_outer_qual<i #test.cmpnd_nested_inner<42 <1, !test.smpla, [5, 6]>>>
 func.func private @qualifiedAttr() attributes {foo = #test.cmpnd_nested_outer_qual<i #test.cmpnd_nested_inner<42 <1, !test.smpla, [5, 6]>>>}

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

private:
   void genWhitespacePrinter(WhitespaceElement *el, FmtContext &ctx,
                             MethodBody &os);
 
   /// The ODS definition of the attribute or type whose format is being used to
   /// generate a parser and printer.
   const AttrOrTypeDef &def;
   /// The list of top-level format elements returned by the assembly format
   /// parser.
   std::vector<FormatElement *> elements;
 
   /// Flags for printing spaces.
   bool shouldEmitSpace = false;
   bool lastWasPunctuation = false;
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
void DefFormat::genParser(MethodBody &os) {
   os.indent();
 
   // Declare variables to store all of the parameters. Allocated parameters
   // such as `ArrayRef` and `StringRef` must provide a `storageType`. Store
   // FailureOr<T> to defer type construction for parameters that are parsed in
   // a loop (parsers return FailureOr anyways).
   ArrayRef<AttrOrTypeParameter> params = def.getParameters();
   for (const AttrOrTypeParameter &param : params) {
+    if (isa<AttributeSelfTypeParameter>(param))
+      continue;
     os << formatv("::mlir::FailureOr<{0}> _result_{1};\n",
                   param.getCppStorageType(), param.getName());
   }
 
   // Store the initial location of the parser.
   ctx.addSubst("_loc", "odsLoc");
   os << tgfmt("::llvm::SMLoc $_loc = $_parser.getCurrentLocation();\n"
               "(void) $_loc;\n",
               &ctx);
 
   // Generate call to each parameter parser.
   for (FormatElement *el : elements)
     genElementParser(el, ctx, os);
 
   // Emit an assert for each mandatory parameter. Triggering an assert means
   // the generated parser is incorrect (i.e. there is a bug in this code).
   for (const AttrOrTypeParameter &param : params) {
-    if (!param.isOptional()) {
-      os << formatv("assert(::mlir::succeeded(_result_{0}));\n",
-                    param.getName());
-    }
+    if (param.isOptional() || isa<AttributeSelfTypeParameter>(param))
+      continue;
+    os << formatv("assert(::mlir::succeeded(_result_{0}));\n", param.getName());
   }
 
   // Generate call to the attribute or type builder. Use the checked getter
   // if one was generated.
   if (def.genVerifyDecl()) {
     os << tgfmt("return $_parser.getChecked<$0>($_loc, $_parser.getContext()",
                 &ctx, def.getCppClassName());
   } else {
     os << tgfmt("return $0::get($_parser.getContext()", &ctx,
                 def.getCppClassName());
   }
   for (const AttrOrTypeParameter &param : params) {
+    os << ",\n    ";
     if (param.isOptional()) {
-      os << formatv(",\n    _result_{0}.getValueOr(", param.getName());
+      os << formatv("_result_{0}.getValueOr(", param.getName());
       if (Optional<StringRef> defaultValue = param.getDefaultValue())
         os << tgfmt(*defaultValue, &ctx);
       else
         os << param.getCppStorageType() << "()";
       os << ")";
+    } else if (isa<AttributeSelfTypeParameter>(param)) {
+      os << tgfmt("$_type", &ctx);
     } else {
-      os << formatv(",\n    *_result_{0}", param.getName());
+      os << formatv("*_result_{0}", param.getName());
     }
   }
   os << ");";
 }
 
 void DefFormat::genElementParser(FormatElement *el, FmtContext &ctx,
                                  MethodBody &os) {
   if (auto *literal = dyn_cast<LiteralElement>(el))

 //===----------------------------------------------------------------------===//
 
 void DefFormat::genPrinter(MethodBody &os) {
   FmtContext ctx;
   ctx.addSubst("_printer", "odsPrinter");
   ctx.addSubst("_ctx", "getContext()");
   os.indent();
 
-  /// Generate printers.
+  // Generate printers.
   shouldEmitSpace = true;
   lastWasPunctuation = false;
   for (FormatElement *el : elements)
     genElementPrinter(el, ctx, os);
 }
 
 void DefFormat::genElementPrinter(FormatElement *el, FmtContext &ctx,
                                   MethodBody &os) {
private:
   BitVector seenParams;
 };
 } // namespace
 
 LogicalResult DefFormatParser::verify(SMLoc loc,
                                       ArrayRef<FormatElement *> elements) {
   // Check that all parameters are referenced in the format.
   for (auto &it : llvm::enumerate(def.getParameters())) {
-    if (!it.value().isOptional() && !seenParams.test(it.index())) {
+    if (it.value().isOptional())
+      continue;
+    if (!seenParams.test(it.index())) {
+      if (isa<AttributeSelfTypeParameter>(it.value()))
+        continue;
       return emitError(loc, "format is missing reference to parameter: " +
                                 it.value().getName());
     }
+    if (isa<AttributeSelfTypeParameter>(it.value())) {
+      return emitError(loc,
+                       "unexpected self type parameter in assembly format");
+    }
   }
   if (elements.empty())
     return success();
   // A `struct` directive that contains optional parameters cannot be followed
   // by a comma literal, which is ambiguous.
   for (auto it : llvm::zip(elements.drop_back(), elements.drop_front())) {
     auto *structEl = dyn_cast<StructDirective>(std::get<0>(it));
     auto *literalEl = dyn_cast<LiteralElement>(std::get<1>(it));