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

[llvmlistbot at llvm.org]

https://github.com/edisongz updated https://github.com/llvm/llvm-project/pull/218495

>From 5d31c306fb667e03ac09f473ed45eec5b511986d Mon Sep 17 00:00:00 2001
From: Yijie Jiang <edisongz123 at gmail.com>
Date: Tue, 25 Aug 2026 03:16:53 +0800
Subject: [PATCH] [MLIR][Arith] Don't narrow shifts whose amount can exceed the
 target width

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
---
 .../Transforms/IntRangeOptimizations.cpp      |  5 ++++
 .../Dialect/Arith/int-range-narrowing.mlir    | 27 +++++++++++++++++++
 2 files changed, 32 insertions(+)

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
 //===----------------------------------------------------------------------===//