[mlir][AttrTypeDefGen] Add support for custom parameter comparators

[rriddle]

Summary

Some parameters to attributes and types rely on special comparison routines other than operator== to ensure equality. This revision adds support for those parameters by allowing them to specify a comparator code block that determines if $_lhs and $_rhs are equal. An example of one of these paramters is APFloat, which requires bitwiseIsEqual for bitwise comparison (which we want for attribute equality).

mlir/docs/OpDefinitions.md

 -   Set `hasCustomStorageConstructor` to generate the TypeStorage class with a
     constructor which is just declared -- no definition -- so you can write it
     yourself.
 -   Use the `TypeParameter` tablegen class instead of the "c++Type" string.
 
 ### TypeParameter tablegen class
 
 This is used to further specify attributes about each of the types parameters.
-It includes documentation (`summary` and `syntax`), the C++ type to use, and a
-custom allocator to use in the storage constructor method.
+It includes documentation (`summary` and `syntax`), the C++ type to use, a
+custom allocator to use in the storage constructor method, and a custom
+comparator to decide if two instances of the parameter type are equal.
 
 ```tablegen
 // DO NOT DO THIS!
 let parameters = (ins "ArrayRef<int>":$dims);
 ```
 
 The default storage constructor blindly copies fields by value. It does not know
 anything about the types. In this case, the ArrayRef<int> requires allocation

 let parameters = (ins ArrayRefIntParam:$dims);
 ```
 
 The `allocator` code block has the following substitutions:
 
 -   `$_allocator` is the TypeStorageAllocator in which to allocate objects.
 -   `$_dst` is the variable in which to place the allocated data.
 
+The `comparator` code block has the following substitutions:
+
+-   `$_lhs` is an instance of the parameter type.
+-   `$_rhs` is an instance of the parameter type.
+
 MLIR includes several specialized classes for common situations:
 
 -   `StringRefParameter<descriptionOfParam>` for StringRefs.
 -   `ArrayRefParameter<arrayOf, descriptionOfParam>` for ArrayRefs of value
     types
 -   `SelfAllocationParameter<descriptionOfParam>` for C++ classes which contain
     a method called `allocateInto(StorageAllocator &allocator)` to allocate
     itself into `allocator`.

mlir/include/mlir/IR/OpBase.td

let predicate = CPred<"$_self.isa<" # dialect.cppNamespace #
                                  "::" # cppClassName # ">()">;
 }
 
 // 'Parameters' should be subclasses of this or simple strings (which is a
 // shorthand for AttrOrTypeParameter<"C++Type">).
 class AttrOrTypeParameter<string type, string desc> {
   // Custom memory allocation code for storage constructor.
   code allocator = ?;
+  // Custom comparator used to compare two instances for equality.
+  code comparator = ?;
   // The C++ type of this parameter.
   string cppType = type;
   // One-line human-readable description of the argument.
   string summary = desc;
   // The format string for the asm syntax (documentation only).
   string syntax = ?;
 }
 class AttrParameter<string type, string desc> : AttrOrTypeParameter<type, desc>;
 class TypeParameter<string type, string desc> : AttrOrTypeParameter<type, desc>;
 
 // For StringRefs, which require allocation.
 class StringRefParameter<string desc = ""> :
     AttrOrTypeParameter<"::llvm::StringRef", desc> {
   let allocator = [{$_dst = $_allocator.copyInto($_self);}];
 }
 
+// For APFloats, which require comparison.
+class APFloatParameter<string desc> :
+    AttrOrTypeParameter<"::llvm::APFloat", desc> {
+  let comparator = "$_lhs.bitwiseIsEqual($_rhs)";
+}
+
 // For standard ArrayRefs, which require allocation.
 class ArrayRefParameter<string arrayOf, string desc = ""> :
     AttrOrTypeParameter<"::llvm::ArrayRef<" # arrayOf # ">", desc> {
   let allocator = [{$_dst = $_allocator.copyInto($_self);}];
 }
 
 // For classes which require allocation and have their own allocateInto method.
 class SelfAllocationParameter<string type, string desc> :

mlir/include/mlir/TableGen/AttrOrTypeDef.h

explicit AttrOrTypeParameter(const llvm::DagInit *def, unsigned index)
       : def(def), index(index) {}
 
   // Get the parameter name.
   StringRef getName() const;
 
   // If specified, get the custom allocator code for this parameter.
   Optional<StringRef> getAllocator() const;
 
+  // If specified, get the custom comparator code for this parameter.
+  Optional<StringRef> getComparator() const;
+
   // Get the C++ type of this parameter.
   StringRef getCppType() const;
 
   // Get a description of this parameter for documentation purposes.
   Optional<StringRef> getSummary() const;
 
   // Get the assembly syntax documentation.
   StringRef getSyntax() const;

mlir/lib/TableGen/AttrOrTypeDef.cpp

 StringRef AttrOrTypeParameter::getName() const {
   return def->getArgName(index)->getValue();
 }
 
 Optional<StringRef> AttrOrTypeParameter::getAllocator() const {
   llvm::Init *parameterType = def->getArg(index);
   if (isa<llvm::StringInit>(parameterType))
     return Optional<StringRef>();
-
-  if (auto *param = dyn_cast<llvm::DefInit>(parameterType)) {
-    llvm::RecordVal *code = param->getDef()->getValue("allocator");
-    if (!code)
-      return Optional<StringRef>();
-    if (llvm::StringInit *ci = dyn_cast<llvm::StringInit>(code->getValue()))
-      return ci->getValue();
-    if (isa<llvm::UnsetInit>(code->getValue()))
-      return Optional<StringRef>();
-
-    llvm::PrintFatalError(
-        param->getDef()->getLoc(),
-        "Record `" + def->getArgName(index)->getValue() +
-            "', field `printer' does not have a code initializer!");
-  }
+  if (auto *param = dyn_cast<llvm::DefInit>(parameterType))
+    return param->getDef()->getValueAsOptionalString("allocator");
+  llvm::PrintFatalError("Parameters DAG arguments must be either strings or "
+                        "defs which inherit from AttrOrTypeParameter\n");
 
+Optional<StringRef> AttrOrTypeParameter::getComparator() const {
+  llvm::Init *parameterType = def->getArg(index);
+  if (isa<llvm::StringInit>(parameterType))
+    return Optional<StringRef>();
+  if (auto *param = dyn_cast<llvm::DefInit>(parameterType))
+    return param->getDef()->getValueAsOptionalString("comparator");
   llvm::PrintFatalError("Parameters DAG arguments must be either strings or "
                         "defs which inherit from AttrOrTypeParameter\n");
 }
 
 StringRef AttrOrTypeParameter::getCppType() const {
   auto *parameterType = def->getArg(index);
   if (auto *stringType = dyn_cast<llvm::StringInit>(parameterType))
     return stringType->getValue();

mlir/test/mlir-tblgen/attrdefs.td

 def B_CompoundAttrA : TestAttr<"CompoundA"> {
   let summary = "A more complex parameterized attribute";
   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,
+      APFloatParameter<"">: $apFloat,
       ArrayRefParameter<"int", "Matrix dimensions">:$dims,
       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 CompoundAAttr getChecked(llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, ::mlir::MLIRContext *context, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, ::llvm::APFloat apFloat, ::llvm::ArrayRef<int> dims, ::mlir::Type inner);
+// DECL: static ::mlir::LogicalResult verify(::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, int widthOfSomething, ::mlir::test::SimpleTypeA exampleTdType, ::llvm::APFloat apFloat, ::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;
+// DECL: ::llvm::APFloat getApFloat() 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-NEXT: if (!(widthOfSomething == std::get<0>(key)))
+// DEF-NEXT:   return false;
+// DEF-NEXT: if (!(exampleTdType == std::get<1>(key)))
+// DEF-NEXT:   return false;
+// DEF-NEXT: if (!(apFloat.bitwiseIsEqual(std::get<2>(key))))
+// DEF-NEXT:   return false;
+// DEF-NEXT: if (!(dims == std::get<3>(key)))
+// DEF-NEXT:   return false;
+// DEF-NEXT: if (!(getType() == std::get<4>(key)))
+// DEF-NEXT:   return false;
+// DEF-NEXT: return true;
 
 // DEF: static CompoundAAttrStorage *construct
 // DEF: return new (allocator.allocate<CompoundAAttrStorage>())
-// DEF-NEXT: CompoundAAttrStorage(widthOfSomething, exampleTdType, exampleCppType, dims, inner);
+// DEF-NEXT: CompoundAAttrStorage(widthOfSomething, exampleTdType, apFloat, dims, inner);
 
 // DEF: ::mlir::Type CompoundAAttr::getInner() const { return getImpl()->getType(); }
 }
 
 def C_IndexAttr : TestAttr<"Index"> {
     let mnemonic = "index";
 
     let parameters = (

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

static ::mlir::LogicalResult generated{0}Printer(
   return ::llvm::TypeSwitch<::mlir::{0}, ::mlir::LogicalResult>(def)
 )";
 
 /// Beginning of storage class.
 /// {0}: Storage class namespace.
 /// {1}: Storage class c++ name.
 /// {2}: Parameters parameters.
 /// {3}: Parameter initializer string.
-/// {4}: Parameter name list.
-/// {5}: Parameter types.
-/// {6}: The name of the base value type, e.g. Attribute or Type.
+/// {4}: Parameter types.
+/// {5}: The name of the base value type, e.g. Attribute or Type.
 static const char *const defStorageClassBeginStr = R"(
 namespace {0} {{
-  struct {1} : public ::mlir::{6}Storage {{
+  struct {1} : public ::mlir::{5}Storage {{
     {1} ({2})
       : {3} {{ }
 
     /// The hash key is a tuple of the parameter types.
-    using KeyTy = std::tuple<{5}>;
-
-    /// Define the comparison function for the key type.
-    bool operator==(const KeyTy &key) const {{
-      return key == KeyTy({4});
-    }
+    using KeyTy = std::tuple<{4}>;
 )";
 
 /// The storage class' constructor template.
 ///
 /// {0}: storage class name.
 /// {1}: The name of the base value type, e.g. Attribute or Type.
 static const char *const defStorageClassConstructorBeginStr = R"(
     /// Define a construction method for creating a new instance of this
void DefGenerator::emitStorageClass(const AttrOrTypeDef &def) {
 
   } 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.
+  // * Emit most of the storage class up until the hashKey body.
   os << formatv(
       defStorageClassBeginStr, def.getStorageNamespace(),
       def.getStorageClassName(),
       ParamCommaFormatter(ParamCommaFormatter::EmitFormat::TypeNamePairs,
                           params, /*prependComma=*/false),
-      paramInitializer, llvm::join(nonStaticParameterNames, ", "),
-      parameterTypeList, valueType);
+      paramInitializer, parameterTypeList, valueType);
+
+  // * Emit the comparison method.
+  os << "  bool operator==(const KeyTy &key) const {\n";
+  for (auto it : llvm::enumerate(params)) {
+    os << "    if (!(";
+
+    // Build the comparator context.
+    bool isSelfType = isa<AttributeSelfTypeParameter>(it.value());
+    FmtContext context;
+    context.addSubst("_lhs", isSelfType ? "getType()" : it.value().getName())
+        .addSubst("_rhs", "std::get<" + Twine(it.index()) + ">(key)");
+
+    // Use the parameter specified comparator if possible, otherwise default to
+    // operator==.
+    Optional<StringRef> comparator = it.value().getComparator();
+    os << tgfmt(comparator ? *comparator : "$_lhs == $_rhs", &context);
+    os << "))\n      return false;\n";
+  }
+  os << "    return true;\n  }\n";
 
-  // 2) Emit the haskKey method.
+  // * 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, params.size()), os,
       [&](unsigned it) { os << "std::get<" << it << ">(key)"; });
   os << ");\n    }\n";
 
-  // 3) Emit the construct method.
+  // * 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);
 
if (def.hasStorageCustomConstructor()) {
 
     // Last, return an allocated copy.
     auto parameterNames = llvm::map_range(
         params, [](const auto &param) { return param.getName(); });
     os << formatv(defStorageClassConstructorEndStr, def.getStorageClassName(),
                   llvm::join(parameterNames, ", "));
   }
 
-  // 4) Emit the parameters as storage class members.
+  // * Emit the parameters as storage class members.
   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";