[ConstraintElim] Add support for decomposing gep nuw

[nikic]

ConstraintElimination currently only supports decomposing gep nusw with non-negative indices (with "non-negative" possibly being enforced via pre-condition).

Add support for gep nuw, which directly gives us the necessary guarantees for the decomposition.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

ConstraintElimination currently only supports decomposing gep nusw with non-negative indices (with "non-negative" possibly being enforced via pre-condition).

Add support for gep nuw, which directly gives us the necessary guarantees for the decomposition.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/ConstraintElimination.cpp (+10-6)
• (modified) llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll (+2-4)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 944be38cb94bc7..8d1e793836c772 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -452,7 +452,9 @@ static Decomposition decomposeGEP(GEPOperator &GEP,
                       "unsigned predicates at the moment.");
   const auto &[BasePtr, ConstantOffset, VariableOffsets, NW] =
       collectOffsets(GEP, DL);
-  if (!BasePtr || !NW.hasNoUnsignedSignedWrap())
+  // We support either plain gep nuw, or gep nusw with non-negative offset,
+  // which implies gep nuw.
+  if (!BasePtr || NW == GEPNoWrapFlags::none())
     return &GEP;
 
   Decomposition Result(ConstantOffset.getSExtValue(), DecompEntry(1, BasePtr));
@@ -461,11 +463,13 @@ static Decomposition decomposeGEP(GEPOperator &GEP,
     IdxResult.mul(Scale.getSExtValue());
     Result.add(IdxResult);
 
-    // If Op0 is signed non-negative, the GEP is increasing monotonically and
-    // can be de-composed.
-    if (!isKnownNonNegative(Index, DL))
-      Preconditions.emplace_back(CmpInst::ICMP_SGE, Index,
-                                 ConstantInt::get(Index->getType(), 0));
+    if (!NW.hasNoUnsignedWrap()) {
+      // Try to prove nuw from nusw and nneg.
+      assert(NW.hasNoUnsignedSignedWrap() && "Must have nusw flag");
+      if (!isKnownNonNegative(Index, DL))
+        Preconditions.emplace_back(CmpInst::ICMP_SGE, Index,
+                                   ConstantInt::get(Index->getType(), 0));
+    }
   }
   return Result;
 }
diff --git a/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll b/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll
index 98bf01ef376317..23e1698136f5e2 100644
--- a/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll
+++ b/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll
@@ -702,8 +702,7 @@ define i1 @test_nuw(ptr %p, i64 %x, i64 %y) {
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP1]])
 ; CHECK-NEXT:    [[GEP_X:%.*]] = getelementptr nuw i8, ptr [[P:%.*]], i64 [[X]]
 ; CHECK-NEXT:    [[GEP_Y:%.*]] = getelementptr nuw i8, ptr [[P]], i64 [[Y]]
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp ugt ptr [[GEP_X]], [[GEP_Y]]
-; CHECK-NEXT:    ret i1 [[CMP2]]
+; CHECK-NEXT:    ret i1 true
 ;
   %cmp1 = icmp ugt i64 %x, %y
   call void @llvm.assume(i1 %cmp1)
@@ -720,8 +719,7 @@ define i1 @test_nuw_nested(ptr %p, i64 %x, i64 %y) {
 ; CHECK-NEXT:    [[GEP_X:%.*]] = getelementptr nuw i8, ptr [[P:%.*]], i64 [[X]]
 ; CHECK-NEXT:    [[GEP_X1:%.*]] = getelementptr nuw i8, ptr [[GEP_X]], i64 1
 ; CHECK-NEXT:    [[GEP_Y:%.*]] = getelementptr nuw i8, ptr [[P]], i64 [[Y]]
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp ugt ptr [[GEP_X1]], [[GEP_Y]]
-; CHECK-NEXT:    ret i1 [[CMP2]]
+; CHECK-NEXT:    ret i1 true
 ;
   %cmp1 = icmp ugt i64 %x, %y
   call void @llvm.assume(i1 %cmp1)

[fhahn]

LGTM, thanks!

[dtcxzyw]

LG. At least it improves compile-time.
Alive2: https://alive2.llvm.org/ce/z/hhriF6