[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:

 Thank you for the review feedback. @dtcxzyw 
	 
 I will address the other feedback to add pre-commit tests with baseline CHECK lines (before the transform) as a separate PR.

My take regarding the suggestion to move the transform to InstCombiner is,
SimplifyCFG already has **mergeConditionalStores** which handles a two-level diamond/triangle pattern where conditional blocks contain stores to the same address. This existing function that does similar motivated me to add the new transform in SimplifyCFG. Our pattern is a natural extension of what mergeConditionalStores already does. The key difference is that our pattern handles the nested diamond case, where the else branch of the outer condition itself branches into two sub-cases. This produces an unconditional select+store instead of a conditional store, which enables loop vectorization.


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