[Mlir-commits] [mlir] [mlir][arith] Fold trivial shifts (`0<<x`, `x>>x`, `-1>>x`) (PR #212159)
Victor Perez
llvmlistbot at llvm.org
Sun Jul 26 17:35:53 PDT 2026
https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/212159
Add value-preserving folds mirroring LLVM's InstructionSimplify for the integer shift ops:
```
shli/shrui/shrsi(0, x) -> 0
shrui/shrsi(x, x) -> 0
shrsi(-1, x) -> -1
```
A shift amount greater than or equal to the bit width yields poison, so refining these cases to a constant is valid. For `shr(x, x)` any in-range amount `v` satisfies `v >> v == 0 (v < 2^v)`.
Folding an out-of-range shift amount (`c >= bit width`) to poison is left as a TODO: it would make the arith dialect depend on the ub dialect to materialize `ub.poison`.
LLVM performs the same simplifications: https://godbolt.org/z/7hsMz9E7Y
Verified with Alive2:
| Fold | shl | lshr | ashr |
| :--- | :---: | :---: | :---: |
| `0 << x`, `0 >> x -> 0` | [✓](https://alive2.llvm.org/ce/z/Sz2fT1) | [✓](https://alive2.llvm.org/ce/z/mlq7k6) | [✓](https://alive2.llvm.org/ce/z/jbeNH9) |
| `x >> x -> 0` | — | [✓](https://alive2.llvm.org/ce/z/JCPBCZ) | [✓](https://alive2.llvm.org/ce/z/YLjo3M) |
| `-1 >> x -> -1` | — | — | [✓](https://alive2.llvm.org/ce/z/sT6pNC) |
---
Code partially generated by Claude Code.
>From 9adcfc04975a72ca21593265597b546ae1d3212d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
<victor.pc.upm at gmail.com>
Date: Sat, 25 Jul 2026 16:06:43 -0700
Subject: [PATCH] [mlir][arith] Fold trivial shifts (0<<x, x>>x, -1>>x)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Add value-preserving folds mirroring LLVM's InstructionSimplify for the
integer shift ops:
shli/shrui/shrsi(0, x) -> 0
shrui/shrsi(x, x) -> 0
shrsi(-1, x) -> -1
A shift amount greater than or equal to the bit width yields poison, so
refining these cases to a constant is valid. For shr(x, x) any in-range
amount v satisfies v >> v == 0 (v < 2^v).
Folding an out-of-range shift amount (c >= bit width) to poison is left as a
TODO: it would make the arith dialect depend on the ub dialect to materialize
ub.poison.
LLVM performs the same simplifications: https://godbolt.org/z/7hsMz9E7Y
Verified with Alive2:
| Fold | shl | lshr | ashr |
| :--- | :---: | :---: | :---: |
| `0 << x`, `0 >> x -> 0` | [✓](https://alive2.llvm.org/ce/z/Sz2fT1) | [✓](https://alive2.llvm.org/ce/z/mlq7k6) | [✓](https://alive2.llvm.org/ce/z/jbeNH9) |
| `x >> x -> 0` | — | [✓](https://alive2.llvm.org/ce/z/JCPBCZ) | [✓](https://alive2.llvm.org/ce/z/YLjo3M) |
| `-1 >> x -> -1` | — | — | [✓](https://alive2.llvm.org/ce/z/sT6pNC) |
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 40 ++++++++++++
.../Conversion/ArithToLLVM/arith-to-llvm.mlir | 8 +--
mlir/test/Dialect/Arith/canonicalize.mlir | 62 +++++++++++++++++++
3 files changed, 106 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 7e4ffeda2bd7e..b7fdb97aba335 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -2975,9 +2975,18 @@ LogicalResult arith::SelectOp::verify() {
//===----------------------------------------------------------------------===//
OpFoldResult arith::ShLIOp::fold(FoldAdaptor adaptor) {
+ // TODO: shli(x, c) -> poison when c is out of range (c >= bit width). An
+ // out-of-range shift amount is undefined behaviour and could fold to poison,
+ // but that would make the arith dialect depend on the ub dialect to
+ // materialize ub.poison; left out for now.
+
// shli(x, 0) -> x
if (matchPattern(adaptor.getRhs(), m_Zero()))
return getLhs();
+ // shli(0, x) -> 0. An out-of-range shift amount yields poison, so refining
+ // it to 0 is valid.
+ if (matchPattern(adaptor.getLhs(), m_Zero()))
+ return getLhs();
// Don't fold if shifting more or equal than the bit width.
bool bounded = false;
auto result = constFoldBinaryOp<IntegerAttr>(
@@ -2993,9 +3002,22 @@ OpFoldResult arith::ShLIOp::fold(FoldAdaptor adaptor) {
//===----------------------------------------------------------------------===//
OpFoldResult arith::ShRUIOp::fold(FoldAdaptor adaptor) {
+ // TODO: shrui(x, c) -> poison when c is out of range (c >= bit width). An
+ // out-of-range shift amount is undefined behaviour and could fold to poison,
+ // but that would make the arith dialect depend on the ub dialect to
+ // materialize ub.poison; left out for now.
+
// shrui(x, 0) -> x
if (matchPattern(adaptor.getRhs(), m_Zero()))
return getLhs();
+ // shrui(0, x) -> 0. An out-of-range shift amount yields poison, so refining
+ // it to 0 is valid.
+ if (matchPattern(adaptor.getLhs(), m_Zero()))
+ return getLhs();
+ // shrui(x, x) -> 0. For any in-range shift amount v < bitwidth, v >> v == 0
+ // (v < 2^v); out-of-range amounts yield poison, so 0 is a valid refinement.
+ if (getLhs() == getRhs())
+ return getIntegerAttrOfType(getType(), 0);
// Don't fold if shifting more or equal than the bit width.
bool bounded = false;
auto result = constFoldBinaryOp<IntegerAttr>(
@@ -3011,9 +3033,27 @@ OpFoldResult arith::ShRUIOp::fold(FoldAdaptor adaptor) {
//===----------------------------------------------------------------------===//
OpFoldResult arith::ShRSIOp::fold(FoldAdaptor adaptor) {
+ // TODO: shrsi(x, c) -> poison when c is out of range (c >= bit width). An
+ // out-of-range shift amount is undefined behaviour and could fold to poison,
+ // but that would make the arith dialect depend on the ub dialect to
+ // materialize ub.poison; left out for now.
+
// shrsi(x, 0) -> x
if (matchPattern(adaptor.getRhs(), m_Zero()))
return getLhs();
+ // shrsi(0, x) -> 0. An out-of-range shift amount yields poison, so refining
+ // it to 0 is valid.
+ if (matchPattern(adaptor.getLhs(), m_Zero()))
+ return getLhs();
+ // shrsi(x, x) -> 0. For any in-range shift amount v < bitwidth, v is a small
+ // non-negative value and v >> v == 0; out-of-range amounts yield poison.
+ if (getLhs() == getRhs())
+ return getIntegerAttrOfType(getType(), 0);
+ // shrsi(-1, x) -> -1. Arithmetic shift of all-ones is all-ones for any
+ // in-range amount; out-of-range amounts yield poison.
+ if (APInt val;
+ matchPattern(adaptor.getLhs(), m_ConstantInt(&val)) && val.isAllOnes())
+ return getLhs();
// Don't fold if shifting more or equal than the bit width.
bool bounded = false;
auto result = constFoldBinaryOp<IntegerAttr>(
diff --git a/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir b/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
index d22677fde774a..34a80eff42fed 100644
--- a/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
+++ b/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
@@ -31,10 +31,10 @@ func.func @vector_ops(%arg0: vector<4xf32>, %arg1: vector<4xi1>, %arg2: vector<4
%11 = arith.xori %arg2, %arg3 : vector<4xi64>
// CHECK-NEXT: %11 = llvm.shl %arg2, %arg2 : vector<4xi64>
%12 = arith.shli %arg2, %arg2 : vector<4xi64>
-// CHECK-NEXT: %12 = llvm.ashr %arg2, %arg2 : vector<4xi64>
- %13 = arith.shrsi %arg2, %arg2 : vector<4xi64>
-// CHECK-NEXT: %13 = llvm.lshr %arg2, %arg2 : vector<4xi64>
- %14 = arith.shrui %arg2, %arg2 : vector<4xi64>
+// CHECK-NEXT: %12 = llvm.ashr %arg2, %arg3 : vector<4xi64>
+ %13 = arith.shrsi %arg2, %arg3 : vector<4xi64>
+// CHECK-NEXT: %13 = llvm.lshr %arg2, %arg3 : vector<4xi64>
+ %14 = arith.shrui %arg2, %arg3 : vector<4xi64>
return %1 : vector<4xf32>
}
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index baf0d859b287f..d6ea02e5508dd 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -3685,6 +3685,68 @@ func.func @nofoldShrs3() -> i64 {
return %r : i64
}
+// CHECK-LABEL: @shli_zero_lhs
+// CHECK: %[[C0:.+]] = arith.constant 0 : i32
+// CHECK: return %[[C0]]
+func.func @shli_zero_lhs(%arg0 : i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %0 = arith.shli %c0, %arg0 : i32
+ return %0 : i32
+}
+
+// CHECK-LABEL: @shrui_zero_lhs_and_self
+// CHECK: %[[C0:.+]] = arith.constant 0 : i32
+// CHECK: return %[[C0]], %[[C0]]
+func.func @shrui_zero_lhs_and_self(%arg0 : i32) -> (i32, i32) {
+ %c0 = arith.constant 0 : i32
+ %0 = arith.shrui %c0, %arg0 : i32
+ %1 = arith.shrui %arg0, %arg0 : i32
+ return %0, %1 : i32, i32
+}
+
+// CHECK-LABEL: @shrsi_zero_lhs_and_self
+// CHECK: %[[C0:.+]] = arith.constant 0 : i32
+// CHECK: return %[[C0]], %[[C0]]
+func.func @shrsi_zero_lhs_and_self(%arg0 : i32) -> (i32, i32) {
+ %c0 = arith.constant 0 : i32
+ %0 = arith.shrsi %c0, %arg0 : i32
+ %1 = arith.shrsi %arg0, %arg0 : i32
+ return %0, %1 : i32, i32
+}
+
+// CHECK-LABEL: @shrsi_all_ones_lhs
+// CHECK: %[[CM1:.+]] = arith.constant -1 : i32
+// CHECK: return %[[CM1]]
+func.func @shrsi_all_ones_lhs(%arg0 : i32) -> i32 {
+ %cm1 = arith.constant -1 : i32
+ %0 = arith.shrsi %cm1, %arg0 : i32
+ return %0 : i32
+}
+
+// CHECK-LABEL: @shift_self_vector
+// CHECK-DAG: %[[C0:.+]] = arith.constant dense<0> : vector<4xi32>
+// CHECK-DAG: %[[CM1:.+]] = arith.constant dense<-1> : vector<4xi32>
+// CHECK: return %[[C0]], %[[C0]], %[[CM1]]
+func.func @shift_self_vector(%arg0 : vector<4xi32>)
+ -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) {
+ %cm1 = arith.constant dense<-1> : vector<4xi32>
+ %0 = arith.shrui %arg0, %arg0 : vector<4xi32>
+ %1 = arith.shrsi %arg0, %arg0 : vector<4xi32>
+ %2 = arith.shrsi %cm1, %arg0 : vector<4xi32>
+ return %0, %1, %2 : vector<4xi32>, vector<4xi32>, vector<4xi32>
+}
+
+// A dynamic shape must bail out of the x>>x constant fold.
+// CHECK-LABEL: @shr_self_dynamic_no_fold
+// CHECK: arith.shrui
+// CHECK: arith.shrsi
+func.func @shr_self_dynamic_no_fold(%arg0 : tensor<?xi32>)
+ -> (tensor<?xi32>, tensor<?xi32>) {
+ %0 = arith.shrui %arg0, %arg0 : tensor<?xi32>
+ %1 = arith.shrsi %arg0, %arg0 : tensor<?xi32>
+ return %0, %1 : tensor<?xi32>, tensor<?xi32>
+}
+
// -----
// CHECK-LABEL: @test_negf(
More information about the Mlir-commits
mailing list