[PR16756] JumpThreading: explicitly update SSA rather than use SSAUpdater.

[mzolotukhin]

Summary

SSAUpdater is often a bottleneck in JumpThreading, and one of the reasons is
that it performs a lot of unnecessary computations (DT/IDF) over and over
again. This patch implements a classic algorithm for PHI-nodes placement and
uses JT-specific properties to optimize it (namely: we only have two blocks with
definitions and these blocks are the same for all instructions we're going to
rewrite).

With this patch the test from PR16756 speeds-up by ~2x, while the time spent in
JumpThreading goes down by ~4x.

Before the patch:

Total Execution Time: 26.6205 seconds (26.6232 wall clock)
 ---User Time---   --System Time--   --User+System--   ---Wall Time---  --- Name ---
13.3016 ( 50.6%)   0.0167 (  4.7%)  13.3183 ( 50.0%)  13.3190 ( 50.0%)  Jump Threading
 5.3226 ( 20.3%)   0.0170 (  4.8%)   5.3397 ( 20.1%)   5.3408 ( 20.1%)  Jump Threading
 1.7753 (  6.8%)   0.0020 (  0.6%)   1.7772 (  6.7%)   1.7772 (  6.7%)  SLP Vectorizer
 1.1617 (  4.4%)   0.1579 ( 44.1%)   1.3195 (  5.0%)   1.3199 (  5.0%)  X86 DAG->DAG Instruction Selection

With the patch:

Total Execution Time: 12.8328 seconds (12.8335 wall clock)
 ---User Time---   --System Time--   --User+System--   ---Wall Time---  --- Name ---
 4.5331 ( 36.3%)   0.0135 (  4.1%)   4.5466 ( 35.4%)   4.5471 ( 35.4%)  Jump Threading
 1.7649 ( 14.1%)   0.0015 (  0.5%)   1.7664 ( 13.8%)   1.7665 ( 13.8%)  SLP Vectorizer
 1.1914 (  9.5%)   0.1594 ( 47.9%)   1.3507 ( 10.5%)   1.3510 ( 10.5%)  X86 DAG->DAG Instruction Selection
 0.8300 (  6.6%)   0.0034 (  1.0%)   0.8333 (  6.5%)   0.8333 (  6.5%)  Jump Threading

Also, some archeology. Some time ago there was a patch for an alternative
SSAUpdater implementation (https://reviews.llvm.org/D28934). It was not
committed back then, but I decided to try it first. Initially, it also showed
huge speed-ups, but then I discovered a couple of bugs, fixes for which ate a
big chunk of the speedups (although the new implementation still was
significantly faster than the existing one). I can upload an updated version of
that patch if there is an interest, but in JumpThreading I decided to pursue
another path: the D28934 implementation, being much faster than the existing
one, still does not scale very well. The algorithm used there is very efficient
for small incremental updates, but doesn't scale well for bulk updates (e.g.
when we have to update big number of instructions), as there is not much to
share across the algorithm invocations. In contrast, the approach I propose here
aims at reusing as much computations as possible.

Compile time impact: no noticable changes on CTMark, a big improvement on the
test from PR16756.

[davide]

I'm going to review this in detail, but I have a general comment about the design.
I'm a little torn about the approach you're using. I think it's a good experiment because you're leveraging the structure of the problem to find a fundamentally more efficient solution. This is, FWIW, not really surprising.
OTOH, this is bad because you're de-facto rewriting a specialized SSA updater. We might need another one for, e.g. LCSSA, and another one for GVN/PRE (until NewGVN is in or we rewrite PRE to not use the updater).
This results in basically a bunch of code duplicated. History shows (and you know this very well) that it turned out to be a long tail of bugs when we did something similar with the dominator (have specialized updates per-pass). If we really want to go this route, we really need to analyze carefully all the implications.

[mzolotukhin]

Yeah, I have such feelings regarding this patch too.

I've been thinking about how to make this more general and reusable: it should be possible to implement something like BulkSSAUpdater (it probably might use a better name), and it will at least won't recompute DT on every iteration. I don't think it would be worth it to expose in its interface a flag showing that all definitions are in the same blocks or something like this, so we'll have to recompute IDF for every instruction we process. It still will be a win compared to the existing implementation, but loss in term of compile time compared to this implementation.

When I submitted this patch, I expected that it would trigger discussions, and that was one of the reasons I wanted to put it out :) I'm eager to suggestions on how to do that better.

Thanks,
Michael

[• dberlin]

So, FWIW:
It's definitely possible to make the SSAUpdater rewrite faster in bulk insertion. I just never bothered because it wasn't used that way :)
You will never make it as fast as exploiting the problem structure, but it could be made a *lot* faster.

As for this one:
There are approaches to reuse a lot of the phi placement computation in bulk. It's also possible to close IDF recomputation under DT update (IE incrementally recompute IDF).

You can see an example of something like that in the TDMSC algorithms here:https://dl.acm.org/citation.cfm?id=1065890

It only relies on the DJ graph, and it's possible, at maximum, to invalidate only the merge sets as DT changes.
(We'd need to get info from DT about what happened during an update, but that seems easy enough)

All depends how far down this rabbit hole you want to go.

As for the pass-specific vs bulk updating, yeah, this is the eternal debate.
JT itself doesn't make too many types of transforms, and i wonder if you could just verify them all working with the ssa updater with all the types of cfgs or something.
(This would be fragile in other ways, but at least would give you confidence the thing works)

I think it also depensd on how much faster you can make the generic bulk updater. For example, if you get it to 0.9x the speed of the impl you have here, i suspect that's plenty good enough.
But i'm also not the one doing the work :)

[mzolotukhin]

I think so far the consensus was that I need to try implementing it in a more general way. I'll do that and come back!

Thanks,
Michael

[hiraditya]

If you want to update using DJ-Graphs and merge sets, I have implemented them in the global-scheduler which I thought could be reused: https://reviews.llvm.org/D32140
Please let me know if you find this useful.

[mzolotukhin]

Thanks! My current plan is to simply move the part that I wrote to a separate class and make it usable from other places. If I manage to keep most of the gains, then it would probably be the easiest way to resolve the issue.

Michael

[mzolotukhin]

• Factor out SSA updating logic to a separate class SSAUpdaterBulk.

[mzolotukhin]

Hi,

I created a separate class for bulk SSA updates. With that implementation, we recompute IDF for every individual variable, but usually the subgraph we're working with is smaller. In the previous implementation we computed IDF once for the union of these subgraphs - that was faster, but also we might accidentally insert unneeded phi-nodes, which we later had to clean-up.

I tried to make the interface close to the existing one, and included an example of how it can be used in JumpThreading (if this change is approved, I'll commit these parts separately). With this change, the speedup on the original test is smaller, but still quite good:

===-------------------------------------------------------------------------===
  Total Execution Time: 15.0158 seconds (15.0163 wall clock)

   ---User Time---   --System Time--   --User+System--   ---Wall Time---  --- Name ---
   4.5962 ( 31.5%)   0.0170 (  3.8%)   4.6132 ( 30.7%)   4.6142 ( 30.7%)  Jump Threading
   3.1789 ( 21.8%)   0.0049 (  1.1%)   3.1838 ( 21.2%)   3.1838 ( 21.2%)  Jump Threading
   1.2104 (  8.3%)   0.1836 ( 41.5%)   1.3940 (  9.3%)   1.3942 (  9.3%)  X86 DAG->DAG Instruction Selection
   1.3558 (  9.3%)   0.0020 (  0.5%)   1.3578 (  9.0%)   1.3577 (  9.0%)  SLP Vectorizer

What do you think?

Thanks,
Michael

PS: Ideas for a better class name are welcome!

[mzolotukhin]

Ping!

[davide]

I don't mind this approach, but as discussed offline we should consider also moving LCSSA to make sure the API makes sense.
In general, what's your plan for this? You want it to replace the SSA updater for all the instances in llvm? If so, we should carefully plan the transition costs.

@zhendongsu could do a round of testing before we decide whether this can go in to shake bugs.

Also, do you know why the second invocation of Jump threading is taking so long?

llvm/lib/Transforms/Utils/SSAUpdaterBulk.cpp

25–26 ↗

(On Diff #138097)

should this be ssaupdaterbulk?

52–54 ↗

(On Diff #138097)

ternary

111 ↗

(On Diff #138097)

auto

113 ↗

(On Diff #138097)

auto

[davide]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2083

typo: s/ee/we/

2102

Why do you need to flush the dominator here? please add a comment.

[mzolotukhin]

• Rebase.
• Address Davide's remarks.

[mzolotukhin]

I don't mind this approach, but as discussed offline we should consider also moving LCSSA to make sure the API makes sense.
In general, what's your plan for this? You want it to replace the SSA updater for all the instances in llvm? If so, we should carefully plan the transition costs.

I looked into LCSSA, and indeed it seems that it also can be improved. I tried direct replacement of the old SSAUpdater with the new one, but that didn't give any benefits. However, I think we can simplify the code in LCSSA by passing LoopInfo to SSAUpdaterBulk, which will then use it to insert a phi node whenever it crosses a loop boundary when rewriting a use. I don't know yet how much work it would take, but I don't think it would require much rewritings - it would probably be an addition to what we currently have in this patch.

Also, do you know why the second invocation of Jump threading is taking so long?

I assume by 'taking so long' you mean 'taking so long compared to the first version of this patch', because even this "slow" version is ~50% faster than what we have in trunk. It is most probably caused by the fact that we recompute IDFs for every rewriting while in the first version we computed it once for a united subgraph. Actually, both approaches have "good" and "bad" cases: I can imagine a set of subgraphs for which computing individual IDFs N times would be faster than computing a united IDF, but seemingly in this particular case we have an opposite example.

Thanks,
Michael

[• dberlin]

I'm ... confused.

[davide]

I have no other comments on this, which looks fine. I'd add some unittest(s), as we have in the SSAUpdater and I think we should be good to do.
@dberlin what do you think?

[• dberlin]

I'm fine with this, we can always improve it more later.

[mzolotukhin]

Thanks! I committed the patch in two parts: r329643 (Add SSAUpdaterBulk) and r329644 (Use SSAUpdaterBulk in JumpThreading).

Michael

[alekseyshl]

JumpThreading pass fails with assertion here:
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-android/builds/9491,
please check it out

[alexcrichton]

I've bisected a recent regression we discovered in Rust to this commit and cc'd some of y'all on the bug there, but if anyone else would like to help take a look at the performance regression there that'd be much appreciated!

llvm/lib/Transforms/Scalar/JumpThreading.cpp

 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/PredIteratorCache.h"
+#include "llvm/Analysis/IteratedDominanceFrontier.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
if (Instruction *Inst = dyn_cast<Instruction>(IV)) {
       if (I != ValueMap.end())
         IV = I->second;
     }
 
     PN.addIncoming(IV, NewPred);
   }
 }
 
+struct RewriteInfo {
+  DenseMap<BasicBlock*, Value*> Defines;
+  SmallPtrSet<Use*, 4> Uses;
+  Instruction *I;
+  RewriteInfo(Instruction *OrigI) : I(OrigI) {}
+};
+
+static void
+ComputeLiveInBlocks(const SmallPtrSetImpl<BasicBlock *> &UsingBlocks,
+                    const SmallPtrSetImpl<BasicBlock *> &DefBlocks,
+                    SmallPtrSetImpl<BasicBlock *> &LiveInBlocks) {
+  // To determine liveness, we must iterate through the predecessors of blocks
+  // where the def is live.  Blocks are added to the worklist if we need to
+  // check their predecessors.  Start with all the using blocks.
+  SmallVector<BasicBlock *, 64> LiveInBlockWorklist(UsingBlocks.begin(),
+                                                    UsingBlocks.end());
+
+  // Now that we have a set of blocks where the phi is live-in, recursively add
+  // their predecessors until we find the full region the value is live.
+  while (!LiveInBlockWorklist.empty()) {
+    BasicBlock *BB = LiveInBlockWorklist.pop_back_val();
+
+    // The block really is live in here, insert it into the set.  If already in
+    // the set, then it has already been processed.
+    if (!LiveInBlocks.insert(BB).second)
+      continue;
+
+    // Since the value is live into BB, it is either defined in a predecessor or
+    // live into it to.  Add the preds to the worklist unless they are a
+    // defining block.
+    for (BasicBlock *P : predecessors(BB)) {
+      // The value is not live into a predecessor if it defines the value.
+      if (DefBlocks.count(P))
+        continue;
+
+      // Otherwise it is, add to the worklist.
+      LiveInBlockWorklist.push_back(P);
+    }
+  }
+}
+
+// Compute value at the given block BB. We either should already know it, or we
+// should be able to recursively reach it going up dominator tree.
+static Value *computeValueAt(BasicBlock *BB,
+                             DenseMap<BasicBlock *, Value *> &KnownValues,
+                             DominatorTree *DT) {
+  if (!KnownValues.count(BB)) {
+    BasicBlock *IDom = DT->getNode(BB)->getIDom()->getBlock();
+    KnownValues[BB] = computeValueAt(IDom, KnownValues, DT);
+  }
+  return KnownValues[BB];
+}
+
+static BasicBlock *getUserBB(Use *U) {
+  Instruction *User = cast<Instruction>(U->getUser());
+
+  if (PHINode *UserPN = dyn_cast<PHINode>(User))
+    return UserPN->getIncomingBlock(*U);
+  else
+    return User->getParent();
+}
+
 /// ThreadEdge - We have decided that it is safe and profitable to factor the
 /// blocks in PredBBs to one predecessor, then thread an edge from it to SuccBB
 /// across BB.  Transform the IR to reflect this change.
 bool JumpThreadingPass::ThreadEdge(BasicBlock *BB,
                                    const SmallVectorImpl<BasicBlock *> &PredBBs,
                                    BasicBlock *SuccBB) {
   // If threading to the same block as we come from, we would infinite loop.
   if (SuccBB == BB) {
bool JumpThreadingPass::ThreadEdge(BasicBlock *BB,
   // Check to see if SuccBB has PHI nodes. If so, we need to add entries to the
   // PHI nodes for NewBB now.
   AddPHINodeEntriesForMappedBlock(SuccBB, BB, NewBB, ValueMapping);
 
   // If there were values defined in BB that are used outside the block, then we
   // now have to update all uses of the value to use either the original value,
   // the cloned value, or some PHI derived value.  This can require arbitrary
   // PHI insertion, of which we are prepared to do, clean these up now.
-  SSAUpdater SSAUpdate;
   SmallVector<Use*, 16> UsesToRename;
+
+  SmallVector<RewriteInfo, 16> Rewrites;
   for (Instruction &I : *BB) {
+    RewriteInfo RewritesForI(&I);
+    UsesToRename.clear();
+
     // Scan all uses of this instruction to see if it is used outside of its
     // block, and if so, record them in UsesToRename.
     for (Use &U : I.uses()) {
       Instruction *User = cast<Instruction>(U.getUser());
       if (PHINode *UserPN = dyn_cast<PHINode>(User)) {
         if (UserPN->getIncomingBlock(U) == BB)
           continue;
       } else if (User->getParent() == BB)
         continue;
 
       UsesToRename.push_back(&U);
+      RewritesForI.Uses.insert(&U);
     }
 
     // If there are no uses outside the block, we're done with this instruction.
     if (UsesToRename.empty())
       continue;
 
-    DEBUG(dbgs() << "JT: Renaming non-local uses of: " << I << "\n");
-
-    // We found a use of I outside of BB.  Rename all uses of I that are outside
-    // its block to be uses of the appropriate PHI node etc.  See ValuesInBlocks
-    // with the two values we know.
-    SSAUpdate.Initialize(I.getType(), I.getName());
-    SSAUpdate.AddAvailableValue(BB, &I);
-    SSAUpdate.AddAvailableValue(NewBB, ValueMapping[&I]);
-
-    while (!UsesToRename.empty())
-      SSAUpdate.RewriteUse(*UsesToRename.pop_back_val());
-    DEBUG(dbgs() << "\n");
+    // We found a use of I outside of BB - ee need to rename all uses of I that
+    // are outside its block to be uses of the appropriate PHI node etc.
+    RewritesForI.Defines[BB] = &I;
+    RewritesForI.Defines[NewBB] = ValueMapping[&I];
+    Rewrites.push_back(RewritesForI);
   }
 
   // Ok, NewBB is good to go.  Update the terminator of PredBB to jump to
   // NewBB instead of BB.  This eliminates predecessors from BB, which requires
   // us to simplify any PHI nodes in BB.
   TerminatorInst *PredTerm = PredBB->getTerminator();
   for (unsigned i = 0, e = PredTerm->getNumSuccessors(); i != e; ++i)
     if (PredTerm->getSuccessor(i) == BB) {
       BB->removePredecessor(PredBB, true);
       PredTerm->setSuccessor(i, NewBB);
     }
 
   DDT->applyUpdates({{DominatorTree::Insert, NewBB, SuccBB},
                      {DominatorTree::Insert, PredBB, NewBB},
                      {DominatorTree::Delete, PredBB, BB}});
 
+  // Now we need to update SSA.
+  // We will use the classic algorithm for placing PHI-nodes: "Efficiently
+  // Computing Static Single Assignment Form and the Control Dependence Graph"
+  // by Cytron et al.
+  //
+  // All the uses we need to rewrite along with the information about available
+  // definitions are stored in Rewrites. In a general case, to properly insert
+  // all necessary PHI-nodes, we need to compute IDF for iterated joins of all
+  // definitions for each instruction - the blocks in IDF will be the blocks
+  // where PHI-nodes need to be inserted.
+  //
+  // In the case of JumpThreading, all definitions for all instructions are
+  // located in just two basic blocks: BB and NewBB. We can exploit this
+  // information to avoid computing IDF for each instruction and compute it only
+  // once. UsingBlocks, which are used to limit the part of CFG for which IDF is
+  // computed, will be a union of UsingBlocks for all instructions. As a side
+  // effect of this, we might insert PHI-nodes without actual uses, but we can
+  // quickly wipe them out in the end.
+  DominatorTree *DT = &DDT->flush();
+
+  ForwardIDFCalculator IDF(*DT);
+  PredIteratorCache PredCache;
+
+  // DefBlocks contain only two blocks: BB and its clone NewBB.
+  SmallPtrSet<BasicBlock *, 2> DefBlocks;
+  DefBlocks.insert(BB);
+  DefBlocks.insert(NewBB);
+  IDF.setDefiningBlocks(DefBlocks);
+
+  // UsingBlocks contain all blocks with uses of instructions we are rewriting.
+  SmallPtrSet<BasicBlock *, 2> UsingBlocks;
+  for (auto R : Rewrites) {
+    for (auto U : R.Uses)
+      UsingBlocks.insert(getUserBB(U));
+  }
+
+  // Compute iterated dominance frontier, IDF.
+  SmallVector<BasicBlock *, 32> IDFBlocks;
+  SmallPtrSet<BasicBlock *, 32> LiveInBlocks;
+  ComputeLiveInBlocks(UsingBlocks, DefBlocks, LiveInBlocks);
+  IDF.resetLiveInBlocks();
+  IDF.setLiveInBlocks(LiveInBlocks);
+  IDF.calculate(IDFBlocks);
+
+  // Now, go through all instructions and perform actual IR updates.
+  for (auto R : Rewrites) {
+    DenseMap<BasicBlock *, Value *> KnownValues;
+
+    DEBUG(dbgs() << "JT: Renaming non-local uses of: " << *R.I << "\n");
+
+    // Fill in known values for the initial blocks.
+    KnownValues[BB] = R.I;
+    KnownValues[NewBB] = ValueMapping[R.I];
+
+    // Insert PHIs into blocks from IDF and record them as known values.
+    SmallVector<PHINode *, 4> InsertedPHIs;
+    for (auto FrontierBB : IDFBlocks) {
+      IRBuilder<> B(FrontierBB, FrontierBB->begin());
+      PHINode *PN = B.CreatePHI(R.I->getType(), 0, R.I->getName());
+      KnownValues[FrontierBB] = PN;
+      InsertedPHIs.push_back(PN);
+    }
+
+    // Fill in arguments of the inserted PHIs.
+    for (auto PN : InsertedPHIs) {
+      BasicBlock *PBB = PN->getParent();
+      for (BasicBlock *Pred : PredCache.get(PBB))
+        PN->addIncoming(computeValueAt(Pred, KnownValues, DT), Pred);
+    }
+
+    // Rewrite actual uses with the inserted definitions.
+    for (auto U : R.Uses) {
+      Value *V = computeValueAt(getUserBB(U), KnownValues, DT);
+      Value *OldVal = U->get();
+      // Notify that users of the existing value that it is being replaced.
+      if (OldVal != V && OldVal->hasValueHandle())
+        ValueHandleBase::ValueIsRAUWd(OldVal, V);
+      U->set(V);
+    }
+
+    // Clean up unused PHIs.
+    for (auto PN : InsertedPHIs) {
+      if (PN->use_empty())
+        PN->eraseFromParent();
+    }
+  }
+  // SSA Update is finished!
+
   // At this point, the IR is fully up to date and consistent.  Do a quick scan
   // over the new instructions and zap any that are constants or dead.  This
   // frequently happens because of phi translation.
   SimplifyInstructionsInBlock(NewBB, TLI);
 
   // Update the edge weight from BB to SuccBB, which should be less than before.
   UpdateBlockFreqAndEdgeWeight(PredBB, BB, NewBB, SuccBB);