[LSR] Allow terminating condition fold on start values by BinaryOperator

[eopXD]

Summary

Orginally only GEP instructions are allowed, this patch allows BinaryOperator
instructions to pass too.

[fhahn]

It looks like this needs additional tests.

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

const SCEV *TermValueS = SE.getAddExpr(
             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))
+      return false;
 
-    // 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 Expander.isSafeToExpand(TermValueS) &&
-           isa<GetElementPtrInst>(PN.getIncomingValueForBlock(LoopPreheader));
+    Value *StartValue = PN.getIncomingValueForBlock(LoopPreheader);
+    // Note: Expanding more types of StartValue requires corresponding change
+    // for `TypeToExpand` below
+    if (isa<GetElementPtrInst>(StartValue) || isa<BinaryOperator>(StartValue))
+      return true;
+    return false;
   };
 
   PHINode *ToFold = nullptr;
   PHINode *ToHelpFold = nullptr;
 
   for (PHINode &PN : L->getHeader()->phis()) {
     if (!SE.isSCEVable(PN.getType())) {
       LLVM_DEBUG(dbgs() << "IV of phi '" << PN
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();
+      Type *TypeToExpand = nullptr;
+      if (isa<GetElementPtrInst>(StartValue)) {
+        TypeToExpand =
+            cast<GetElementPtrInst>(StartValue)->getPointerOperandType();
+      } else if (isa<BinaryOperator>(StartValue)) {
+        TypeToExpand = StartValue->getType();
+      } else {
+        llvm_unreachable("Unhandled type should be guarded by IsToHelpFold "
+                         "under canFoldTermCondOfLoop");
+      }
 
       const SCEV *BECount = SE.getBackedgeTakenCount(L);
       const SCEVAddRecExpr *AddRec =
           cast<SCEVAddRecExpr>(SE.getSCEV(ToHelpFold));
 
       // TermValue = Start + Stride * (BackedgeCount + 1)
       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()));
 
       Changed = true;
       NumTermFold++;
 
-      Value *TermValue = Expander.expandCodeFor(TermValueS, PtrTy,
+      Value *TermValue = Expander.expandCodeFor(TermValueS, TypeToExpand,
                                                 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