[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

6789

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

6799

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

[reames]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6789

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.

6799

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.

[reames]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6789

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

6789

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

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

for (User *U : V->users()) {
       }
     }
     return VToPN && VToTermCond;
   };
 
   // If this is an IV which we could replace the terminating condition, return
   // the final value of the alternative IV on the last iteration.
   auto getAlternateIVEnd = [&](PHINode &PN) -> const SCEV * {
-    // FIXME: This does not properly account for overflow.
     const SCEVAddRecExpr *AddRec = cast<SCEVAddRecExpr>(SE.getSCEV(&PN));
     const SCEV *BECount = SE.getBackedgeTakenCount(L);
+
+    // 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)
+      return nullptr;
+
     const SCEV *TermValueS = 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())),
             AddRec->getOperand(0)->getType()));
     const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
     SCEVExpander Expander(SE, DL, "lsr_fold_term_cond");
     if (!Expander.isSafeToExpand(TermValueS)) {
       LLVM_DEBUG(
           dbgs() << "Is not safe to expand terminating value for phi node" << PN
                  << "\n");
       return nullptr;
     }

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 ]