[SimplifyIndVar] Replace IVUsers with loop invariant if possible

[efriedma]

lib/Transforms/Utils/SimplifyIndVar.cpp

551 ↗

(On Diff #117177)

Could you use SCEVExpander rather than writing out the cast logic here? I'm not sure this covers all the relevant cases.

557 ↗

(On Diff #117177)

Does it actually matter whether the value here is a Constant?

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

551 ↗

(On Diff #117177)

Ok, will do

557 ↗

(On Diff #117177)

You means we gain something if we can "simplify" the IV user from a SCEVAddRec to a SCEVUnknown?*

SE->getSCEV(I); may directly return an SCEUnknown, so we need to compare the new LLVM Value and the original one.
We are not going to do the replace if they are the same

[efriedma]

lib/Transforms/Utils/SimplifyIndVar.cpp

557 ↗

(On Diff #117177)

Rather than try to check for equality, you can just check whether the SCEVUnknown is loop-invariant.

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

557 ↗

(On Diff #117177)

This sounds great

[etherzhhb]

Address reviewer's comments

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

553–554 ↗

(On Diff #117260)

Should we only expand SCEVConstant and SCEVUnknown?

If we should expand all SCEV that is Loop Invariant, should we make the SCEVExpander global, such that we have a global "InsertedExpressions" map (in the SCEVExpander) to remember the SCEV expression that we expanded?

[etherzhhb]

Looks like I need to pass the SCEVExpander from the IndVarSimpilfy pass to coordinate SCEV rewrite, otherwise we will fail test/Transforms/IndVarSimplify/pr20680.ll

[etherzhhb]

Update format

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

552 ↗

(On Diff #117271)

We may only want to do the replace if S is invariant of the IV's Loop, which may be different from "L" which is the innermost loop that contains the IV user

[efriedma]

lib/Transforms/Utils/SimplifyIndVar.cpp

552 ↗

(On Diff #117271)

Yes; otherwise, the transform isn't doing anything useful.

We might need to call isHighCostExpansion here, so we don't generate something ridiculous?

555 ↗

(On Diff #117271)

"I" is probably not the right place the put any expansion code, given it's inside the loop.

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

555 ↗

(On Diff #117271)

SCEVExpander figure out the correct place by itself (at least in test4 of test/Transforms/IndVarSimplify/replace-iv-with-loop-invariant.ll), otherwise we need to find the preheader (and fallback to I if there is no preheader)

test/Transforms/IndVarSimplify/replace-iv-with-loop-invariant.ll

84 ↗

(On Diff #117271)

Here the SCEVExpander expand the expression before the loop (before the phi)

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

552 ↗

(On Diff #117271)

Yes, this make sense.

I did some experiments and we may want to actually write:

if (!SE->isLoopInvariant(S, IV's Loop) || !SE->isLoopInvariant(S, IV User's Loop))

If IV's Loop do not contains IV User's Loop

[efriedma]

lib/Transforms/Utils/SimplifyIndVar.cpp

552 ↗

(On Diff #117271)

Given a function in LCSSA form, I don't think this relevant? But feel free to prove me wrong. :)

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

552 ↗

(On Diff #117271)

Not sure if we assume LCSSA (or SCEV could see through it), I hit some crash when I just test 'if (!SE->isLoopInvariant(S, IV's Loop))' and it turned out that IV's Loop do not contain IV's user. I will find those cases

[etherzhhb]

Partially address reviewer's comment. This is a WIP patch

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

553 ↗

(On Diff #117542)

Maybe we should not push IV users in pushIVUsers at all if IV users do not in the IV loop (e.g. LCSSA phi)

[efriedma]

lib/Transforms/Utils/SimplifyIndVar.cpp

553 ↗

(On Diff #117542)

Yeah, that makes sense; we have special handling for exit values in IndVarSimplify::rewriteLoopExitValues.

[etherzhhb]

Address reviewer comments

[etherzhhb]

Ping?

[hfinkel]

lib/Transforms/Utils/SimplifyIndVar.cpp

562 ↗

(On Diff #117618)

End comments with a period (here and also below).

[etherzhhb]

Update comments

[efriedma]

lib/Transforms/Utils/SimplifyIndVar.cpp

557 ↗

(On Diff #118312)

SE->getSCEVAtScope(SE->getSCEV(I), LI->getLoopFor(I->getParent())) should be equivalent to SE->getSCEV(I). The point of getSCEVAtScope is to compute the value of a SCEV outside of its associated loop (the exit value).

Sorry, I should have spotted this earlier.

[etherzhhb]

lib/Transforms/Utils/SimplifyIndVar.cpp

557 ↗

(On Diff #118312)

Thanks for pointing this out, I will update the patch

[etherzhhb]

Delete the redundant getSCEVAtScope

[efriedma]

LGTM.

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

 namespace llvm {
 
 class CastInst;
 class DominatorTree;
 class Loop;
 class LoopInfo;
 class PHINode;
 class ScalarEvolution;
+class SCEVExpander;
 
 /// Interface for visiting interesting IV users that are recognized but not
 /// simplified by this utility.
 class IVVisitor {
 protected:
   const DominatorTree *DT = nullptr;
 
   virtual void anchor();
 
 public:
   IVVisitor() = default;
   virtual ~IVVisitor() = default;
 
   const DominatorTree *getDomTree() const { return DT; }
   virtual void visitCast(CastInst *Cast) = 0;
 };
 
 /// simplifyUsersOfIV - Simplify instructions that use this induction variable
 /// by using ScalarEvolution to analyze the IV's recurrence.
 bool simplifyUsersOfIV(PHINode *CurrIV, ScalarEvolution *SE, DominatorTree *DT,
                        LoopInfo *LI, SmallVectorImpl<WeakTrackingVH> &Dead,
-                       IVVisitor *V = nullptr);
+                       SCEVExpander &Rewriter, IVVisitor *V = nullptr);
 
 /// SimplifyLoopIVs - Simplify users of induction variables within this
 /// loop. This does not actually change or add IVs.
 bool simplifyLoopIVs(Loop *L, ScalarEvolution *SE, DominatorTree *DT,
                      LoopInfo *LI, SmallVectorImpl<WeakTrackingVH> &Dead);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_SIMPLIFYINDVAR_H

llvm/trunk/lib/Transforms/Scalar/IndVarSimplify.cpp

while (!LoopPhis.empty()) {
     // evaluation of sign/zero extend evaluation until needed, and avoid running
     // other SCEV based analysis prior to simplifyAndExtend.
     do {
       PHINode *CurrIV = LoopPhis.pop_back_val();
 
       // Information about sign/zero extensions of CurrIV.
       IndVarSimplifyVisitor Visitor(CurrIV, SE, TTI, DT);
 
-      Changed |= simplifyUsersOfIV(CurrIV, SE, DT, LI, DeadInsts, &Visitor);
+      Changed |=
+          simplifyUsersOfIV(CurrIV, SE, DT, LI, DeadInsts, Rewriter, &Visitor);
 
       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, HasGuards);

llvm/trunk/lib/Transforms/Utils/SimplifyIndVar.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/SimplifyIndVar.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
-#include "llvm/Analysis/ScalarEvolutionExpressions.h"
+#include "llvm/Analysis/ScalarEvolutionExpander.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
namespace {
   /// based on ScalarEvolution. It is the primary instrument of the
   /// IndvarSimplify pass, but it may also be directly invoked to cleanup after
   /// other loop passes that preserve SCEV.
   class SimplifyIndvar {
     Loop             *L;
     LoopInfo         *LI;
     ScalarEvolution  *SE;
     DominatorTree    *DT;
-
+    SCEVExpander     &Rewriter;
     SmallVectorImpl<WeakTrackingVH> &DeadInsts;
 
     bool Changed;
 
   public:
     SimplifyIndvar(Loop *Loop, ScalarEvolution *SE, DominatorTree *DT,
-                   LoopInfo *LI, SmallVectorImpl<WeakTrackingVH> &Dead)
-        : L(Loop), LI(LI), SE(SE), DT(DT), DeadInsts(Dead), Changed(false) {
+                   LoopInfo *LI, SCEVExpander &Rewriter,
+                   SmallVectorImpl<WeakTrackingVH> &Dead)
+        : L(Loop), LI(LI), SE(SE), DT(DT), Rewriter(Rewriter), DeadInsts(Dead),
+          Changed(false) {
       assert(LI && "IV simplification requires LoopInfo");
     }
 
     bool hasChanged() const { return Changed; }
 
     /// Iteratively perform simplification on a worklist of users of the
     /// specified induction variable. This is the top-level driver that applies
     /// all simplifications to users of an IV.
     void simplifyUsers(PHINode *CurrIV, IVVisitor *V = nullptr);
 
     Value *foldIVUser(Instruction *UseInst, Instruction *IVOperand);
 
     bool eliminateIdentitySCEV(Instruction *UseInst, Instruction *IVOperand);
-    bool foldConstantSCEV(Instruction *UseInst);
+    bool replaceIVUserWithLoopInvariant(Instruction *UseInst);
 
     bool eliminateOverflowIntrinsic(CallInst *CI);
     bool eliminateIVUser(Instruction *UseInst, Instruction *IVOperand);
     void eliminateIVComparison(ICmpInst *ICmp, Value *IVOperand);
     void simplifyIVRemainder(BinaryOperator *Rem, Value *IVOperand,
                              bool IsSigned);
     void replaceRemWithNumerator(BinaryOperator *Rem);
     void replaceRemWithNumeratorOrZero(BinaryOperator *Rem);
if (eliminateOverflowIntrinsic(CI))
       return true;
 
   if (eliminateIdentitySCEV(UseInst, IVOperand))
     return true;
 
   return false;
 }
 
-/// Replace the UseInst with a constant if possible
-bool SimplifyIndvar::foldConstantSCEV(Instruction *I) {
+static Instruction *GetLoopInvariantInsertPosition(Loop *L, Instruction *Hint) {
+  if (auto *BB = L->getLoopPreheader())
+    return BB->getTerminator();
+
+  return Hint;
+}
+
+/// Replace the UseInst with a constant if possible.
+bool SimplifyIndvar::replaceIVUserWithLoopInvariant(Instruction *I) {
   if (!SE->isSCEVable(I->getType()))
     return false;
 
   // Get the symbolic expression for this instruction.
   const SCEV *S = SE->getSCEV(I);
 
-  const Loop *L = LI->getLoopFor(I->getParent());
-  S = SE->getSCEVAtScope(S, L);
-  auto *C = dyn_cast<SCEVConstant>(S);
-
-  if (!C)
+  if (!SE->isLoopInvariant(S, L))
     return false;
 
-  Constant *V = C->getValue();
-  // The SCEV will have a different type than the instruction if the instruction
-  // has a pointer type. Skip the replacement
-  // TODO: Replace ConstantInt Zero by ConstantPointerNull
-  if (V->getType() != I->getType())
+  // Do not generate something ridiculous even if S is loop invariant.
+  if (Rewriter.isHighCostExpansion(S, L, I))
     return false;
 
-  I->replaceAllUsesWith(V);
-  DEBUG(dbgs() << "INDVARS: Replace IV user: " << *I << " with constant: " << *C
-               << '\n');
+  auto *IP = GetLoopInvariantInsertPosition(L, I);
+  auto *Invariant = Rewriter.expandCodeFor(S, I->getType(), IP);
+
+  I->replaceAllUsesWith(Invariant);
+  DEBUG(dbgs() << "INDVARS: Replace IV user: " << *I
+               << " with loop invariant: " << *S << '\n');
   ++NumFoldedUser;
   Changed = true;
   DeadInsts.emplace_back(I);
   return true;
 }
 
 /// Eliminate any operation that SCEV can prove is an identity function.
 bool SimplifyIndvar::eliminateIdentitySCEV(Instruction *UseInst,
if (BO->getOpcode() == Instruction::Shl) {
     return Changed;
   }
 
   return false;
 }
 
 /// Add all uses of Def to the current IV's worklist.
 static void pushIVUsers(
-  Instruction *Def,
+  Instruction *Def, Loop *L,
   SmallPtrSet<Instruction*,16> &Simplified,
   SmallVectorImpl< std::pair<Instruction*,Instruction*> > &SimpleIVUsers) {
 
   for (User *U : Def->users()) {
     Instruction *UI = cast<Instruction>(U);
 
     // Avoid infinite or exponential worklist processing.
     // Also ensure unique worklist users.
     // If Def is a LoopPhi, it may not be in the Simplified set, so check for
     // self edges first.
-    if (UI != Def && Simplified.insert(UI).second)
-      SimpleIVUsers.push_back(std::make_pair(UI, Def));
+    if (UI == Def)
+      continue;
+
+    // Only change the current Loop, do not change the other parts (e.g. other
+    // Loops).
+    if (!L->contains(UI))
+      continue;
+
+    // Do not push the same instruction more than once.
+    if (!Simplified.insert(UI).second)
+      continue;
+
   }
 }
 
 /// Return true if this instruction generates a simple SCEV
 /// expression in terms of that IV.
 ///
 /// This is similar to IVUsers' isInteresting() but processes each instruction
 /// non-recursively when the operand is already known to be a simpleIVUser.
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;
 
   // 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, Simplified, SimpleIVUsers);
+  pushIVUsers(CurrIV, L, Simplified, SimpleIVUsers);
 
   while (!SimpleIVUsers.empty()) {
     std::pair<Instruction*, Instruction*> UseOper =
       SimpleIVUsers.pop_back_val();
     Instruction *UseInst = UseOper.first;
 
     // Bypass back edges to avoid extra work.
     if (UseInst == CurrIV) continue;
 
-    // Try to replace UseInst with a constant before any other simplifications
-    if (foldConstantSCEV(UseInst))
+    // Try to replace UseInst with a loop invariant before any other
+    // simplifications.
+    if (replaceIVUserWithLoopInvariant(UseInst))
       continue;
 
     Instruction *IVOperand = UseOper.second;
     for (unsigned N = 0; IVOperand; ++N) {
       assert(N <= Simplified.size() && "runaway iteration");
 
       Value *NewOper = foldIVUser(UseOper.first, IVOperand);
       if (!NewOper)
         break; // done folding
       IVOperand = dyn_cast<Instruction>(NewOper);
     }
     if (!IVOperand)
       continue;
 
     if (eliminateIVUser(UseOper.first, IVOperand)) {
-      pushIVUsers(IVOperand, Simplified, SimpleIVUsers);
+      pushIVUsers(IVOperand, L, Simplified, SimpleIVUsers);
       continue;
     }
 
     if (BinaryOperator *BO = dyn_cast<BinaryOperator>(UseOper.first)) {
       if ((isa<OverflowingBinaryOperator>(BO) &&
            strengthenOverflowingOperation(BO, IVOperand)) ||
           (isa<ShlOperator>(BO) && strengthenRightShift(BO, IVOperand))) {
         // re-queue uses of the now modified binary operator and fall
         // through to the checks that remain.
-        pushIVUsers(IVOperand, Simplified, SimpleIVUsers);
+        pushIVUsers(IVOperand, L, Simplified, SimpleIVUsers);
       }
     }
 
     CastInst *Cast = dyn_cast<CastInst>(UseOper.first);
     if (V && Cast) {
       V->visitCast(Cast);
       continue;
     }
     if (isSimpleIVUser(UseOper.first, L, SE)) {
-      pushIVUsers(UseOper.first, Simplified, SimpleIVUsers);
+      pushIVUsers(UseOper.first, L, Simplified, SimpleIVUsers);
     }
   }
 }
 
 namespace llvm {
 
 void IVVisitor::anchor() { }
 
 /// Simplify instructions that use this induction variable
 /// by using ScalarEvolution to analyze the IV's recurrence.
 bool simplifyUsersOfIV(PHINode *CurrIV, ScalarEvolution *SE, DominatorTree *DT,
                        LoopInfo *LI, SmallVectorImpl<WeakTrackingVH> &Dead,
-                       IVVisitor *V) {
-  SimplifyIndvar SIV(LI->getLoopFor(CurrIV->getParent()), SE, DT, LI, Dead);
+                       SCEVExpander &Rewriter, IVVisitor *V) {
+  SimplifyIndvar SIV(LI->getLoopFor(CurrIV->getParent()), SE, DT, LI, Rewriter,
+                     Dead);
   SIV.simplifyUsers(CurrIV, V);
   return SIV.hasChanged();
 }
 
 /// Simplify users of induction variables within this
 /// loop. This does not actually change or add IVs.
 bool simplifyLoopIVs(Loop *L, ScalarEvolution *SE, DominatorTree *DT,
                      LoopInfo *LI, SmallVectorImpl<WeakTrackingVH> &Dead) {
+  SCEVExpander Rewriter(*SE, SE->getDataLayout(), "indvars");
+#ifndef NDEBUG
+  Rewriter.setDebugType(DEBUG_TYPE);
+#endif
   bool Changed = false;
   for (BasicBlock::iterator I = L->getHeader()->begin(); isa<PHINode>(I); ++I) {
-    Changed |= simplifyUsersOfIV(cast<PHINode>(I), SE, DT, LI, Dead);
+    Changed |= simplifyUsersOfIV(cast<PHINode>(I), SE, DT, LI, Dead, Rewriter);
   }
   return Changed;
 }
 
 } // namespace llvm

llvm/trunk/test/Transforms/IndVarSimplify/constant-fold-1.ll

-; RUN: opt < %s -indvars -S | FileCheck %s
-
-target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
-target triple = "x86_64-unknown-linux-gnu"
-
-define void @test(i64* %arg) unnamed_addr align 2 {
-bb:
-  switch i64 undef, label %bb1 [
-  ]
-
-bb1:                                              ; preds = %bb
-  br label %bb2
-
-bb2:                                              ; preds = %bb6, %bb1
-  %tmp = phi i64* [%arg, %bb1 ], [ %tmp7, %bb6 ]
-  switch i2 undef, label %bb6 [
-    i2 1, label %bb5
-    i2 -2, label %bb3
-    i2 -1, label %bb3
-  ]
-
-bb3:                                              ; preds = %bb2, %bb2
-  %tmp4 = call fastcc i32* @wobble(i64* nonnull %tmp, i32* null)
-  %tmp5 = load i32, i32* %tmp4 , align 8
-  br label %bb6
-
-bb5:                                              ; preds = %bb2
-  unreachable
-
-bb6:                                              ; preds = %bb3, %bb2
-  %tmp7 = load i64*, i64** undef, align 8
-  br label %bb2
-}
-
-declare i32* @wobble(i64*, i32* returned)
-
-; Should not fail when SCEV is fold to ConstantPointerNull
-; CHECK-LABEL: void @test
-; CHECK:         load i32, i32* %{{[a-zA-Z$._0-9]+}}

llvm/trunk/test/Transforms/IndVarSimplify/replace-iv-with-loop-invariant.ll

+; RUN: opt < %s -indvars -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@G = external global i32
+
+define void @test0(i64* %arg) {
+bb:
+  br label %bb2
+
+bb2:
+  %tmp = phi i64* [%arg, %bb ], [ %tmp7, %bb2 ]
+  %tmp4 = call i32* @wobble(i64* nonnull %tmp, i32* null)
+  %tmp5 = load i32, i32* %tmp4, align 8
+  %tmp7 = load i64*, i64** undef, align 8
+  br label %bb2
+}
+
+; CHECK-LABEL: void @test0
+; CHECK:         load i32, i32* null
+
+define void @test1(i64* %arg) {
+bb:
+  br label %bb2
+
+bb2:
+  %tmp = phi i64* [%arg, %bb ], [ %tmp7, %bb2 ]
+  %tmp4 = call i32* @wobble(i64* nonnull %tmp, i32* inttoptr (i64 4 to i32*))
+  %tmp5 = load i32, i32* %tmp4
+  %tmp7 = load i64*, i64** undef, align 8
+  br label %bb2
+}
+
+; CHECK-LABEL: void @test1
+; CHECK:         load i32, i32* inttoptr (i64 4 to i32*)
+
+define void @test2(i64* %arg) {
+bb:
+  br label %bb2
+
+bb2:
+  %tmp = phi i64* [%arg, %bb ], [ %tmp7, %bb2 ]
+  %tmp4 = call i32* @wobble(i64* nonnull %tmp, i32* @G)
+  %tmp5 = load i32, i32* %tmp4
+  %tmp7 = load i64*, i64** undef, align 8
+  br label %bb2
+}
+
+; CHECK-LABEL: void @test2
+; CHECK:         load i32, i32* @G
+
+
+define void @test3(i64* %arg, i32* %loop.invariant) {
+bb:
+  br label %bb2
+
+bb2:
+  %tmp = phi i64* [%arg, %bb ], [ %tmp7, %bb2 ]
+  %tmp4 = call i32* @wobble(i64* nonnull %tmp, i32* %loop.invariant)
+  %tmp5 = load i32, i32* %tmp4
+  %tmp7 = load i64*, i64** undef, align 8
+  br label %bb2
+}
+
+; CHECK-LABEL: void @test3
+; CHECK:         load i32, i32* %loop.invariant
+
+define void @test4(i64* %arg, i32* %loop.invariant, i64 %N) {
+bb:
+  br label %bb2
+
+bb2:
+  %tmp = phi i64* [%arg, %bb ], [ %tmp7, %bb2 ]
+  %mul = mul nsw i64 %N, 64
+  %ptr = getelementptr inbounds i32, i32* %loop.invariant, i64 %mul 
+  %tmp4 = call i32* @wobble(i64* nonnull %tmp, i32* %ptr)
+  %tmp5 = load i32, i32* %tmp4
+  %tmp7 = load i64*, i64** undef, align 8
+  br label %bb2
+}
+
+; CHECK-LABEL: void @test4
+; CHECK:         [[P:%[a-zA-Z$._0-9]+]] = getelementptr i32, i32* %loop.invariant
+; CHECK:         phi
+; CHECK:         load i32, i32* [[P]]
+
+declare i32* @wobble(i64*, i32* returned)