[PlaceSafepoints] Reduce dominator tree recalculation

[reames]

Summary

Reduce recalculation of the dominator tree by identifying all sites that will need a safepoint poll before doing any of the insertion. This allows us to invalidate the dominator info once, rather than once per safepoint poll inserted.

While I'm at it, update findLocationForEntrySafepoint to properly update the dom tree now that the interface has been made easy. When first written, it wasn't per comment in the code.

[sanjoy]

LGTM with comments.

lib/Transforms/Scalar/PlaceSafepoints.cpp

661 ↗

(On Diff #25649)

I think a SmallVector is more appropriate here.

666 ↗

(On Diff #25649)

I don't think this is the right way to avoid using PollsNeeded. As a later change, we should factor PlaceSafepoints::runOnFunction into smaller functions such that PollsNeeded is statically not in scope of the functions that should not be using it.

[reames]

lib/Transforms/Scalar/PlaceSafepoints.cpp

661 ↗

(On Diff #25649)

We can't currently create a small vector of CallSites. I haven't dug into why not, but it fails to compile.

666 ↗

(On Diff #25649)

Agreed. I'll submit as is, but plan to refactor in this direction.

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

void getAnalysisUsage(AnalysisUsage &AU) const override {
     // if that was worth doing
   }
 };
 }
 
 // Insert a safepoint poll immediately before the given instruction.  Does
 // not handle the parsability of state at the runtime call, that's the
 // callers job.
-static void
-InsertSafepointPoll(DominatorTree &DT, Instruction *after,
+static void 
+InsertSafepointPoll(Instruction *after,
                     std::vector<CallSite> &ParsePointsNeeded /*rval*/);
 
 static bool isGCLeafFunction(const CallSite &CS);
 
 static bool needsStatepoint(const CallSite &CS) {
   if (isGCLeafFunction(CS))
     return false;
   if (CS.isCall()) {
static Instruction *findLocationForEntrySafepoint(Function &F,
   assert((hasNextInstruction(cursor) || cursor->isTerminator()) &&
          "either we stopped because of a call, or because of terminator");
 
   if (cursor->isTerminator()) {
     return cursor;
   }
 
   BasicBlock *BB = cursor->getParent();
-  SplitBlock(BB, cursor, nullptr);
-
-  // Note: SplitBlock modifies the DT.  Simply passing a Pass (which is a
-  // module pass) is not enough.
-  DT.recalculate(F);
+  SplitBlock(BB, cursor, &DT);
 
   // SplitBlock updates the DT
   DEBUG(DT.verifyDomTree());
 
   return BB->getTerminator();
 }
 
 /// Identify the list of call sites which need to be have parseable state
bool PlaceSafepoints::runOnFunction(Function &F) {
   // and reachablity queries return non-sensical results.  Thus, we preprocess
   // the function to ensure these properties hold.
   modified |= removeUnreachableBlocks(F);
 
   // STEP 1 - Insert the safepoint polling locations.  We do not need to
   // actually insert parse points yet.  That will be done for all polls and
   // calls in a single pass.
 
-  // Note: With the migration, we need to recompute this for each 'pass'.  Once
-  // we merge these, we'll do it once before the analysis
   DominatorTree DT;
+  DT.recalculate(F);
 
+  SmallVector<Instruction *, 16> PollsNeeded;
   std::vector<CallSite> ParsePointNeeded;
 
   if (enableBackedgeSafepoints(F)) {
     // Construct a pass manager to run the LoopPass backedge logic.  We
     // need the pass manager to handle scheduling all the loop passes
     // appropriately.  Doing this by hand is painful and just not worth messing
     // with for the moment.
     legacy::FunctionPassManager FPM(F.getParent());
PollLocations.erase(std::unique(PollLocations.begin(),
                         PollLocations.end());
 
     // Insert a poll at each point the analysis pass identified
     // The poll location must be the terminator of a loop latch block.
     for (TerminatorInst *Term : PollLocations) {
       // We are inserting a poll, the function is modified
       modified = true;
       
-      std::vector<CallSite> ParsePoints;
       if (SplitBackedge) {
         // Split the backedge of the loop and insert the poll within that new
         // basic block.  This creates a loop with two latches per original
         // latch (which is non-ideal), but this appears to be easier to
         // optimize in practice than inserting the poll immediately before the
         // latch test.
 
         // Since this is a latch, at least one of the successors must dominate
for (TerminatorInst *Term : PollLocations) {
         assert(!Headers.empty() && "poll location is not a loop latch?");
 
         // The split loop structure here is so that we only need to recalculate
         // the dominator tree once.  Alternatively, we could just keep it up to
         // date and use a more natural merged loop.
         SetVector<BasicBlock *> SplitBackedges;
         for (BasicBlock *Header : Headers) {
           BasicBlock *NewBB = SplitEdge(Term->getParent(), Header, &DT);
-          
-          std::vector<CallSite> RuntimeCalls;
-          InsertSafepointPoll(DT, NewBB->getTerminator(), RuntimeCalls);
+          PollsNeeded.push_back(NewBB->getTerminator());
           NumBackedgeSafepoints++;
-          ParsePointNeeded.insert(ParsePointNeeded.end(), RuntimeCalls.begin(),
-                                  RuntimeCalls.end());
         }
-
       } else {
         // Split the latch block itself, right before the terminator.
-        std::vector<CallSite> RuntimeCalls;
-        InsertSafepointPoll(DT, Term, RuntimeCalls);
+        PollsNeeded.push_back(Term);
         NumBackedgeSafepoints++;
-        ParsePointNeeded.insert(ParsePointNeeded.end(), RuntimeCalls.begin(),
-                                RuntimeCalls.end());
       }
-
-      // Record the parse points for later use
-      ParsePointNeeded.insert(ParsePointNeeded.end(), ParsePoints.begin(),
-                              ParsePoints.end());
     }
   }
 
   if (enableEntrySafepoints(F)) {
-    DT.recalculate(F);
-    Instruction *term = findLocationForEntrySafepoint(F, DT);
-    if (!term) {
+    Instruction *Location = findLocationForEntrySafepoint(F, DT);
+    if (!Location) {
       // policy choice not to insert?
     } else {
-      std::vector<CallSite> RuntimeCalls;
-      InsertSafepointPoll(DT, term, RuntimeCalls);
+      PollsNeeded.push_back(Location);
       modified = true;
       NumEntrySafepoints++;
-      ParsePointNeeded.insert(ParsePointNeeded.end(), RuntimeCalls.begin(),
-                              RuntimeCalls.end());
-    }
-  }
+    }
+  }
+
+  // Now that we've identified all the needed safepoint poll locations, insert
+  // safepoint polls themselves.
+  for (Instruction *PollLocation : PollsNeeded) {
+    std::vector<CallSite> RuntimeCalls;
+    InsertSafepointPoll(PollLocation, RuntimeCalls);
+  PollsNeeded.clear(); // make sure we don't accidentally use
+  // The dominator tree has been invalidated by the inlining performed in the
+  // above loop.  TODO: Teach the inliner how to update the dom tree?
+  DT.recalculate(F);
   
   if (enableCallSafepoints(F)) {
-    DT.recalculate(F);
     std::vector<CallSite> Calls;
     findCallSafepoints(F, Calls);
     NumCallSafepoints += Calls.size();
     ParsePointNeeded.insert(ParsePointNeeded.end(), Calls.begin(), Calls.end());
   }
 
   // Unique the vectors since we can end up with duplicates if we scan the call
   // site for call safepoints after we add it for entry or backedge.  The
   // only reason we need tracking at all is that some functions might have
   // polls but not call safepoints and thus we might miss marking the runtime
   // calls for the polls. (This is useful in test cases!)
   unique_unsorted(ParsePointNeeded);
 
   // Any parse point (no matter what source) will be handled here
-  DT.recalculate(F); // Needed?
 
   // We're about to start modifying the function
   if (!ParsePointNeeded.empty())
     modified = true;
 
   // Now run through and insert the safepoints, but do _NOT_ update or remove
   // any existing uses.  We have references to live variables that need to
   // survive to the last iteration of this loop.
bool isLeaf =
       F->getFnAttribute("gc-leaf-function").getValueAsString().equals("true");
   if (isLeaf) {
     return true;
   }
   return false;
 }
 
 static void
-InsertSafepointPoll(DominatorTree &DT, Instruction *term,
+InsertSafepointPoll(Instruction *term,
                     std::vector<CallSite> &ParsePointsNeeded /*rval*/) {
   Module *M = term->getParent()->getParent()->getParent();
   assert(M);
 
   // Inline the safepoint poll implementation - this will get all the branch,
   // control flow, etc..  Most importantly, it will introduce the actual slow
   // path call - where we need to insert a safepoint (parsepoint).
   FunctionType *ftype =
InsertSafepointPoll(Instruction *term,
   bool isBegin(false);
   if (before == term->getParent()->begin()) {
     isBegin = true;
   } else {
     before--;
   }
   after++;
   assert(after != poll->getParent()->end() && "must have successor");
-  assert(DT.dominates(before, after) && "trivially true");
 
   // do the actual inlining
   InlineFunctionInfo IFI;
   bool inlineStatus = InlineFunction(poll, IFI);
   assert(inlineStatus && "inline must succeed");
   (void)inlineStatus; // suppress warning in release-asserts
 
   // Check post conditions
InsertSafepointPoll(Instruction *term,
   }
 
   // If your poll function includes an unreachable at the end, that's not
   // valid.  Bugpoint likes to create this, so check for it.
   assert(isPotentiallyReachable(&*start, &*after, nullptr, nullptr) &&
          "malformed poll function");
 
   scanInlinedCode(&*(start), &*(after), calls, BBs);
-
-  // Recompute since we've invalidated cached data.  Conceptually we
-  // shouldn't need to do this, but implementation wise we appear to.  Needed
-  // so we can insert safepoints correctly.
-  // TODO: update more cheaply
-  DT.recalculate(*after->getParent()->getParent());
-
   assert(!calls.empty() && "slow path not found for safepoint poll");
 
   // Record the fact we need a parsable state at the runtime call contained in
   // the poll function.  This is required so that the runtime knows how to
   // parse the last frame when we actually take  the safepoint (i.e. execute
   // the slow path)
   assert(ParsePointsNeeded.empty());
   for (size_t i = 0; i < calls.size(); i++) {