[DSE] Preserve basic block ordering using OrderedBasicBlock.

[fhahn]

Summary

By extending OrderedBB to allow removing and replacing cached
instructions, we can preserve OrderedBBs in DSE easily. This eliminates
one source of quadratic compile time in DSE.

Fixes PR38829.

[efriedma]

llvm/lib/Analysis/OrderedBasicBlock.cpp

92

Can you just set "LastInstFound = BB->end()", instead of actually numbering the new first instruction?

[fhahn]

Thanks Eli. In the case LastInstFound == BB->begin(), we have to renumber all
instructions after removing LastInstFound, so there is no need to manually number
the next instruction with 0.

[efriedma]

LGTM

llvm/include/llvm/Analysis/OrderedBasicBlock.h

public:
   OrderedBasicBlock(const BasicBlock *BasicB);
 
   /// Find out whether \p A dominates \p B, meaning whether \p A
   /// comes before \p B in \p BB. This is a simplification that considers
   /// cached instruction positions and ignores other basic blocks, being
   /// only relevant to compare relative instructions positions inside \p BB.
   /// Returns false for A == B.
   bool dominates(const Instruction *A, const Instruction *B);
+
+  /// Remove \p from the ordering, if it is present.
+  void eraseInstruction(const Instruction *I);
+
+  /// Replace \p Old with \p New in the ordering. \p New is assigned the
+  /// numbering of \p Old, so it must be inserted at the same position in the
+  /// IR.
+  void replaceInstruction(const Instruction *Old, const Instruction *New);
 };
 
 } // End llvm namespace
 
 #endif

llvm/lib/Analysis/OrderedBasicBlock.cpp

if (NAI != NumberedInsts.end() && NBI != NumberedInsts.end())
     return NAI->second < NBI->second;
   if (NAI != NumberedInsts.end())
     return true;
   if (NBI != NumberedInsts.end())
     return false;
 
   return comesBefore(A, B);
 }
+
+void OrderedBasicBlock::eraseInstruction(const Instruction *I) {
+  if (LastInstFound != BB->end() && I == &*LastInstFound) {
+    if (LastInstFound == BB->begin()) {
+      LastInstFound++;
+      if (LastInstFound != BB->end())
+        NumberedInsts[&*LastInstFound] = 0;
+    } else
+      LastInstFound--;
+  }
+
+  NumberedInsts.erase(I);
+}
+
+void OrderedBasicBlock::replaceInstruction(const Instruction *Old,
+                                           const Instruction *New) {
+  auto OI = NumberedInsts.find(Old);
+  if (OI == NumberedInsts.end())
+    return;
+
+  NumberedInsts[New] = OI->second;
+  if (LastInstFound != BB->end() && Old == &*LastInstFound)
+    LastInstFound = New->getIterator();
+  NumberedInsts.erase(Old);
+}

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MemoryDependenceAnalysis.h"
 #include "llvm/Analysis/MemoryLocation.h"
+#include "llvm/Analysis/OrderedBasicBlock.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
-#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"

 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
 #include <map>
 #include <utility>
 

 
 /// Delete this instruction.  Before we do, go through and zero out all the
 /// operands of this instruction.  If any of them become dead, delete them and
 /// the computation tree that feeds them.
 /// If ValueSet is non-null, remove any deleted instructions from it as well.
 static void
 deleteDeadInstruction(Instruction *I, BasicBlock::iterator *BBI,
                       MemoryDependenceResults &MD, const TargetLibraryInfo &TLI,
-                      InstOverlapIntervalsTy &IOL,
-                      DenseMap<Instruction*, size_t> *InstrOrdering,
+                      InstOverlapIntervalsTy &IOL, OrderedBasicBlock &OBB,
                       SmallSetVector<Value *, 16> *ValueSet = nullptr) {
   SmallVector<Instruction*, 32> NowDeadInsts;
 
   NowDeadInsts.push_back(I);
   --NumFastOther;
 
   // Keeping the iterator straight is a pain, so we let this routine tell the
   // caller what the next instruction is after we're done mucking about.
for (unsigned op = 0, e = DeadInst->getNumOperands(); op != e; ++op) {
       if (!Op->use_empty()) continue;
 
       if (Instruction *OpI = dyn_cast<Instruction>(Op))
         if (isInstructionTriviallyDead(OpI, &TLI))
           NowDeadInsts.push_back(OpI);
     }
 
     if (ValueSet) ValueSet->remove(DeadInst);
-    InstrOrdering->erase(DeadInst);
     IOL.erase(DeadInst);
+    OBB.eraseInstruction(DeadInst);
 
     if (NewIter == DeadInst->getIterator())
       NewIter = DeadInst->eraseFromParent();
     else
       DeadInst->eraseFromParent();
   } while (!NowDeadInsts.empty());
   *BBI = NewIter;
 }
static void findUnconditionalPreds(SmallVectorImpl<BasicBlock *> &Blocks,
   }
 }
 
 /// Handle frees of entire structures whose dependency is a store
 /// to a field of that structure.
 static bool handleFree(CallInst *F, AliasAnalysis *AA,
                        MemoryDependenceResults *MD, DominatorTree *DT,
                        const TargetLibraryInfo *TLI,
-                       InstOverlapIntervalsTy &IOL,
-                       DenseMap<Instruction*, size_t> *InstrOrdering) {
+                       InstOverlapIntervalsTy &IOL, OrderedBasicBlock &OBB) {
   bool MadeChange = false;
 
   MemoryLocation Loc = MemoryLocation(F->getOperand(0));
   SmallVector<BasicBlock *, 16> Blocks;
   Blocks.push_back(F->getParent());
   const DataLayout &DL = F->getModule()->getDataLayout();
 
   while (!Blocks.empty()) {
while (Dep.isDef() || Dep.isClobber()) {
         break;
 
       LLVM_DEBUG(
           dbgs() << "DSE: Dead Store to soon to be freed memory:\n  DEAD: "
                  << *Dependency << '\n');
 
       // DCE instructions only used to calculate that store.
       BasicBlock::iterator BBI(Dependency);
-      deleteDeadInstruction(Dependency, &BBI, *MD, *TLI, IOL, InstrOrdering);
+      deleteDeadInstruction(Dependency, &BBI, *MD, *TLI, IOL, OBB);
       ++NumFastStores;
       MadeChange = true;
 
       // Inst's old Dependency is now deleted. Compute the next dependency,
       // which may also be dead, as in
       //    s[0] = 0;
       //    s[1] = 0; // This has just been deleted.
       //    free(s);

 
 /// Remove dead stores to stack-allocated locations in the function end block.
 /// Ex:
 /// %A = alloca i32
 /// ...
 /// store i32 1, i32* %A
 /// ret void
 static bool handleEndBlock(BasicBlock &BB, AliasAnalysis *AA,
-                             MemoryDependenceResults *MD,
-                             const TargetLibraryInfo *TLI,
-                             InstOverlapIntervalsTy &IOL,
-                             DenseMap<Instruction*, size_t> *InstrOrdering) {
+                           OrderedBasicBlock &OBB) {
   bool MadeChange = false;
 
   // Keep track of all of the stack objects that are dead at the end of the
   // function.
   SmallSetVector<Value*, 16> DeadStackObjects;
 
   // Find all of the alloca'd pointers in the entry block.
   BasicBlock &Entry = BB.getParent()->front();
if (hasAnalyzableMemoryWrite(&*BBI, *TLI) && isRemovable(&*BBI)) {
                         I != E; ++I) {
                      dbgs() << **I;
                      if (std::next(I) != E)
                        dbgs() << ", ";
                    } dbgs()
                    << '\n');
 
         // DCE instructions only used to calculate that store.
-        deleteDeadInstruction(Dead, &BBI, *MD, *TLI, IOL, InstrOrdering, &DeadStackObjects);
+        deleteDeadInstruction(Dead, &BBI, *MD, *TLI, IOL, OBB,
+                              &DeadStackObjects);
         ++NumFastStores;
         MadeChange = true;
         continue;
       }
     }
 
     // Remove any dead non-memory-mutating instructions.
     if (isInstructionTriviallyDead(&*BBI, TLI)) {
       LLVM_DEBUG(dbgs() << "DSE: Removing trivially dead instruction:\n  DEAD: "
                         << *&*BBI << '\n');
-      deleteDeadInstruction(&*BBI, &BBI, *MD, *TLI, IOL, InstrOrdering, &DeadStackObjects);
+      deleteDeadInstruction(&*BBI, &BBI, *MD, *TLI, IOL, OBB,
+                            &DeadStackObjects);
       ++NumFastOther;
       MadeChange = true;
       continue;
     }
 
     if (isa<AllocaInst>(BBI)) {
       // Remove allocas from the list of dead stack objects; there can't be
       // any references before the definition.
static bool removePartiallyOverlappedStores(AliasAnalysis *AA,
   return Changed;
 }
 
 static bool eliminateNoopStore(Instruction *Inst, BasicBlock::iterator &BBI,
                                AliasAnalysis *AA, MemoryDependenceResults *MD,
                                const DataLayout &DL,
                                const TargetLibraryInfo *TLI,
                                InstOverlapIntervalsTy &IOL,
-                               DenseMap<Instruction*, size_t> *InstrOrdering) {
+                               OrderedBasicBlock &OBB) {
   // Must be a store instruction.
   StoreInst *SI = dyn_cast<StoreInst>(Inst);
   if (!SI)
     return false;
 
   // If we're storing the same value back to a pointer that we just loaded from,
   // then the store can be removed.
   if (LoadInst *DepLoad = dyn_cast<LoadInst>(SI->getValueOperand())) {
     if (SI->getPointerOperand() == DepLoad->getPointerOperand() &&
         isRemovable(SI) && memoryIsNotModifiedBetween(DepLoad, SI, AA)) {
 
       LLVM_DEBUG(
           dbgs() << "DSE: Remove Store Of Load from same pointer:\n  LOAD: "
                  << *DepLoad << "\n  STORE: " << *SI << '\n');
 
-      deleteDeadInstruction(SI, &BBI, *MD, *TLI, IOL, InstrOrdering);
+      deleteDeadInstruction(SI, &BBI, *MD, *TLI, IOL, OBB);
       ++NumRedundantStores;
       return true;
     }
   }
 
   // Remove null stores into the calloc'ed objects
   Constant *StoredConstant = dyn_cast<Constant>(SI->getValueOperand());
   if (StoredConstant && StoredConstant->isNullValue() && isRemovable(SI)) {
     Instruction *UnderlyingPointer =
         dyn_cast<Instruction>(GetUnderlyingObject(SI->getPointerOperand(), DL));
 
     if (UnderlyingPointer && isCallocLikeFn(UnderlyingPointer, TLI) &&
         memoryIsNotModifiedBetween(UnderlyingPointer, SI, AA)) {
       LLVM_DEBUG(
           dbgs() << "DSE: Remove null store to the calloc'ed object:\n  DEAD: "
                  << *Inst << "\n  OBJECT: " << *UnderlyingPointer << '\n');
 
-      deleteDeadInstruction(SI, &BBI, *MD, *TLI, IOL, InstrOrdering);
+      deleteDeadInstruction(SI, &BBI, *MD, *TLI, IOL, OBB);
       ++NumRedundantStores;
       return true;
     }
   }
   return false;
 }
 
 static bool eliminateDeadStores(BasicBlock &BB, AliasAnalysis *AA,
                                 MemoryDependenceResults *MD, DominatorTree *DT,
                                 const TargetLibraryInfo *TLI) {
   const DataLayout &DL = BB.getModule()->getDataLayout();
   bool MadeChange = false;
 
-  // FIXME: Maybe change this to use some abstraction like OrderedBasicBlock?
-  // The current OrderedBasicBlock can't deal with mutation at the moment.
-  size_t LastThrowingInstIndex = 0;
-  DenseMap<Instruction*, size_t> InstrOrdering;
-  size_t InstrIndex = 1;
+  OrderedBasicBlock OBB(&BB);
+  Instruction *LastThrowing = nullptr;
 
   // A map of interval maps representing partially-overwritten value parts.
   InstOverlapIntervalsTy IOL;
 
   // Do a top-down walk on the BB.
   for (BasicBlock::iterator BBI = BB.begin(), BBE = BB.end(); BBI != BBE; ) {
     // Handle 'free' calls specially.
     if (CallInst *F = isFreeCall(&*BBI, TLI)) {
-      MadeChange |= handleFree(F, AA, MD, DT, TLI, IOL, &InstrOrdering);
+      MadeChange |= handleFree(F, AA, MD, DT, TLI, IOL, OBB);
       // Increment BBI after handleFree has potentially deleted instructions.
       // This ensures we maintain a valid iterator.
       ++BBI;
       continue;
     }
 
     Instruction *Inst = &*BBI++;
 
-    size_t CurInstNumber = InstrIndex++;
-    InstrOrdering.insert(std::make_pair(Inst, CurInstNumber));
     if (Inst->mayThrow()) {
-      LastThrowingInstIndex = CurInstNumber;
+      LastThrowing = Inst;
       continue;
     }
 
     // Check to see if Inst writes to memory.  If not, continue.
     if (!hasAnalyzableMemoryWrite(Inst, *TLI))
       continue;
 
     // eliminateNoopStore will update in iterator, if necessary.
-    if (eliminateNoopStore(Inst, BBI, AA, MD, DL, TLI, IOL, &InstrOrdering)) {
+    if (eliminateNoopStore(Inst, BBI, AA, MD, DL, TLI, IOL, OBB)) {
       MadeChange = true;
       continue;
     }
 
     // If we find something that writes memory, get its memory dependence.
-    MemDepResult InstDep = MD->getDependency(Inst);
+    MemDepResult InstDep = MD->getDependency(Inst, &OBB);
 
     // Ignore any store where we can't find a local dependence.
     // FIXME: cross-block DSE would be fun. :)
     if (!InstDep.isDef() && !InstDep.isClobber())
       continue;
 
     // Figure out what location is being stored to.
     MemoryLocation Loc = getLocForWrite(Inst);
while (InstDep.isDef() || InstDep.isClobber()) {
       // Make sure we don't look past a call which might throw. This is an
       // issue because MemoryDependenceAnalysis works in the wrong direction:
       // it finds instructions which dominate the current instruction, rather than
       // instructions which are post-dominated by the current instruction.
       //
       // If the underlying object is a non-escaping memory allocation, any store
       // to it is dead along the unwind edge. Otherwise, we need to preserve
       // the store.
-      size_t DepIndex = InstrOrdering.lookup(DepWrite);
-      assert(DepIndex && "Unexpected instruction");
-      if (DepIndex <= LastThrowingInstIndex) {
+      if (LastThrowing && OBB.dominates(DepWrite, LastThrowing)) {
         const Value* Underlying = GetUnderlyingObject(DepLoc.Ptr, DL);
         bool IsStoreDeadOnUnwind = isa<AllocaInst>(Underlying);
         if (!IsStoreDeadOnUnwind) {
             // We're looking for a call to an allocation function
             // where the allocation doesn't escape before the last
             // throwing instruction; PointerMayBeCaptured
             // reasonably fast approximation.
             IsStoreDeadOnUnwind = isAllocLikeFn(Underlying, TLI) &&
while (InstDep.isDef() || InstDep.isClobber()) {
         OverwriteResult OR = isOverwrite(Loc, DepLoc, DL, *TLI, DepWriteOffset,
                                          InstWriteOffset, DepWrite, IOL, *AA,
                                          BB.getParent());
         if (OR == OW_Complete) {
           LLVM_DEBUG(dbgs() << "DSE: Remove Dead Store:\n  DEAD: " << *DepWrite
                             << "\n  KILLER: " << *Inst << '\n');
 
           // Delete the store and now-dead instructions that feed it.
-          deleteDeadInstruction(DepWrite, &BBI, *MD, *TLI, IOL, &InstrOrdering);
+          deleteDeadInstruction(DepWrite, &BBI, *MD, *TLI, IOL, OBB);
           ++NumFastStores;
           MadeChange = true;
 
           // We erased DepWrite; start over.
-          InstDep = MD->getDependency(Inst);
+          InstDep = MD->getDependency(Inst, &OBB);
           continue;
         } else if ((OR == OW_End && isShortenableAtTheEnd(DepWrite)) ||
                    ((OR == OW_Begin &&
                      isShortenableAtTheBeginning(DepWrite)))) {
           assert(!EnablePartialOverwriteTracking && "Do not expect to perform "
                                                     "when partial-overwrite "
                                                     "tracking is enabled");
           // The overwrite result is known, so these must be known, too.
while (InstDep.isDef() || InstDep.isClobber()) {
             unsigned MDToKeep[] = {LLVMContext::MD_dbg, LLVMContext::MD_tbaa,
                                    LLVMContext::MD_alias_scope,
                                    LLVMContext::MD_noalias,
                                    LLVMContext::MD_nontemporal};
             SI->copyMetadata(*DepWrite, MDToKeep);
             ++NumModifiedStores;
 
             // Remove earlier, wider, store
-            size_t Idx = InstrOrdering.lookup(DepWrite);
-            InstrOrdering.erase(DepWrite);
-            InstrOrdering.insert(std::make_pair(SI, Idx));
+            OBB.replaceInstruction(DepWrite, SI);
 
             // Delete the old stores and now-dead instructions that feed them.
-            deleteDeadInstruction(Inst, &BBI, *MD, *TLI, IOL, &InstrOrdering);
-            deleteDeadInstruction(DepWrite, &BBI, *MD, *TLI, IOL,
-                                  &InstrOrdering);
+            deleteDeadInstruction(Inst, &BBI, *MD, *TLI, IOL, OBB);
+            deleteDeadInstruction(DepWrite, &BBI, *MD, *TLI, IOL, OBB);
             MadeChange = true;
 
             // We erased DepWrite and Inst (Loc); start over.
             break;
           }
         }
       }
 
static bool eliminateDeadStores(BasicBlock &BB, AliasAnalysis *AA,
   }
 
   if (EnablePartialOverwriteTracking)
     MadeChange |= removePartiallyOverlappedStores(AA, DL, IOL);
 
   // If this block ends in a return, unwind, or unreachable, all allocas are
   // dead at its end, which means stores to them are also dead.
   if (BB.getTerminator()->getNumSuccessors() == 0)
-    MadeChange |= handleEndBlock(BB, AA, MD, TLI, IOL, &InstrOrdering);
+    MadeChange |= handleEndBlock(BB, AA, MD, TLI, IOL, OBB);
 
   return MadeChange;
 }
 
 static bool eliminateDeadStores(Function &F, AliasAnalysis *AA,
                                 MemoryDependenceResults *MD, DominatorTree *DT,
                                 const TargetLibraryInfo *TLI) {
   bool MadeChange = false;