[MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer #135772

Jezurko · 2025-04-15T10:15:06Z

Add APIntParameter with custom implementation for comparison and use it in llvm.constant_range attribute. This is necessary because the default equality operator of APInt asserts when the bit widths of the compared APInts differ. The comparison is used by StorageUniquer when hashes of two ranges with different bit widths collide.

This PR adds the bitwidth parameter to the constant range to allow for comparing of two instances of constant range. This fixes a crash in storage uniquer when two ranges with different bitwidths hashed to the same value and then the comparison triggered an assert in APInt because of the different bitwidths.

llvmbot · 2025-04-15T10:15:42Z

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir-llvm

Author: Robert Konicar (Jezurko)

Changes

This PR adds the bitwidth parameter to the constant range to allow for comparing of two instances of constant range. This fixes a crash in storage uniquer when two ranges with different bitwidths hashed to the same value and then the comparison triggered an assert in APInt because of the different bitwidths.

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

3 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td (+6-2)
(modified) mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp (+8-3)
(added) mlir/test/Dialect/LLVMIR/range-attr.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 690243525ede4..69376061bac72 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -1095,6 +1095,7 @@ def LLVM_TBAATagArrayAttr
 //===----------------------------------------------------------------------===//
 def LLVM_ConstantRangeAttr : LLVM_Attr<"ConstantRange", "constant_range"> {
   let parameters = (ins
+    "uint32_t":$width,
     "::llvm::APInt":$lower,
     "::llvm::APInt":$upper
   );
@@ -1110,13 +1111,16 @@ def LLVM_ConstantRangeAttr : LLVM_Attr<"ConstantRange", "constant_range"> {
 
     Syntax:
     ```
-    `<` `i`(width($lower)) $lower `,` $upper `>`
+    `<` `i`(width) $lower `,` $upper `>`
     ```
   }];
 
   let builders = [
     AttrBuilder<(ins "uint32_t":$bitWidth, "int64_t":$lower, "int64_t":$upper), [{
-      return $_get($_ctxt, ::llvm::APInt(bitWidth, lower), ::llvm::APInt(bitWidth, upper));
+      return $_get($_ctxt, bitWidth, ::llvm::APInt(bitWidth, lower), ::llvm::APInt(bitWidth, upper));
+    }]>,
+    AttrBuilder<(ins "::llvm::APInt":$lower, "::llvm::APInt":$upper), [{
+      return $_get($_ctxt, lower.getBitWidth(), lower, upper);
     }]>
   ];
 
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index e4f9d6f987401..6975c593d7f7e 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -278,13 +278,18 @@ Attribute ConstantRangeAttr::parse(AsmParser &parser, Type odsType) {
 }
 
 void ConstantRangeAttr::print(AsmPrinter &printer) const {
-  printer << "<i" << getLower().getBitWidth() << ", " << getLower() << ", "
-          << getUpper() << ">";
+  printer << "<i" << getWidth() << ", " << getLower() << ", " << getUpper()
+          << ">";
 }
 
 LogicalResult
 ConstantRangeAttr::verify(llvm::function_ref<InFlightDiagnostic()> emitError,
-                          APInt lower, APInt upper) {
+                          uint32_t width, llvm::APInt lower,
+                          llvm::APInt upper) {
+  if (width != lower.getBitWidth())
+    return emitError()
+           << "expected type and value to have matching bitwidths but got "
+           << width << " vs. " << lower.getBitWidth();
   if (lower.getBitWidth() != upper.getBitWidth())
     return emitError()
            << "expected lower and upper to have matching bitwidths but got "
diff --git a/mlir/test/Dialect/LLVMIR/range-attr.mlir b/mlir/test/Dialect/LLVMIR/range-attr.mlir
new file mode 100644
index 0000000000000..5f2b67609743b
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/range-attr.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s -o - | FileCheck %s
+
+// CHECK: #llvm.constant_range<i32, 0, 12>
+llvm.func external @foo1(!llvm.ptr, i64) -> (i32 {llvm.range = #llvm.constant_range<i32, 0, 12>})
+// CHECK: #llvm.constant_range<i8, 1, 10>
+llvm.func external @foo2(!llvm.ptr, i64) -> (i8 {llvm.range = #llvm.constant_range<i8, 1, 10>})
+// CHECK: #llvm.constant_range<i64, 0, 2147483648>
+llvm.func external @foo3(!llvm.ptr, i64) -> (i64 {llvm.range = #llvm.constant_range<i64, 0, 2147483648>})
+// CHECK: #llvm.constant_range<i32, 1, -2147483648>
+llvm.func external @foo4(!llvm.ptr, i64) -> (i32 {llvm.range = #llvm.constant_range<i32, 1, -2147483648>})

llvmbot · 2025-04-15T10:15:43Z

@llvm/pr-subscribers-mlir

Author: Robert Konicar (Jezurko)

Changes

This PR adds the bitwidth parameter to the constant range to allow for comparing of two instances of constant range. This fixes a crash in storage uniquer when two ranges with different bitwidths hashed to the same value and then the comparison triggered an assert in APInt because of the different bitwidths.

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

3 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td (+6-2)
(modified) mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp (+8-3)
(added) mlir/test/Dialect/LLVMIR/range-attr.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 690243525ede4..69376061bac72 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -1095,6 +1095,7 @@ def LLVM_TBAATagArrayAttr
 //===----------------------------------------------------------------------===//
 def LLVM_ConstantRangeAttr : LLVM_Attr<"ConstantRange", "constant_range"> {
   let parameters = (ins
+    "uint32_t":$width,
     "::llvm::APInt":$lower,
     "::llvm::APInt":$upper
   );
@@ -1110,13 +1111,16 @@ def LLVM_ConstantRangeAttr : LLVM_Attr<"ConstantRange", "constant_range"> {
 
     Syntax:
     ```
-    `<` `i`(width($lower)) $lower `,` $upper `>`
+    `<` `i`(width) $lower `,` $upper `>`
     ```
   }];
 
   let builders = [
     AttrBuilder<(ins "uint32_t":$bitWidth, "int64_t":$lower, "int64_t":$upper), [{
-      return $_get($_ctxt, ::llvm::APInt(bitWidth, lower), ::llvm::APInt(bitWidth, upper));
+      return $_get($_ctxt, bitWidth, ::llvm::APInt(bitWidth, lower), ::llvm::APInt(bitWidth, upper));
+    }]>,
+    AttrBuilder<(ins "::llvm::APInt":$lower, "::llvm::APInt":$upper), [{
+      return $_get($_ctxt, lower.getBitWidth(), lower, upper);
     }]>
   ];
 
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index e4f9d6f987401..6975c593d7f7e 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -278,13 +278,18 @@ Attribute ConstantRangeAttr::parse(AsmParser &parser, Type odsType) {
 }
 
 void ConstantRangeAttr::print(AsmPrinter &printer) const {
-  printer << "<i" << getLower().getBitWidth() << ", " << getLower() << ", "
-          << getUpper() << ">";
+  printer << "<i" << getWidth() << ", " << getLower() << ", " << getUpper()
+          << ">";
 }
 
 LogicalResult
 ConstantRangeAttr::verify(llvm::function_ref<InFlightDiagnostic()> emitError,
-                          APInt lower, APInt upper) {
+                          uint32_t width, llvm::APInt lower,
+                          llvm::APInt upper) {
+  if (width != lower.getBitWidth())
+    return emitError()
+           << "expected type and value to have matching bitwidths but got "
+           << width << " vs. " << lower.getBitWidth();
   if (lower.getBitWidth() != upper.getBitWidth())
     return emitError()
            << "expected lower and upper to have matching bitwidths but got "
diff --git a/mlir/test/Dialect/LLVMIR/range-attr.mlir b/mlir/test/Dialect/LLVMIR/range-attr.mlir
new file mode 100644
index 0000000000000..5f2b67609743b
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/range-attr.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s -o - | FileCheck %s
+
+// CHECK: #llvm.constant_range<i32, 0, 12>
+llvm.func external @foo1(!llvm.ptr, i64) -> (i32 {llvm.range = #llvm.constant_range<i32, 0, 12>})
+// CHECK: #llvm.constant_range<i8, 1, 10>
+llvm.func external @foo2(!llvm.ptr, i64) -> (i8 {llvm.range = #llvm.constant_range<i8, 1, 10>})
+// CHECK: #llvm.constant_range<i64, 0, 2147483648>
+llvm.func external @foo3(!llvm.ptr, i64) -> (i64 {llvm.range = #llvm.constant_range<i64, 0, 2147483648>})
+// CHECK: #llvm.constant_range<i32, 1, -2147483648>
+llvm.func external @foo4(!llvm.ptr, i64) -> (i32 {llvm.range = #llvm.constant_range<i32, 1, -2147483648>})

Jezurko · 2025-04-15T10:17:35Z

Fwiw, I will add a test that triggers this with the constant seed used in llvm::hash for release builds when I figure out the correct constants. At the moment I can reproduce the issue in the already added test on debug build with the "random" seed (enabled by allowing ABI breaks in cmake config) in some executions.
I wanted to create the PR to get feedback on the solution and if this approach is fine or I should take different approach.

joker-eph · 2025-04-15T11:06:18Z

Isn't the underlying issue in the hash function of the APInt that does not take the bit width into account?

Jezurko · 2025-04-15T11:21:07Z

Unless I'm missing something, I believe it does take it into account: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Support/APInt.cpp#L590

gysit · 2025-04-15T12:38:33Z

hmm could it be that this is related to the operator== implementation:

llvm-project/llvm/include/llvm/ADT/APInt.h

Line 1056 in 9a6c001

bool operator==(const APInt &RHS) const {

For some reason APInt seems to expect that the operator== is only called on APInts of the same bit width.

I am not entirely sure but I believe that the storage uniquer use the equality operator to avoid has value collisions. So maybe this is the root of the problem?

With regards to the approach, I would first like to fully understand the problem before moving forward. In theory APIInt should work with the storage uniquer. If it doesn't, there may be more problematic attributes.

Jezurko · 2025-04-15T12:51:52Z

Yes, the core of the issue is that two APInts with different bit widths can not be compared. I believe in other attributes the issue will be often avoided by the fact, that they contain some information about the type and the comparison short-circuits before comparing the APInts.

I have a stack trace from when the assert is triggered:
constant_range_crash.log

I will try to extract the hash seed that causes it with the input I added to the tests.

xlauko · 2025-04-15T13:00:30Z

A alternative solution is to add custom StorageClass for this: trailofbits@83a8b2b
with

  /// The hash key for this storage is a pair of the integer and type params.
  using KeyTy = std::pair<llvm::APInt, llvm::APInt>;

  /// Define the comparison function for the key type.
  bool operator==(const KeyTy &key) const {
    if (lower.getBitWidth() != key.first.getBitWidth() ||
        upper.getBitWidth() != key.second.getBitWidth()) {
      return false;
    }
    return lower == key.first && upper == key.second;
  }

gysit · 2025-04-15T13:13:17Z

A alternative solution is to add custom StorageClass for this

Do you know if it is possibly to provide a storage class for all APInts to overwrite the strange equality operator? That way the problem could be solved for all attributes that use APInt.

I also wonder why APInt has such a strange implementation of the equality operator?

With regards to the two workarounds the question is probably a trade-off between simplicity and storing four additional bits? I have a slight tendency for using the custom storage class approach, but no strong opinion.

xlauko · 2025-04-15T13:26:54Z

Hacky solution is to add special case to handle APInt bitwidth in

llvm-project/mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

Line 611 in e1382b3

"bool", "operator==", MethodParameter("const KeyTy &", "tblgenKey"));

which is not nice too

Jezurko · 2025-04-15T13:27:17Z

For debugging purpose: making the get_execution_seed() function in llvm/include/llvm/ADT/Hashing.h return 0x5556587D0350ULLmakes the test-case trigger the assert every time.

Jezurko · 2025-04-15T13:36:45Z

Another way I see would be defining APIntParameter with correct comparator in mlir/include/mlir/IR/AttrTypeBase.td.

llvm-project/mlir/include/mlir/IR/AttrTypeBase.td

Line 387 in 535eaa4

class APFloatParameter<string desc> :

For example APFloat is already handled there. But that doesn't fix the StorageClass for users of the raw APInt.

gysit · 2025-04-15T13:49:33Z

Another way I see would be defining APIntParameter with correct comparator in mlir/include/mlir/IR/AttrTypeBase.td.

Yeah I think this is my preferred solution so far.

River707

Adding an APInt parameter looks reasonable to me, any other uses of APInt as a parameter in-tree that should be updated?

River707 · 2025-04-15T16:58:25Z

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

+  let parameters = (ins APIntParameter<"">:$lower,
+    APIntParameter<"">:$upper


Suggested change

let parameters = (ins APIntParameter<"">:$lower,

APIntParameter<"">:$upper

let parameters = (ins

APIntParameter<"">:$lower,

APIntParameter<"">:$upper

I found only:

https://github.com/trail-of-forks/llvm-project/blob/bba2507c19ff678c5d7b18e0b220406be87451fe/mlir/include/mlir/IR/BuiltinAttributes.td#L683

llvm-project/mlir/docs/DefiningDialects/AttributesAndTypes.md

Line 127 in a955399

def My_IntegerAttr : MyDialect_Attr<"Integer", "int"> {

We discussed with @gysit in side channel that it might be worthwhile to add a check to tablegen to warn on APInt use and suggest APIntParameter instead.

One conceptual alternative to that I suppose, would be to detect APInt and use APIntParameter equivalent logic instead. Either way, would be nice to remove a footgun.

gysit

Thanks for fixing, I like this solution.

LGTM modulo nit comments.

If we can add guardrails in tablegen that would be great. But that would be something for a separate PR.

gysit · 2025-04-16T06:08:28Z

mlir/include/mlir/IR/AttrTypeBase.td

@@ -383,6 +383,12 @@ class StringRefParameter<string desc = "", string value = ""> :
  let defaultValue = value;
 }

+// For APInts, which require comparison over different bitwidths


Suggested change

// For APInts, which require comparison over different bitwidths

// For APInts, which require comparison supporting different bitwidths. The default

// APInt comparison operator asserts when the bitwidths differ, so a custom

// implementation is necessary.

nit: Let's maybe expand a bit why this is necessary.

Jezurko · 2025-04-16T09:32:11Z

I have integrated the suggested changes and updated the PR description to match the current solution.
I will add the warning to tablegen and update the remaining uses of llvm::APInt in a second PR.

If it's okay like this, can someone merge it for me, please? I do not have write access yet.

…vm#135772) Add APIntParameter with custom implementation for comparison and use it in llvm.constant_range attribute. This is necessary because the default equality operator of APInt asserts when the bit widths of the compared APInts differ. The comparison is used by StorageUniquer when hashes of two ranges with different bit widths collide.

llvmbot added mlir:llvm mlir labels Apr 15, 2025

xlauko requested review from gysit and ftynse April 15, 2025 10:40

llvmbot added mlir:core MLIR Core Infrastructure mlir:ods labels Apr 15, 2025

fixup! [MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer

de30d07

Jezurko force-pushed the mlir-llvm-constant-range branch from 9bb9bdc to de30d07 Compare April 15, 2025 16:56

fixup! [MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer

8b59fe5

River707 reviewed Apr 15, 2025

View reviewed changes

gysit approved these changes Apr 16, 2025

View reviewed changes

joker-eph approved these changes Apr 16, 2025

View reviewed changes

fixup! [MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer

d8b04e6

joker-eph merged commit b9ce185 into llvm:main Apr 16, 2025
11 checks passed

Jezurko mentioned this pull request Apr 16, 2025

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

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer #135772

[MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer #135772

Jezurko commented Apr 15, 2025 •

edited

Loading

llvmbot commented Apr 15, 2025 •

edited

Loading

llvmbot commented Apr 15, 2025

Jezurko commented Apr 15, 2025 •

edited

Loading

joker-eph commented Apr 15, 2025

Jezurko commented Apr 15, 2025 •

edited

Loading

gysit commented Apr 15, 2025

Jezurko commented Apr 15, 2025

xlauko commented Apr 15, 2025

gysit commented Apr 15, 2025 •

edited

Loading

xlauko commented Apr 15, 2025

Jezurko commented Apr 15, 2025

Jezurko commented Apr 15, 2025 •

edited

Loading

gysit commented Apr 15, 2025

River707 left a comment

River707 Apr 15, 2025

xlauko Apr 15, 2025

River707 Apr 15, 2025

gysit left a comment

gysit Apr 16, 2025

Jezurko commented Apr 16, 2025

		let parameters = (ins APIntParameter<"">:$lower,
		APIntParameter<"">:$upper

-// For APInts, which require comparison over different bitwidths
+// For APInts, which require comparison supporting different bitwidths. The default
+// APInt comparison operator asserts when the bitwidths differ, so a custom
+// implementation is necessary.

[MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer #135772

[MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer #135772

Conversation

Jezurko commented Apr 15, 2025 • edited Loading

llvmbot commented Apr 15, 2025 • edited Loading

llvmbot commented Apr 15, 2025

Jezurko commented Apr 15, 2025 • edited Loading

joker-eph commented Apr 15, 2025

Jezurko commented Apr 15, 2025 • edited Loading

gysit commented Apr 15, 2025

Jezurko commented Apr 15, 2025

xlauko commented Apr 15, 2025

gysit commented Apr 15, 2025 • edited Loading

xlauko commented Apr 15, 2025

Jezurko commented Apr 15, 2025

Jezurko commented Apr 15, 2025 • edited Loading

gysit commented Apr 15, 2025

River707 left a comment

Choose a reason for hiding this comment

River707 Apr 15, 2025

Choose a reason for hiding this comment

xlauko Apr 15, 2025

Choose a reason for hiding this comment

River707 Apr 15, 2025

Choose a reason for hiding this comment

gysit left a comment

Choose a reason for hiding this comment

gysit Apr 16, 2025

Choose a reason for hiding this comment

Jezurko commented Apr 16, 2025

Jezurko commented Apr 15, 2025 •

edited

Loading

llvmbot commented Apr 15, 2025 •

edited

Loading

Jezurko commented Apr 15, 2025 •

edited

Loading

Jezurko commented Apr 15, 2025 •

edited

Loading

gysit commented Apr 15, 2025 •

edited

Loading

Jezurko commented Apr 15, 2025 •

edited

Loading