[MLIR][TableGen] Error on APInt parameter without custom comparator

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-warn.td

@@ -0,0 +1,17 @@
+// RUN: 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-warn.td:11:5: warning: 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()) {
+ PrintWarning(
+ 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());
+ }
 }
 
 //===----------------------------------------------------------------------===//