[llvm] [ConstraintElimination] Scale comparisons against exact right shifts (PR #223212)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Alina Sbirlea (alinas)

<details>
<summary>Changes</summary>

When forming constraints for relational comparisons where an operand is an exact right shift (e.g. X < (Sub >> k)), scale the comparison by multiplying the opposite side by (1 << k), yielding (1 << k) * X < Sub. If both operands share the same exact right shift, decompose the unshifted sub-expressions directly.

Added lit tests in llvm/test/Transforms/ConstraintElimination/exact-shr-scaling.ll verifying comparisons with exact shifts on RHS, LHS, and both sides.

Gemini assisted changes.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/ConstraintElimination.cpp (+34) 
- (added) llvm/test/Transforms/ConstraintElimination/exact-shr-scaling.ll (+59) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index ed95686e9e5d9..0c613593cfe14 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -796,6 +796,40 @@ ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
                         IsSigned, DL);
   auto BDec = decompose(Op1->stripPointerCastsSameRepresentation(), *this,
                         IsSigned, DL);
+
+  // For relational comparisons where an operand is an exact right shift
+  // (e.g. X < (Sub >> k)), scale the comparison by multiplying the opposite
+  // side by (1 << k), yielding (1 << k) * X < Sub. If both operands share the
+  // same exact right shift, decompose the unshifted sub-expressions directly.
+  Value *Sub0 = nullptr, *Sub1 = nullptr;
+  ConstantInt *Shift0 = nullptr, *Shift1 = nullptr;
+  bool IsShift0 =
+      match(Op0, m_Exact(m_Shr(m_Value(Sub0), m_ConstantInt(Shift0)))) &&
+      canUseSExt(Shift0);
+  bool IsShift1 =
+      match(Op1, m_Exact(m_Shr(m_Value(Sub1), m_ConstantInt(Shift1)))) &&
+      canUseSExt(Shift1);
+  if (IsShift0 && IsShift1 && Shift0->getValue() == Shift1->getValue()) {
+    ADec = decompose(Sub0->stripPointerCastsSameRepresentation(), *this,
+                     IsSigned, DL);
+    BDec = decompose(Sub1->stripPointerCastsSameRepresentation(), *this,
+                     IsSigned, DL);
+  } else if (IsShift1 && Shift1->getSExtValue() > 0 &&
+             Shift1->getSExtValue() < 63) {
+    int64_t Factor = int64_t(1) << Shift1->getSExtValue();
+    if (BDec.Vars.size() <= 1 && !ADec.mul(Factor)) {
+      BDec = decompose(Sub1->stripPointerCastsSameRepresentation(), *this,
+                       IsSigned, DL);
+    }
+  } else if (IsShift0 && Shift0->getSExtValue() > 0 &&
+             Shift0->getSExtValue() < 63) {
+    int64_t Factor = int64_t(1) << Shift0->getSExtValue();
+    if (ADec.Vars.size() <= 1 && !BDec.mul(Factor)) {
+      ADec = decompose(Sub0->stripPointerCastsSameRepresentation(), *this,
+                       IsSigned, DL);
+    }
+  }
+
   int64_t Offset1 = ADec.Offset;
   int64_t Offset2 = BDec.Offset;
   if (MulOverflow(Offset1, int64_t(-1), Offset1))
diff --git a/llvm/test/Transforms/ConstraintElimination/exact-shr-scaling.ll b/llvm/test/Transforms/ConstraintElimination/exact-shr-scaling.ll
new file mode 100644
index 0000000000000..9958f90cb97e7
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/exact-shr-scaling.ll
@@ -0,0 +1,59 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=constraint-elimination -S %s | FileCheck %s
+
+declare void @llvm.assume(i1 noundef)
+
+; When one operand is an exact right shift (Sub >> k), scale comparison by multiplying
+; the opposite side by (1 << k), proving iv < (Sub >> 2) when (4 * iv) < Sub.
+define i1 @test_scale_rhs_exact_shift(i64 %iv, i64 %sub) {
+; CHECK-LABEL: define i1 @test_scale_rhs_exact_shift(
+; CHECK-SAME: i64 [[IV:%.*]], i64 [[SUB:%.*]]) {
+; CHECK-NEXT:    [[SCALED:%.*]] = mul nuw i64 [[IV]], 4
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i64 [[SCALED]], [[SUB]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
+; CHECK-NEXT:    [[SIZE:%.*]] = ashr exact i64 [[SUB]], 2
+; CHECK-NEXT:    ret i1 true
+;
+  %scaled = mul nuw i64 %iv, 4
+  %c = icmp ult i64 %scaled, %sub
+  call void @llvm.assume(i1 %c)
+  %size = ashr exact i64 %sub, 2
+  %check = icmp ult i64 %iv, %size
+  ret i1 %check
+}
+
+; Both operands have matching exact right shift: Sub0 < Sub1 implies (Sub0 >> 2) < (Sub1 >> 2).
+define i1 @test_scale_both_exact_shifts(i64 %sub0, i64 %sub1) {
+; CHECK-LABEL: define i1 @test_scale_both_exact_shifts(
+; CHECK-SAME: i64 [[SUB0:%.*]], i64 [[SUB1:%.*]]) {
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i64 [[SUB0]], [[SUB1]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
+; CHECK-NEXT:    [[SH0:%.*]] = ashr exact i64 [[SUB0]], 2
+; CHECK-NEXT:    [[SH1:%.*]] = ashr exact i64 [[SUB1]], 2
+; CHECK-NEXT:    ret i1 true
+;
+  %c = icmp ult i64 %sub0, %sub1
+  call void @llvm.assume(i1 %c)
+  %sh0 = ashr exact i64 %sub0, 2
+  %sh1 = ashr exact i64 %sub1, 2
+  %check = icmp ult i64 %sh0, %sh1
+  ret i1 %check
+}
+
+; Exact shift on LHS: (Sub >> 2) <= iv implies Sub <= (4 * iv).
+define i1 @test_scale_lhs_exact_shift(i64 %iv, i64 %sub) {
+; CHECK-LABEL: define i1 @test_scale_lhs_exact_shift(
+; CHECK-SAME: i64 [[IV:%.*]], i64 [[SUB:%.*]]) {
+; CHECK-NEXT:    [[SIZE:%.*]] = ashr exact i64 [[SUB]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i64 [[SIZE]], [[IV]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
+; CHECK-NEXT:    [[SCALED:%.*]] = mul nuw i64 [[IV]], 4
+; CHECK-NEXT:    ret i1 true
+;
+  %size = ashr exact i64 %sub, 2
+  %c = icmp ule i64 %size, %iv
+  call void @llvm.assume(i1 %c)
+  %scaled = mul nuw i64 %iv, 4
+  %check = icmp ule i64 %sub, %scaled
+  ret i1 %check
+}

``````````

</details>


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