[Mlir-commits] [mlir] [MLIR][LLVM] Avoid narrowing wide constant shift amounts (PR #220777)

[Oscar Priego]

https://github.com/Opriego created https://github.com/llvm/llvm-project/pull/220777

LLVM::ShlOp::fold() used APInt::getZExtValue() to compare a constant shift
amount against the operand bit width.

For wide integer constants whose active value exceeds 64 bits, such as an i128
shift amount of 2^100, getZExtValue() asserts before the fold can determine
that the shift is out of range.

Compare the APInt directly with the operand bit width using uge() instead. This
preserves the existing behavior for out-of-range shifts, which are currently
left unfolded, while avoiding the unsafe narrowing.

Add regression coverage for:
- a wide out-of-range shift amount
- the exact bit-width boundary
- the largest valid shift amount

Fixes #220610.

>From ffebda1162dafca4aecd6644fc8515ad9e648021 Mon Sep 17 00:00:00 2001
From: Oscar Priego Verdugo <oscar.priegov at gmail.com>
Date: Wed, 2 Sep 2026 19:00:26 -0600
Subject: [PATCH] [MLIR][LLVM] Avoid narrowing wide constant shift amounts

---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp    |  3 +-
 .../test/Dialect/LLVMIR/constant-folding.mlir | 34 +++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 1e1ec23791775..d4005c9dde500 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -4107,8 +4107,7 @@ OpFoldResult LLVM::ShlOp::fold(FoldAdaptor adaptor) {
   if (!rhs)
     return {};
 
-  if (rhs.getValue().getZExtValue() >=
-      getLhs().getType().getIntOrFloatBitWidth())
+  if (rhs.getValue().uge(getLhs().getType().getIntOrFloatBitWidth()))
     return {}; // TODO: Fold into poison.
 
   auto lhs = dyn_cast_or_null<IntegerAttr>(adaptor.getLhs());
diff --git a/mlir/test/Dialect/LLVMIR/constant-folding.mlir b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
index 777213e067c26..0a6227bd2607b 100644
--- a/mlir/test/Dialect/LLVMIR/constant-folding.mlir
+++ b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
@@ -40,6 +40,40 @@ llvm.func @shl_multiple() -> i32 {
   llvm.return %2 : i32
 }
 
+// CHECK-LABEL: llvm.func @shl_wide_out_of_range_shift
+llvm.func @shl_wide_out_of_range_shift() -> i128 {
+  // CHECK-DAG: %[[LHS:.*]] = llvm.mlir.constant(1 : i128) : i128
+  %lhs = llvm.mlir.constant(1 : i128) : i128
+  // CHECK-DAG: %[[RHS:.*]] = llvm.mlir.constant(1267650600228229401496703205376 : i128) : i128
+  %rhs = llvm.mlir.constant(1267650600228229401496703205376 : i128) : i128
+  // CHECK: %[[RESULT:.*]] = llvm.shl %[[LHS]], %[[RHS]] : i128
+  %result = llvm.shl %lhs, %rhs : i128
+  // CHECK: llvm.return %[[RESULT]] : i128
+  llvm.return %result : i128
+}
+
+// CHECK-LABEL: llvm.func @shl_boundary_out_of_range_shift
+llvm.func @shl_boundary_out_of_range_shift() -> i128 {
+  // CHECK-DAG: %[[LHS:.*]] = llvm.mlir.constant(1 : i128) : i128
+  %lhs = llvm.mlir.constant(1 : i128) : i128
+  // CHECK-DAG: %[[RHS:.*]] = llvm.mlir.constant(128 : i128) : i128
+  %rhs = llvm.mlir.constant(128 : i128) : i128
+  // CHECK: %[[RESULT:.*]] = llvm.shl %[[LHS]], %[[RHS]] : i128
+  %result = llvm.shl %lhs, %rhs : i128
+  // CHECK: llvm.return %[[RESULT]] : i128
+  llvm.return %result : i128
+}
+
+// CHECK-LABEL: llvm.func @shl_largest_valid_shift
+llvm.func @shl_largest_valid_shift() -> i128 {
+  %lhs = llvm.mlir.constant(1 : i128) : i128
+  %rhs = llvm.mlir.constant(127 : i128) : i128
+  %result = llvm.shl %lhs, %rhs : i128
+  // CHECK: %[[RES:.*]] = llvm.mlir.constant(-170141183460469231731687303715884105728 : i128) : i128
+  // CHECK: llvm.return %[[RES]] : i128
+  llvm.return %result : i128
+}
+
 // -----
 
 // CHECK-LABEL: llvm.func @or_basic