[ConstraintElim] Use SCEV to check for multiples (#76925)

When adding constraints for induction variables, if the step is not one,
we need to make sure that (end-start) is a multiple of step, otherwise
we might step over the end value.

Currently this only supports one specific pattern for pointers, where
the end is a gep of the start with an appropriate offset.

Generalize this by using SCEV to check for multiples, which also makes
this work for integer IVs.

2 files changed, 9 insertions, 27 deletions

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index cc93c86..5e57aa7 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -968,23 +968,10 @@ void State::addInfoForInductions(BasicBlock &BB) {
     return;
 
   if (!StepOffset.isOne()) {
-    auto *UpperGEP = dyn_cast<GetElementPtrInst>(B);
-    if (!UpperGEP || UpperGEP->getPointerOperand() != StartValue ||
-        !UpperGEP->isInBounds())
-      return;
-
-    MapVector<Value *, APInt> UpperVariableOffsets;
-    APInt UpperConstantOffset(StepOffset.getBitWidth(), 0);
-    const DataLayout &DL = BB.getModule()->getDataLayout();
-    if (!UpperGEP->collectOffset(DL, StepOffset.getBitWidth(),
-                                 UpperVariableOffsets, UpperConstantOffset))
-      return;
-    // All variable offsets and the constant offset have to be a multiple of the
-    // step.
-    if (!UpperConstantOffset.urem(StepOffset).isZero() ||
-        any_of(UpperVariableOffsets, [&StepOffset](const auto &P) {
-          return !P.second.urem(StepOffset).isZero();
-        }))
+    // Check whether B-Start is known to be a multiple of StepOffset.
+    const SCEV *BMinusStart = SE.getMinusSCEV(SE.getSCEV(B), StartSCEV);
+    if (isa<SCEVCouldNotCompute>(BMinusStart) ||
+        !SE.getConstantMultiple(BMinusStart).urem(StepOffset).isZero())
       return;
   }
 
diff --git a/llvm/test/Transforms/ConstraintElimination/monotonic-int-phis-multiples.ll b/llvm/test/Transforms/ConstraintElimination/monotonic-int-phis-multiples.ll
index 035cea4..a952c40 100644
--- a/llvm/test/Transforms/ConstraintElimination/monotonic-int-phis-multiples.ll
+++ b/llvm/test/Transforms/ConstraintElimination/monotonic-int-phis-multiples.ll
@@ -13,8 +13,7 @@ define void @multiple_pow2(i64 %count) {
 ; CHECK-NEXT:    [[CMP_I_NOT:%.*]] = icmp eq i64 [[IV]], [[END]]
 ; CHECK-NEXT:    br i1 [[CMP_I_NOT]], label [[EXIT:%.*]], label [[LOOP_LATCH]]
 ; CHECK:       loop.latch:
-; CHECK-NEXT:    [[CMP2_I_I:%.*]] = icmp ult i64 [[IV]], [[END]]
-; CHECK-NEXT:    br i1 [[CMP2_I_I]], label [[LOOP]], label [[EXIT]]
+; CHECK-NEXT:    br i1 true, label [[LOOP]], label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
@@ -48,8 +47,7 @@ define void @multiple_pow2_larger_than_needed(i64 %count) {
 ; CHECK-NEXT:    [[CMP_I_NOT:%.*]] = icmp eq i64 [[IV]], [[END]]
 ; CHECK-NEXT:    br i1 [[CMP_I_NOT]], label [[EXIT:%.*]], label [[LOOP_LATCH]]
 ; CHECK:       loop.latch:
-; CHECK-NEXT:    [[CMP2_I_I:%.*]] = icmp ult i64 [[IV]], [[END]]
-; CHECK-NEXT:    br i1 [[CMP2_I_I]], label [[LOOP]], label [[EXIT]]
+; CHECK-NEXT:    br i1 true, label [[LOOP]], label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
@@ -119,8 +117,7 @@ define void @multiple_pow2_start_offset(i64 %count) {
 ; CHECK-NEXT:    [[CMP_I_NOT:%.*]] = icmp eq i64 [[IV]], [[END]]
 ; CHECK-NEXT:    br i1 [[CMP_I_NOT]], label [[EXIT]], label [[LOOP_LATCH]]
 ; CHECK:       loop.latch:
-; CHECK-NEXT:    [[CMP2_I_I:%.*]] = icmp ult i64 [[IV]], [[END]]
-; CHECK-NEXT:    br i1 [[CMP2_I_I]], label [[LOOP]], label [[EXIT]]
+; CHECK-NEXT:    br i1 true, label [[LOOP]], label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
@@ -194,8 +191,7 @@ define void @multiple_pow2_start_offset_dynamic(i64 %count) {
 ; CHECK-NEXT:    [[CMP_I_NOT:%.*]] = icmp eq i64 [[IV]], [[END]]
 ; CHECK-NEXT:    br i1 [[CMP_I_NOT]], label [[EXIT]], label [[LOOP_LATCH]]
 ; CHECK:       loop.latch:
-; CHECK-NEXT:    [[CMP2_I_I:%.*]] = icmp ult i64 [[IV]], [[END]]
-; CHECK-NEXT:    br i1 [[CMP2_I_I]], label [[LOOP]], label [[EXIT]]
+; CHECK-NEXT:    br i1 true, label [[LOOP]], label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
@@ -231,8 +227,7 @@ define void @multiple_non_pow2_nuw(i64 %count) {
 ; CHECK-NEXT:    [[CMP_I_NOT:%.*]] = icmp eq i64 [[IV]], [[END]]
 ; CHECK-NEXT:    br i1 [[CMP_I_NOT]], label [[EXIT:%.*]], label [[LOOP_LATCH]]
 ; CHECK:       loop.latch:
-; CHECK-NEXT:    [[CMP2_I_I:%.*]] = icmp ult i64 [[IV]], [[END]]
-; CHECK-NEXT:    br i1 [[CMP2_I_I]], label [[LOOP]], label [[EXIT]]
+; CHECK-NEXT:    br i1 true, label [[LOOP]], label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;