Experimental late jump threading pass

[kparzysz]

Summary

It's not invoked anywhere by default.

[kparzysz]

Reduced code duplication between the legacy (old PM) passes.

[bmakam]

lib/Transforms/Scalar/JumpThreading.cpp

1824

Should this be
if (!IsLate && (LoopHeaders.count(BB) || LoopHeaders.count(SuccBB))) {

[kparzysz]

lib/Transforms/Scalar/JumpThreading.cpp

1824

Jump threading over loop headers was never allowed. You can try it and see what happens, but I'm not sure if it would be a good idea to enable it for all target.

[bmakam]

lib/Transforms/Scalar/JumpThreading.cpp

1824

I'm testing with a similar patch and I used this check. Wasn't JT over loop headers never allowed because it would turn the loop into irreducible loop, if that is the case this shouldn't be a concern for LateJumpThreading? However, I admit there is no difference in spec2017/perlbench with this.

[kparzysz]

lib/Transforms/Scalar/JumpThreading.cpp

1824

It could be a problem for the CodeGen passes that use MachineLoopInfo, since some loops would become non-optimizable.

include/llvm/InitializePasses.h

 void initializeInstructionSelectPass(PassRegistry&);
 void initializeInterleavedAccessPass(PassRegistry&);
 void initializeInternalizeLegacyPassPass(PassRegistry&);
 void initializeIntervalPartitionPass(PassRegistry&);
 void initializeJumpThreadingPass(PassRegistry&);
 void initializeLCSSAVerificationPassPass(PassRegistry&);
 void initializeLCSSAWrapperPassPass(PassRegistry&);
 void initializeLateCFGSimplifyPassPass(PassRegistry&);
+void initializeLateJumpThreadingPass(PassRegistry&);
 void initializeLazyBlockFrequencyInfoPassPass(PassRegistry&);
 void initializeLazyBranchProbabilityInfoPassPass(PassRegistry&);
 void initializeLazyMachineBlockFrequencyInfoPassPass(PassRegistry&);
 void initializeLazyValueInfoPrinterPass(PassRegistry&);
 void initializeLazyValueInfoWrapperPassPass(PassRegistry&);
 void initializeLegacyLICMPassPass(PassRegistry&);
 void initializeLegacyLoopSinkPassPass(PassRegistry&);
 void initializeLegalizerPass(PassRegistry&);

include/llvm/Transforms/Scalar.h

 
 //===----------------------------------------------------------------------===//
 //
 // JumpThreading - Thread control through mult-pred/multi-succ blocks where some
 // preds always go to some succ. Thresholds other than minus one override the
 // internal BB duplication default threshold.
 //
 FunctionPass *createJumpThreadingPass(int Threshold = -1);
+FunctionPass *createLateJumpThreadingPass(int Threshold = -1);
 
 //===----------------------------------------------------------------------===//
 //
 // CFGSimplification - Merge basic blocks, eliminate unreachable blocks,
 // simplify terminator instructions, etc...
 //
 FunctionPass *createCFGSimplificationPass(
     int Threshold = -1, std::function<bool(const Function &)> Ftor = nullptr);

include/llvm/Transforms/Scalar/JumpThreading.h

 class JumpThreadingPass : public PassInfoMixin<JumpThreadingPass> {
   TargetLibraryInfo *TLI;
   LazyValueInfo *LVI;
   AliasAnalysis *AA;
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   bool HasProfileData = false;
   bool HasGuards = false;
+  bool IsLate = false;
 #ifdef NDEBUG
   SmallPtrSet<const BasicBlock *, 16> LoopHeaders;
 #else
   SmallSet<AssertingVH<const BasicBlock>, 16> LoopHeaders;
 #endif
   DenseSet<std::pair<Value *, BasicBlock *>> RecursionSet;
 
   unsigned BBDupThreshold;
 
   // RAII helper for updating the recursion stack.
   struct RecursionSetRemover {
     DenseSet<std::pair<Value *, BasicBlock *>> &TheSet;
     std::pair<Value *, BasicBlock *> ThePair;
 
     RecursionSetRemover(DenseSet<std::pair<Value *, BasicBlock *>> &S,
                         std::pair<Value *, BasicBlock *> P)
         : TheSet(S), ThePair(P) {}
 
     ~RecursionSetRemover() { TheSet.erase(ThePair); }
   };
 
 public:
-  JumpThreadingPass(int T = -1);
+  JumpThreadingPass(int T = -1, bool Late = false);
 
   // Glue for old PM.
   bool runImpl(Function &F, TargetLibraryInfo *TLI_, LazyValueInfo *LVI_,
                AliasAnalysis *AA_, bool HasProfileData_,
                std::unique_ptr<BlockFrequencyInfo> BFI_,
                std::unique_ptr<BranchProbabilityInfo> BPI_);
 
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
private:
   BasicBlock *SplitBlockPreds(BasicBlock *BB, ArrayRef<BasicBlock *> Preds,
                               const char *Suffix);
   void UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB, BasicBlock *BB,
                                     BasicBlock *NewBB, BasicBlock *SuccBB);
   /// Check if the block has profile metadata for its outgoing edges.
   bool doesBlockHaveProfileData(BasicBlock *BB);
 };
 
+class LateJumpThreadingPass : public JumpThreadingPass {
+public:
+  LateJumpThreadingPass(int T = -1) : JumpThreadingPass(T, true) {}
+};
+
 } // end namespace llvm
 
 #endif

lib/Passes/PassRegistry.def

 FUNCTION_PASS("loop-sink", LoopSinkPass())
 FUNCTION_PASS("lowerinvoke", LowerInvokePass())
 FUNCTION_PASS("mem2reg", PromotePass())
 FUNCTION_PASS("memcpyopt", MemCpyOptPass())
 FUNCTION_PASS("mldst-motion", MergedLoadStoreMotionPass())
 FUNCTION_PASS("nary-reassociate", NaryReassociatePass())
 FUNCTION_PASS("newgvn", NewGVNPass())
 FUNCTION_PASS("jump-threading", JumpThreadingPass())
+FUNCTION_PASS("late-jump-threading", LateJumpThreadingPass())
 FUNCTION_PASS("partially-inline-libcalls", PartiallyInlineLibCallsPass())
 FUNCTION_PASS("lcssa", LCSSAPass())
 FUNCTION_PASS("loop-data-prefetch", LoopDataPrefetchPass())
 FUNCTION_PASS("loop-load-elim", LoopLoadEliminationPass())
 FUNCTION_PASS("loop-distribute", LoopDistributePass())
 FUNCTION_PASS("loop-vectorize", LoopVectorizePass())
 FUNCTION_PASS("pgo-memop-opt", PGOMemOPSizeOpt())
 FUNCTION_PASS("print", PrintFunctionPass(dbgs()))

lib/Transforms/Scalar/JumpThreading.cpp

namespace {
   ///   if () { ...
   ///     X = 4;
   ///   }
   ///   if (X < 3) {
   ///
   /// In this case, the unconditional branch at the end of the first if can be
   /// revectored to the false side of the second if.
   ///
-  class JumpThreading : public FunctionPass {
+  class JumpThreadingBase : public FunctionPass {
     JumpThreadingPass Impl;
 
   public:
-    static char ID; // Pass identification
-    JumpThreading(int T = -1) : FunctionPass(ID), Impl(T) {
-      initializeJumpThreadingPass(*PassRegistry::getPassRegistry());
-    }
+    JumpThreadingBase(int T, bool Late, char &ID)
+      : FunctionPass(ID), Impl(T, false) {}
 
     bool runOnFunction(Function &F) override;
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       if (PrintLVIAfterJumpThreading)
         AU.addRequired<DominatorTreeWrapperPass>();
       AU.addRequired<AAResultsWrapperPass>();
       AU.addRequired<LazyValueInfoWrapperPass>();
       AU.addPreserved<GlobalsAAWrapperPass>();
       AU.addRequired<TargetLibraryInfoWrapperPass>();
     }
 
     void releaseMemory() override { Impl.releaseMemory(); }
   };
+
+  class JumpThreading : public JumpThreadingBase {
+  public:
+    static char ID; // Pass identification
+    JumpThreading(int T = -1) : JumpThreadingBase(T, false, ID) {
+      initializeJumpThreadingPass(*PassRegistry::getPassRegistry());
+    }
+  };
+
+  class LateJumpThreading : public JumpThreadingBase {
+    JumpThreadingPass Impl;
+
+  public:
+    static char ID; // Pass identification
+    LateJumpThreading(int T = -1) : JumpThreadingBase(T, false, ID) {
+      initializeLateJumpThreadingPass(*PassRegistry::getPassRegistry());
+    }
+  };
 }
 
 char JumpThreading::ID = 0;
 INITIALIZE_PASS_BEGIN(JumpThreading, "jump-threading",
                 "Jump Threading", false, false)
 INITIALIZE_PASS_DEPENDENCY(LazyValueInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_END(JumpThreading, "jump-threading",
                 "Jump Threading", false, false)
 
+char LateJumpThreading::ID = 0;
+INITIALIZE_PASS_BEGIN(LateJumpThreading, "late-jump-threading",
+                "Late Jump Threading", false, false)
+INITIALIZE_PASS_DEPENDENCY(LazyValueInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+INITIALIZE_PASS_END(LateJumpThreading, "late-jump-threading",
+                "Late Jump Threading", false, false)
+
 // Public interface to the Jump Threading pass
-FunctionPass *llvm::createJumpThreadingPass(int Threshold) { return new JumpThreading(Threshold); }
+FunctionPass *llvm::createJumpThreadingPass(int Threshold) {
+  return new JumpThreading(Threshold);
+}
+FunctionPass *llvm::createLateJumpThreadingPass(int Threshold) {
+  return new LateJumpThreading(Threshold);
+}
 
-JumpThreadingPass::JumpThreadingPass(int T) {
+JumpThreadingPass::JumpThreadingPass(int T, bool Late) : IsLate(Late) {
   BBDupThreshold = (T == -1) ? BBDuplicateThreshold : unsigned(T);
 }
 
 // Update branch probability information according to conditional
 // branch probablity. This is usually made possible for cloned branches
 // in inline instances by the context specific profile in the caller.
 // For instance,
 //
for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
     PredBr->setMetadata(LLVMContext::MD_prof,
                         MDBuilder(PredBr->getParent()->getContext())
                             .createBranchWeights(Weights));
   }
 }
 
 /// runOnFunction - Toplevel algorithm.
 ///
-bool JumpThreading::runOnFunction(Function &F) {
+bool JumpThreadingBase::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
   auto TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
   auto LVI = &getAnalysis<LazyValueInfoWrapperPass>().getLVI();
   auto AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   bool HasProfileData = F.getEntryCount().hasValue();
bool JumpThreadingPass::ThreadEdge(BasicBlock *BB,
   if (SuccBB == BB) {
     DEBUG(dbgs() << "  Not threading across BB '" << BB->getName()
           << "' - would thread to self!\n");
     return false;
   }
 
   // If threading this would thread across a loop header, don't thread the edge.
   // See the comments above FindLoopHeaders for justifications and caveats.
-  if (LoopHeaders.count(BB) || LoopHeaders.count(SuccBB)) {
+  if (LoopHeaders.count(BB) || (!IsLate && LoopHeaders.count(SuccBB))) {
     DEBUG({
       bool BBIsHeader = LoopHeaders.count(BB);
       bool SuccIsHeader = LoopHeaders.count(SuccBB);
       dbgs() << "  Not threading across "
           << (BBIsHeader ? "loop header BB '" : "block BB '") << BB->getName()
           << "' to dest " << (SuccIsHeader ? "loop header BB '" : "block BB '")
           << SuccBB->getName() << "' - it might create an irreducible loop!\n";
     });

lib/Transforms/Scalar/Scalar.cpp

void llvm::initializeScalarOpts(PassRegistry &Registry) {
   initializeEarlyCSEMemSSALegacyPassPass(Registry);
   initializeGVNHoistLegacyPassPass(Registry);
   initializeGVNSinkLegacyPassPass(Registry);
   initializeFlattenCFGPassPass(Registry);
   initializeInductiveRangeCheckEliminationPass(Registry);
   initializeIndVarSimplifyLegacyPassPass(Registry);
   initializeInferAddressSpacesPass(Registry);
   initializeJumpThreadingPass(Registry);
+  initializeLateJumpThreadingPass(Registry);
   initializeLegacyLICMPassPass(Registry);
   initializeLegacyLoopSinkPassPass(Registry);
   initializeLoopDataPrefetchLegacyPassPass(Registry);
   initializeLoopDeletionLegacyPassPass(Registry);
   initializeLoopAccessLegacyAnalysisPass(Registry);
   initializeLoopInstSimplifyLegacyPassPass(Registry);
   initializeLoopInterchangePass(Registry);
   initializeLoopPredicationLegacyPassPass(Registry);