[Mlir-commits] [mlir] 44a4dbf - [MLIR][Arith] Don't narrow shifts whose amount can exceed the target width (#218495)

[llvmlistbot at llvm.org]

Author: edisongz
Date: 2026-09-07T22:58:09+02:00
New Revision: 44a4dbf32a6b54de32e4d064756d7a6ba5b9e808

URL: https://github.com/llvm/llvm-project/commit/44a4dbf32a6b54de32e4d064756d7a6ba5b9e808
DIFF: https://github.com/llvm/llvm-project/commit/44a4dbf32a6b54de32e4d064756d7a6ba5b9e808.diff

LOG: [MLIR][Arith] Don't narrow shifts whose amount can exceed the target width (#218495)

A shift by an amount >= the operand bitwidth is poison. Narrowing
changes that bitwidth, so an amount valid at the original width (e.g. 32
for i64) can become poison at the narrowed one (>= 32 for i32),
miscompiling `4 >> 32` from 0 to 4.

Only narrow shifts when the shift amount is provably below the target
width.

Fixes #218191

Added: 
    

Modified: 
    mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
    mlir/test/Dialect/Arith/int-range-narrowing.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 298c0dc2f3bda..10f436a8df1be 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -395,6 +395,11 @@ struct NarrowElementwise final : OpTraitRewritePattern<OpTrait::Elementwise> {
       castKind = mergeCastKinds(castKind, castKindForOp);
       if (castKind == CastKind::None)
         continue;
+      // A shift by an amount >= the bitwidth is poison, so only narrow shifts
+      // when the shift amount (second operand) stays below the target width.
+      if (isa<arith::ShLIOp, arith::ShRSIOp, arith::ShRUIOp>(op) &&
+          !ranges[1].umax().ult(targetBitwidth))
+        continue;
       Type targetType = getTargetType(srcType, targetBitwidth);
       if (targetType == srcType)
         continue;

diff  --git a/mlir/test/Dialect/Arith/int-range-narrowing.mlir b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
index 5938f24868cb9..06ca5d19de9c9 100644
--- a/mlir/test/Dialect/Arith/int-range-narrowing.mlir
+++ b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
@@ -321,6 +321,33 @@ func.func @unsigned_ops_out_of_narrowed_signed_range() -> (i64, i64, i64, i64, i
   return %3, %4, %5, %6, %7, %8 : i64, i64, i64, i64, i64, i64
 }
 
+//===----------------------------------------------------------------------===//
+// Shift ops must not be narrowed when the shift amount can reach the target
+// bitwidth, since a shift >= the bitwidth is poison. Here the value fits in
+// i32 but the amount [0, 63] does not, so narrowing must be skipped.
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @shrsi_amount_out_of_narrowed_range
+// CHECK: arith.shrsi {{.*}} : i64
+// CHECK-NOT: arith.shrsi {{.*}} : i32
+func.func @shrsi_amount_out_of_narrowed_range() -> i64 {
+  %0 = test.with_bounds { umin = 0 : i64, umax = 4 : i64, smin = 0 : i64, smax = 4 : i64 } : i64
+  %1 = test.with_bounds { umin = 0 : i64, umax = 63 : i64, smin = 0 : i64, smax = 63 : i64 } : i64
+  %2 = arith.shrsi %0, %1 : i64
+  return %2 : i64
+}
+
+// When the amount is provably below the target bitwidth, narrowing is safe.
+//
+// CHECK-LABEL: func.func @shrui_amount_in_narrowed_range
+// CHECK: arith.shrui {{.*}} : i32
+func.func @shrui_amount_in_narrowed_range() -> i64 {
+  %0 = test.with_bounds { umin = 0 : i64, umax = 4 : i64, smin = 0 : i64, smax = 4 : i64 } : i64
+  %1 = test.with_bounds { umin = 0 : i64, umax = 31 : i64, smin = 0 : i64, smax = 31 : i64 } : i64
+  %2 = arith.shrui %0, %1 : i64
+  return %2 : i64
+}
+
 //===----------------------------------------------------------------------===//
 // arith.muli
 //===----------------------------------------------------------------------===//