[llvm] [ConstraintElim] Add support for decomposing gep nuw (PR #118639)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/118639

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.

>From 7ae6f3a9c5881550744834dfaae2cfaaef189e86 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Wed, 4 Dec 2024 13:15:32 +0100
Subject: [PATCH] [ConstraintElim] Add support for decomposing gep nuw

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.
---
 .../Transforms/Scalar/ConstraintElimination.cpp  | 16 ++++++++++------
 .../ConstraintElimination/gep-arithmetic.ll      |  6 ++----
 2 files changed, 12 insertions(+), 10 deletions(-)

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)