[MLIR] Validate APInt bitwidth in IntegerAttr::get(Type, APInt)

[joker-eph]

IntegerAttr::get(Type, APInt) did not validate that the APInt's bit width matched the expected bit width for the given type. For integer types, the APInt width must equal the integer type's width. For index types, the APInt width must equal IndexType::kInternalStorageBitWidth (64 bits).

Passing an APInt with the wrong bit width could cause a non-deterministic crash in StorageUniquer when comparing two IntegerAttr instances for the same type but with different APInt widths.

This commit adds assertions in the get(Type, APInt) builder to catch such misuse early in debug builds, providing a clear error message at the call site rather than a cryptic crash in the storage uniquer.

Fixes #56401

Assisted-by: Claude Code

[llvmbot]

@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir-ods

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

IntegerAttr::get(Type, APInt) did not validate that the APInt's bit width matched the expected bit width for the given type. For integer types, the APInt width must equal the integer type's width. For index types, the APInt width must equal IndexType::kInternalStorageBitWidth (64 bits).

Passing an APInt with the wrong bit width could cause a non-deterministic crash in StorageUniquer when comparing two IntegerAttr instances for the same type but with different APInt widths.

This commit adds assertions in the get(Type, APInt) builder to catch such misuse early in debug builds, providing a clear error message at the call site rather than a cryptic crash in the storage uniquer.

Fixes #56401

Assisted-by: Claude Code

---

Full diff: https://github.com/llvm/llvm-project/pull/188725.diff

2 Files Affected:

• (modified) mlir/include/mlir/IR/BuiltinAttributes.td (+9)
• (modified) mlir/unittests/IR/AttributeTest.cpp (+42)

diff --git a/mlir/include/mlir/IR/BuiltinAttributes.td b/mlir/include/mlir/IR/BuiltinAttributes.td
index 0cc556ef5d852..a0d72d082b5fb 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.td
+++ b/mlir/include/mlir/IR/BuiltinAttributes.td
@@ -720,6 +720,15 @@ def Builtin_IntegerAttr : Builtin_Attr<"Integer", "integer",
                                         "const APInt &":$value), [{
       if (type.isSignlessInteger(1))
         return BoolAttr::get(type.getContext(), value.getBoolValue());
+      // Validate that the APInt has the correct bit width for the given type.
+      if (auto intTy = ::llvm::dyn_cast<IntegerType>(type)) {
+        assert(value.getBitWidth() == intTy.getWidth() &&
+               "IntegerAttr::get: APInt bit width must match integer type width");
+      } else if (::llvm::isa<IndexType>(type)) {
+        assert(value.getBitWidth() == IndexType::kInternalStorageBitWidth &&
+               "IntegerAttr::get: APInt bit width must match IndexType internal "
+               "storage bit width");
+      }
       return $_get(type.getContext(), type, value);
     }]>,
     AttrBuilder<(ins "const APSInt &":$value), [{
diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index 404aa8c0dcf3d..900cacabd592e 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -523,4 +523,46 @@ TEST(CopyCountAttr, PrintStripped) {
   EXPECT_EQ(str, "|#test.copy_count<hello>|[copy_count<hello>]");
 }
 
+//===----------------------------------------------------------------------===//
+// IntegerAttr
+//===----------------------------------------------------------------------===//
+
+TEST(IntegerAttrTest, CorrectBitWidths) {
+  MLIRContext context;
+
+  // Correct APInt width for i32.
+  IntegerType i32 = IntegerType::get(&context, 32);
+  IntegerAttr attr32 = IntegerAttr::get(i32, APInt(32, 42));
+  EXPECT_EQ(attr32.getType(), i32);
+  EXPECT_EQ(attr32.getValue().getBitWidth(), 32u);
+  EXPECT_EQ(attr32.getInt(), 42);
+
+  // Correct APInt width for index type.
+  IndexType indexTy = IndexType::get(&context);
+  IntegerAttr attrIdx =
+      IntegerAttr::get(indexTy, APInt(IndexType::kInternalStorageBitWidth, 5));
+  EXPECT_EQ(attrIdx.getType(), indexTy);
+  EXPECT_EQ(attrIdx.getValue().getBitWidth(),
+            (unsigned)IndexType::kInternalStorageBitWidth);
+}
+
+#ifndef NDEBUG
+TEST(IntegerAttrDeathTest, WrongBitWidthForIntegerType) {
+  MLIRContext context;
+  IntegerType i32 = IntegerType::get(&context, 32);
+  // APInt(8, 1) has bit width 8, but i32 requires 32.
+  EXPECT_DEATH(IntegerAttr::get(i32, APInt(8, 1)),
+               "APInt bit width must match integer type width");
+}
+
+TEST(IntegerAttrDeathTest, WrongBitWidthForIndexType) {
+  MLIRContext context;
+  IndexType indexTy = IndexType::get(&context);
+  // APInt(1, 1) has bit width 1, but index type requires 64.
+  EXPECT_DEATH(
+      IntegerAttr::get(indexTy, APInt(1, 1)),
+      "APInt bit width must match IndexType internal storage bit width");
+}
+#endif // NDEBUG
+
 } // namespace

[github-actions]

🐧 Linux x64 Test Results

• 7731 tests passed
• 603 tests skipped

✅ The build succeeded and all tests passed.

[github-actions]

🪟 Windows x64 Test Results

• 3584 tests passed
• 415 tests skipped

✅ The build succeeded and all tests passed.

[hockyy]

Feeling like cast narrowing is definitely not OK, but cast widening can be handled automatically i guess. Thinking of APInt API(s) can be called from different call sites, bailing out with assert seems to be too aggressive?

Maybe something like this?
But considering the type can be variative then an assert is also fine. Just be careful when not in assert build.

    AttrBuilderWithInferredContext<(ins "Type":$type,
                                        "const APInt &":$value), [{
      if (type.isSignlessInteger(1))
        return BoolAttr::get(type.getContext(), value.getBoolValue());
      if (auto intTy = ::llvm::dyn_cast<IntegerType>(type)) {
        unsigned w = intTy.getWidth();
        unsigned vw = value.getBitWidth();
        APInt canon = value;
        if (vw < w) {
          if (intTy.isSignedInteger())
            canon = value.sext(w);
          else
            canon = value.zext(w);
        } else {
          assert(vw == w && "IntegerAttr::get: APInt wider than integer type");
        }
        return $_get(type.getContext(), type, canon);
      }
      if (::llvm::isa<IndexType>(type)) {
        unsigned w = IndexType::kInternalStorageBitWidth;
        unsigned vw = value.getBitWidth();
        APInt canon = value;
        if (vw < w)
          canon = value.zext(w);
        else
          assert(vw == w &&
                 "IntegerAttr::get: APInt wider than index internal storage");
        return $_get(type.getContext(), type, canon);
      }
      return $_get(type.getContext(), type, value);
    }]>,

[joker-eph]

bailing out with assert seems to be too aggressive?

There isn't a single use-case upstream, so I'd say it seem reasonable...

[ftynse]

I have had a crash from this downstream. Makes sense to have and looks consistent with the APFloat/FloatAttr counterpart that does verify flt semantics.