Fix a bug in dead store elimination. Make sure we consider SILargument when we

invalidate base.

Also small change on how BlockState is allocated.

Author: trentxintong

lib/SILOptimizer/Transforms/DeadStoreElimination.cpp

@@ -98,8 +98,8 @@ enum class DSEKind : unsigned {
 
 /// Return the deallocate stack instructions corresponding to the given
 /// AllocStackInst.
-static llvm::SmallVector<SILInstruction *, 1> findDeallocStackInst(
-                                                   AllocStackInst *ASI) {
+static llvm::SmallVector<SILInstruction *, 1>
+findDeallocStackInst(AllocStackInst *ASI) {
   llvm::SmallVector<SILInstruction *, 1> DSIs;
   for (auto UI = ASI->use_begin(), E = ASI->use_end(); UI != E; ++UI) {
     if (DeallocStackInst *D = dyn_cast<DeallocStackInst>(UI->getUser())) {
@@ -249,13 +249,16 @@ class BlockState {
   llvm::DenseMap<SILValue, SILValue> LiveStores;
 
   /// Constructors.
-  BlockState(SILBasicBlock *B) : BB(B) {}
+  BlockState(SILBasicBlock *B, unsigned LocationNum, bool Optimistic) 
+      : BB(B), LocationNum(LocationNum) {
+    init(LocationNum, Optimistic);
+  }
 
   /// Return the current basic block.
   SILBasicBlock *getBB() const { return BB; }
 
   /// Initialize the bitvectors for the current basic block.
-  void init(DSEContext &Ctx, bool Optimistic);
+  void init(unsigned LocationNum, bool Optimistic);
 
   /// Check whether the BBWriteSetIn has changed. If it does, we need to rerun
   /// the data flow on this block's predecessors to reach fixed point.
@@ -298,11 +301,13 @@ bool BlockState::isTrackingLocation(llvm::SmallBitVector &BV, unsigned i) {
 namespace {
 
 class DSEContext {
-  enum class ProcessKind {
-    ProcessOptimistic = 0,
-    ProcessPessimistic = 1,
-    ProcessNone = 2,
-  }; 
+/// How to process the current function.
+enum class ProcessKind {
+  ProcessOptimistic = 0,
+  ProcessPessimistic = 1,
+  ProcessNone = 2,
+}; 
+
 private:
   /// The module we are currently processing.
   SILModule *Mod;
@@ -322,17 +327,14 @@ class DSEContext {
   /// Type Expansion Analysis.
   TypeExpansionAnalysis *TE;
 
+  /// The allocator we are using.
+  llvm::BumpPtrAllocator &BPA;
+
   /// The epilogue release matcher we are using.
   ConsumedArgToEpilogueReleaseMatcher& ERM;
 
-  /// Allocator.
-  llvm::BumpPtrAllocator BPA;
-
   /// Map every basic block to its location state.
-  llvm::DenseMap<SILBasicBlock *, BlockState *> BBToLocState;
-
-  /// Keeps the actual BlockStates.
-  std::vector<BlockState> BlockStates;
+  llvm::SmallDenseMap<SILBasicBlock *, BlockState *> BBToLocState;
 
   /// Keeps all the locations for the current function. The BitVector in each
   /// BlockState is then laid on top of it to keep track of which LSLocation
@@ -371,14 +373,13 @@ class DSEContext {
 
   /// There is a read to a location, expand the location into individual fields
   /// before processing them.
-  void processRead(SILInstruction *Inst, BlockState *S, SILValue M, DSEKind K);
+  void processRead(SILInstruction *Inst, SILValue M, DSEKind K);
   void processReadForGenKillSet(BlockState *S, unsigned bit);
   void processReadForDSE(BlockState *S, unsigned Bit);
 
   /// There is a write to a location, expand the location into individual fields
   /// before processing them.
-  void processWrite(SILInstruction *Inst, BlockState *S, SILValue V, SILValue M,
-                    DSEKind K);
+  void processWrite(SILInstruction *Inst, SILValue V, SILValue M, DSEKind K);
   void processWriteForGenKillSet(BlockState *S, unsigned bit);
   bool processWriteForDSE(BlockState *S, unsigned bit);
 
@@ -395,7 +396,8 @@ class DSEContext {
   void processDebugValueAddrInstForGenKillSet(SILInstruction *I);
   void processDebugValueAddrInstForDSE(SILInstruction *I);
 
-  /// Process instructions. Extract locations from unknown memory inst.
+  /// Process unknown read instructions. Extract locations from unknown memory
+  /// inst.
   void processUnknownReadInst(SILInstruction *Inst, DSEKind Kind);
   void processUnknownReadInstForGenKillSet(SILInstruction *Inst);
   void processUnknownReadInstForDSE(SILInstruction *Inst);
@@ -421,9 +423,9 @@ class DSEContext {
   /// If not, on the second iteration, the intersection of the successors of
   /// the loop basic block will have store to x.a and therefore x.a = 13 can now
   /// be considered dead.
-  void invalidateLSLocationBase(SILInstruction *Inst, DSEKind Kind);
-  void invalidateLSLocationBaseForGenKillSet(SILInstruction *Inst);
-  void invalidateLSLocationBaseForDSE(SILInstruction *Inst);
+  void invalidateBase(SILValue B, BlockState *S, DSEKind Kind);
+  void invalidateBaseForGenKillSet(SILValue B, BlockState *S);
+  void invalidateBaseForDSE(SILValue B, BlockState *S);
 
   /// Get the bit representing the location in the LocationVault.
   unsigned getLocationBit(const LSLocation &L);
@@ -432,8 +434,9 @@ class DSEContext {
   /// Constructor.
   DSEContext(SILFunction *F, SILModule *M, SILPassManager *PM,
              AliasAnalysis *AA, EscapeAnalysis *EA, TypeExpansionAnalysis *TE,
+             llvm::BumpPtrAllocator &BPA,
              ConsumedArgToEpilogueReleaseMatcher &ERM)
-      : Mod(M), F(F), PM(PM), AA(AA), EA(EA), TE(TE), ERM(ERM) {}
+      : Mod(M), F(F), PM(PM), AA(AA), EA(EA), TE(TE), BPA(BPA), ERM(ERM) {}
 
   /// Entry point for dead store elimination.
   bool run();
@@ -444,12 +447,10 @@ class DSEContext {
   /// Returns the escape analysis we use.
   EscapeAnalysis *getEA() { return EA; }
 
-  /// Returns the function currently being processing.
-  SILFunction *getFn() { return F; }
-
   /// Returns the location vault of the current function.
   std::vector<LSLocation> &getLocationVault() { return LocationVault; }
 
+  /// Returns the epilogue release matcher we are using.
   ConsumedArgToEpilogueReleaseMatcher &getERM() const { return ERM; };
 
   /// Use a set of ad hoc rules to tell whether we should run a pessimistic
@@ -476,9 +477,7 @@ class DSEContext {
 
 } // end anonymous namespace
 
-void BlockState::init(DSEContext &Ctx, bool Optimistic) {
-  std::vector<LSLocation> &LV = Ctx.getLocationVault();
-  LocationNum = LV.size();
+void BlockState::init(unsigned LocationNum, bool Optimistic) {
   // For function that requires just 1 iteration of the data flow to converge
   // we set the initial state of BBWriteSetIn to 0.
   //
@@ -506,7 +505,6 @@ void BlockState::init(DSEContext &Ctx, bool Optimistic) {
   // MaxStoreSet is optimistically set to true initially.
   BBMaxStoreSet.resize(LocationNum, true);
 
-
   // DeallocateLocation initially empty.
   BBDeallocateLocation.resize(LocationNum, false);
 }
@@ -604,6 +602,12 @@ void DSEContext::processBasicBlockForGenKillSet(SILBasicBlock *BB) {
     // Compute the genset and killset for this instruction.
     processInstruction(&(*I), DSEKind::BuildGenKillSet);
   }
+
+  // Handle SILArgument for base invalidation.
+  ArrayRef<SILArgument *> Args = BB->getBBArgs();
+  for(auto &X : Args) {
+    invalidateBase(X, BBState, DSEKind::BuildGenKillSet);
+  }
 }
 
 bool DSEContext::processBasicBlockWithGenKillSet(SILBasicBlock *BB) {
@@ -620,8 +624,7 @@ bool DSEContext::processBasicBlockWithGenKillSet(SILBasicBlock *BB) {
   return S->updateBBWriteSetIn(S->BBWriteSetMid);
 }
 
-void DSEContext::processBasicBlockForDSE(SILBasicBlock *BB,
-                                         bool Optimistic) {
+void DSEContext::processBasicBlockForDSE(SILBasicBlock *BB, bool Optimistic) {
   // If we know this is not a one iteration function which means its
   // its BBWriteSetIn and BBWriteSetOut have been computed and converged, 
   // and this basic block does not even have StoreInsts, there is no point
@@ -643,6 +646,12 @@ void DSEContext::processBasicBlockForDSE(SILBasicBlock *BB,
     processInstruction(&(*I), DSEKind::PerformDSE);
   }
 
+  // Handle SILArgument for base invalidation.
+  ArrayRef<SILArgument *> Args = BB->getBBArgs();
+  for(auto &X : Args) {
+    invalidateBase(X, S, DSEKind::BuildGenKillSet);
+  }
+
   S->BBWriteSetIn = S->BBWriteSetMid;
 }
 
@@ -697,40 +706,38 @@ void DSEContext::mergeSuccessorLiveIns(SILBasicBlock *BB) {
   C->BBWriteSetOut |= C->BBDeallocateLocation;
 }
 
-void DSEContext::invalidateLSLocationBaseForGenKillSet(SILInstruction *I) {
-  BlockState *S = getBlockState(I);
+void DSEContext::invalidateBaseForGenKillSet(SILValue B, BlockState *S) {
   for (unsigned i = 0; i < S->LocationNum; ++i) {
-    if (LocationVault[i].getBase() != I)
+    if (LocationVault[i].getBase() != B)
       continue;
     S->startTrackingLocation(S->BBKillSet, i);
     S->stopTrackingLocation(S->BBGenSet, i);
   }
 }
 
-void DSEContext::invalidateLSLocationBaseForDSE(SILInstruction *I) {
-  BlockState *S = getBlockState(I);
+void DSEContext::invalidateBaseForDSE(SILValue B, BlockState *S) {
   for (unsigned i = 0; i < S->LocationNum; ++i) {
     if (!S->BBWriteSetMid.test(i))
       continue;
-    if (LocationVault[i].getBase() != I)
+    if (LocationVault[i].getBase() != B)
       continue;
     S->stopTrackingLocation(S->BBWriteSetMid, i);
   }
 }
 
-void DSEContext::invalidateLSLocationBase(SILInstruction *I, DSEKind Kind) {
+void DSEContext::invalidateBase(SILValue B, BlockState *S, DSEKind Kind) {
   // If this instruction defines the base of a location, then we need to
   // invalidate any locations with the same base.
   //
   // Are we building genset and killset.
   if (isBuildingGenKillSet(Kind)) {
-    invalidateLSLocationBaseForGenKillSet(I);
+    invalidateBaseForGenKillSet(B, S);
     return;
   }
 
   // Are we performing dead store elimination.
   if (isPerformingDSE(Kind)) {
-    invalidateLSLocationBaseForDSE(I);
+    invalidateBaseForDSE(B, S);
     return;
   }
 
@@ -771,8 +778,8 @@ void DSEContext::processReadForGenKillSet(BlockState *S, unsigned bit) {
   }
 }
 
-void DSEContext::processRead(SILInstruction *I, BlockState *S, SILValue Mem,
-                             DSEKind Kind) {
+void DSEContext::processRead(SILInstruction *I, SILValue Mem, DSEKind Kind) {
+  auto *S = getBlockState(I);
   // Construct a LSLocation to represent the memory read by this instruction.
   // NOTE: The base will point to the actual object this inst is accessing,
   // not this particular field.
@@ -856,8 +863,9 @@ void DSEContext::processWriteForMaxStoreSet(BlockState *S, unsigned bit) {
   S->startTrackingLocation(S->BBMaxStoreSet, bit);
 }
 
-void DSEContext::processWrite(SILInstruction *I, BlockState *S, SILValue Val,
-                              SILValue Mem, DSEKind Kind) {
+void DSEContext::processWrite(SILInstruction *I, SILValue Val, SILValue Mem,
+                              DSEKind Kind) {
+  auto *S = getBlockState(I);
   // Construct a LSLocation to represent the memory read by this instruction.
   // NOTE: The base will point to the actual object this inst is accessing,
   // not this particular field.
@@ -979,12 +987,12 @@ void DSEContext::processWrite(SILInstruction *I, BlockState *S, SILValue Val,
 }
 
 void DSEContext::processLoadInst(SILInstruction *I, DSEKind Kind) {
-  processRead(I, getBlockState(I), cast<LoadInst>(I)->getOperand(), Kind);
+  processRead(I, cast<LoadInst>(I)->getOperand(), Kind);
 }
 
 void DSEContext::processStoreInst(SILInstruction *I, DSEKind Kind) {
   auto *SI = cast<StoreInst>(I);
-  processWrite(I, getBlockState(I), SI->getSrc(), SI->getDest(), Kind);
+  processWrite(I, SI->getSrc(), SI->getDest(), Kind);
 }
 
 void DSEContext::processDebugValueAddrInstForGenKillSet(SILInstruction *I) {
@@ -1091,7 +1099,7 @@ void DSEContext::processInstruction(SILInstruction *I, DSEKind Kind) {
   }  
 
   // Check whether this instruction will invalidate any other locations.
-  invalidateLSLocationBase(I, Kind);
+  invalidateBase(I, getBlockState(I), Kind);
 }
 
 void DSEContext::runIterativeDSE() {
@@ -1132,46 +1140,30 @@ void DSEContext::runIterativeDSE() {
 }
 
 bool DSEContext::run() {
-  // Is this a one iteration function.
-  auto *PO = PM->getAnalysis<PostOrderAnalysis>()->get(F);
-
   std::pair<int, int> LSCount = std::make_pair(0, 0);
   // Walk over the function and find all the locations accessed by
   // this function.
-  LSLocation::enumerateLSLocations(*F, LocationVault,
-                                   LocToBitIndex,
+  LSLocation::enumerateLSLocations(*F, LocationVault, LocToBitIndex,
                                    BaseToLocIndex, TE, LSCount);
 
   // Check how to optimize this function.
   ProcessKind Kind = getProcessFunctionKind(LSCount.second);
 
   // We do not optimize this function at all.
   if (Kind == ProcessKind::ProcessNone)
-      return false;
+    return false;
 
   // Do we run a pessimistic data flow ?
   bool Optimistic = Kind == ProcessKind::ProcessOptimistic ? true : false;
 
   // For all basic blocks in the function, initialize a BB state.
   //
-  // DenseMap has a minimum size of 64, while many functions do not have more
-  // than 64 basic blocks. Therefore, allocate the BlockState in a vector and
-  // use pointer in BBToLocState to access them.
-  for (auto &B : *F) {
-    BlockStates.push_back(BlockState(&B));
-    // Since we know all the locations accessed in this function, we can resize
-    // the bit vector to the appropriate size.
-    BlockStates.back().init(*this, Optimistic);
-  }
-
   // Initialize the BBToLocState mapping.
-  for (auto &S : BlockStates) {
-    BBToLocState[S.getBB()] = &S;
-  }
-
-  // Initialize the store bit state at the end of each basic block.
-  for (auto &X : BlockStates) {
-    X.initStoreSetAtEndOfBlock(*this);
+  unsigned LocationNum = this->getLocationVault().size();
+  for (auto &B : *F) {
+    auto *State = new (BPA) BlockState(&B, LocationNum, Optimistic);
+    BBToLocState[&B] = State;
+    State->initStoreSetAtEndOfBlock(*this);
   }
 
   // We perform dead store elimination in the following phases.
@@ -1198,6 +1190,8 @@ bool DSEContext::run() {
 
   // The data flow has stabilized, run one last iteration over all the basic
   // blocks and try to remove dead stores.
+  // Is this a one iteration function.
+  auto *PO = PM->getAnalysis<PostOrderAnalysis>()->get(F);
   for (SILBasicBlock *B : PO->getPostOrder()) {
     processBasicBlockForDSE(B, Optimistic);
   }
@@ -1245,9 +1239,13 @@ class DeadStoreElimination : public SILFunctionTransform {
     auto *TE = PM->getAnalysis<TypeExpansionAnalysis>();
     auto *RCFI = PM->getAnalysis<RCIdentityAnalysis>()->get(F);
 
+    // The allocator we are using.
+    llvm::BumpPtrAllocator BPA;
+
+    // The epilogue release matcher we are using.
     ConsumedArgToEpilogueReleaseMatcher ERM(RCFI, F);
 
-    DSEContext DSE(F, &F->getModule(), PM, AA, EA, TE, ERM);
+    DSEContext DSE(F, &F->getModule(), PM, AA, EA, TE, BPA, ERM);
     if (DSE.run()) {
       invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
     }

test/SILOptimizer/dead_store_elim.sil

@@ -1340,3 +1340,30 @@ bb0(%0 : $*Int):
 
 sil @closure : $@convention(thin) (Bool, @in Int) -> Int
 
+// Make sure we invalidate the store in bb1 when we process the SILargument. If not
+// we could remove the store in bb1 which is incorrect.
+//
+// CHECK-LABEL: invalidate_base_for_silargument
+// CHECK: bb1([[A0:%.+]] : $foo):
+// CHECK: {{ store}}
+// CHECK: bb2:
+// CHECK: {{ store}}
+sil hidden @invalidate_base_for_silargument : $@convention(thin) () -> () {
+bb0:
+  %1 = alloc_ref $foo
+  %4 = integer_literal $Builtin.Int64, 12
+  %5 = struct $Int (%4 : $Builtin.Int64)
+  br bb1(%1 : $foo)
+
+bb1(%2 : $foo):
+  %6 = ref_element_addr %2 : $foo, #foo.a
+  store %5 to %6 : $*Int
+  %9 = alloc_ref $foo
+  cond_br undef, bb1(%9: $foo), bb2
+
+bb2:
+  %7 = ref_element_addr %2 : $foo, #foo.a
+  store %5 to %7 : $*Int
+  %14 = tuple ()
+  return %14 : $()
+}