[MLIR][TableGen] Error on APInt parameter without custom comparator (llvm#135970)

The error is triggered when an attribute or type uses an APInt typed
parameter with the generated equality operator. If the compared APInts
have different bit widths the equality operator triggers an assert. This
is dangerous, since `StorageUniquer` for types and attributes uses the
equality operator when a hash collision appears. As such, it is
necessary to use custom provided comarator or `APIntParameter` that
already has it.
This commit also replaces uses of the raw `APInt` parameter with the
`APIntParameter` and removes the no longer necessary custom StorageClass
for the `BitVectorAttr` from the SMT dialect that was a workaround for
the described issue.

---------

Co-authored-by: Tobias Gysi <tobias.gysi@nextsilicon.com>

Author: Jezurko

mlir/include/mlir/Dialect/SMT/IR/SMTAttributes.td

@@ -45,21 +45,11 @@ def BitVectorAttr : AttrDef<SMTDialect, "BitVector", [
     present).
   }];
 
-  let parameters = (ins "llvm::APInt":$value);
+  let parameters = (ins APIntParameter<"">:$value);
 
   let hasCustomAssemblyFormat = true;
   let genVerifyDecl = true;
 
-  // We need to manually define the storage class because the generated one is
-  // buggy (because the APInt asserts matching bitwidth in the `==` operator and
-  // the generated storage uses that directly.
-  // Alternatively: add a type parameter to redundantly store the bitwidth of
-  // of the attribute type, it it's in the order before the 'value' it will be
-  // checked before the APInt equality (this is the reason it works for the
-  // builtin integer attribute), but would be more fragile (and we'd store
-  // duplicate data).
-  let genStorageClass = false;
-
   let builders = [
     AttrBuilder<(ins "llvm::StringRef":$value)>,
     AttrBuilder<(ins "uint64_t":$value, "unsigned":$width)>,

mlir/include/mlir/IR/BuiltinAttributes.td

@@ -700,7 +700,7 @@ def Builtin_IntegerAttr : Builtin_Attr<"Integer", "integer",
     false // A bool, i.e. i1, value.
     ```
   }];
-  let parameters = (ins AttributeSelfTypeParameter<"">:$type, "APInt":$value);
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type, APIntParameter<"">:$value);
   let builders = [
     AttrBuilderWithInferredContext<(ins "Type":$type,
                                         "const APInt &":$value), [{

mlir/include/mlir/TableGen/AttrOrTypeDef.h

@@ -68,7 +68,11 @@ class AttrOrTypeParameter {
   /// If specified, get the custom allocator code for this parameter.
   std::optional<StringRef> getAllocator() const;
 
-  /// If specified, get the custom comparator code for this parameter.
+  /// Return true if user defined comparator is specified.
+  bool hasCustomComparator() const;
+
+  /// Get the custom comparator code for this parameter or fallback to the
+  /// default.
   StringRef getComparator() const;
 
   /// Get the C++ type of this parameter.

mlir/lib/Dialect/SMT/IR/SMTAttributes.cpp

@@ -21,42 +21,6 @@ using namespace mlir::smt;
 // BitVectorAttr
 //===----------------------------------------------------------------------===//
 
-namespace mlir {
-namespace smt {
-namespace detail {
-struct BitVectorAttrStorage : public mlir::AttributeStorage {
-  using KeyTy = APInt;
-  BitVectorAttrStorage(APInt value) : value(std::move(value)) {}
-
-  KeyTy getAsKey() const { return value; }
-
-  // NOTE: the implementation of this operator is the reason we need to define
-  // the storage manually. The auto-generated version would just do the direct
-  // equality check of the APInt, but that asserts the bitwidth of both to be
-  // the same, leading to a crash. This implementation, therefore, checks for
-  // matching bit-width beforehand.
-  bool operator==(const KeyTy &key) const {
-    return (value.getBitWidth() == key.getBitWidth() && value == key);
-  }
-
-  static llvm::hash_code hashKey(const KeyTy &key) {
-    return llvm::hash_value(key);
-  }
-
-  static BitVectorAttrStorage *
-  construct(mlir::AttributeStorageAllocator &allocator, KeyTy &&key) {
-    return new (allocator.allocate<BitVectorAttrStorage>())
-        BitVectorAttrStorage(std::move(key));
-  }
-
-  APInt value;
-};
-} // namespace detail
-} // namespace smt
-} // namespace mlir
-
-APInt BitVectorAttr::getValue() const { return getImpl()->value; }
-
 LogicalResult BitVectorAttr::verify(
     function_ref<InFlightDiagnostic()> emitError,
     APInt value) { // NOLINT(performance-unnecessary-value-param)

mlir/lib/TableGen/AttrOrTypeDef.cpp

@@ -278,6 +278,10 @@ std::optional<StringRef> AttrOrTypeParameter::getAllocator() const {
   return getDefValue<StringInit>("allocator");
 }
 
+bool AttrOrTypeParameter::hasCustomComparator() const {
+  return getDefValue<StringInit>("comparator").has_value();
+}
+
 StringRef AttrOrTypeParameter::getComparator() const {
   return getDefValue<StringInit>("comparator").value_or("$_lhs == $_rhs");
 }

mlir/test/mlir-tblgen/apint-param-error.td

@@ -0,0 +1,17 @@
+// RUN: not mlir-tblgen -gen-attrdef-decls -I %S/../../include %s 2>&1 | FileCheck %s
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/OpBase.td"
+
+def Test_Dialect: Dialect {
+  let name = "TestDialect";
+  let cppNamespace = "::test";
+}
+
+def RawAPIntAttr : AttrDef<Test_Dialect, "RawAPInt"> {
+  let mnemonic = "raw_ap_int";
+  let parameters = (ins "APInt":$value);
+  let hasCustomAssemblyFormat = 1;
+}
+
+// CHECK: apint-param-error.td:11:5: error: Using a raw APInt parameter

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

@@ -678,8 +678,18 @@ void DefGen::emitStorageClass() {
   emitConstruct();
   // Emit the storage class members as public, at the very end of the struct.
   storageCls->finalize();
-  for (auto &param : params)
+  for (auto &param : params) {
+    if (param.getCppType().contains("APInt") && !param.hasCustomComparator()) {
+      PrintFatalError(
+          def.getLoc(),
+          "Using a raw APInt parameter without a custom comparator is "
+          "not supported because an assert in the equality operator is "
+          "triggered when the two APInts have different bit widths. This can "
+          "lead to unexpected crashes. Use an `APIntParameter` or "
+          "provide a custom comparator.");
+    }
     storageCls->declare<Field>(param.getCppType(), param.getName());
+  }
 }
 
 //===----------------------------------------------------------------------===//