Backport fixes for JumpThreading/assume bug

include/llvm/Analysis/LazyValueInfo.h

@@ -113,6 +113,13 @@ class LazyValueInfo {
   /// in LVI, so we need to pass it here as an argument.
   void printLVI(Function &F, DominatorTree &DTree, raw_ostream &OS);
 
+  /// Disables use of the DominatorTree within LVI.
+  void disableDT();
+
+  /// Enables use of the DominatorTree within LVI. Does nothing if the class
+  /// instance was initialized without a DT pointer.
+  void enableDT();
+
   // For old PM pass. Delete once LazyValueInfoWrapperPass is gone.
   void releaseMemory();
 

include/llvm/IR/Dominators.h

@@ -290,6 +290,93 @@ class DominatorTreeWrapperPass : public FunctionPass {
   void print(raw_ostream &OS, const Module *M = nullptr) const override;
 };
 
+//===-------------------------------------
+/// \brief Class to defer updates to a DominatorTree.
+///
+/// Definition: Applying updates to every edge insertion and deletion is
+/// expensive and not necessary. When one needs the DominatorTree for analysis
+/// they can request a flush() to perform a larger batch update. This has the
+/// advantage of the DominatorTree inspecting the set of updates to find
+/// duplicates or unnecessary subtree updates.
+///
+/// The scope of DeferredDominance operates at a Function level.
+///
+/// It is not necessary for the user to scrub the updates for duplicates or
+/// updates that point to the same block (Delete, BB_A, BB_A). Performance
+/// can be gained if the caller attempts to batch updates before submitting
+/// to applyUpdates(ArrayRef) in cases where duplicate edge requests will
+/// occur.
+///
+/// It is required for the state of the LLVM IR to be applied *before*
+/// submitting updates. The update routines must analyze the current state
+/// between a pair of (From, To) basic blocks to determine if the update
+/// needs to be queued.
+/// Example (good):
+///     TerminatorInstructionBB->removeFromParent();
+///     DDT->deleteEdge(BB, Successor);
+/// Example (bad):
+///     DDT->deleteEdge(BB, Successor);
+///     TerminatorInstructionBB->removeFromParent();
+class DeferredDominance {
+public:
+  DeferredDominance(DominatorTree &DT_) : DT(DT_) {}
+
+  /// \brief Queues multiple updates and discards duplicates.
+  void applyUpdates(ArrayRef<DominatorTree::UpdateType> Updates);
+
+  /// \brief Helper method for a single edge insertion. It's almost always
+  /// better to batch updates and call applyUpdates to quickly remove duplicate
+  /// edges. This is best used when there is only a single insertion needed to
+  /// update Dominators.
+  void insertEdge(BasicBlock *From, BasicBlock *To);
+
+  /// \brief Helper method for a single edge deletion. It's almost always better
+  /// to batch updates and call applyUpdates to quickly remove duplicate edges.
+  /// This is best used when there is only a single deletion needed to update
+  /// Dominators.
+  void deleteEdge(BasicBlock *From, BasicBlock *To);
+
+  /// \brief Delays the deletion of a basic block until a flush() event.
+  void deleteBB(BasicBlock *DelBB);
+
+  /// \brief Returns true if DelBB is awaiting deletion at a flush() event.
+  bool pendingDeletedBB(BasicBlock *DelBB);
+
+  /// \brief Returns true if pending DT updates are queued for a flush() event.
+  bool pending();
+
+  /// \brief Flushes all pending updates and block deletions. Returns a
+  /// correct DominatorTree reference to be used by the caller for analysis.
+  DominatorTree &flush();
+
+  /// \brief Drops all internal state and forces a (slow) recalculation of the
+  /// DominatorTree based on the current state of the LLVM IR in F. This should
+  /// only be used in corner cases such as the Entry block of F being deleted.
+  void recalculate(Function &F);
+
+  /// \brief Debug method to help view the state of pending updates.
+  LLVM_DUMP_METHOD void dump() const;
+
+private:
+  DominatorTree &DT;
+  SmallVector<DominatorTree::UpdateType, 16> PendUpdates;
+  SmallPtrSet<BasicBlock *, 8> DeletedBBs;
+
+  /// Apply an update (Kind, From, To) to the internal queued updates. The
+  /// update is only added when determined to be necessary. Checks for
+  /// self-domination, unnecessary updates, duplicate requests, and balanced
+  /// pairs of requests are all performed. Returns true if the update is
+  /// queued and false if it is discarded.
+  bool applyUpdate(DominatorTree::UpdateKind Kind, BasicBlock *From,
+                   BasicBlock *To);
+
+  /// Performs all pending basic block deletions. We have to defer the deletion
+  /// of these blocks until after the DominatorTree updates are applied. The
+  /// internal workings of the DominatorTree code expect every update's From
+  /// and To blocks to exist and to be a member of the same Function.
+  bool flushDelBB();
+};
+
 } // end namespace llvm
 
 #endif // LLVM_IR_DOMINATORS_H

include/llvm/Transforms/Scalar/JumpThreading.h

@@ -34,6 +34,7 @@ class BinaryOperator;
 class BranchInst;
 class CmpInst;
 class Constant;
+class DeferredDominance;
 class Function;
 class Instruction;
 class IntrinsicInst;
@@ -77,6 +78,7 @@ class JumpThreadingPass : public PassInfoMixin<JumpThreadingPass> {
   TargetLibraryInfo *TLI;
   LazyValueInfo *LVI;
   AliasAnalysis *AA;
+  DeferredDominance *DDT;
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   bool HasProfileData = false;
@@ -107,8 +109,8 @@ class JumpThreadingPass : public PassInfoMixin<JumpThreadingPass> {
 
   // Glue for old PM.
   bool runImpl(Function &F, TargetLibraryInfo *TLI_, LazyValueInfo *LVI_,
-               AliasAnalysis *AA_, bool HasProfileData_,
-               std::unique_ptr<BlockFrequencyInfo> BFI_,
+               AliasAnalysis *AA_, DeferredDominance *DDT_,
+               bool HasProfileData_, std::unique_ptr<BlockFrequencyInfo> BFI_,
                std::unique_ptr<BranchProbabilityInfo> BPI_);
 
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

include/llvm/Transforms/Utils/BasicBlockUtils.h

@@ -27,6 +27,7 @@ namespace llvm {
 
 class BlockFrequencyInfo;
 class BranchProbabilityInfo;
+class DeferredDominance;
 class DominatorTree;
 class Function;
 class Instruction;
@@ -38,7 +39,7 @@ class TargetLibraryInfo;
 class Value;
 
 /// Delete the specified block, which must have no predecessors.
-void DeleteDeadBlock(BasicBlock *BB);
+void DeleteDeadBlock(BasicBlock *BB, DeferredDominance *DDT = nullptr);
 
 /// We know that BB has one predecessor. If there are any single-entry PHI nodes
 /// in it, fold them away. This handles the case when all entries to the PHI

include/llvm/Transforms/Utils/Local.h

@@ -117,7 +117,8 @@ struct SimplifyCFGOptions {
 /// conditions and indirectbr addresses this might make dead if
 /// DeleteDeadConditions is true.
 bool ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions = false,
-                            const TargetLibraryInfo *TLI = nullptr);
+                            const TargetLibraryInfo *TLI = nullptr,
+                            DeferredDominance *DDT = nullptr);
 
 //===----------------------------------------------------------------------===//
 //  Local dead code elimination.
@@ -171,18 +172,21 @@ bool SimplifyInstructionsInBlock(BasicBlock *BB,
 ///
 /// .. and delete the predecessor corresponding to the '1', this will attempt to
 /// recursively fold the 'and' to 0.
-void RemovePredecessorAndSimplify(BasicBlock *BB, BasicBlock *Pred);
+void RemovePredecessorAndSimplify(BasicBlock *BB, BasicBlock *Pred,
+                                  DeferredDominance *DDT = nullptr);
 
 /// BB is a block with one predecessor and its predecessor is known to have one
 /// successor (BB!). Eliminate the edge between them, moving the instructions in
 /// the predecessor into BB. This deletes the predecessor block.
-void MergeBasicBlockIntoOnlyPred(BasicBlock *BB, DominatorTree *DT = nullptr);
+void MergeBasicBlockIntoOnlyPred(BasicBlock *BB, DominatorTree *DT = nullptr,
+                                 DeferredDominance *DDT = nullptr);
 
 /// BB is known to contain an unconditional branch, and contains no instructions
 /// other than PHI nodes, potential debug intrinsics and the branch. If
 /// possible, eliminate BB by rewriting all the predecessors to branch to the
 /// successor block and return true. If we can't transform, return false.
-bool TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB);
+bool TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
+                                             DeferredDominance *DDT = nullptr);
 
 /// Check for and eliminate duplicate PHI nodes in this block. This doesn't try
 /// to be clever about PHI nodes which differ only in the order of the incoming
@@ -382,7 +386,8 @@ unsigned removeAllNonTerminatorAndEHPadInstructions(BasicBlock *BB);
 /// Insert an unreachable instruction before the specified
 /// instruction, making it and the rest of the code in the block dead.
 unsigned changeToUnreachable(Instruction *I, bool UseLLVMTrap,
-                             bool PreserveLCSSA = false);
+                             bool PreserveLCSSA = false,
+                             DeferredDominance *DDT = nullptr);
 
 /// Convert the CallInst to InvokeInst with the specified unwind edge basic
 /// block.  This also splits the basic block where CI is located, because
@@ -397,12 +402,13 @@ BasicBlock *changeToInvokeAndSplitBasicBlock(CallInst *CI,
 ///
 /// \param BB  Block whose terminator will be replaced.  Its terminator must
 ///            have an unwind successor.
-void removeUnwindEdge(BasicBlock *BB);
+void removeUnwindEdge(BasicBlock *BB, DeferredDominance *DDT = nullptr);
 
 /// Remove all blocks that can not be reached from the function's entry.
 ///
 /// Returns true if any basic block was removed.
-bool removeUnreachableBlocks(Function &F, LazyValueInfo *LVI = nullptr);
+bool removeUnreachableBlocks(Function &F, LazyValueInfo *LVI = nullptr,
+                             DeferredDominance *DDT = nullptr);
 
 /// Combine the metadata of two instructions so that K can replace J
 ///

lib/Analysis/LazyValueInfo.cpp

@@ -401,6 +401,7 @@ namespace {
     AssumptionCache *AC;  ///< A pointer to the cache of @llvm.assume calls.
     const DataLayout &DL; ///< A mandatory DataLayout
     DominatorTree *DT;    ///< An optional DT pointer.
+    DominatorTree *DisabledDT; ///< Stores DT if it's disabled.
 
   ValueLatticeElement getBlockValue(Value *Val, BasicBlock *BB);
   bool getEdgeValue(Value *V, BasicBlock *F, BasicBlock *T,
@@ -463,13 +464,30 @@ namespace {
       TheCache.eraseBlock(BB);
     }
 
+    /// Disables use of the DominatorTree within LVI.
+    void disableDT() {
+      if (DT) {
+        assert(!DisabledDT && "Both DT and DisabledDT are not nullptr!");
+        std::swap(DT, DisabledDT);
+      }
+    }
+
+    /// Enables use of the DominatorTree within LVI. Does nothing if the class
+    /// instance was initialized without a DT pointer.
+    void enableDT() {
+      if (DisabledDT) {
+        assert(!DT && "Both DT and DisabledDT are not nullptr!");
+        std::swap(DT, DisabledDT);
+      }
+    }
+
     /// This is the update interface to inform the cache that an edge from
     /// PredBB to OldSucc has been threaded to be from PredBB to NewSucc.
     void threadEdge(BasicBlock *PredBB,BasicBlock *OldSucc,BasicBlock *NewSucc);
 
     LazyValueInfoImpl(AssumptionCache *AC, const DataLayout &DL,
                        DominatorTree *DT = nullptr)
-        : AC(AC), DL(DL), DT(DT) {}
+        : AC(AC), DL(DL), DT(DT), DisabledDT(nullptr) {}
   };
 } // end anonymous namespace
 
@@ -1791,6 +1809,16 @@ void LazyValueInfo::printLVI(Function &F, DominatorTree &DTree, raw_ostream &OS)
   }
 }
 
+void LazyValueInfo::disableDT() {
+  if (PImpl)
+    getImpl(PImpl, AC, DL, DT).disableDT();
+}
+
+void LazyValueInfo::enableDT() {
+  if (PImpl)
+    getImpl(PImpl, AC, DL, DT).enableDT();
+}
+
 // Print the LVI for the function arguments at the start of each basic block.
 void LazyValueInfoAnnotatedWriter::emitBasicBlockStartAnnot(
     const BasicBlock *BB, formatted_raw_ostream &OS) {