[LSR] Fix wrapping bug in lsr-term-fold logic

[reames]

Summary

The existing logic was unsound, in two ways.

First, due to wrapping on the trip count computation, it could compute a value which convert a loop which exiting on iteration 256, to one which exited at 255. (With i8 trip counts.)

Second, it allowed rewriting when the trip count implies wrapping around the alternate IV. As a trivial example, it allowed rewriting an i128 exit test in terms of an i64 IV. This is obviously wrong.

Note that the test change is fairly minimal - i.e. only the targeted test - but that's only because I precommitted a change which switched the test from 32 to 64 bit pointers. For 32 bit point architectures with 32 bit primary inductions, this transform is almost always unsound to perform.

[nikic]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6743–6744

Side note, but this code looks like evaluateAtIteration().

6762

Hm, can't we do something like check the nowrap (as in <nw>) flag on the post-inc addrec here?

[reames]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6743–6744

I'd noticed that too, I'm just trying to fix this code one bit at a time. It also looks somewhat like LFTR's genLoopLimit and I have yet dug into the three enough to be really sure they compute the same values under all edge cases.

6762

I think you might be on to something here. Let me give this a bit more thought and see if I can find a counter example.

[reames]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6762

I tried this out, and on simple tests (e.g. const_tripcount for this transform), we appear to not be inferring no-self-wrap. I think we should, but it looks like we've got some inference problems here.

Mind if I go forward with the current scheme? I'd really like to get this transform into a somewhat sound state.

[nikic]

LGTM

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6762

Yeah, I'm fine with that, thanks for checking.

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
   // If that IV isn't dead after we rewrite the exit condition in terms of
   // another IV, there's no point in doing the transform.
   if (!isAlmostDeadIV(ToFold, LoopLatch, TermCond))
     return std::nullopt;
 
   const SCEV *BECount = SE.getBackedgeTakenCount(L);
   const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
   SCEVExpander Expander(SE, DL, "lsr_fold_term_cond");
 
   PHINode *ToHelpFold = nullptr;
   const SCEV *TermValueS = nullptr;
   for (PHINode &PN : L->getHeader()->phis()) {
     if (ToFold == &PN)
       continue;
 
     if (!SE.isSCEVable(PN.getType())) {
       LLVM_DEBUG(dbgs() << "IV of phi '" << PN
                         << "' is not SCEV-able, not qualified for the "
                            "terminating condition folding.\n");
       continue;
     }
     const SCEVAddRecExpr *AddRec = dyn_cast<SCEVAddRecExpr>(SE.getSCEV(&PN));
     // Only speculate on affine AddRec
     if (!AddRec || !AddRec->isAffine()) {
       LLVM_DEBUG(dbgs() << "SCEV of phi '" << PN
                         << "' is not an affine add recursion, not qualified "
                            "for the terminating condition folding.\n");
       continue;
     }
 
-    // FIXME: This does not properly account for overflow.
+    // Check that we can compute the value of AddRec on the exiting iteration
+    // without soundness problems.  There are two cases to be worried about:
+    // 1) BECount could be 255 with type i8.  Simply adding one would be
+    //    incorrect.  We may need one extra bit to represent the unsigned
+    //    trip count.
+    // 2) The multiplication of stride by TC may wrap around.  This is subtle
+    //    because computing the result accounting for wrap is insufficient.
+    //    In order to use the result in an exit test, we must also know that
+    //    AddRec doesn't take the same value on any previous iteration.
+    //    The simplest case to consider is a candidate IV which is narrower
+    //    than the trip count (and thus original IV), but this can also
+    //    happen due to non-unit strides on the candidate IVs.
+    ConstantRange StepCR = SE.getSignedRange(AddRec->getStepRecurrence(SE));
+    ConstantRange BECountCR = SE.getUnsignedRange(BECount);
+    unsigned NoOverflowBitWidth = BECountCR.getActiveBits() + 1 + StepCR.getMinSignedBits();
+    unsigned ARBitWidth = SE.getTypeSizeInBits(AddRec->getType());
+    if (NoOverflowBitWidth > ARBitWidth)
+      continue;
+
     const SCEV *TermValueSLocal = SE.getAddExpr(
         AddRec->getOperand(0),
         SE.getTruncateOrZeroExtend(
             SE.getMulExpr(
                 AddRec->getOperand(1),
                 SE.getTruncateOrZeroExtend(
                     SE.getAddExpr(BECount, SE.getOne(BECount->getType())),
                     AddRec->getOperand(1)->getType())),

llvm/test/Transforms/LoopStrengthReduce/lsr-term-fold.ll

 
 for.end:                                          ; preds = %for.body
   ret void
 }
 
 ; In this case, the integer IV has a larger bitwidth than the pointer IV.
 ; This means that the smaller IV may wrap around multiple times before
 ; the original loop exit is taken.
-; FIXME: miscompile
 define void @iv_size(ptr %a, i128 %N) {
 ; CHECK-LABEL: @iv_size(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = trunc i128 [[N:%.*]] to i64
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i64 [[TMP0]], 2
-; CHECK-NEXT:    [[SCEVGEP:%.*]] = getelementptr i8, ptr [[A:%.*]], i64 [[TMP1]]
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
-; CHECK-NEXT:    [[LSR_IV1:%.*]] = phi ptr [ [[UGLYGEP2:%.*]], [[FOR_BODY]] ], [ [[A]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[LSR_IV1:%.*]] = phi ptr [ [[UGLYGEP2:%.*]], [[FOR_BODY]] ], [ [[A:%.*]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[LSR_IV:%.*]] = phi i128 [ [[LSR_IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[N:%.*]], [[ENTRY]] ]
 ; CHECK-NEXT:    store i32 1, ptr [[LSR_IV1]], align 4
+; CHECK-NEXT:    [[LSR_IV_NEXT]] = add nsw i128 [[LSR_IV]], -1
 ; CHECK-NEXT:    [[UGLYGEP2]] = getelementptr i8, ptr [[LSR_IV1]], i64 4
-; CHECK-NEXT:    [[LSR_FOLD_TERM_COND_REPLACED_TERM_COND:%.*]] = icmp eq ptr [[UGLYGEP2]], [[SCEVGEP]]
-; CHECK-NEXT:    br i1 [[LSR_FOLD_TERM_COND_REPLACED_TERM_COND]], label [[FOR_END:%.*]], label [[FOR_BODY]]
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i128 [[LSR_IV_NEXT]], 0
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %for.body, %entry
   %lsr.iv1 = phi ptr [ %uglygep2, %for.body ], [ %a, %entry ]