[IndVarSimplify] Extend trip count instead of truncating IV in LFTR, when original IV does not overflow

[amehsan]

Summary

When Original IV, before widening, does not overflow, extending trip count in the loop control logic generates better code compared to truncating IV. This is because

(1) extending trip count is a loop invariant operation (see genLoopLimit where we prove trip count is loop invariant).
(2) Scalar Evolution seems to have problems understanding trunc. So removing them allows better analysis performed in Scalar Evolution. (In particular this fixes PR 28363 which is the motivation for this change).

So far I have tested this with unit-tests and test suite on pwr. I am planning to do more functional testing and also add some unit-tests. But I thought to post this here, to get some feedback.

I am not going to perform any performance test. Any degradation caused by this should be an indication of a bug elsewhere.

[mehdi_amini]

Please reopen with llvm-commits as a subscriber since the beginning.

[amehsan]

@mehdi_amini done

[amehsan]

abandoning this if you don't mind @mehdi_amini

lib/Transforms/Scalar/IndVarSimplify.cpp

cl::values(clEnumValN(NeverRepl, "never", "never replace exit value"),
                           "only replace exit value when the cost is cheap"),
                clEnumValN(AlwaysRepl, "always",
                           "always replace exit value whenever possible"),
                clEnumValEnd));
 
 namespace {
 struct RewritePhi;
 
+struct OriginalIVInfo {
+  bool HasNUW;
+  bool HasNSW;
+  bool IsSigned;
+  Type *OriginalType;
+
+  OriginalIVInfo(bool HasNUW, bool HasNSW,
+                 bool IsSigned, Type *OriginalType) :
+                 HasNUW(HasNUW), HasNSW(HasNSW),
+                 IsSigned(IsSigned),
+                 OriginalType(OriginalType){};
+
+  OriginalIVInfo() {};
+
+};
+
+typedef ValueMap<Value *, std::unique_ptr<OriginalIVInfo>> WidenedIVMapType;
+
 class IndVarSimplify {
   LoopInfo *LI;
   ScalarEvolution *SE;
   DominatorTree *DT;
   const DataLayout &DL;
   TargetLibraryInfo *TLI;
   const TargetTransformInfo *TTI;
 
+  WidenedIVMapType WidenedIVMap;
+
   SmallVector<WeakVH, 16> DeadInsts;
   bool Changed = false;
 
   bool isValidRewrite(Value *FromVal, Value *ToVal);
 
   void handleFloatingPointIV(Loop *L, PHINode *PH);
   void rewriteNonIntegerIVs(Loop *L);
 
WidenIV(const WideIVInfo &WI, LoopInfo *LInfo,
     DT(DTree),
     WidePhi(nullptr),
     WideInc(nullptr),
     WideIncExpr(nullptr),
     DeadInsts(DI) {
     assert(L->getHeader() == OrigPhi->getParent() && "Phi must be an IV");
   }
 
-  PHINode *createWideIV(SCEVExpander &Rewriter);
+  PHINode *createWideIV(SCEVExpander &Rewriter, WidenedIVMapType &WIMap);
 
 protected:
   Value *createExtendInst(Value *NarrowOper, Type *WideType, bool IsSigned,
                           Instruction *Use);
 
   Instruction *cloneIVUser(NarrowIVDefUse DU, const SCEVAddRecExpr *WideAR);
   Instruction *cloneArithmeticIVUser(NarrowIVDefUse DU,
                                      const SCEVAddRecExpr *WideAR);

 /// wide induction variable that evaluates to the same recurrence as the
 /// original narrow IV. Then use a worklist to forward traverse the narrow IV's
 /// def-use chain. After widenIVUse has processed all interesting IV users, the
 /// narrow IV will be isolated for removal by DeleteDeadPHIs.
 ///
 /// It would be simpler to delete uses as they are processed, but we must avoid
 /// invalidating SCEV expressions.
 ///
-PHINode *WidenIV::createWideIV(SCEVExpander &Rewriter) {
+PHINode *WidenIV::createWideIV(SCEVExpander &Rewriter, WidenedIVMapType &WIMap) {
   // Is this phi an induction variable?
   const SCEVAddRecExpr *AddRec = dyn_cast<SCEVAddRecExpr>(SE->getSCEV(OrigPhi));
   if (!AddRec)
     return nullptr;
 
   // Widen the induction variable expression.
   const SCEV *WideIVExpr = IsSigned ?
     SE->getSignExtendExpr(AddRec, WideType) :
PHINode *WidenIV::createWideIV(SCEVExpander &Rewriter, WidenedIVMapType &WIMap) {
 
   // The rewriter provides a value for the desired IV expression. This may
   // either find an existing phi or materialize a new one. Either way, we
   // expect a well-formed cyclic phi-with-increments. i.e. any operand not part
   // of the phi-SCC dominates the loop entry.
   Instruction *InsertPt = &L->getHeader()->front();
   WidePhi = cast<PHINode>(Rewriter.expandCodeFor(AddRec, WideType, InsertPt));
 
+  WIMap[WidePhi] = std::unique_ptr<OriginalIVInfo>
+                        (new OriginalIVInfo(AddRec->hasNoUnsignedWrap(),
+                                            AddRec->hasNoSignedWrap(),
+                                            IsSigned,
+                                            AddRec->getType()));
+
   // Remembering the WideIV increment generated by SCEVExpander allows
   // widenIVUse to reuse it when widening the narrow IV's increment. We don't
   // employ a general reuse mechanism because the call above is the only call to
   // SCEVExpander. Henceforth, we produce 1-to-1 narrow to wide uses.
   if (BasicBlock *LatchBlock = L->getLoopLatch()) {
     WideInc =
       cast<Instruction>(WidePhi->getIncomingValueForBlock(LatchBlock));
     WideIncExpr = SE->getSCEV(WideInc);
do {
 
       if (Visitor.WI.WidestNativeType) {
         WideIVs.push_back(Visitor.WI);
       }
     } while(!LoopPhis.empty());
 
     for (; !WideIVs.empty(); WideIVs.pop_back()) {
       WidenIV Widener(WideIVs.back(), LI, SE, DT, DeadInsts);
-      if (PHINode *WidePhi = Widener.createWideIV(Rewriter)) {
+      if (PHINode *WidePhi = Widener.createWideIV(Rewriter, WidenedIVMap)) {
         Changed = true;
         LoopPhis.push_back(WidePhi);
       }
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
if (isa<SCEVConstant>(ARStart) && isa<SCEVConstant>(IVCount)) {
       if (cast<SCEVConstant>(ARStep)->getValue()->isNegative())
         NewLimit = Start - Count;
       else
         NewLimit = Start + Count;
       ExitCnt = ConstantInt::get(CmpIndVar->getType(), NewLimit);
 
       DEBUG(dbgs() << "  Widen RHS:\t" << *ExitCnt << "\n");
     } else {
+      bool extended = false;
+      if (WidenedIVMap.count(IndVar) &&
+         (WidenedIVMap[IndVar]->OriginalType == ExitCnt->getType())) {
+
+        if (WidenedIVMap[IndVar]->HasNSW && WidenedIVMap[IndVar]->IsSigned) {
+          extended = true;
+          ExitCnt = Builder.CreateSExt(ExitCnt, IndVar->getType(),
+                                       "wide.trip.count");
+        }
+        if (WidenedIVMap[IndVar]->HasNUW && !WidenedIVMap[IndVar]->IsSigned) {
+          extended = true;
+          ExitCnt = Builder.CreateZExt(ExitCnt, IndVar->getType(),
+                                       "wide.trip.count");
+        }
+      }
+      if (!extended) {
         CmpIndVar = Builder.CreateTrunc(CmpIndVar, ExitCnt->getType(),
                                         "lftr.wideiv");
       }
     }
+  }
   Value *Cond = Builder.CreateICmp(P, CmpIndVar, ExitCnt, "exitcond");
   Value *OrigCond = BI->getCondition();
   // It's tempting to use replaceAllUsesWith here to fully replace the old
   // comparison, but that's not immediately safe, since users of the old
   // comparison may not be dominated by the new comparison. Instead, just
   // update the branch to use the new comparison; in the common case this
   // will make old comparison dead.
   BI->setCondition(Cond);
if (IndVar) {
       // FIXME: SCEV expansion has no way to bail out, so the caller must
       // explicitly check any assumptions made by SCEV. Brittle.
       const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(BackedgeTakenCount);
       if (!AR || AR->getLoop()->getLoopPreheader())
         (void)linearFunctionTestReplace(L, BackedgeTakenCount, IndVar,
                                         Rewriter);
     }
   }
+  WidenedIVMap.clear();
   // Clear the rewriter cache, because values that are in the rewriter's cache
   // can be deleted in the loop below, causing the AssertingVH in the cache to
   // trigger.
   Rewriter.clear();
 
   // Now that we're done iterating through lists, clean up any instructions
   // which are now dead.
   while (!DeadInsts.empty())

test/Transforms/IndVarSimplify/elim-extend.ll

-; RUN: opt < %s -indvars -S | FileCheck %s
+; RUN: opt < %s -indvars -S | grep -v wide.trip.count| FileCheck %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
 
 ; IV with constant start, preinc and postinc sign extends, with and without NSW.
 ; IV rewrite only removes one sext. WidenIVs removes all three.
 define void @postincConstIV(i8* %base, i32 %limit) nounwind {
 entry:
   br label %loop

test/Transforms/IndVarSimplify/preserve-signed-wrap.ll

-; RUN: opt < %s -indvars -S | FileCheck %s
+; RUN: opt < %s -indvars -S | grep -v wide.trip.count | FileCheck %s
 
 ; Indvars should insert a 64-bit induction variable to eliminate the
 ; sext for the addressing, however it shouldn't eliminate the sext
 ; on the other phi, since that value undergoes signed wrapping.
 
 ; Provide legal integer types.
 target datalayout = "n8:16:32:64"
 

test/Transforms/IndVarSimplify/promote-iv-to-eliminate-casts.ll

-; RUN: opt < %s -indvars -S | not grep sext
+; RUN: opt < %s -indvars -S | grep -v wide.trip.count | not grep sext
 
 ; Provide legal integer types.
 target datalayout = "n8:16:32:64"
 
 
 define i64 @test(i64* nocapture %first, i32 %count) nounwind readonly {
 entry:
   %t0 = icmp sgt i32 %count, 0    ; <i1> [#uses=1]

test/Transforms/IndVarSimplify/signed-trip-count.ll

-; RUN: opt < %s -indvars -S > %t
+; RUN: opt < %s -indvars -S | grep -v wide.trip.count > %t
 ; RUN: not grep sext %t
 ; RUN: grep phi %t | count 1
 
 ; Provide legal integer types.
 target datalayout = "n8:16:32:64"
 
 
 define void @foo(i64* nocapture %x, i32 %n) nounwind {