[Mlir-commits] [mlir] [mlir][arith] Fix crash when narrowing or truncating to i0 (PR #179894)

[Shoaib Khan]

https://github.com/shoaibkhan-sde updated https://github.com/llvm/llvm-project/pull/179894

>From 097b6567ca549c244a18f4cce0b005244b17f67e Mon Sep 17 00:00:00 2001
From: shoaibkhan-sde <shoaib.cse.engineer at gmail.com>
Date: Thu, 5 Feb 2026 14:41:50 +0530
Subject: [PATCH 1/3] [mlir][arith] Fix crash when narrowing/truncating to i0

Added guards to prevent integer range inference and narrowing patterns from processing zero-bitwidth types, which previously caused an APInt shift assertion failure.
---
 mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp   | 5 +++++
 mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
index 49f89e1bd17f3..900d986e79c9a 100644
--- a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
+++ b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
@@ -245,6 +245,11 @@ void arith::TruncIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                         SetIntRangeFn setResultRange) {
   unsigned destWidth =
       ConstantIntRanges::getStorageBitwidth(getResult().getType());
+
+  // Shoaib's Ultimate Guard: If width is 0, don't calculate ranges.
+  if (destWidth == 0)
+    return;
+
   setResultRange(getResult(), truncRange(argRanges[0], destWidth));
 }
 
diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index fefbba989b996..bad12c0a010f0 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -347,6 +347,9 @@ struct NarrowElementwise final : OpTraitRewritePattern<OpTrait::Elementwise> {
           op, "no operands or operand types don't match result type");
 
     for (unsigned targetBitwidth : targetBitwidths) {
+      // Shoaib's Fix: Bitwidth must be at least 1 to be valid in MLIR.
+      if (targetBitwidth == 0)
+        continue;
       CastKind castKind = CastKind::Both;
       for (const ConstantIntRanges &range : ranges) {
         castKind = mergeCastKinds(castKind,

>From 91a80670e76bb384dcad62a1e9ebb19eaae759b6 Mon Sep 17 00:00:00 2001
From: shoaibkhan-sde <shoaib.cse.engineer at gmail.com>
Date: Thu, 5 Feb 2026 15:26:05 +0530
Subject: [PATCH 2/3] [mlir][arith] Add test case for i0 truncation crash

---
 mlir/test/Dialect/Arith/int-range-opts.mlir | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/mlir/test/Dialect/Arith/int-range-opts.mlir b/mlir/test/Dialect/Arith/int-range-opts.mlir
index e6e48d30cece5..ef7ff3683edfc 100644
--- a/mlir/test/Dialect/Arith/int-range-opts.mlir
+++ b/mlir/test/Dialect/Arith/int-range-opts.mlir
@@ -148,3 +148,14 @@ func.func @analysis_crash(%arg0: i32, %arg1: tensor<128xi1>) -> tensor<128xi64>
   %2 = arith.extsi %1 : tensor<128xi32> to tensor<128xi64>
   return %2 : tensor<128xi64>
 }
+
+// -----
+
+// CHECK-LABEL: func @test_i0_trunc_crash
+func.func @test_i0_trunc_crash() {
+  %c1_i1 = arith.constant true
+  // Shoaib's Fix: Ensure trunci to i0 doesn't crash during range inference.
+  // The test passes if mlir-opt processes this without an assertion failure.
+  %0 = arith.trunci %c1_i1 : i1 to i0
+  return
+}
\ No newline at end of file

>From df94032ba1b05d850aed9a8cd5779cb303b60531 Mon Sep 17 00:00:00 2001
From: shoaibkhan-sde <shoaib.cse.engineer at gmail.com>
Date: Thu, 5 Feb 2026 15:32:16 +0530
Subject: [PATCH 3/3] [mlir][arith] Add guard to IntRangeNarrowing pattern for
 i0

---
 mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index bad12c0a010f0..51d5db5d6ac8c 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -742,4 +742,4 @@ void mlir::arith::populateControlFlowValuesNarrowingPatterns(
 
 std::unique_ptr<Pass> mlir::arith::createIntRangeOptimizationsPass() {
   return std::make_unique<IntRangeOptimizationsPass>();
-}
+}
\ No newline at end of file