[Mlir-commits] [mlir] [MLIR][Arith] Fix crash in TruncIOp::inferResultRanges for i0 destination (PR #183564)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When `arith.trunci` truncates to `i0`, `getStorageBitwidth` returns 0. The subsequent call to `truncRange(range, 0)` then computes `range.smin().ashr(destWidth - 1)` where `destWidth - 1` wraps to a huge unsigned value, triggering the assertion
`ShiftAmt <= BitWidth && "Invalid shift amount"` in `APInt::ashrInPlace`.

Guard against this in `TruncIOp::inferResultRanges`: skip calling `setResultRange` when `destWidth == 0` since APInt arithmetic is undefined for zero-width types.

Fixes #<!-- -->177822

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp (+3) 
- (modified) mlir/test/Dialect/Arith/int-range-narrowing.mlir (+9) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
index 49f89e1bd17f3..c515d92e24ff5 100644
--- a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
+++ b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
@@ -245,6 +245,9 @@ void arith::TruncIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                         SetIntRangeFn setResultRange) {
   unsigned destWidth =
       ConstantIntRanges::getStorageBitwidth(getResult().getType());
+  // i0 has no bits; APInt arithmetic is undefined for width 0, skip it.
+  if (destWidth == 0)
+    return;
   setResultRange(getResult(), truncRange(argRanges[0], destWidth));
 }
 
diff --git a/mlir/test/Dialect/Arith/int-range-narrowing.mlir b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
index e2cd9b50f6736..f527f8c4bb7ee 100644
--- a/mlir/test/Dialect/Arith/int-range-narrowing.mlir
+++ b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
@@ -417,3 +417,12 @@ func.func @narrow_loop_bounds() {
   }
   return
 }
+
+// Verify that truncating to i0 does not crash (issue #177822).
+// The pass cannot narrow i0 types; the trunci should be preserved.
+// CHECK-LABEL: @trunci_to_i0_no_crash
+// CHECK: arith.trunci %arg0 : i32 to i0
+func.func @trunci_to_i0_no_crash(%arg0: i32) -> i0 {
+  %0 = arith.trunci %arg0 : i32 to i0
+  return %0 : i0
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/183564