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

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.

Co-authored-by: Robert Konicar <[email protected]>

Author: frabert

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

@@ -1095,8 +1095,8 @@ def LLVM_TBAATagArrayAttr
 //===----------------------------------------------------------------------===//
 def LLVM_ConstantRangeAttr : LLVM_Attr<"ConstantRange", "constant_range"> {
   let parameters = (ins
-    "::llvm::APInt":$lower,
-    "::llvm::APInt":$upper
+    APIntParameter<"">:$lower,
+    APIntParameter<"">:$upper
   );
   let summary = "A range of two integers, corresponding to LLVM's ConstantRange";
   let description = [{

mlir/include/mlir/IR/AttrTypeBase.td

@@ -383,6 +383,14 @@ class StringRefParameter<string desc = "", string value = ""> :
   let defaultValue = value;
 }
 
+// For APInts, which require comparison supporting different bitwidths. The
+// default APInt comparison operator asserts when the bitwidths differ, so
+// a custom implementation is necessary.
+class APIntParameter<string desc> :
+    AttrOrTypeParameter<"::llvm::APInt", desc> {
+  let comparator = "$_lhs.getBitWidth() == $_rhs.getBitWidth() && $_lhs == $_rhs";
+}
+
 // For APFloats, which require comparison.
 class APFloatParameter<string desc> :
     AttrOrTypeParameter<"::llvm::APFloat", desc> {

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