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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Oscar Priego (Opriego)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/220777.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp (+1-2) 
- (modified) mlir/test/Dialect/LLVMIR/constant-folding.mlir (+34) 


``````````diff
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

``````````

</details>


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