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

This reverts commit 4bcc414.

Author: raikonenfnu

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

@@ -45,11 +45,21 @@ def BitVectorAttr : AttrDef<SMTDialect, "BitVector", [
     present).
   }];
 
-  let parameters = (ins APIntParameter<"">:$value);
+  let parameters = (ins "llvm::APInt":$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, APIntParameter<"">:$value);
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type, "APInt":$value);
   let builders = [
     AttrBuilderWithInferredContext<(ins "Type":$type,
                                         "const APInt &":$value), [{

mlir/include/mlir/TableGen/AttrOrTypeDef.h

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

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

@@ -21,6 +21,42 @@ 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,10 +278,6 @@ 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

@@ -678,18 +678,8 @@ void DefGen::emitStorageClass() {
   emitConstruct();
   // Emit the storage class members as public, at the very end of the struct.
   storageCls->finalize();
-  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.");
-    }
+  for (auto &param : params)
     storageCls->declare<Field>(param.getCppType(), param.getName());
-  }
 }
 
 //===----------------------------------------------------------------------===//