[NFCI][InstCombine] Move store merging from `visitStoreInst()` into `visitUnconditionalBranchInst()`

Summary:
As @nikic is pointing out in https://bugs.llvm.org/show_bug.cgi?id=46680#c5,
InstCombine should not have forward instruction scans,
so let's move this transform into the proper place.

This is pretty much NFCI.

Reviewers: nikic, spatel

Reviewed By: nikic

Subscribers: hiraditya, llvm-commits, nikic

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D83670

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

@@ -440,6 +440,7 @@ public:
   Instruction *visitLoadInst(LoadInst &LI);
   Instruction *visitStoreInst(StoreInst &SI);
   Instruction *visitAtomicRMWInst(AtomicRMWInst &SI);
+  Instruction *visitUnconditionalBranchInst(BranchInst &BI);
   Instruction *visitBranchInst(BranchInst &BI);
   Instruction *visitFenceInst(FenceInst &FI);
   Instruction *visitSwitchInst(SwitchInst &SI);

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

@@ -1425,34 +1425,6 @@ Instruction *InstCombiner::visitStoreInst(StoreInst &SI) {
   if (isa<UndefValue>(Val))
     return eraseInstFromFunction(SI);
 
-  auto IsNoopInstrForStoreMerging = [](BasicBlock::iterator BBI) {
-    return isa<DbgInfoIntrinsic>(BBI) ||
-           (isa<BitCastInst>(BBI) && BBI->getType()->isPointerTy());
-  };
-
-  // If this store is the second-to-last instruction in the basic block
-  // (excluding debug info and bitcasts of pointers) and if the block ends with
-  // an unconditional branch, try to move the store to the successor block.
-  BBI = SI.getIterator();
-  do {
-    ++BBI;
-  } while (IsNoopInstrForStoreMerging(BBI));
-
-  if (BranchInst *BI = dyn_cast<BranchInst>(BBI))
-    if (BI->isUnconditional())
-      if (mergeStoreIntoSuccessor(SI)) {
-        // Okay, we've managed to do that. Now, let's see if now-second-to-last
-        // instruction is also a store that we can also sink.
-        BasicBlock::iterator FirstInstr = BBI->getParent()->begin();
-        do {
-          if (BBI != FirstInstr)
-            --BBI;
-        } while (BBI != FirstInstr && IsNoopInstrForStoreMerging(BBI));
-        if (StoreInst *PrevStore = dyn_cast<StoreInst>(BBI))
-          Worklist.add(PrevStore);
-        return nullptr;
-      }
-
   return nullptr;
 }
 
@@ -1462,8 +1434,8 @@ Instruction *InstCombiner::visitStoreInst(StoreInst &SI) {
 ///   *P = v1; if () { *P = v2; }
 /// into a phi node with a store in the successor.
 bool InstCombiner::mergeStoreIntoSuccessor(StoreInst &SI) {
-  assert(SI.isUnordered() &&
-         "This code has not been audited for volatile or ordered store case.");
+  if (!SI.isUnordered())
+    return false; // This code has not been audited for volatile/ordered case.
 
   // Check if the successor block has exactly 2 incoming edges.
   BasicBlock *StoreBB = SI.getParent();

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -2734,10 +2734,38 @@ Instruction *InstCombiner::visitReturnInst(ReturnInst &RI) {
   return nullptr;
 }
 
+Instruction *InstCombiner::visitUnconditionalBranchInst(BranchInst &BI) {
+  assert(BI.isUnconditional() && "Only for unconditional branches.");
+
+  // If this store is the second-to-last instruction in the basic block
+  // (excluding debug info and bitcasts of pointers) and if the block ends with
+  // an unconditional branch, try to move the store to the successor block.
+
+  auto GetLastSinkableStore = [](BasicBlock::iterator BBI) {
+    auto IsNoopInstrForStoreMerging = [](BasicBlock::iterator BBI) {
+      return isa<DbgInfoIntrinsic>(BBI) ||
+             (isa<BitCastInst>(BBI) && BBI->getType()->isPointerTy());
+    };
+
+    BasicBlock::iterator FirstInstr = BBI->getParent()->begin();
+    do {
+      if (BBI != FirstInstr)
+        --BBI;
+    } while (BBI != FirstInstr && IsNoopInstrForStoreMerging(BBI));
+
+    return dyn_cast<StoreInst>(BBI);
+  };
+
+  if (StoreInst *SI = GetLastSinkableStore(BasicBlock::iterator(BI)))
+    if (mergeStoreIntoSuccessor(*SI))
+      return &BI;
+
+  return nullptr;
+}
+
 Instruction *InstCombiner::visitBranchInst(BranchInst &BI) {
-  // Nothing to do about unconditional branches.
   if (BI.isUnconditional())
-    return nullptr;
+    return visitUnconditionalBranchInst(BI);
 
   // Change br (not X), label True, label False to: br X, label False, True
   Value *X = nullptr;