[LoopUnroll] Move visitBB and cost computation to UnrolledInstAnalyzer, for reuse (WIP)

[fhahn]

Summary

This is a rough patch moving some infrastructure to estimate the size of a
BB after unrolling to UnrolledInstAnalyzer. The interface still needs
improvement and maybe better naming as well.

llvm/include/llvm/Analysis/LoopUnrollAnalyzer.h

 // potential effects that loop unrolling might have, such as enabling constant
 // propagation and other optimizations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_LOOPUNROLLANALYZER_H
 #define LLVM_ANALYSIS_LOOPUNROLLANALYZER_H
 
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/InstVisitor.h"
 
 // This class is used to get an estimate of the optimization effects that we
 // could get from complete loop unrolling. It comes from the fact that some
 // loads might be replaced with concrete constant values and that could trigger
 // a chain of instruction simplifications.
 //
 // E.g. we might have:
 //   int a[] = {0, 1, 0};
 //   v = 0;
 //   for (i = 0; i < 3; i ++)
 //     v += b[i]*a[i];
 // If we completely unroll the loop, we would get:
 //   v = b[0]*a[0] + b[1]*a[1] + b[2]*a[2]
 // Which then will be simplified to:
 //   v = b[0]* 0 + b[1]* 1 + b[2]* 0
 // And finally:
 //   v = b[1]
 namespace llvm {
+
+namespace {
+/// A struct to densely store the state of an instruction after unrolling at
+/// each iteration.
+///
+/// This is designed to work like a tuple of <Instruction *, int> for the
+/// purposes of hashing and lookup, but to be able to associate two boolean
+/// states with each key.
+struct UnrolledInstState {
+  Instruction *I;
+  int Iteration : 30;
+  unsigned IsFree : 1;
+  unsigned IsCounted : 1;
+};
+
+/// Hashing and equality testing for a set of the instruction states.
+struct UnrolledInstStateKeyInfo {
+  using PtrInfo = DenseMapInfo<Instruction *>;
+  using PairInfo = DenseMapInfo<std::pair<Instruction *, int>>;
+
+  static inline UnrolledInstState getEmptyKey() {
+    return {PtrInfo::getEmptyKey(), 0, 0, 0};
+  }
+
+  static inline UnrolledInstState getTombstoneKey() {
+    return {PtrInfo::getTombstoneKey(), 0, 0, 0};
+  }
+
+  static inline unsigned getHashValue(const UnrolledInstState &S) {
+    return PairInfo::getHashValue({S.I, S.Iteration});
+  }
+
+  static inline bool isEqual(const UnrolledInstState &LHS,
+                             const UnrolledInstState &RHS) {
+    return PairInfo::isEqual({LHS.I, LHS.Iteration}, {RHS.I, RHS.Iteration});
+  }
+};
+
+struct EstimatedUnrollCost {
+  /// The estimated cost after unrolling.
+  unsigned UnrolledCost;
+
+  /// The estimated dynamic cost of executing the instructions in the
+  /// rolled form.
+  unsigned RolledDynamicCost;
+};
+
+} // namespace
+
 class UnrolledInstAnalyzer : private InstVisitor<UnrolledInstAnalyzer, bool> {
   typedef InstVisitor<UnrolledInstAnalyzer, bool> Base;
   friend class InstVisitor<UnrolledInstAnalyzer, bool>;
   struct SimplifiedAddress {
     Value *Base = nullptr;
     ConstantInt *Offset = nullptr;
   };
 
 public:
-  UnrolledInstAnalyzer(unsigned Iteration,
-                       DenseMap<Value *, Constant *> &SimplifiedValues,
-                       ScalarEvolution &SE, const Loop *L)
-      : SimplifiedValues(SimplifiedValues), SE(SE), L(L) {
-      IterationNumber = SE.getConstant(APInt(64, Iteration));
+  UnrolledInstAnalyzer(
+      unsigned Iteration, DenseMap<Value *, Constant *> &SimplifiedValues,
+      const SmallPtrSetImpl<const Value *> &EphValues,
+      DenseSet<UnrolledInstState, UnrolledInstStateKeyInfo> &InstCostMap,
+      unsigned &RolledDynamicCost, unsigned &UnrolledCost,
+      unsigned MaxUnrolledLoopSize, ScalarEvolution &SE,
+      const TargetTransformInfo &TTI, const Loop *L)
+      : SimplifiedValues(SimplifiedValues), EphValues(EphValues),
+        InstCostMap(InstCostMap), RolledDynamicCost(RolledDynamicCost),
+        UnrolledCost(UnrolledCost), MaxUnrolledLoopSize(MaxUnrolledLoopSize),
+        SE(SE), TTI(TTI), L(L) {
   }
 
   // Allow access to the initial visit method.
   using Base::visit;
 
+  bool visitBB(BasicBlock &B);
+
+  void AddCostRecursively(Instruction &RootI, int Iteration);
+
 private:
   /// A cache of pointer bases and constant-folded offsets corresponding
   /// to GEP (or derived from GEP) instructions.
   ///
   /// In order to find the base pointer one needs to perform non-trivial
   /// traversal of the corresponding SCEV expression, so it's good to have the
   /// results saved.
   DenseMap<Value *, SimplifiedAddress> SimplifiedAddresses;
 
   /// SCEV expression corresponding to number of currently simulated
   /// iteration.
   const SCEV *IterationNumber;
 
   /// A Value->Constant map for keeping values that we managed to
   /// constant-fold on the given iteration.
   ///
   /// While we walk the loop instructions, we build up and maintain a mapping
   /// of simplified values specific to this iteration.  The idea is to propagate
   /// any special information we have about loads that can be replaced with
   /// constants after complete unrolling, and account for likely simplifications
   /// post-unrolling.
   DenseMap<Value *, Constant *> &SimplifiedValues;
+  const SmallPtrSetImpl<const Value *> &EphValues;
+  DenseSet<UnrolledInstState, UnrolledInstStateKeyInfo> &InstCostMap;
+
+  unsigned &RolledDynamicCost;
+  unsigned &UnrolledCost;
+  unsigned MaxUnrolledLoopSize;
 
   ScalarEvolution &SE;
+  const TargetTransformInfo &TTI;
   const Loop *L;
 
   bool simplifyInstWithSCEV(Instruction *I);
 
   bool visitInstruction(Instruction &I) { return simplifyInstWithSCEV(&I); }
   bool visitBinaryOperator(BinaryOperator &I);
   bool visitLoad(LoadInst &I);
   bool visitCastInst(CastInst &I);
   bool visitCmpInst(CmpInst &I);
   bool visitPHINode(PHINode &PN);
 };
 }
 #endif

llvm/lib/Analysis/LoopUnrollAnalyzer.cpp

 // propagation and other optimizations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/LoopUnrollAnalyzer.h"
 
 using namespace llvm;
 
+#define DEBUG_TYPE "loop-unroll"
+
+// Helper function to accumulate cost for instructions in the loop.
+void UnrolledInstAnalyzer::AddCostRecursively(Instruction &RootI,
+                                              int Iteration) {
+  // A small worklist used to accumulate cost of instructions from each
+  // observable and reached root in the loop.
+  SmallVector<Instruction *, 16> CostWorklist;
+
+  // PHI-used worklist used between iterations while accumulating cost.
+  SmallVector<Instruction *, 4> PHIUsedList;
+
+  assert(Iteration >= 0 && "Cannot have a negative iteration!");
+  assert(CostWorklist.empty() && "Must start with an empty cost list");
+  assert(PHIUsedList.empty() && "Must start with an empty phi used list");
+  CostWorklist.push_back(&RootI);
+  for (;; --Iteration) {
+    do {
+      Instruction *I = CostWorklist.pop_back_val();
+
+      // InstCostMap only uses I and Iteration as a key, the other two values
+      // don't matter here.
+      auto CostIter = InstCostMap.find({I, Iteration, 0, 0});
+      if (CostIter == InstCostMap.end())
+        // If an input to a PHI node comes from a dead path through the loop
+        // we may have no cost data for it here. What that actually means is
+        // that it is free.
+        continue;
+      auto &Cost = *CostIter;
+      if (Cost.IsCounted)
+        // Already counted this instruction.
+        continue;
+
+      // Mark that we are counting the cost of this instruction now.
+      Cost.IsCounted = true;
+
+      // If this is a PHI node in the loop header, just add it to the PHI set.
+      if (auto *PhiI = dyn_cast<PHINode>(I))
+        if (PhiI->getParent() == L->getHeader()) {
+          assert(Cost.IsFree && "Loop PHIs shouldn't be evaluated as they "
+                                "inherently simplify during unrolling.");
+          if (Iteration == 0)
+            continue;
+
+          // Push the incoming value from the backedge into the PHI used list
+          // if it is an in-loop instruction. We'll use this to populate the
+          // cost worklist for the next iteration (as we count backwards).
+          if (auto *OpI = dyn_cast<Instruction>(
+                  PhiI->getIncomingValueForBlock(L->getLoopLatch())))
+            if (L->contains(OpI))
+              PHIUsedList.push_back(OpI);
+          continue;
+        }
+
+      // First accumulate the cost of this instruction.
+      if (!Cost.IsFree) {
+        UnrolledCost += TTI.getUserCost(I);
+        LLVM_DEBUG(dbgs() << "Adding cost of instruction (iteration "
+                          << Iteration << "): ");
+        LLVM_DEBUG(I->dump());
+      }
+
+      // We must count the cost of every operand which is not free,
+      // recursively. If we reach a loop PHI node, simply add it to the set
+      // to be considered on the next iteration (backwards!).
+      for (Value *Op : I->operands()) {
+        // Check whether this operand is free due to being a constant or
+        // outside the loop.
+        auto *OpI = dyn_cast<Instruction>(Op);
+        if (!OpI || !L->contains(OpI))
+          continue;
+
+        // Otherwise accumulate its cost.
+        CostWorklist.push_back(OpI);
+      }
+    } while (!CostWorklist.empty());
+
+    if (PHIUsedList.empty())
+      // We've exhausted the search.
+      break;
+
+    assert(Iteration > 0 &&
+           "Cannot track PHI-used values past the first iteration!");
+    CostWorklist.append(PHIUsedList.begin(), PHIUsedList.end());
+    PHIUsedList.clear();
+  }
+};
+
 /// Try to simplify instruction \param I using its SCEV expression.
 ///
 /// The idea is that some AddRec expressions become constants, which then
 /// could trigger folding of other instructions. However, that only happens
 /// for expressions whose start value is also constant, which isn't always the
 /// case. In another common and important case the start value is just some
 /// address (i.e. SCEVUnknown) - in this case we compute the offset and save
 /// it along with the base address instead.
bool UnrolledInstAnalyzer::visitPHINode(PHINode &PN) {
   // Run base visitor first. This way we can gather some useful for later
   // analysis information.
   if (Base::visitPHINode(PN))
     return true;
 
   // The loop induction PHI nodes are definitionally free.
   return PN.getParent() == L->getHeader();
 }
+
+bool UnrolledInstAnalyzer::visitBB(BasicBlock &BB) {
+  int Iteration =
+      cast<SCEVConstant>(IterationNumber)->getValue()->getZExtValue();
+  for (Instruction &I : BB) {
+    // These won't get into the final code - don't even try calculating the
+    // cost for them.
+    if (isa<DbgInfoIntrinsic>(I) || EphValues.count(&I))
+      continue;
+
+    // Track this instruction's expected baseline cost when executing the
+    // rolled loop form.
+    RolledDynamicCost += TTI.getUserCost(&I);
+
+    // Visit the instruction to analyze its loop cost after unrolling,
+    // and if the visitor returns true, mark the instruction as free after
+    // unrolling and continue.
+    bool IsFree = visit(I);
+    bool Inserted = InstCostMap
+                        .insert({&I, (int)Iteration, (unsigned)IsFree,
+                                 /*IsCounted*/ false})
+                        .second;
+    (void)Inserted;
+    assert(Inserted && "Cannot have a state for an unvisited instruction!");
+
+    if (IsFree)
+      continue;
+
+    // Can't properly model a cost of a call.
+    // FIXME: With a proper cost model we should be able to do it.
+    if (auto *CI = dyn_cast<CallInst>(&I)) {
+      const Function *Callee = CI->getCalledFunction();
+      if (!Callee || TTI.isLoweredToCall(Callee)) {
+        LLVM_DEBUG(dbgs() << "Can't analyze cost of loop with call\n");
+        return false;
+      }
+    }
+
+    // If the instruction might have a side-effect recursively account for
+    // the cost of it and all the instructions leading up to it.
+    if (I.mayHaveSideEffects())
+      AddCostRecursively(I, Iteration);
+
+    // If unrolled body turns out to be too big, bail out.
+    if (UnrolledCost > MaxUnrolledLoopSize) {
+      LLVM_DEBUG(dbgs() << "  Exceeded threshold.. exiting.\n"
+                        << "  UnrolledCost: " << UnrolledCost
+                        << ", MaxUnrolledLoopSize: " << MaxUnrolledLoopSize
+                        << "\n");
+      return false;
+    }
+  }
+  return true;
+}

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

TargetTransformInfo::UnrollingPreferences llvm::gatherUnrollingPreferences(
   if (UserAllowPeeling.hasValue())
     UP.AllowPeeling = *UserAllowPeeling;
 
   return UP;
 }
 
 namespace {
 
-/// A struct to densely store the state of an instruction after unrolling at
-/// each iteration.
-///
-/// This is designed to work like a tuple of <Instruction *, int> for the
-/// purposes of hashing and lookup, but to be able to associate two boolean
-/// states with each key.
-struct UnrolledInstState {
-  Instruction *I;
-  int Iteration : 30;
-  unsigned IsFree : 1;
-  unsigned IsCounted : 1;
-};
-
-/// Hashing and equality testing for a set of the instruction states.
-struct UnrolledInstStateKeyInfo {
-  using PtrInfo = DenseMapInfo<Instruction *>;
-  using PairInfo = DenseMapInfo<std::pair<Instruction *, int>>;
-
-  static inline UnrolledInstState getEmptyKey() {
-    return {PtrInfo::getEmptyKey(), 0, 0, 0};
-  }
-
-  static inline UnrolledInstState getTombstoneKey() {
-    return {PtrInfo::getTombstoneKey(), 0, 0, 0};
-  }
-
-  static inline unsigned getHashValue(const UnrolledInstState &S) {
-    return PairInfo::getHashValue({S.I, S.Iteration});
-  }
-
-  static inline bool isEqual(const UnrolledInstState &LHS,
-                             const UnrolledInstState &RHS) {
-    return PairInfo::isEqual({LHS.I, LHS.Iteration}, {RHS.I, RHS.Iteration});
-  }
-};
-
-struct EstimatedUnrollCost {
-  /// The estimated cost after unrolling.
-  unsigned UnrolledCost;
-
-  /// The estimated dynamic cost of executing the instructions in the
-  /// rolled form.
-  unsigned RolledDynamicCost;
-};
-
 } // end anonymous namespace
 
 /// Figure out if the loop is worth full unrolling.
 ///
 /// Complete loop unrolling can make some loads constant, and we need to know
 /// if that would expose any further optimization opportunities.  This routine
 /// estimates this optimization.  It computes cost of unrolled loop
 /// (UnrolledCost) and dynamic cost of the original loop (RolledDynamicCost). By
for (Instruction &I : *L->getHeader()) {
         SimplifiedInputValues.push_back({PHI, C});
     }
 
     // Now clear and re-populate the map for the next iteration.
     SimplifiedValues.clear();
     while (!SimplifiedInputValues.empty())
       SimplifiedValues.insert(SimplifiedInputValues.pop_back_val());
 
-    UnrolledInstAnalyzer Analyzer(Iteration, SimplifiedValues, SE, L);
+    UnrolledInstAnalyzer Analyzer(Iteration, SimplifiedValues, EphValues,
+                                  InstCostMap, RolledDynamicCost, UnrolledCost,
+                                  MaxUnrolledLoopSize, SE, TTI, L);
 
     BBWorklist.clear();
     BBWorklist.insert(L->getHeader());
     // Note that we *must not* cache the size, this loop grows the worklist.
     for (unsigned Idx = 0; Idx != BBWorklist.size(); ++Idx) {
       BasicBlock *BB = BBWorklist[Idx];
 
       // Visit all instructions in the given basic block and try to simplify
       // it.  We don't change the actual IR, just count optimization
       // opportunities.
-      for (Instruction &I : *BB) {
-        // These won't get into the final code - don't even try calculating the
-        // cost for them.
-        if (isa<DbgInfoIntrinsic>(I) || EphValues.count(&I))
-          continue;
-
-        // Track this instruction's expected baseline cost when executing the
-        // rolled loop form.
-        RolledDynamicCost += TTI.getUserCost(&I);
-
-        // Visit the instruction to analyze its loop cost after unrolling,
-        // and if the visitor returns true, mark the instruction as free after
-        // unrolling and continue.
-        bool IsFree = Analyzer.visit(I);
-        bool Inserted = InstCostMap.insert({&I, (int)Iteration,
-                                           (unsigned)IsFree,
-                                           /*IsCounted*/ false}).second;
-        (void)Inserted;
-        assert(Inserted && "Cannot have a state for an unvisited instruction!");
-
-        if (IsFree)
-          continue;
-
-        // Can't properly model a cost of a call.
-        // FIXME: With a proper cost model we should be able to do it.
-        if (auto *CI = dyn_cast<CallInst>(&I)) {
-          const Function *Callee = CI->getCalledFunction();
-          if (!Callee || TTI.isLoweredToCall(Callee)) {
-            LLVM_DEBUG(dbgs() << "Can't analyze cost of loop with call\n");
-            return None;
-          }
-        }
-
-        // If the instruction might have a side-effect recursively account for
-        // the cost of it and all the instructions leading up to it.
-        if (I.mayHaveSideEffects())
-          AddCostRecursively(I, Iteration);
-
-        // If unrolled body turns out to be too big, bail out.
-        if (UnrolledCost > MaxUnrolledLoopSize) {
-          LLVM_DEBUG(dbgs() << "  Exceeded threshold.. exiting.\n"
-                            << "  UnrolledCost: " << UnrolledCost
-                            << ", MaxUnrolledLoopSize: " << MaxUnrolledLoopSize
-                            << "\n");
-          return None;
-        }
-      }
+      if (!Analyzer.visitBB(*BB))
 
       Instruction *TI = BB->getTerminator();
 
       // Add in the live successors by first checking whether we have terminator
       // that may be simplified based on the values simplified by this call.
       BasicBlock *KnownSucc = nullptr;
       if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
         if (BI->isConditional()) {
for (unsigned Idx = 0; Idx != BBWorklist.size(); ++Idx) {
       }
 
       // Add BB's successors to the worklist.
       for (BasicBlock *Succ : successors(BB))
         if (L->contains(Succ))
           BBWorklist.insert(Succ);
         else
           ExitWorklist.insert({BB, Succ});
-      AddCostRecursively(*TI, Iteration);
+      Analyzer.AddCostRecursively(*TI, Iteration);
     }
 
     // If we found no optimization opportunities on the first iteration, we
     // won't find them on later ones too.
     if (UnrolledCost == RolledDynamicCost) {
       LLVM_DEBUG(dbgs() << "  No opportunities found.. exiting.\n"
                         << "  UnrolledCost: " << UnrolledCost << "\n");
       return None;

llvm/unittests/Analysis/UnrollAnalyzerTest.cpp

 static unsigned TripCount = 0;
 
 namespace {
 struct UnrollAnalyzerTest : public FunctionPass {
   static char ID;
   bool runOnFunction(Function &F) override {
     LoopInfo *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     ScalarEvolution *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
+    const TargetTransformInfo &TTI =
+        getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
     Function::iterator FI = F.begin();
     FI++; // First basic block is entry - skip it.
     BasicBlock *Header = &*FI++;
     Loop *L = LI->getLoopFor(Header);
     BasicBlock *Exiting = L->getExitingBlock();
 
+    DenseSet<UnrolledInstState, UnrolledInstStateKeyInfo> InstCostMap;
+    unsigned UnrolledCost = 0;
+    unsigned RolledDynamicCost = 0;
+
+    SmallPtrSet<const Value *, 32> EphValues;
     SimplifiedValuesVector.clear();
     TripCount = SE->getSmallConstantTripCount(L, Exiting);
     for (unsigned Iteration = 0; Iteration < TripCount; Iteration++) {
       DenseMap<Value *, Constant *> SimplifiedValues;
-      UnrolledInstAnalyzer Analyzer(Iteration, SimplifiedValues, *SE, L);
+      UnrolledInstAnalyzer Analyzer(Iteration, SimplifiedValues, EphValues, InstCostMap, RolledDynamicCost, UnrolledCost, 0,*SE, TTI, L);
       for (auto *BB : L->getBlocks())
         for (Instruction &I : *BB)
           Analyzer.visit(I);
       SimplifiedValuesVector.push_back(SimplifiedValues);
     }
     return false;
   }
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.setPreservesAll();
   }
   UnrollAnalyzerTest() : FunctionPass(ID) {
     initializeUnrollAnalyzerTestPass(*PassRegistry::getPassRegistry());
   }
TEST(UnrollAnalyzerTest, CastSimplifications) {
   EXPECT_TRUE(I3 != SimplifiedValuesVector[5].end());
   EXPECT_EQ(cast<ConstantInt>((*I3).second)->getZExtValue(), 3U);
 }
 
 } // end namespace llvm
 
 INITIALIZE_PASS_BEGIN(UnrollAnalyzerTest, "unrollanalyzertestpass",
                       "unrollanalyzertestpass", false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
 INITIALIZE_PASS_END(UnrollAnalyzerTest, "unrollanalyzertestpass",
                     "unrollanalyzertestpass", false, false)