[SimplifyIndVar] Push more users to worklist for simplifyUsers

Instead of considering only users that are inside the loop,
consider all the users with parent dominated by the loop header.

Author: v01dxyz

llvm/include/llvm/Transforms/Utils/SimplifyIndVar.h

@@ -52,12 +52,11 @@ class IVVisitor {
 /// where the first entry indicates that the function makes changes and the
 /// second entry indicates that it introduced new opportunities for loop
 /// unswitching.
-std::pair<bool, bool> simplifyUsersOfIV(PHINode *CurrIV, ScalarEvolution *SE,
-                                        DominatorTree *DT, LoopInfo *LI,
-                                        const TargetTransformInfo *TTI,
-                                        SmallVectorImpl<WeakTrackingVH> &Dead,
-                                        SCEVExpander &Rewriter,
-                                        IVVisitor *V = nullptr);
+std::pair<bool, bool>
+simplifyUsersOfIV(PHINode *CurrIV, ScalarEvolution *SE, DominatorTree *DT,
+                  LoopInfo *LI, const TargetTransformInfo *TTI,
+                  SmallVectorImpl<WeakTrackingVH> &Dead, SCEVExpander &Rewriter,
+                  unsigned MaxDepthOutOfLoop = 1, IVVisitor *V = nullptr);
 
 /// SimplifyLoopIVs - Simplify users of induction variables within this
 /// loop. This does not actually change or add IVs.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

@@ -124,6 +124,12 @@ static cl::opt<bool>
 AllowIVWidening("indvars-widen-indvars", cl::Hidden, cl::init(true),
                 cl::desc("Allow widening of indvars to eliminate s/zext"));
 
+static cl::opt<unsigned> MaxDepthOutOfLoop(
+    "indvars-max-depth-out-of-loop", cl::Hidden, cl::init(1),
+    cl::desc(
+        "Strict upper bound for the number of successive out-of-loop blocks "
+        "when traversing use-def chains. 0 enables full traversal"));
+
 namespace {
 
 class IndVarSimplify {
@@ -624,8 +630,9 @@ bool IndVarSimplify::simplifyAndExtend(Loop *L,
       // Information about sign/zero extensions of CurrIV.
       IndVarSimplifyVisitor Visitor(CurrIV, SE, TTI, DT);
 
-      const auto &[C, U] = simplifyUsersOfIV(CurrIV, SE, DT, LI, TTI, DeadInsts,
-                                             Rewriter, &Visitor);
+      const auto &[C, U] =
+          simplifyUsersOfIV(CurrIV, SE, DT, LI, TTI, DeadInsts, Rewriter,
+                            MaxDepthOutOfLoop, &Visitor);
 
       Changed |= C;
       RunUnswitching |= U;

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

@@ -62,13 +62,18 @@ namespace {
     bool Changed = false;
     bool RunUnswitching = false;
 
+    // When following the def-use chains, it can go outside the loop.
+    // Strict upper bound on number of traversed out-of-loop blocks.
+    unsigned MaxDepthOutOfLoop;
+
   public:
     SimplifyIndvar(Loop *Loop, ScalarEvolution *SE, DominatorTree *DT,
                    LoopInfo *LI, const TargetTransformInfo *TTI,
                    SCEVExpander &Rewriter,
-                   SmallVectorImpl<WeakTrackingVH> &Dead)
+                   SmallVectorImpl<WeakTrackingVH> &Dead,
+                   unsigned MaxDepthOutOfLoop = 1)
         : L(Loop), LI(LI), SE(SE), DT(DT), TTI(TTI), Rewriter(Rewriter),
-          DeadInsts(Dead) {
+          DeadInsts(Dead), MaxDepthOutOfLoop(MaxDepthOutOfLoop) {
       assert(LI && "IV simplification requires LoopInfo");
     }
 
@@ -509,8 +514,8 @@ bool SimplifyIndvar::eliminateTrunc(TruncInst *TI) {
         !DT->isReachableFromEntry(cast<Instruction>(U)->getParent()))
       continue;
     ICmpInst *ICI = dyn_cast<ICmpInst>(U);
-    if (!ICI) return false;
-    assert(L->contains(ICI->getParent()) && "LCSSA form broken?");
+    if (!ICI)
+      return false;
     if (!(ICI->getOperand(0) == TI && L->isLoopInvariant(ICI->getOperand(1))) &&
         !(ICI->getOperand(1) == TI && L->isLoopInvariant(ICI->getOperand(0))))
       return false;
@@ -839,10 +844,12 @@ bool SimplifyIndvar::strengthenRightShift(BinaryOperator *BO,
 }
 
 /// Add all uses of Def to the current IV's worklist.
-static void pushIVUsers(
-  Instruction *Def, Loop *L,
-  SmallPtrSet<Instruction*,16> &Simplified,
-  SmallVectorImpl< std::pair<Instruction*,Instruction*> > &SimpleIVUsers) {
+static void
+pushIVUsers(Instruction *Def, Loop *L, DominatorTree *DT,
+            SmallPtrSet<Instruction *, 16> &Simplified,
+            SmallVectorImpl<std::tuple<Instruction *, Instruction *, unsigned>>
+                &SimpleIVUsers,
+            unsigned OutOfLoopChainCounter, unsigned MaxOutOfLoopChainCounter) {
 
   for (User *U : Def->users()) {
     Instruction *UI = cast<Instruction>(U);
@@ -854,16 +861,23 @@ static void pushIVUsers(
     if (UI == Def)
       continue;
 
-    // Only change the current Loop, do not change the other parts (e.g. other
-    // Loops).
-    if (!L->contains(UI))
+    // Avoid adding Defs that SCEV expand to themselves, e.g. the LoopPhis
+    // of the outter loops.
+    if (!DT->dominates(L->getHeader(), UI->getParent()))
       continue;
 
     // Do not push the same instruction more than once.
     if (!Simplified.insert(UI).second)
       continue;
 
-    SimpleIVUsers.push_back(std::make_pair(UI, Def));
+    unsigned Counter =
+        L->contains(UI)
+            ? 0 // reset depth if we go back inside the loop.
+            : OutOfLoopChainCounter + (UI->getParent() != Def->getParent());
+
+    if (!MaxOutOfLoopChainCounter || Counter < MaxOutOfLoopChainCounter) {
+      SimpleIVUsers.push_back(std::make_tuple(UI, Def, Counter));
+    }
   }
 }
 
@@ -908,17 +922,17 @@ void SimplifyIndvar::simplifyUsers(PHINode *CurrIV, IVVisitor *V) {
   SmallPtrSet<Instruction*,16> Simplified;
 
   // Use-def pairs if IV users waiting to be processed for CurrIV.
-  SmallVector<std::pair<Instruction*, Instruction*>, 8> SimpleIVUsers;
+  SmallVector<std::tuple<Instruction *, Instruction *, unsigned>, 8>
+      SimpleIVUsers;
 
   // Push users of the current LoopPhi. In rare cases, pushIVUsers may be
   // called multiple times for the same LoopPhi. This is the proper thing to
   // do for loop header phis that use each other.
-  pushIVUsers(CurrIV, L, Simplified, SimpleIVUsers);
+  pushIVUsers(CurrIV, L, DT, Simplified, SimpleIVUsers, 0, MaxDepthOutOfLoop);
 
   while (!SimpleIVUsers.empty()) {
-    std::pair<Instruction*, Instruction*> UseOper =
-      SimpleIVUsers.pop_back_val();
-    Instruction *UseInst = UseOper.first;
+    auto [UseInst, IVOperand, OutOfLoopChainCounter] =
+        SimpleIVUsers.pop_back_val();
 
     // If a user of the IndVar is trivially dead, we prefer just to mark it dead
     // rather than try to do some complex analysis or transformation (such as
@@ -942,11 +956,11 @@ void SimplifyIndvar::simplifyUsers(PHINode *CurrIV, IVVisitor *V) {
     if ((isa<PtrToIntInst>(UseInst)) || (isa<TruncInst>(UseInst)))
       for (Use &U : UseInst->uses()) {
         Instruction *User = cast<Instruction>(U.getUser());
-        if (replaceIVUserWithLoopInvariant(User))
+        if (DT->dominates(L->getHeader(), User->getParent()) &&
+            replaceIVUserWithLoopInvariant(User))
           break; // done replacing
       }
 
-    Instruction *IVOperand = UseOper.second;
     for (unsigned N = 0; IVOperand; ++N) {
       assert(N <= Simplified.size() && "runaway iteration");
       (void) N;
@@ -960,22 +974,25 @@ void SimplifyIndvar::simplifyUsers(PHINode *CurrIV, IVVisitor *V) {
       continue;
 
     if (eliminateIVUser(UseInst, IVOperand)) {
-      pushIVUsers(IVOperand, L, Simplified, SimpleIVUsers);
+      pushIVUsers(IVOperand, L, DT, Simplified, SimpleIVUsers,
+                  OutOfLoopChainCounter, MaxDepthOutOfLoop);
       continue;
     }
 
     if (BinaryOperator *BO = dyn_cast<BinaryOperator>(UseInst)) {
       if (strengthenBinaryOp(BO, IVOperand)) {
         // re-queue uses of the now modified binary operator and fall
         // through to the checks that remain.
-        pushIVUsers(IVOperand, L, Simplified, SimpleIVUsers);
+        pushIVUsers(IVOperand, L, DT, Simplified, SimpleIVUsers,
+                    OutOfLoopChainCounter, MaxDepthOutOfLoop);
       }
     }
 
     // Try to use integer induction for FPToSI of float induction directly.
     if (replaceFloatIVWithIntegerIV(UseInst)) {
       // Re-queue the potentially new direct uses of IVOperand.
-      pushIVUsers(IVOperand, L, Simplified, SimpleIVUsers);
+      pushIVUsers(IVOperand, L, DT, Simplified, SimpleIVUsers,
+                  OutOfLoopChainCounter, MaxDepthOutOfLoop);
       continue;
     }
 
@@ -985,7 +1002,8 @@ void SimplifyIndvar::simplifyUsers(PHINode *CurrIV, IVVisitor *V) {
       continue;
     }
     if (isSimpleIVUser(UseInst, L, SE)) {
-      pushIVUsers(UseInst, L, Simplified, SimpleIVUsers);
+      pushIVUsers(UseInst, L, DT, Simplified, SimpleIVUsers,
+                  OutOfLoopChainCounter, MaxDepthOutOfLoop);
     }
   }
 }
@@ -999,13 +1017,13 @@ void IVVisitor::anchor() { }
 ///  Returns a pair where the first entry indicates that the function makes
 ///  changes and the second entry indicates that it introduced new opportunities
 ///  for loop unswitching.
-std::pair<bool, bool> simplifyUsersOfIV(PHINode *CurrIV, ScalarEvolution *SE,
-                                        DominatorTree *DT, LoopInfo *LI,
-                                        const TargetTransformInfo *TTI,
-                                        SmallVectorImpl<WeakTrackingVH> &Dead,
-                                        SCEVExpander &Rewriter, IVVisitor *V) {
+std::pair<bool, bool>
+simplifyUsersOfIV(PHINode *CurrIV, ScalarEvolution *SE, DominatorTree *DT,
+                  LoopInfo *LI, const TargetTransformInfo *TTI,
+                  SmallVectorImpl<WeakTrackingVH> &Dead, SCEVExpander &Rewriter,
+                  unsigned MaxDepthOutOfLoop, IVVisitor *V) {
   SimplifyIndvar SIV(LI->getLoopFor(CurrIV->getParent()), SE, DT, LI, TTI,
-                     Rewriter, Dead);
+                     Rewriter, Dead, MaxDepthOutOfLoop);
   SIV.simplifyUsers(CurrIV, V);
   return {SIV.hasChanged(), SIV.runUnswitching()};
 }