[llvm] [SimplifyCFG] Add foldCondStoreToSelect optimization (PR #207654)

[Karthika Devi C via llvm-commits]

================
@@ -4292,6 +4292,287 @@ static StoreInst *findUniqueStoreInBlocks(BasicBlock *BB1, BasicBlock *BB2) {
   return S;
 }
 
+/// Return true if two pointer-producing values refer to the same address.
+/// Strips pointer casts and matches identical GEPs with the same base and
+/// indices.
+static bool isSameGEPAddress(Value *A, Value *B) {
+  if (A == B)
+    return true;
+  A = A->stripPointerCasts();
+  B = B->stripPointerCasts();
+  if (A == B)
+    return true;
+  auto *GA = dyn_cast<GetElementPtrInst>(A);
+  auto *GB = dyn_cast<GetElementPtrInst>(B);
+  if (!GA || !GB)
+    return false;
+  if (GA->getPointerOperand()->stripPointerCasts() !=
+      GB->getPointerOperand()->stripPointerCasts())
+    return false;
+  if (GA->getSourceElementType() != GB->getSourceElementType())
+    return false;
+  if (GA->getNumIndices() != GB->getNumIndices())
+    return false;
+  return std::equal(GA->idx_begin(), GA->idx_end(), GB->idx_begin());
+}
+
+/// Find the unique simple store in \p BB. Returns nullptr if there is no
+/// store, if there are multiple stores, or if the store is not simple.
+static StoreInst *findUniqueSimpleStoreInBlock(BasicBlock *BB) {
+  StoreInst *Found = nullptr;
+  for (Instruction &I : *BB) {
+    auto *SI = dyn_cast<StoreInst>(&I);
+    if (!SI)
+      continue;
+    if (!SI->isSimple() || Found)
+      return nullptr;
+    Found = SI;
+  }
+  return Found;
+}
+
+/// Return true if none of the non-terminator, non-store instructions in \p BB
+/// are too expensive to speculate. Uses TTI::isExpensiveToSpeculativelyExecute
+/// which is the established LLVM norm for this check (e.g. CodeGenPrepare).
+/// \p FoldedStore is the store being folded and is considered free.
+static bool isBlockCheapToSpeculate(BasicBlock *BB, StoreInst *FoldedStore,
+                                    const TargetTransformInfo &TTI) {
+  for (Instruction &I : *BB) {
+    if (I.isTerminator())
+      continue;
+    if (&I == FoldedStore)
+      continue;
+    if (TTI.isExpensiveToSpeculativelyExecute(&I))
+      return false;
+  }
+  return true;
+}
+
+/// Fold a triangle+diamond CFG pattern where all three branches store to the
+/// same address into a single select+store in the head block.
+///
+///   HeadBB: br i1 %cond1, ThenBB, ElseBB
+///   ThenBB: store VThen, addr;  br MergeBB
+///   ElseBB: [cheap instrs]; br i1 %cond2, ElseThenBB, ElseElseBB
+///   ElseThenBB:  store VElseThen, addr;  br MergeBB
+///   ElseElseBB:  store VElseElse, addr;  br MergeBB
+///
+/// Transforms to:
+///   HeadBB: [hoisted ElseBB instrs]
+///           %sel = select cond1, VThen, (select cond2, VElseThen, VElseElse)
+///           store %sel, addr
+///           br MergeBB
+///
+/// The profitability check uses TTI::isExpensiveToSpeculativelyExecute on the
+/// four speculated blocks (ThenBB, ElseBB, ElseThenBB, ElseElseBB) — HeadBB
+/// instructions already execute unconditionally so they are not checked.
+static bool foldCondStoreToSelectImpl(CondBrInst *BI, DomTreeUpdater *DTU,
+                                      const TargetTransformInfo &TTI) {
+  BasicBlock *HeadBB = BI->getParent();
+  BasicBlock *ThenBB = BI->getSuccessor(0);
+  BasicBlock *ElseBB = BI->getSuccessor(1);
+
+  // ThenBB: single predecessor (HeadBB), one simple store, unconditional
+  // branch to MergeBB.
+  if (ThenBB->getSinglePredecessor() != HeadBB)
+    return false;
+  UncondBrInst *ThenTerm = dyn_cast<UncondBrInst>(ThenBB->getTerminator());
+  if (!ThenTerm)
+    return false;
+  BasicBlock *MergeBB = ThenTerm->getSuccessor(0);
+
+  StoreInst *ThenStore = findUniqueSimpleStoreInBlock(ThenBB);
+  if (!ThenStore)
+    return false;
+
+  // ElseBB: single predecessor (HeadBB), conditional branch to two leaf
+  // blocks. May contain cheap instructions before the branch (e.g. icmp).
+  if (ElseBB->getSinglePredecessor() != HeadBB)
+    return false;
+  CondBrInst *ElseTerm = dyn_cast<CondBrInst>(ElseBB->getTerminator());
+  if (!ElseTerm)
+    return false;
+
+  // Collect non-terminator instructions in ElseBB to hoist into HeadBB.
+  // They must be side-effect-free and not read/write memory.
+  SmallVector<Instruction *, 4> ElseBBInstsToHoist;
+  for (Instruction &I : *ElseBB) {
+    if (&I == ElseTerm)
+      break;
+    if (I.mayHaveSideEffects() || I.mayReadOrWriteMemory())
+      return false;
+    ElseBBInstsToHoist.push_back(&I);
+  }
+
+  BasicBlock *ElseThenBB = ElseTerm->getSuccessor(0);
+  BasicBlock *ElseElseBB = ElseTerm->getSuccessor(1);
+
+  // Each leaf block: single predecessor (ElseBB), one simple store,
+  // unconditional branch to MergeBB.
+  auto CheckLeafBlock = [&](BasicBlock *BB, StoreInst *&Store) -> bool {
+    if (BB->getSinglePredecessor() != ElseBB)
+      return false;
+    UncondBrInst *Term = dyn_cast<UncondBrInst>(BB->getTerminator());
+    if (!Term || Term->getSuccessor(0) != MergeBB)
+      return false;
+    Store = findUniqueSimpleStoreInBlock(BB);
+    return Store != nullptr;
+  };
+
+  StoreInst *ElseThenStore, *ElseElseStore;
+  if (!CheckLeafBlock(ElseThenBB, ElseThenStore) ||
+      !CheckLeafBlock(ElseElseBB, ElseElseStore))
+    return false;
+
+  // Bail out if any of the blocks contain PHI nodes. These are single-
+  // predecessor blocks so PHIs here are degenerate and not expected in
+  // well-optimized IR.
+  if (!ThenBB->phis().empty() || !ElseBB->phis().empty() ||
+      !ElseThenBB->phis().empty() || !ElseElseBB->phis().empty())
+    return false;
+
+  // All three stores must write to the same logical address and have the same
+  // value type.
+  Value *Addr = ThenStore->getPointerOperand();
+  if (!isSameGEPAddress(Addr, ElseThenStore->getPointerOperand()) ||
+      !isSameGEPAddress(Addr, ElseElseStore->getPointerOperand()))
+    return false;
+
+  Type *StoreTy = ThenStore->getValueOperand()->getType();
+  if (ElseThenStore->getValueOperand()->getType() != StoreTy ||
+      ElseElseStore->getValueOperand()->getType() != StoreTy)
+    return false;
+
+  // Profitability: reject if any of the four speculated blocks contains an
+  // instruction that is expensive to execute speculatively (e.g. fdiv).
+  // HeadBB is excluded — its instructions already run unconditionally.
+  if (!isBlockCheapToSpeculate(ThenBB, ThenStore, TTI) ||
+      !isBlockCheapToSpeculate(ElseBB, nullptr, TTI) ||
+      !isBlockCheapToSpeculate(ElseThenBB, ElseThenStore, TTI) ||
+      !isBlockCheapToSpeculate(ElseElseBB, ElseElseStore, TTI))
+    return false;
+
+  // MergeBB PHI nodes: all three leaf predecessors must supply the same value
+  // (since we're collapsing them into HeadBB).
+  for (PHINode &PHI : MergeBB->phis()) {
+    Value *Val = nullptr;
+    for (BasicBlock *Pred : {ThenBB, ElseThenBB, ElseElseBB}) {
+      int Idx = PHI.getBasicBlockIndex(Pred);
+      if (Idx < 0)
+        continue;
+      Value *InVal = PHI.getIncomingValue(Idx);
+      if (!Val)
+        Val = InVal;
+      else if (Val != InVal)
+        return false;
+    }
+  }
+
+  // All checks passed. Build select+store in HeadBB.
+  errs() << "[foldCondStore] firing in function: "
----------------
kartcq wrote:

Hi @mjulian31 , I have added the test case as suggested. Thanks for the comment.

https://github.com/llvm/llvm-project/pull/207654