[mlir][ods] Allow specifying return types of builders

[Mogball]

Summary

This patch allows custom attribute and type builders to return
something other than the C++ type of the attribute or type.

This is useful for attributes or types that may perform extra work during
construction (e.g. canonicalization) that could result in a different
kind of attribute or type being returned.

[jpienaar]

Could you also add rationale to the description? (Especially as one now has a attribute with a C++ type that you'd want to build using a different C++ type, makes me wonder, why not just define the attribute that way to begin with)

[Mogball]

Added a justification.

[lattner]

Nice, this is super useful for attributes whose get() methods can fold them into simpler operators, like these in CIRCT:
https://github.com/llvm/circt/blob/main/include/circt/Dialect/HW/HWAttributes.td#L258

@darthscsi fyi

[rriddle]

This is missing documentation in AttributesAndTypes.md

mlir/lib/TableGen/AttrOrTypeDef.cpp

87

This looks unrelated, why was this dropped? AFAICT this commit doesn't address the TODO here.

[Mogball]

Update docs and remove extraneous diff

[rriddle]

LG, nice work!

[Mogball]

rebase

mlir/docs/AttributesAndTypes.md

TypeBuilder<(ins CArg<"int", "0">:$intParam), [{
       // Write the body of the `get` builder inline here.
       return Base::get($_ctxt, intParam);
     }]>,
     TypeBuilderWithInferredContext<(ins "Type":$typeParam), [{
       // This builder states that it can infer an MLIRContext instance from
       // its arguments.
       return Base::get(typeParam.getContext(), ...);
     }]>,
+    TypeBuilder<(ins "int":$intParam), [{}], "IntegerType">,
   ];
 }
 ```
 
 In this example, we provide several different convenience builders that are
 useful in different scenarios. The `ins` prefix is common to many function
 declarations in ODS, which use a TableGen [`dag`](#tablegen-syntax). What
 follows is a comma-separated list of types (quoted string or `CArg`) and names

 ```
 
 In this builder example, the main difference from the third builder example
 there is that the `MLIRContext` parameter is no longer added. This is because
 the builder used `TypeBuilderWithInferredContext` implies that the context
 parameter is not necessary as it can be inferred from the arguments to the
 builder.
 
+The fifth builder will generate the declaration of a builder method with a
+custom return type, like:
+
+```tablegen
+  let builders = [
+    TypeBuilder<(ins "int":$intParam), [{}], "IntegerType">,
+  ]
+```
+
+```c++
+class MyType : /*...*/ {
+  /*...*/
+  static IntegerType get(::mlir::MLIRContext *context, int intParam);
+
+};
+```
+
+This generates a builder declaration the same as the first three examples, but
+the return type of the builder is user-specified instead of the attribute or
+type class. This is useful for defining builders of attributes and types that
+may fold or canonicalize on construction.
+
 ### Parsing and Printing
 
 If a mnemonic was specified, the `hasCustomAssemblyFormat` and `assemblyFormat`
 fields may be used to specify the assembly format of an attribute or type. Attributes
 and Types with no parameters need not use either of these fields, in which case
 the syntax for the Attribute or Type is simply the mnemonic.
 
 For each dialect, two "dispatch" functions will be created: one for parsing and

mlir/include/mlir/IR/AttrTypeBase.td

 // AttrOrTypeBuilder<(ins "int":$integerArg, CArg<"float", "3.0f">:$floatArg), [{
 //   return $_get($_ctxt, integerArg, floatArg);
 // }]>
 // ```
 //
 // This is necessary because the `body` is also used to generate `getChecked`
 // methods, which have a different underlying `Base::get*` call.
 //
-class AttrOrTypeBuilder<dag parameters, code bodyCode = ""> {
+class AttrOrTypeBuilder<dag parameters, code bodyCode = "",
+                        string returnTypeStr = ""> {
   dag dagParams = parameters;
   code body = bodyCode;
 
+  // Change the return type of the builder. By default, it is the type of the
+  // attribute or type.
+  string returnType = returnTypeStr;
+
   // The context parameter can be inferred from one of the other parameters and
   // is not implicitly added to the parameter list.
   bit hasInferredContextParam = 0;
 }
-class AttrBuilder<dag parameters, code bodyCode = "">
-  : AttrOrTypeBuilder<parameters, bodyCode>;
-class TypeBuilder<dag parameters, code bodyCode = "">
-  : AttrOrTypeBuilder<parameters, bodyCode>;
+class AttrBuilder<dag parameters, code bodyCode = "", string returnType = "">
+  : AttrOrTypeBuilder<parameters, bodyCode, returnType>;
+class TypeBuilder<dag parameters, code bodyCode = "", string returnType = "">
+  : AttrOrTypeBuilder<parameters, bodyCode, returnType>;
 
 // A class of AttrOrTypeBuilder that is able to infer the MLIRContext parameter
 // from one of the other builder parameters. Instances of this builder do not
 // have `MLIRContext *` implicitly added to the parameter list.
-class AttrOrTypeBuilderWithInferredContext<dag parameters, code bodyCode = "">
-  : TypeBuilder<parameters, bodyCode> {
+class AttrOrTypeBuilderWithInferredContext<dag parameters, code bodyCode = "",
+                                           string returnType = "">
+  : TypeBuilder<parameters, bodyCode, returnType> {
   let hasInferredContextParam = 1;
 }
-class AttrBuilderWithInferredContext<dag parameters, code bodyCode = "">
-  : AttrOrTypeBuilderWithInferredContext<parameters, bodyCode>;
-class TypeBuilderWithInferredContext<dag parameters, code bodyCode = "">
-  : AttrOrTypeBuilderWithInferredContext<parameters, bodyCode>;
+class AttrBuilderWithInferredContext<dag parameters, code bodyCode = "",
+                                     string returnType = "">
+  : AttrOrTypeBuilderWithInferredContext<parameters, bodyCode, returnType>;
+class TypeBuilderWithInferredContext<dag parameters, code bodyCode = "",
+                                     string returnType = "">
+  : AttrOrTypeBuilderWithInferredContext<parameters, bodyCode, returnType>;
 
 //===----------------------------------------------------------------------===//
 // Definitions
 //===----------------------------------------------------------------------===//
 
 // Define a new attribute or type, named `name`, that inherits from the given
 // C++ base class.
 class AttrOrTypeDef<string valueType, string name, list<Trait> defTraits,

mlir/include/mlir/IR/OpImplementation.h

public:
   // Attribute/Type Parsing
   //===--------------------------------------------------------------------===//
 
   /// Invoke the `getChecked` method of the given Attribute or Type class, using
   /// the provided location to emit errors in the case of failure. Note that
   /// unlike `OpBuilder::getType`, this method does not implicitly insert a
   /// context parameter.
   template <typename T, typename... ParamsT>
-  T getChecked(SMLoc loc, ParamsT &&...params) {
+  auto getChecked(SMLoc loc, ParamsT &&...params) {
     return T::getChecked([&] { return emitError(loc); },
                          std::forward<ParamsT>(params)...);
   }
   /// A variant of `getChecked` that uses the result of `getNameLoc` to emit
   /// errors.
   template <typename T, typename... ParamsT>
-  T getChecked(ParamsT &&...params) {
+  auto getChecked(ParamsT &&...params) {
     return T::getChecked([&] { return emitError(getNameLoc()); },
                          std::forward<ParamsT>(params)...);
   }
 
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/TableGen/AttrOrTypeDef.h

 // AttrOrTypeBuilder
 //===----------------------------------------------------------------------===//
 
 /// Wrapper class that represents a Tablegen AttrOrTypeBuilder.
 class AttrOrTypeBuilder : public Builder {
 public:
   using Builder::Builder;
 
+  /// Returns an optional builder return type.
+  Optional<StringRef> getReturnType() const;
+
   /// Returns true if this builder is able to infer the MLIRContext parameter.
   bool hasInferredContextParameter() const;
 };
 
 //===----------------------------------------------------------------------===//
 // AttrOrTypeParameter
 //===----------------------------------------------------------------------===//
 

mlir/lib/TableGen/AttrOrTypeDef.cpp

 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 //===----------------------------------------------------------------------===//
 // AttrOrTypeBuilder
 //===----------------------------------------------------------------------===//
 
-/// Returns true if this builder is able to infer the MLIRContext parameter.
+Optional<StringRef> AttrOrTypeBuilder::getReturnType() const {
+  Optional<StringRef> type = def->getValueAsOptionalString("returnType");
+  return type && !type->empty() ? type : llvm::None;
+}
+
 bool AttrOrTypeBuilder::hasInferredContextParameter() const {
   return def->getValueAsBit("hasInferredContextParam");
 }
 
 //===----------------------------------------------------------------------===//
 // AttrOrTypeDef
 //===----------------------------------------------------------------------===//
 
AttrOrTypeDef::AttrOrTypeDef(const llvm::Record *def) : def(def) {
   } else if (hasCppFormat || hasDeclarativeFormat) {
     PrintFatalError(getLoc(),
                     "'assemblyFormat' or 'hasCustomAssemblyFormat' can only be "
                     "used when 'mnemonic' is set");
   }
   // Assembly format parser requires builders with the same prototype
   // as the default-builders.
   // TODO: attempt to detect when a custom builder matches the prototype.
   if (hasDeclarativeFormat && skipDefaultBuilders()) {
     PrintWarning(getLoc(),
                  "using 'assemblyFormat' with 'skipDefaultBuilders=1' may "
                  "result in C++ compilation errors");
   }
   // Assembly format printer requires accessors to be generated.
   if (hasDeclarativeFormat && !genAccessors()) {
     PrintFatalError(getLoc(),
                     "'assemblyFormat' requires 'genAccessors' to be true");

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

 def TestAttrSelfTypeParameterFormat
     : Test_Attr<"TestAttrSelfTypeParameterFormat"> {
   let parameters = (ins "int":$a, AttributeSelfTypeParameter<"">:$type);
 
   let mnemonic = "attr_self_type_format";
   let assemblyFormat = "`<` $a `>`";
 }
 
+// Test overridding attribute builders with a custom builder.
+def TestOverrideBuilderAttr : Test_Attr<"TestOverrideBuilder"> {
+  let mnemonic = "override_builder";
+  let parameters = (ins "int":$a);
+  let assemblyFormat = "`<` $a `>`";
+
+  let skipDefaultBuilders = 1;
+  let genVerifyDecl = 1;
+  let builders = [AttrBuilder<(ins "int":$a), [{
+    return ::mlir::IntegerAttr::get(::mlir::IndexType::get($_ctxt), a);
+  }], "::mlir::Attribute">];
+}
+
 // Test simple extern 1D vector using ElementsAttrInterface.
 def TestExtern1DI64ElementsAttr : Test_Attr<"TestExtern1DI64Elements", [
     ElementsAttrInterface
   ]> {
   let mnemonic = "e1di64_elements";
   let parameters = (ins
     AttributeSelfTypeParameter<"", "::mlir::ShapedType">:$type,
     ResourceHandleParameter<"TestExternalElementsDataHandle">:$handle

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

 // ATTR:   if (odsParser.parseComma())
 // ATTR:     return {};
 // ATTR:   _result_complex = ::parseAttrParamA(odsParser, odsType);
 // ATTR:   if (::mlir::failed(_result_complex))
 // ATTR:     return {};
 // ATTR:   if (odsParser.parseRParen())
 // ATTR:     return {};
 // ATTR:   return TestAAttr::get(odsParser.getContext(),
-// ATTR:                         (*_result_value),
-// ATTR:                         (*_result_complex));
+// ATTR:                         IntegerAttr((*_result_value)),
+// ATTR:                         TestParamA((*_result_complex)));
 // ATTR: }
 
 // ATTR: void TestAAttr::print(::mlir::AsmPrinter &odsPrinter) const {
 // ATTR:   odsPrinter << ' ' << "hello";
 // ATTR:   odsPrinter << ' ' << "=";
 // ATTR:   odsPrinter << ' ';
 // ATTR:   odsPrinter.printStrippedAttrOrType(getValue());
 // ATTR:   odsPrinter << ",";

 // ATTR:     StringRef _paramKey;
 // ATTR:     if (odsParser.parseKeyword(&_paramKey))
 // ATTR:       return {};
 // ATTR:     if (!_loop_body(_paramKey)) return {};
 // ATTR:     if ((odsStructIndex != 2 - 1) && odsParser.parseComma())
 // ATTR:       return {};
 // ATTR:   }
 // ATTR:   return TestBAttr::get(odsParser.getContext(),
-// ATTR:                         (*_result_v0),
-// ATTR:                         (*_result_v1));
+// ATTR:                         TestParamA((*_result_v0)),
+// ATTR:                         TestParamB((*_result_v1)));
 // ATTR: }
 
 // ATTR: void TestBAttr::print(::mlir::AsmPrinter &odsPrinter) const {
 // ATTR:   odsPrinter << "v0 = ";
 // ATTR:   ::printAttrParamA(odsPrinter, getV0());
 // ATTR:   odsPrinter << ", ";
 // ATTR:   odsPrinter << "v1 = ";
 // ATTR:   ::printAttrB(odsPrinter, getV1());

 // ATTR:   if (::mlir::failed(_result_v0))
 // ATTR:     return {};
 // ATTR:   if (odsParser.parseComma())
 // ATTR:     return {};
 // ATTR:   _result_v1 = ::mlir::FieldParser<int>::parse(odsParser);
 // ATTR:   if (::mlir::failed(_result_v1))
 // ATTR:     return {};
 // ATTR:   return TestFAttr::get(odsParser.getContext(),
-// ATTR:     (*_result_v0),
-// ATTR:     (*_result_v1));
+// ATTR:     int((*_result_v0)),
+// ATTR:     int((*_result_v1)));
 // ATTR: }
 
 def AttrC : TestAttr<"TestF"> {
   let parameters = (ins "int":$v0, "int":$v1);
 
   let mnemonic = "attr_c";
   let assemblyFormat = "params";
 }

 // TYPE:     if (!_loop_body(_paramKey)) return {};
 // TYPE:     if ((odsStructIndex != 2 - 1) && odsParser.parseComma())
 // TYPE:       return {};
 // TYPE:   }
 // TYPE:   _result_v3 = someFcnCall();
 // TYPE:   if (::mlir::failed(_result_v3))
 // TYPE:     return {};
 // TYPE:   return TestDType::get(odsParser.getContext(),
-// TYPE:                         (*_result_v0),
-// TYPE:                         (*_result_v1),
-// TYPE:                         (*_result_v2),
-// TYPE:                         (*_result_v3));
+// TYPE:                         TestParamC((*_result_v0)),
+// TYPE:                         TestParamD((*_result_v1)),
+// TYPE:                         TestParamC((*_result_v2)),
+// TYPE:                         TestParamD((*_result_v3)));
 // TYPE: }
 
 // TYPE: void TestDType::print(::mlir::AsmPrinter &odsPrinter) const {
 // TYPE:   odsPrinter << getV0();
 // TYPE:   myPrinter(getV1());
 // TYPE:   odsPrinter << "v2 = ";
 // TYPE:   odsPrinter << getV2();
 // TYPE:   myPrinter(getV3());

 // TYPE:     StringRef _paramKey;
 // TYPE:     if (odsParser.parseKeyword(&_paramKey))
 // TYPE:       return {};
 // TYPE:     if (!_loop_body(_paramKey)) return {};
 // TYPE:     if ((odsStructIndex != 2 - 1) && odsParser.parseComma())
 // TYPE:       return {};
 // TYPE:   }
 // TYPE:   return TestEType::get(odsParser.getContext(),
-// TYPE:     (*_result_v0),
-// TYPE:     (*_result_v1),
-// TYPE:     (*_result_v2),
-// TYPE:     (*_result_v3));
+// TYPE:     IntegerAttr((*_result_v0)),
+// TYPE:     IntegerAttr((*_result_v1)),
+// TYPE:     IntegerAttr((*_result_v2)),
+// TYPE:     IntegerAttr((*_result_v3)));
 // TYPE: }
 
 // TYPE: void TestEType::print(::mlir::AsmPrinter &odsPrinter) const {
 // TYPE:   odsPrinter << "v0 = ";
 // TYPE:   odsPrinter.printStrippedAttrOrType(getV0());
 // TYPE:   odsPrinter << ", ";
 // TYPE:   odsPrinter << "v2 = ";
 // TYPE:   odsPrinter.printStrippedAttrOrType(getV2());

mlir/test/mlir-tblgen/attrdefs.td

 // DEF: return ::mlir::AsmParser::KeywordSwitch<::mlir::OptionalParseResult>(parser)
 // DEF: .Case(::test::CompoundAAttr::getMnemonic()
 // DEF-NEXT: value = ::test::CompoundAAttr::parse(parser, type);
 // DEF-NEXT: return ::mlir::success(!!value);
 // DEF-NEXT: })
 // DEF-NEXT: .Case(::test::IndexAttr::getMnemonic()
 // DEF-NEXT:   value = ::test::IndexAttr::parse(parser, type);
 // DEF-NEXT:   return ::mlir::success(!!value);
-// DEF: .Default([&](llvm::StringRef keyword, 
+// DEF: .Default([&](llvm::StringRef keyword,
 // DEF-NEXT:   *mnemonic = keyword;
-// DEF-NEXT:   return llvm::None; 
+// DEF-NEXT:   return llvm::None;
 
 def Test_Dialect: Dialect {
 // DECL-NOT: TestDialect
 // DEF-NOT: TestDialect
     let name = "TestDialect";
     let cppNamespace = "::test";
 }
 
def F_ParamWithAccessorTypeAttr : TestAttr<"ParamWithAccessorType"> {
   let parameters = (ins AttrParameter<"std::string", "", "StringRef">:$param);
 }
 
 // DECL-LABEL: class ParamWithAccessorTypeAttr
 // DECL: StringRef getParam()
 // DEF: ParamWithAccessorTypeAttrStorage
 // DEF: ParamWithAccessorTypeAttrStorage(std::string param)
 // DEF: StringRef ParamWithAccessorTypeAttr::getParam()
+
+def G_BuilderWithReturnTypeAttr : TestAttr<"BuilderWithReturnType"> {
+  let parameters = (ins "int":$a);
+  let genVerifyDecl = 1;
+  let builders = [AttrBuilder<(ins), [{ return {}; }], "::mlir::Attribute">];
+}
+
+// DECL-LABEL: class BuilderWithReturnTypeAttr
+// DECL: ::mlir::Attribute get(
+// DECL: ::mlir::Attribute getChecked(

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_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]>>>}
+
+// CHECK-LABEL: @overriddenAttr
+// CHECK-SAME: foo = 5 : index
+func.func private @overriddenAttr() attributes {
+  foo = #test.override_builder<5>
+}

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

builderParams.emplace_back(param.getCppType(), *param.getName(),
                                param.getDefaultValue());
   }
   return builderParams;
 }
 
 void DefGen::emitCustomBuilder(const AttrOrTypeBuilder &builder) {
   // Don't emit a body if there isn't one.
   auto props = builder.getBody() ? Method::Static : Method::StaticDeclaration;
-  Method *m = defCls.addMethod(def.getCppClassName(), "get", props,
+  StringRef returnType = def.getCppClassName();
+  if (Optional<StringRef> builderReturnType = builder.getReturnType())
+    returnType = *builderReturnType;
+  Method *m = defCls.addMethod(returnType, "get", props,
                                getCustomBuilderParams({}, builder));
   if (!builder.getBody())
     return;
 
   // Format the body and emit it.
   FmtContext ctx;
   ctx.addSubst("_get", "Base::get");
   if (!builder.hasInferredContextParameter())
     ctx.addSubst("_ctxt", "context");
   std::string bodyStr = tgfmt(*builder.getBody(), &ctx);
   m->body().indent().getStream().printReindented(bodyStr);
 }
 
 /// Replace all instances of 'from' to 'to' in `str` and return the new string.
 static std::string replaceInStr(std::string str, StringRef from, StringRef to) {
   size_t pos = 0;
   while ((pos = str.find(from.data(), pos, from.size())) != std::string::npos)
     str.replace(pos, from.size(), to.data(), to.size());
   return str;
 }
 
 void DefGen::emitCheckedCustomBuilder(const AttrOrTypeBuilder &builder) {
   // Don't emit a body if there isn't one.
   auto props = builder.getBody() ? Method::Static : Method::StaticDeclaration;
+  StringRef returnType = def.getCppClassName();
+  if (Optional<StringRef> builderReturnType = builder.getReturnType())
+    returnType = *builderReturnType;
   Method *m = defCls.addMethod(
-      def.getCppClassName(), "getChecked", props,
+      returnType, "getChecked", props,
       getCustomBuilderParams(
           {{"::llvm::function_ref<::mlir::InFlightDiagnostic()>", "emitError"}},
           builder));
   if (!builder.getBody())
     return;
 
   // Format the body and emit it. Replace $_get(...) with
   // Base::getChecked(emitError, ...)

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

if (param.isOptional()) {
       else
         selfOs << param.getCppStorageType() << "()";
       selfOs << "))";
     } else if (isa<AttributeSelfTypeParameter>(param)) {
       selfOs << tgfmt("$_type", &ctx);
     } else {
       selfOs << formatv("(*_result_{0})", param.getName());
     }
-    os << tgfmt(param.getConvertFromStorage(), &ctx.withSelf(selfOs.str()));
+    os << param.getCppType() << "("
+       << tgfmt(param.getConvertFromStorage(), &ctx.withSelf(selfOs.str()))
+       << ")";
   }
   os << ");";
 }
 
 void DefFormat::genElementParser(FormatElement *el, FmtContext &ctx,
                                  MethodBody &os) {
   if (auto *literal = dyn_cast<LiteralElement>(el))
     return genLiteralParser(literal->getSpelling(), ctx, os);