[ConstraintElim] Add support for decomposing gep nuw (#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.

Author: nikic

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;
 }

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)