[Mlir-commits] [mlir] [MLIR][LLVM] Fix #llvm.constant_range crashing in storage uniquer (PR #135772)

[Robert Konicar]

https://github.com/Jezurko created https://github.com/llvm/llvm-project/pull/135772

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.

>From 9a1bb3e165c1642ac9979d3258bb6e6ad25f60a2 Mon Sep 17 00:00:00 2001
From: Robert Konicar <rkonicar at mail.muni.cz>
Date: Mon, 14 Apr 2025 18:18:31 +0200
Subject: [PATCH] [MLIR][LLVM] Fix #llvm.constant_range crashing in storage
 uniquer

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.
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td |  8 ++++++--
 mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp         | 11 ++++++++---
 mlir/test/Dialect/LLVMIR/range-attr.mlir         | 10 ++++++++++
 3 files changed, 24 insertions(+), 5 deletions(-)
 create mode 100644 mlir/test/Dialect/LLVMIR/range-attr.mlir

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>})