[mlir][AttrDefGen] Add support for specifying the value type of an attribute

[rriddle]

Summary

The value type of the attribute can be specified by either overriding the typeBuilder field on the AttrDef, or by providing a parameter of type AttributeSelfTypeParameter. This removes the need to define custom storage class constructors for attributes that have a value type other than NoneType.

Depends On D97589

[bondhugula]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

488

Missing doc comment.

[mehdi_amini]

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

488

(Please fix)

mlir/include/mlir/IR/OpBase.td

 // from the given C++ base class.
 class AttrDef<Dialect dialect, string name,
               string baseCppClass = "::mlir::Attribute">
     : DialectAttr<dialect, CPred<"">, /*descr*/"">,
       AttrOrTypeDef<"Attr", name, baseCppClass> {
   // The name of the C++ Attribute class.
   string cppClassName = name # "Attr";
 
+  // A code block used to build the value 'Type' of an Attribute when
+  // initializing its storage instance. This field is optional, and if not
+  // present the attribute will have its value type set to `NoneType`. This code
+  // block may reference any of the attributes parameters via
+  // `$_<parameter-name`. If one of the parameters of the attribute is of type
+  // `AttributeSelfTypeParameter`, this field is ignored.
+  code typeBuilder = ?;
+
   // The predicate for when this def is used as a constraint.
   let predicate = CPred<"$_self.isa<" # dialect.cppNamespace #
                                  "::" # cppClassName # ">()">;
 }
 
 // Define a new type, named `name`, belonging to `dialect` that inherits from
 // the given C++ base class.
 class TypeDef<Dialect dialect, string name,
class ArrayRefOfSelfAllocationParameter<string arrayOf, string desc> :
   let allocator = [{
     llvm::SmallVector<}] # arrayOf # [{, 4> tmpFields;
     for (size_t i = 0, e = $_self.size(); i < e; ++i)
       tmpFields.push_back($_self[i].allocateInto($_allocator));
     $_dst = $_allocator.copyInto(ArrayRef<}] # arrayOf # [{>(tmpFields));
   }];
 }
 
+// This is a special parameter used for AttrDefs that represents a `mlir::Type`
+// that is also used as the value `Type` of the attribute. Only one parameter
+// of the attribute may be of this type.
+class AttributeSelfTypeParameter<string desc> :
+    AttrOrTypeParameter<"::mlir::Type", desc> {}
+
 #endif // OP_BASE

mlir/include/mlir/TableGen/AttrOrTypeDef.h

public:
   bool operator==(const AttrOrTypeDef &other) const;
 
   // Compares two AttrOrTypeDefs by comparing the names of the dialects.
   bool operator<(const AttrOrTypeDef &other) const;
 
   // Returns whether the AttrOrTypeDef is defined.
   operator bool() const { return def != nullptr; }
 
-private:
+  // Return the underlying def.
+  const llvm::Record *getDef() const { return def; }
+
+protected:
   const llvm::Record *def;
 
   // The builders of this type definition.
   SmallVector<AttrOrTypeBuilder> builders;
 };
 
 //===----------------------------------------------------------------------===//
 // AttrDef
 //===----------------------------------------------------------------------===//
 
 /// This class represents a wrapper around a tablegen AttrDef record.
 class AttrDef : public AttrOrTypeDef {
 public:
   using AttrOrTypeDef::AttrOrTypeDef;
+
+  // Returns the attributes value type builder code block, or None if it doesn't
+  // have one.
+  Optional<StringRef> getTypeBuilder() const;
+
+  static bool classof(const AttrOrTypeDef *def);
 };
 
 //===----------------------------------------------------------------------===//
 // TypeDef
 //===----------------------------------------------------------------------===//
 
 /// This class represents a wrapper around a tablegen TypeDef record.
 class TypeDef : public AttrOrTypeDef {
public:
   StringRef getCppType() const;
 
   // Get a description of this parameter for documentation purposes.
   Optional<StringRef> getSummary() const;
 
   // Get the assembly syntax documentation.
   StringRef getSyntax() const;
 
+  // Return the underlying def of this parameter.
+  const llvm::Init *getDef() const;
+
 private:
   /// The underlying tablegen parameter list this parameter is a part of.
   const llvm::DagInit *def;
   /// The index of the parameter within the parameter list (`def`).
   unsigned index;
 };
 
+//===----------------------------------------------------------------------===//
+// AttributeSelfTypeParameter
+//===----------------------------------------------------------------------===//
+
+// A wrapper class for the AttributeSelfTypeParameter tblgen class. This
+// represents a parameter of mlir::Type that is the value type of an AttrDef.
+class AttributeSelfTypeParameter : public AttrOrTypeParameter {
+public:
+  static bool classof(const AttrOrTypeParameter *param);
+};
+
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_ATTRORTYPEDEF_H

mlir/lib/TableGen/AttrOrTypeDef.cpp

bool AttrOrTypeDef::operator==(const AttrOrTypeDef &other) const {
   return def == other.def;
 }
 
 bool AttrOrTypeDef::operator<(const AttrOrTypeDef &other) const {
   return getName() < other.getName();
 }
 
 //===----------------------------------------------------------------------===//
+// AttrDef
+//===----------------------------------------------------------------------===//
+
+Optional<StringRef> AttrDef::getTypeBuilder() const {
+  return def->getValueAsOptionalString("typeBuilder");
+}
+
+bool AttrDef::classof(const AttrOrTypeDef *def) {
+  return def->getDef()->isSubClassOf("AttrDef");
+}
+
+//===----------------------------------------------------------------------===//
 // AttrOrTypeParameter
 //===----------------------------------------------------------------------===//
 
 StringRef AttrOrTypeParameter::getName() const {
   return def->getArgName(index)->getValue();
 }
 
 Optional<StringRef> AttrOrTypeParameter::getAllocator() const {
if (auto *param = dyn_cast<llvm::DefInit>(parameterType)) {
     const auto *syntax = param->getDef()->getValue("syntax");
     if (syntax && isa<llvm::StringInit>(syntax->getValue()))
       return cast<llvm::StringInit>(syntax->getValue())->getValue();
     return getCppType();
   }
   llvm::PrintFatalError("Parameters DAG arguments must be either strings or "
                         "defs which inherit from AttrOrTypeParameter");
 }
+
+const llvm::Init *AttrOrTypeParameter::getDef() const {
+  return def->getArg(index);
+}
+
+//===----------------------------------------------------------------------===//
+// AttributeSelfTypeParameter
+//===----------------------------------------------------------------------===//
+
+bool AttributeSelfTypeParameter::classof(const AttrOrTypeParameter *param) {
+  const llvm::Init *paramDef = param->getDef();
+  if (auto *paramDefInit = dyn_cast<llvm::DefInit>(paramDef))
+    return paramDefInit->getDef()->isSubClassOf("AttributeSelfTypeParameter");
+  return false;
+}

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

let parameters = (
     // constructor.
     ArrayRefParameter<
       "int", // The parameter C++ type.
       "An example of an array of ints" // Parameter description.
       >: $arrayOfInts
   );
 }
 
+// An attribute testing AttributeSelfTypeParameter.
+def AttrWithSelfTypeParam : Test_Attr<"AttrWithSelfTypeParam"> {
+  let mnemonic = "attr_with_self_type_param";
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type);
+}
+
+// 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()";
+}
+
 #endif // TEST_ATTRDEFS

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

 #include "mlir/IR/Types.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::test;
 
+//===----------------------------------------------------------------------===//
+// AttrWithSelfTypeParamAttr
+//===----------------------------------------------------------------------===//
+
+Attribute AttrWithSelfTypeParamAttr::parse(MLIRContext *context,
+                                           DialectAsmParser &parser,
+                                           Type type) {
+  Type selfType;
+  if (parser.parseType(selfType))
+    return Attribute();
+  return get(context, selfType);
+}
+
+void AttrWithSelfTypeParamAttr::print(DialectAsmPrinter &printer) const {
+  printer << "attr_with_self_type_param " << getType();
+}
+
+//===----------------------------------------------------------------------===//
+// AttrWithTypeBuilderAttr
+//===----------------------------------------------------------------------===//
+
+Attribute AttrWithTypeBuilderAttr::parse(MLIRContext *context,
+                                         DialectAsmParser &parser, Type type) {
+  IntegerAttr element;
+  if (parser.parseAttribute(element))
+    return Attribute();
+  return get(context, element);
+}
+
+void AttrWithTypeBuilderAttr::print(DialectAsmPrinter &printer) const {
+  printer << "attr_with_type_builder " << getAttr();
+}
+
+//===----------------------------------------------------------------------===//
+// CompoundAAttr
+//===----------------------------------------------------------------------===//
+
 Attribute CompoundAAttr::parse(MLIRContext *context, DialectAsmParser &parser,
                                Type type) {
   int widthOfSomething;
   Type oneType;
   SmallVector<int, 4> arrayOfInts;
   if (parser.parseLess() || parser.parseInteger(widthOfSomething) ||
       parser.parseComma() || parser.parseType(oneType) || parser.parseComma() ||
       parser.parseLSquare())

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

def OperandZeroAndResultHaveSameElementType : TEST_Op<
   let results = (outs AnyType:$res);
 }
 
 def OperandsHaveSameType :
     TEST_Op<"operands_have_same_type", [AllTypesMatch<["x", "y"]>]> {
   let arguments = (ins AnyType:$x, AnyType:$y);
 }
 
+def ResultHasSameTypeAsAttr :
+    TEST_Op<"result_has_same_type_as_attr",
+            [AllTypesMatch<["attr", "result"]>]> {
+  let arguments = (ins AnyAttr:$attr);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "$attr `->` type($result) attr-dict";
+}
+
 def OperandZeroAndResultHaveSameType :
     TEST_Op<"operand0_and_result_have_same_type",
             [AllTypesMatch<["x", "res"]>]> {
   let arguments = (ins AnyType:$x, AnyType:$y);
   let results = (outs AnyType:$res);
 }
 
 def OperandsHaveSameRank :

mlir/test/mlir-tblgen/attrdefs.td

 // DEF: ::mlir::test::IndexAttr,
 // DEF: ::mlir::test::SingleParameterAttr
 
 // DEF-LABEL: ::mlir::Attribute generatedAttributeParser(::mlir::MLIRContext *context,
 // DEF-NEXT: ::mlir::DialectAsmParser &parser,
 // DEF-NEXT: ::llvm::StringRef mnemonic, ::mlir::Type type) {
 // DEF: if (mnemonic == ::mlir::test::CompoundAAttr::getMnemonic()) return ::mlir::test::CompoundAAttr::parse(context, parser, type);
 // DEF-NEXT: if (mnemonic == ::mlir::test::IndexAttr::getMnemonic()) return ::mlir::test::IndexAttr::parse(context, parser, type);
-// DEF-NEXT: return ::mlir::Attribute();
+// DEF: return ::mlir::Attribute();
 
 def Test_Dialect: Dialect {
 // DECL-NOT: TestDialect
 // DEF-NOT: TestDialect
     let name = "TestDialect";
     let cppNamespace = "::mlir::test";
 }
 
def B_CompoundAttrA : TestAttr<"CompoundA"> {
   let description = "This attribute is to test a reasonably complex attribute";
   let mnemonic = "cmpnd_a";
   let parameters = (
       ins
       "int":$widthOfSomething,
       "::mlir::test::SimpleTypeA": $exampleTdType,
       "SomeCppStruct": $exampleCppType,
       ArrayRefParameter<"int", "Matrix dimensions">:$dims,
-      "::mlir::Type":$inner
+      AttributeSelfTypeParameter<"">:$inner
   );
 
   let genVerifyDecl = 1;
 
 // DECL-LABEL: class CompoundAAttr : public ::mlir::Attribute
 // DECL: static CompoundAAttr getChecked(llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, ::mlir::MLIRContext *context, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, SomeCppStruct exampleCppType, ::llvm::ArrayRef<int> dims, ::mlir::Type inner);
 // DECL: static ::mlir::LogicalResult verify(::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, SomeCppStruct exampleCppType, ::llvm::ArrayRef<int> dims, ::mlir::Type inner);
 // DECL: static constexpr ::llvm::StringLiteral getMnemonic() {
 // DECL:   return ::llvm::StringLiteral("cmpnd_a");
 // DECL: }
 // DECL: static ::mlir::Attribute parse(::mlir::MLIRContext *context,
 // DECL-NEXT: ::mlir::DialectAsmParser &parser, ::mlir::Type type);
 // DECL: void print(::mlir::DialectAsmPrinter &printer) const;
 // DECL: int getWidthOfSomething() const;
 // DECL: ::mlir::test::SimpleTypeA getExampleTdType() const;
 // DECL: SomeCppStruct getExampleCppType() const;
+
+// Check that AttributeSelfTypeParameter is handled properly.
+// DEF-LABEL: struct CompoundAAttrStorage
+// DEF: CompoundAAttrStorage (
+// DEF-NEXT: : ::mlir::AttributeStorage(inner),
+
+// DEF: bool operator==(const KeyTy &key) const {
+// DEF-NEXT: return key == KeyTy(widthOfSomething, exampleTdType, exampleCppType, dims, getType());
+
+// DEF: static CompoundAAttrStorage *construct
+// DEF: return new (allocator.allocate<CompoundAAttrStorage>())
+// DEF-NEXT: CompoundAAttrStorage(widthOfSomething, exampleTdType, exampleCppType, dims, inner);
+
+// DEF: ::mlir::Type CompoundAAttr::getInner() const { return getImpl()->getType(); }
 }
 
 def C_IndexAttr : TestAttr<"Index"> {
     let mnemonic = "index";
 
     let parameters = (
       ins
       StringRefParameter<"Label for index">:$label
def D_SingleParameterAttr : TestAttr<"SingleParameter"> {
   let parameters = (
     ins
     "int": $num
   );
 // DECL-LABEL: struct SingleParameterAttrStorage;
 // DECL-LABEL: class SingleParameterAttr
 // DECL-NEXT:                   detail::SingleParameterAttrStorage
 }
+
+// An attribute testing AttributeSelfTypeParameter.
+def E_AttrWithTypeBuilder : TestAttr<"AttrWithTypeBuilder"> {
+  let mnemonic = "attr_with_type_builder";
+  let parameters = (ins "::mlir::IntegerAttr":$attr);
+  let typeBuilder = "$_attr.getType()";
+}
+
+// DEF-LABEL: struct AttrWithTypeBuilderAttrStorage
+// DEF: AttrWithTypeBuilderAttrStorage (::mlir::IntegerAttr attr)
+// DEF-NEXT: : ::mlir::AttributeStorage(attr.getType()), attr(attr)

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 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
+%b = test.result_has_same_type_as_attr #test<"attr_with_type_builder 10 : i16"> -> i16

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

for (AttrOrTypeParameter &parameter : parameters) {
     if (Optional<StringRef> allocCode = parameter.getAllocator()) {
       fmtCtxt.withSelf(parameter.getName());
       fmtCtxt.addSubst("_dst", parameter.getName());
       os << "      " << tgfmt(*allocCode, &fmtCtxt) << "\n";
     }
   }
 }
 
+/// Builds a code block that initializes the attribute storage of 'def'.
+/// Attribute initialization is separated from Type initialization given that
+/// the Attribute also needs to initialize its self-type, which has multiple
+/// means of initialization.
+static std::string buildAttributeStorageParamInitializer(
+    const AttrOrTypeDef &def, ArrayRef<AttrOrTypeParameter> parameters) {
+  std::string paramInitializer;
+  llvm::raw_string_ostream paramOS(paramInitializer);
+  paramOS << "::mlir::AttributeStorage(";
+
+  // If this is an attribute, we need to check for value type initialization.
+  Optional<size_t> selfParamIndex;
+  for (auto it : llvm::enumerate(parameters)) {
+    const auto *selfParam = dyn_cast<AttributeSelfTypeParameter>(&it.value());
+    if (!selfParam)
+      continue;
+    if (selfParamIndex) {
+      llvm::PrintFatalError(def.getLoc(),
+                            "Only one attribute parameter can be marked as "
+                            "AttributeSelfTypeParameter");
+    }
+    paramOS << selfParam->getName();
+    selfParamIndex = it.index();
+  }
+
+  // If we didn't find a self param, but the def has a type builder we use that
+  // to construct the type.
+  if (!selfParamIndex) {
+    const AttrDef &attrDef = cast<AttrDef>(def);
+    if (Optional<StringRef> typeBuilder = attrDef.getTypeBuilder()) {
+      FmtContext fmtContext;
+      for (const AttrOrTypeParameter &param : parameters)
+        fmtContext.addSubst(("_" + param.getName()).str(), param.getName());
+      paramOS << tgfmt(*typeBuilder, &fmtContext);
+    }
+  }
+  paramOS << ")";
+
+  // Append the parameters to the initializer.
+  for (auto it : llvm::enumerate(parameters))
+    if (it.index() != selfParamIndex)
+      paramOS << llvm::formatv(", {0}({0})", it.value().getName());
+
+  return paramOS.str();
+}
+
 void DefGenerator::emitStorageClass(const AttrOrTypeDef &def) {
-  SmallVector<AttrOrTypeParameter, 4> parameters;
-  def.getParameters(parameters);
+  SmallVector<AttrOrTypeParameter, 4> params;
+  def.getParameters(params);
 
-  // Collect the parameter names and types.
-  auto parameterNames =
-      map_range(parameters, [](AttrOrTypeParameter parameter) {
-        return parameter.getName();
-      });
+  // Collect the parameter types.
   auto parameterTypes =
-      map_range(parameters, [](AttrOrTypeParameter parameter) {
+      llvm::map_range(params, [](const AttrOrTypeParameter &parameter) {
         return parameter.getCppType();
       });
-  auto parameterList = join(parameterNames, ", ");
-  auto parameterTypeList = join(parameterTypes, ", ");
+  std::string parameterTypeList = llvm::join(parameterTypes, ", ");
+
+  // Collect the parameter initializer.
+  std::string paramInitializer;
+  if (isAttrGenerator) {
+    paramInitializer = buildAttributeStorageParamInitializer(def, params);
+
+  } else {
+    llvm::raw_string_ostream initOS(paramInitializer);
+    llvm::interleaveComma(params, initOS, [&](const AttrOrTypeParameter &it) {
+      initOS << llvm::formatv("{0}({0})", it.getName());
+    });
+  }
+
+  // Construct the parameter list that is used when a concrete instance of the
+  // storage exists.
+  auto nonStaticParameterNames = llvm::map_range(params, [](const auto &param) {
+    return isa<AttributeSelfTypeParameter>(param) ? "getType()"
+                                                  : param.getName();
+  });
 
   // 1) Emit most of the storage class up until the hashKey body.
   os << formatv(
       defStorageClassBeginStr, def.getStorageNamespace(),
       def.getStorageClassName(),
       ParamCommaFormatter(ParamCommaFormatter::EmitFormat::TypeNamePairs,
-                          parameters, /*prependComma=*/false),
-      ParamCommaFormatter(ParamCommaFormatter::EmitFormat::TypeNameInitializer,
-                          parameters, /*prependComma=*/false),
-      parameterList, parameterTypeList, valueType);
+                          params, /*prependComma=*/false),
+      paramInitializer, llvm::join(nonStaticParameterNames, ", "),
+      parameterTypeList, valueType);
 
   // 2) Emit the haskKey method.
   os << "  static ::llvm::hash_code hashKey(const KeyTy &key) {\n";
 
   // Extract each parameter from the key.
   os << "      return ::llvm::hash_combine(";
   llvm::interleaveComma(
-      llvm::seq<unsigned>(0, parameters.size()), os,
+      llvm::seq<unsigned>(0, params.size()), os,
       [&](unsigned it) { os << "std::get<" << it << ">(key)"; });
   os << ");\n    }\n";
 
   // 3) Emit the construct method.
 
   // If user wants to build the storage constructor themselves, declare it
   // here and then they can write the definition elsewhere.
   if (def.hasStorageCustomConstructor()) {
     os << llvm::formatv("    static {0} *construct(::mlir::{1}StorageAllocator "
                         "&allocator, const KeyTy &key);\n",
                         def.getStorageClassName(), valueType);
 
     // Otherwise, generate one.
   } else {
     // First, unbox the parameters.
     os << formatv(defStorageClassConstructorBeginStr, def.getStorageClassName(),
                   valueType);
-    for (unsigned i = 0, e = parameters.size(); i < e; ++i) {
+    for (unsigned i = 0, e = params.size(); i < e; ++i) {
       os << formatv("      auto {0} = std::get<{1}>(key);\n",
-                    parameters[i].getName(), i);
+                    params[i].getName(), i);
     }
 
     // Second, reassign the parameter variables with allocation code, if it's
     // specified.
     emitStorageParameterAllocation(def, os);
 
     // Last, return an allocated copy.
+    auto parameterNames = llvm::map_range(
+        params, [](const auto &param) { return param.getName(); });
     os << formatv(defStorageClassConstructorEndStr, def.getStorageClassName(),
-                  parameterList);
+                  llvm::join(parameterNames, ", "));
   }
 
   // 4) Emit the parameters as storage class members.
-  for (auto parameter : parameters) {
+  for (const AttrOrTypeParameter &parameter : params) {
+    // Attribute value types are not stored as fields in the storage.
+    if (!isa<AttributeSelfTypeParameter>(parameter))
       os << "      " << parameter.getCppType() << " " << parameter.getName()
          << ";\n";
   }
   os << "  };\n";
 
   os << "} // namespace " << def.getStorageNamespace() << "\n";
 }
 
 void DefGenerator::emitParsePrint(const AttrOrTypeDef &def) {
   // Emit the printer code, if specified.
if (!parameters.empty()) {
 
     // Emit the builders for this def.
     emitBuilderDefs(def, os, parameters);
 
     // Generate accessor definitions only if we also generate the storage class.
     // Otherwise, let the user define the exact accessor definition.
     if (def.genAccessors() && def.genStorageClass()) {
       for (const AttrOrTypeParameter &parameter : parameters) {
+        StringRef paramStorageName = isa<AttributeSelfTypeParameter>(parameter)
+                                         ? "getType()"
+                                         : parameter.getName();
+
         SmallString<16> name = parameter.getName();
         name[0] = llvm::toUpper(name[0]);
         os << formatv("{0} {3}::get{1}() const {{ return getImpl()->{2}; }\n",
-                      parameter.getCppType(), name, parameter.getName(),
+                      parameter.getCppType(), name, paramStorageName,
                       def.getCppClassName());
       }
     }
   }
 
   // If mnemonic is specified maybe print definitions for the parser and printer
   // code, if they're specified.
   if (def.getMnemonic())