[IVUsers] Move preheader check into SCEVExpander

[nikic]

Summary

Rather than checking for loop nest preheaders upfront in IVUsers, move this requirement into isSafeToExpand() from SCEVExpander, which seems like logically the right place for this check.

This is a followup to https://reviews.llvm.org/D111493#3055286.

[mkazantsev]

That's amazing that no test shows difference. Do you expect it to be a 100% NFC?

[mkazantsev]

If I'm reading this correctly, isPreheaderLoopNest may find a loop that is dominating, but not containing a given block, and then bail if it doesn't have a preheader. This logic is kind of strange, but I don't really understand its implications.

[reames]

LGTM.

Have to admit, this is slightly more delicate than I'd realized. Not sure I'd have pushed for this if I realized at the time it had to be an expansion safety predicate as opposed to simply handled (with possibly sub-optimal codegen) during expansion. The benefit seems smaller. Happy to have this land, but would also be fine abandoning this. I'll leave it to your judgement.

[nikic]

Make the check more precise: We can expand affine addrecs in canonical mode without a preheader. This addresses the PowerPC test failure.

[reames]

On the updated revision, my earlier LGTM still stands, but I'm less and less sure this is worth doing. As previously, I'll defer to your judgement.

llvm/include/llvm/Analysis/IVUsers.h

public:
   }
 
   void releaseMemory();
 
   void print(raw_ostream &OS, const Module * = nullptr) const;
 
   /// dump - This method is used for debugging.
   void dump() const;
-
-protected:
-  bool AddUsersImpl(Instruction *I, SmallPtrSetImpl<Loop*> &SimpleLoopNests);
 };
 
 Pass *createIVUsersPass();
 
 class IVUsersWrapperPass : public LoopPass {
   std::unique_ptr<IVUsers> IU;
 
 public:

llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/InstructionCost.h"
 
 namespace llvm {
 extern cl::opt<unsigned> SCEVCheapExpansionBudget;
 
 /// Return true if the given expression is safe to expand in the sense that
 /// all materialized values are safe to speculate anywhere their operands are
-/// defined.
-bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE);
+/// defined, and the expander is capable of expanding the expression.
+/// CanonicalMode indicates whether the expander will be used in canonical mode.
+bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE,
+                    bool CanonicalMode = true);
 
 /// Return true if the given expression is safe to expand in the sense that
 /// all materialized values are defined and safe to speculate at the specified
 /// location and their operands are defined at this location.
 bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint,
                       ScalarEvolution &SE);
 
 /// struct for holding enough information to help calculate the cost of the

llvm/lib/Analysis/IVUsers.cpp

for (const auto *Op : Add->operands())
       }
     return AnyInterestingYet;
   }
 
   // Nothing else is interesting here.
   return false;
 }
 
-/// Return true if all loop headers that dominate this block have a preheader.
-static bool isPreheaderLoopNest(BasicBlock *BB, const DominatorTree *DT,
-                                const LoopInfo *LI,
-                                SmallPtrSetImpl<Loop*> &PreheaderLoopNests) {
-  Loop *NearestLoop = nullptr;
-  for (DomTreeNode *Rung = DT->getNode(BB);
-       Rung; Rung = Rung->getIDom()) {
-    BasicBlock *DomBB = Rung->getBlock();
-    Loop *DomLoop = LI->getLoopFor(DomBB);
-    if (DomLoop && DomLoop->getHeader() == DomBB) {
-      // If we have already checked this loop nest, stop checking.
-      if (PreheaderLoopNests.count(DomLoop))
-        break;
-      // If the domtree walk reaches a loop with no preheader, return false.
-      if (!DomLoop->getLoopPreheader())
-        return false;
-      // If we have not already checked this loop nest, remember the loop
-      // header nearest to BB. The nearest loop may not contain BB.
-      if (!NearestLoop)
-        NearestLoop = DomLoop;
-    }
-  }
-  if (NearestLoop)
-    PreheaderLoopNests.insert(NearestLoop);
-  return true;
-}
-
 /// IVUseShouldUsePostIncValue - We have discovered a "User" of an IV expression
 /// and now we need to decide whether the user should use the preinc or post-inc
 /// value.  If this user should use the post-inc version of the IV, return true.
 ///
 /// Choosing wrong here can break dominance properties (if we choose to use the
 /// post-inc value when we cannot) or it can end up adding extra live-ranges to
 /// the loop, resulting in reg-reg copies (if we use the pre-inc value when we
 /// should use the post-inc value).
if (PN->getIncomingValue(i) == Operand &&
         !DT->dominates(LatchBlock, PN->getIncomingBlock(i)))
       return false;
 
   // Okay, all uses of Operand by PN are in predecessor blocks that really are
   // dominated by the latch block.  Use the post-incremented value.
   return true;
 }
 
-/// AddUsersImpl - Inspect the specified instruction.  If it is a
-/// reducible SCEV, recursively add its users to the IVUsesByStride set and
-/// return true.  Otherwise, return false.
-bool IVUsers::AddUsersImpl(Instruction *I,
-                           SmallPtrSetImpl<Loop*> &PreheaderLoopNests) {
+/// Inspect the specified instruction.  If it is a reducible SCEV, recursively
+/// add its users to the IVUsesByStride set and return true.  Otherwise, return
+/// false.
+bool IVUsers::AddUsersIfInteresting(Instruction *I) {
   const DataLayout &DL = I->getModule()->getDataLayout();
 
   // Add this IV user to the Processed set before returning false to ensure that
   // all IV users are members of the set. See IVUsers::isIVUserOrOperand.
   if (!Processed.insert(I).second)
     return true;    // Instruction already handled.
 
   if (!SE->isSCEVable(I->getType()))
for (Use &U : I->uses()) {
     Instruction *User = cast<Instruction>(U.getUser());
     if (!UniqueUsers.insert(User).second)
       continue;
 
     // Do not infinitely recurse on PHI nodes.
     if (isa<PHINode>(User) && Processed.count(User))
       continue;
 
-    // Only consider IVUsers that are dominated by simplified loop
-    // headers. Otherwise, SCEVExpander will crash.
-    BasicBlock *UseBB = User->getParent();
-    // A phi's use is live out of its predecessor block.
-    if (PHINode *PHI = dyn_cast<PHINode>(User)) {
-      unsigned OperandNo = U.getOperandNo();
-      unsigned ValNo = PHINode::getIncomingValueNumForOperand(OperandNo);
-      UseBB = PHI->getIncomingBlock(ValNo);
-    }
-    if (!isPreheaderLoopNest(UseBB, DT, LI, PreheaderLoopNests))
-      return false;
-
     // Descend recursively, but not into PHI nodes outside the current loop.
     // It's important to see the entire expression outside the loop to get
     // choices that depend on addressing mode use right, although we won't
     // consider references outside the loop in all cases.
     // If User is already in Processed, we don't want to recurse into it again,
     // but do want to record a second reference in the same instruction.
     bool AddUserToIVUsers = false;
     if (LI->getLoopFor(User->getParent()) != L) {
       if (isa<PHINode>(User) || Processed.count(User) ||
-          !AddUsersImpl(User, PreheaderLoopNests)) {
+          !AddUsersIfInteresting(User)) {
         LLVM_DEBUG(dbgs() << "FOUND USER in other loop: " << *User << '\n'
                           << "   OF SCEV: " << *ISE << '\n');
         AddUserToIVUsers = true;
       }
-    } else if (Processed.count(User) ||
-               !AddUsersImpl(User, PreheaderLoopNests)) {
+    } else if (Processed.count(User) || !AddUsersIfInteresting(User)) {
       LLVM_DEBUG(dbgs() << "FOUND USER: " << *User << '\n'
                         << "   OF SCEV: " << *ISE << '\n');
       AddUserToIVUsers = true;
     }
 
     if (AddUserToIVUsers) {
       // Okay, we found a user that we cannot reduce.
       IVStrideUse &NewUse = AddUser(User, I);
if (AddUserToIVUsers) {
       }
       LLVM_DEBUG(if (SE->getSCEV(I) != ISE) dbgs()
                  << "   NORMALIZED TO: " << *ISE << '\n');
     }
   }
   return true;
 }
 
-bool IVUsers::AddUsersIfInteresting(Instruction *I) {
-  // SCEVExpander can only handle users that are dominated by loops with
-  // preheaders. Keep track of all loops that are only dominated by preheader
-  // loops so we don't traverse the domtree for each user.
-  SmallPtrSet<Loop*,16> PreheaderLoopNests;
-
-  return AddUsersImpl(I, PreheaderLoopNests);
-}
-
 IVStrideUse &IVUsers::AddUser(Instruction *User, Value *Operand) {
   IVUses.push_back(new IVStrideUse(this, User, Operand));
   return IVUses.back();
 }
 
 IVUsers::IVUsers(Loop *L, AssumptionCache *AC, LoopInfo *LI, DominatorTree *DT,
                  ScalarEvolution *SE)
     : L(L), AC(AC), LI(LI), DT(DT), SE(SE), IVUses() {

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

void LSRInstance::CollectFixupsAndInitialFormulae() {
   LLVM_DEBUG(print_fixups(dbgs()));
 }
 
 /// Insert a formula for the given expression into the given use, separating out
 /// loop-variant portions from loop-invariant and loop-computable portions.
 void
 LSRInstance::InsertInitialFormula(const SCEV *S, LSRUse &LU, size_t LUIdx) {
   // Mark uses whose expressions cannot be expanded.
-  if (!isSafeToExpand(S, SE))
+  if (!isSafeToExpand(S, SE, /*CanonicalMode*/ false))
     LU.RigidFormula = true;
 
   Formula F;
   F.initialMatch(S, L, SE);
   bool Inserted = InsertFormula(LU, LUIdx, F);
   assert(Inserted && "Initial formula already exists!"); (void)Inserted;
 }
 
if (auto *PN = dyn_cast<PHINode>(U.getUser())) {
        if (isa<FuncletPadInst>(FirstNonPHI) ||
            isa<CatchSwitchInst>(FirstNonPHI))
          for (BasicBlock *PredBB : PN->blocks())
            if (isa<CatchSwitchInst>(PredBB->getFirstNonPHI()))
              return;
     }
   }
 
-#ifndef NDEBUG
-  // All dominating loops must have preheaders, or SCEVExpander may not be able
-  // to materialize an AddRecExpr whose Start is an outer AddRecExpr.
-  //
-  // IVUsers analysis should only create users that are dominated by simple loop
-  // headers. Since this loop should dominate all of its users, its user list
-  // should be empty if this loop itself is not within a simple loop nest.
-  for (DomTreeNode *Rung = DT.getNode(L->getLoopPreheader());
-       Rung; Rung = Rung->getIDom()) {
-    BasicBlock *BB = Rung->getBlock();
-    const Loop *DomLoop = LI.getLoopFor(BB);
-    if (DomLoop && DomLoop->getHeader() == BB) {
-      assert(DomLoop->getLoopPreheader() && "LSR needs a simplified loop nest");
-    }
-  }
-#endif // DEBUG
-
   LLVM_DEBUG(dbgs() << "\nLSR on loop ";
              L->getHeader()->printAsOperand(dbgs(), /*PrintType=*/false);
              dbgs() << ":\n");
 
   // First, perform some low-level loop optimizations.
   OptimizeShadowIV();
   OptimizeLoopTermCond();
 

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

 // We cannot generally expand recurrences unless the step dominates the loop
 // header. The expander handles the special case of affine recurrences by
 // scaling the recurrence outside the loop, but this technique isn't generally
 // applicable. Expanding a nested recurrence outside a loop requires computing
 // binomial coefficients. This could be done, but the recurrence has to be in a
 // perfectly reduced form, which can't be guaranteed.
 struct SCEVFindUnsafe {
   ScalarEvolution &SE;
+  bool CanonicalMode;
   bool IsUnsafe;
 
-  SCEVFindUnsafe(ScalarEvolution &se): SE(se), IsUnsafe(false) {}
+  SCEVFindUnsafe(ScalarEvolution &SE, bool CanonicalMode)
+      : SE(SE), CanonicalMode(CanonicalMode), IsUnsafe(false) {}
 
   bool follow(const SCEV *S) {
     if (const SCEVUDivExpr *D = dyn_cast<SCEVUDivExpr>(S)) {
       const SCEVConstant *SC = dyn_cast<SCEVConstant>(D->getRHS());
       if (!SC || SC->getValue()->isZero()) {
         IsUnsafe = true;
         return false;
       }
     }
     if (const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(S)) {
       const SCEV *Step = AR->getStepRecurrence(SE);
       if (!AR->isAffine() && !SE.dominates(Step, AR->getLoop()->getHeader())) {
         IsUnsafe = true;
         return false;
       }
+
+      // For non-affine addrecs or in non-canonical mode we need a preheader
+      // to insert into.
+      if (!AR->getLoop()->getLoopPreheader() &&
+          (!CanonicalMode || !AR->isAffine())) {
+        IsUnsafe = true;
+        return false;
+      }
     }
     return true;
   }
   bool isDone() const { return IsUnsafe; }
 };
 }
 
 namespace llvm {
-bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE) {
-  SCEVFindUnsafe Search(SE);
+bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE, bool CanonicalMode) {
+  SCEVFindUnsafe Search(SE, CanonicalMode);
   visitAll(S, Search);
   return !Search.IsUnsafe;
 }
 
 bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint,
                       ScalarEvolution &SE) {
   if (!isSafeToExpand(S, SE))
     return false;