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

[Ramkumar Ramachandra via llvm-commits]

================
@@ -4292,6 +4292,302 @@ 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.
+/// Also rejects blocks containing non-store instructions that read/write
+/// memory or have side effects, as they cannot be safely speculated.
+static StoreInst *findUniqueSimpleStoreInBlock(BasicBlock *BB) {
+  StoreInst *Found = nullptr;
+  unsigned NumOther = 0;
+  for (Instruction &I : *BB) {
+    if (I.isTerminator())
+      continue;
+    auto *SI = dyn_cast<StoreInst>(&I);
+    if (SI) {
+      if (!SI->isSimple() || Found)
+        return nullptr;
+      Found = SI;
+      continue;
+    }
+    // Reject if any non-store instruction reads/writes memory or has
+    // side effects — it cannot be safely speculated.
+    if (I.mayReadOrWriteMemory() || I.mayHaveSideEffects())
+      return nullptr;
+    // Allow at most one non-store, non-terminator instruction. This ensures
+    // that HoistIfNeeded (which moves a single instruction) is sufficient —
+    // with only one instruction, its operands must come from outside the block.
+    if (++NumOther > 1)
+      return nullptr;
+  }
+  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
----------------
artagnon wrote:

I'm confused why this is specific to stores: wouldn't any instruction potentially benefit from this select-folding?

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