[LSR] Generalize one aspect of terminator folding (recently introduced in D132443)

[reames]

Summary

There's no need to require the start value to come directly from the loop predecessor.

[nikic]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6714

The bool return value can be dropped -- check nullptr return instead.

[reames]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6714

I plan to do this in a followup NFC. There's a bunch of other simplification that follows, and some of the naming could stand cleaned up too.

[nikic]

LG

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

6714

Sounds good.

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

154

This is a miscompile, but you are already know that...

[reames]

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

154

I actually hadn't noticed. But yes, that's one of the issued I'd called out on the original review. :)

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

if (!L->isLoopSimplifyForm()) {
     return std::nullopt;
   }
 
   if (!SE.hasLoopInvariantBackedgeTakenCount(L)) {
     LLVM_DEBUG(dbgs() << "Cannot fold on backedge that is loop variant\n");
     return std::nullopt;
   }
 
-  BasicBlock *LoopPreheader = L->getLoopPreheader();
   BasicBlock *LoopLatch = L->getLoopLatch();
 
   // TODO: Can we do something for greater than and less than?
   // Terminating condition is foldable when it is an eq/ne icmp
   BranchInst *BI = cast<BranchInst>(LoopLatch->getTerminator());
   if (BI->isUnconditional())
     return std::nullopt;
   Value *TermCond = BI->getCondition();
auto IsToHelpFold = [&](PHINode &PN) -> std::pair<bool, const SCEV *> {
     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 {false, nullptr};
     }
-    // TODO: Right now we limit the phi node to help the folding be of a start
-    // value of getelementptr. We can extend to any kinds of IV as long as it is
-    // an affine AddRec. Add a switch to cover more types of instructions here
-    // and down in the actual transformation.
-    return {isa<GetElementPtrInst>(PN.getIncomingValueForBlock(LoopPreheader)),
-            TermValueS};
+    return {true, TermValueS};
   };
 
   PHINode *ToFold = nullptr;
   PHINode *ToHelpFold = nullptr;
   const SCEV *TermValueS = nullptr;
 
   for (PHINode &PN : L->getHeader()->phis()) {
     if (!SE.isSCEVable(PN.getType())) {
if (CanFoldTerminatingCondition) {
       Value *LoopValue = ToHelpFold->getIncomingValueForBlock(LoopLatch);
 
       // SCEVExpander for both use in preheader and latch
       const DataLayout &DL = L->getHeader()->getModule()->getDataLayout();
       SCEVExpander Expander(SE, DL, "lsr_fold_term_cond");
       SCEVExpanderCleaner ExpCleaner(Expander);
 
       // Create new terminating value at loop header
-      GetElementPtrInst *StartValueGEP = cast<GetElementPtrInst>(StartValue);
-      Type *PtrTy = StartValueGEP->getPointerOperand()->getType();
-
       const SCEV *TermValueS = CanFoldTerminatingCondition->second.second;
       assert(
           Expander.isSafeToExpand(TermValueS) &&
           "Terminating value was checked safe in canFoldTerminatingCondition");
 
-      Value *TermValue = Expander.expandCodeFor(TermValueS, PtrTy,
+      Value *TermValue = Expander.expandCodeFor(TermValueS, ToHelpFold->getType(),
                                                 LoopPreheader->getTerminator());
 
       LLVM_DEBUG(dbgs() << "Start value of new term-cond phi-node:\n"
                         << *StartValue << "\n"
                         << "Terminating value of new term-cond phi-node:\n"
                         << *TermValue << "\n");
 
       // Create new terminating condition at loop latch

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

for.body4:                                        ; preds = %for.body, %for.body4
   %j.014 = phi i32 [ %mul, %for.body4 ], [ %sub, %for.body ]
   %cnt.113 = phi i32 [ %inc, %for.body4 ], [ %cnt.016, %for.body ]
   %inc = add nsw i32 %cnt.113, 1
   %mul = shl nsw i32 %j.014, 1
   %cmp2 = icmp slt i32 %mul, %m
   br i1 %cmp2, label %for.body4, label %for.cond.cleanup3
 }
 
-; The terminating condition folding transformation cannot find the ptr IV
-; because it checks if the value comes from the LoopPreheader. %mark is from
-; the function argument, so it is not qualified for the transformation.
-define void @no_iv_to_help(ptr %mark, i32 signext %length) {
-; CHECK: Cannot find other AddRec IV to help folding
-entry:
-  %tobool.not3 = icmp eq i32 %length, 0
-  br i1 %tobool.not3, label %for.cond.cleanup, label %for.body
-
-for.cond.cleanup:                                 ; preds = %for.body, %entry
-  ret void
-
-for.body:                                         ; preds = %entry, %for.body
-  %i.05 = phi i32 [ %dec, %for.body ], [ %length, %entry ]
-  %dst.04 = phi ptr [ %incdec.ptr, %for.body ], [ %mark, %entry ]
-  %0 = load ptr, ptr %dst.04, align 8
-  call ptr @foo(ptr %0)
-  %incdec.ptr = getelementptr inbounds ptr, ptr %dst.04, i64 1
-  %dec = add nsw i32 %i.05, -1
-  %tobool.not = icmp eq i32 %dec, 0
-  br i1 %tobool.not, label %for.cond.cleanup, label %for.body
-}
-
 declare void @foo(ptr)
 
 define void @NonAddRecIV(ptr %a) {
 ; CHECK: SCEV of phi '  %lsr.iv = phi i32 [ %lsr.iv.next, %for.body ], [ 1, %entry ]'
 ; CHECK: is not an affine add recursion, not qualified for the terminating condition folding.
 entry:
   %uglygep = getelementptr i8, ptr %a, i32 84
   br label %for.body

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

for.body:                                         ; preds = %entry, %for.body
   %tobool.not = icmp eq i32 %dec, 0
   br i1 %tobool.not, label %for.end, label %for.body
 
 for.end:                                 ; preds = %for.body
   ret void
 }
 
 declare void @foo(ptr)
+
+define void @iv_start_non_preheader(ptr %mark, i32 signext %length) {
+; CHECK-LABEL: @iv_start_non_preheader(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL_NOT3:%.*]] = icmp eq i32 [[LENGTH:%.*]], 0
+; CHECK-NEXT:    br i1 [[TOBOOL_NOT3]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY_PREHEADER:%.*]]
+; CHECK:       for.body.preheader:
+; CHECK-NEXT:    [[TMP0:%.*]] = shl i32 [[LENGTH]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[MARK:%.*]], i32 [[TMP0]]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup.loopexit:
+; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    ret void
+; CHECK:       for.body:
+; CHECK-NEXT:    [[DST_04:%.*]] = phi ptr [ [[INCDEC_PTR:%.*]], [[FOR_BODY]] ], [ [[MARK]], [[FOR_BODY_PREHEADER]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = load ptr, ptr [[DST_04]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = call ptr @foo(ptr [[TMP1]])
+; CHECK-NEXT:    [[INCDEC_PTR]] = getelementptr inbounds ptr, ptr [[DST_04]], i64 1
+; CHECK-NEXT:    [[LSR_FOLD_TERM_COND_REPLACED_TERM_COND:%.*]] = icmp eq ptr [[INCDEC_PTR]], [[UGLYGEP]]
+; CHECK-NEXT:    br i1 [[LSR_FOLD_TERM_COND_REPLACED_TERM_COND]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], label [[FOR_BODY]]
+;
+entry:
+  %tobool.not3 = icmp eq i32 %length, 0
+  br i1 %tobool.not3, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+
+for.body:                                         ; preds = %entry, %for.body
+  %i.05 = phi i32 [ %dec, %for.body ], [ %length, %entry ]
+  %dst.04 = phi ptr [ %incdec.ptr, %for.body ], [ %mark, %entry ]
+  %0 = load ptr, ptr %dst.04, align 8
+  call ptr @foo(ptr %0)
+  %incdec.ptr = getelementptr inbounds ptr, ptr %dst.04, i64 1
+  %dec = add nsw i32 %i.05, -1
+  %tobool.not = icmp eq i32 %dec, 0
+  br i1 %tobool.not, label %for.cond.cleanup, label %for.body
+}