[DSE] Add variant of conditions implying value established by all predecessors

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

@@ -1109,9 +1109,11 @@ struct DSEState {
   /// try folding it into a call to calloc.
   bool tryFoldIntoCalloc(MemoryDef *Def, const Value *DefUO);
 
-  // Check if there is a dominating condition, that implies that the value
-  // being stored in a ptr is already present in the ptr.
-  bool dominatingConditionImpliesValue(MemoryDef *Def);
+  // If there is a dominating condition that implies the value being stored in a
+  // pointer, and such a condition appears either in its idom or in a node that
+  // strictly dominates the store, then the store may be redundant as long as
+  // no write occurs in between.
+  bool dominatingConditionOrPredsImplyValue(StoreInst *SI, MemoryDef *Def);
 
   /// \returns true if \p Def is a no-op store, either because it
   /// directly stores back a loaded value or stores zero to a calloced object.
@@ -2245,53 +2247,104 @@ bool DSEState::tryFoldIntoCalloc(MemoryDef *Def, const Value *DefUO) {
   return true;
 }
 
-bool DSEState::dominatingConditionImpliesValue(MemoryDef *Def) {
-  auto *StoreI = cast<StoreInst>(Def->getMemoryInst());
-  BasicBlock *StoreBB = StoreI->getParent();
-  Value *StorePtr = StoreI->getPointerOperand();
-  Value *StoreVal = StoreI->getValueOperand();
+bool DSEState::dominatingConditionOrPredsImplyValue(StoreInst *SI,
+                                                    MemoryDef *Def) {
+  BasicBlock *StoreBB = SI->getParent();
+  Value *StorePtr = SI->getPointerOperand();
+  Value *StoreVal = SI->getValueOperand();
 
-  DomTreeNode *IDom = DT.getNode(StoreBB)->getIDom();
-  if (!IDom)
-    return false;
+  auto MatchCondition = [&](BasicBlock *BB)
+      -> std::optional<std::pair<BasicBlock *, Instruction *>> {
+    auto *BI = dyn_cast<BranchInst>(BB->getTerminator());
+    if (!BI || !BI->isConditional())
+      return std::nullopt;
 
-  auto *BI = dyn_cast<BranchInst>(IDom->getBlock()->getTerminator());
-  if (!BI || !BI->isConditional())
-    return false;
+    // In case both blocks are the same, it is not possible to determine
+    // if optimization is possible. (We would not want to optimize a store
+    // in the FalseBB if condition is true and vice versa.)
+    if (BI->getSuccessor(0) == BI->getSuccessor(1))
+      return std::nullopt;
 
-  // In case both blocks are the same, it is not possible to determine
-  // if optimization is possible. (We would not want to optimize a store
-  // in the FalseBB if condition is true and vice versa.)
-  if (BI->getSuccessor(0) == BI->getSuccessor(1))
-    return false;
+    Instruction *ICmpL;
+    CmpPredicate Pred;
+    if (!match(BI->getCondition(),
+               m_c_ICmp(Pred,
+                        m_CombineAnd(m_Load(m_Specific(StorePtr)),
+                                     m_Instruction(ICmpL)),
+                        m_Specific(StoreVal))) ||
+        !ICmpInst::isEquality(Pred))
+      return std::nullopt;
 
-  Instruction *ICmpL;
-  CmpPredicate Pred;
-  if (!match(BI->getCondition(),
-             m_c_ICmp(Pred,
-                      m_CombineAnd(m_Load(m_Specific(StorePtr)),
-                                   m_Instruction(ICmpL)),
-                      m_Specific(StoreVal))) ||
-      !ICmpInst::isEquality(Pred))
-    return false;
+    unsigned ImpliedIdx = (Pred == ICmpInst::ICMP_EQ) ? 0 : 1;
+    return {{BI->getSuccessor(ImpliedIdx), ICmpL}};
+  };
+
+  // Walk up the dominator tree looking for dominating conditions, up to limit.
+  static constexpr unsigned Limit = 4;
+  BasicBlock *Node = StoreBB;
+  for (unsigned Depth = 0; Depth < Limit; ++Depth) {
+    DomTreeNode *IDomNode = DT.getNode(Node)->getIDom();
+    if (!IDomNode)
+      break;
+    Node = IDomNode->getBlock();
+
+    auto Match = MatchCondition(Node);
+    if (!Match)
+      continue;
+
+    const auto &[ImpliedSucc, LI] = *Match;
+    if (!DT.dominates(BasicBlockEdge(Node, ImpliedSucc), StoreBB))
+      continue;
+
+    // Found a dominating condition. Make sure there does not exist any
+    // clobbering access between the load and the potential redundant store.
+    MemoryAccess *LoadAcc = MSSA.getMemoryAccess(LI);
+    MemoryAccess *ClobAcc =
+        MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(Def, BatchAA);
 
-  // In case the else blocks also branches to the if block or the other way
-  // around it is not possible to determine if the optimization is possible.
-  if (Pred == ICmpInst::ICMP_EQ &&
-      !DT.dominates(BasicBlockEdge(BI->getParent(), BI->getSuccessor(0)),
-                    StoreBB))
+    if (MSSA.dominates(ClobAcc, LoadAcc))
+      return true;
+    break;
+  }
+
+  // If we haven't found any dominating condition, try evaluating whether all
+  // the predecessors establish a condition implying the value being stored.
+  unsigned NumPreds = pred_size(StoreBB);
+  if (NumPreds < 2 || NumPreds > 4)
     return false;
 
-  if (Pred == ICmpInst::ICMP_NE &&
-      !DT.dominates(BasicBlockEdge(BI->getParent(), BI->getSuccessor(1)),
-                    StoreBB))
+  SmallDenseMap<BasicBlock *, Instruction *, 4> PredToLoad;
+  for (BasicBlock *PredBB : predecessors(StoreBB)) {
+    auto Match = MatchCondition(PredBB);
+    if (!Match)
+      return false;
+
+    const auto &[ImpliedSucc, LI] = *Match;
+    if (ImpliedSucc != StoreBB)
+      return false;
+
+    PredToLoad[PredBB] = LI;
+  }
+  assert(PredToLoad.size() == NumPreds);
+
+  // If we are not merging the memory reads from the predecessors, the memory
+  // location may be clobbered.
+  auto *MPhi = dyn_cast<MemoryPhi>(Def->getDefiningAccess());
+  if (!MPhi || MPhi->getBlock() != StoreBB)
     return false;
 
-  MemoryAccess *LoadAcc = MSSA.getMemoryAccess(ICmpL);
-  MemoryAccess *ClobAcc =
-      MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(Def, BatchAA);
+  MemoryLocation StoreLoc = MemoryLocation::get(SI);
+  for (unsigned I = 0; I < MPhi->getNumIncomingValues(); ++I) {
+    MemoryAccess *LoadAccess =
+        MSSA.getMemoryAccess(PredToLoad[MPhi->getIncomingBlock(I)]);
+    MemoryAccess *ClobberingAccess =
+        MSSA.getWalker()->getClobberingMemoryAccess(MPhi->getIncomingValue(I),
+                                                    StoreLoc, BatchAA);
+    if (!MSSA.dominates(ClobberingAccess, LoadAccess))
+      return false;
+  }
 
-  return MSSA.dominates(ClobAcc, LoadAcc);
+  return true;
 }
 
 bool DSEState::storeIsNoop(MemoryDef *Def, const Value *DefUO) {
@@ -2322,7 +2375,7 @@ bool DSEState::storeIsNoop(MemoryDef *Def, const Value *DefUO) {
   if (!Store)
     return false;
 
-  if (dominatingConditionImpliesValue(Def))
+  if (dominatingConditionOrPredsImplyValue(Store, Def))
     return true;
 
   if (auto *LoadI = dyn_cast<LoadInst>(Store->getOperand(0))) {