[MergedLoadStoreMotion] Sink stores to BB with more than 2 predecessors

[dendibakh]

Summary

If we have:

bb5:
  br i1 %arg3, label %bb6, label %bb7

bb6:
  %tmp = getelementptr inbounds i32, i32* %arg1, i64 2
  store i32 3, i32* %tmp, align 4
  br label %bb9

bb7:
  %tmp8 = getelementptr inbounds i32, i32* %arg1, i64 2
  store i32 3, i32* %tmp8, align 4
  br label %bb9

bb9:	; preds = %bb4, %bb6, %bb7
  ...

We can't sink stores directly into bb9.
This patch creates new BB that is successor of %bb6 and %bb7
and sinks stores into that block.

[lebedev.ri]

Please improve patch description, current one is not too useful.

[dendibakh]

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

364

Maybe we can invalidate only if new BB was created.

394

Maybe we can invalidate only if new BB was created.

[xbolva00]

Can you use utils/update_test_checks.py on this new IR file, please?

[lebedev.ri]

llvm/test/Transforms/InstMerge/st_sink_split_bb.ll

43–45

This test is hard to read, can you use utils/update_test_checks.py and precommit it first?

[dendibakh]

Ran update_test_checks.py

[bjope]

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

20

Could probably add something here that a new tail/footer block may be insterted if the tail/footer block has more predecessors (not only the two predecessors that are forming the diamond).

286

coding rules: asserts should have a string

289

coding rules: asserts should have a string

[dendibakh]

Fixed Bjorn's comments.

[bjope]

Code changes basically looks good to me, but I have no idea about:

• the expected gain
• if this is handled elsewhere (or if it should be)
• if there is a (high) cost related to "not preserving the CFG"

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

275–276

nit: Afaict it seems like a limitation in the current implementation to only sink when we find a diamond (there is a todo on line 72 about generalizing to other code structures). I guess the transform done by this method could be valid for any pair of joining flows, so it is a little bit weird to base it on a diamond head. The function would be more general if it for example took Pred0 and Pred1 as input (asserting that they both have identical single successors, and that Pred0!=Pred1), and then moving the logic about how we detect the predecessor pair to the caller.

[dendibakh]

I introduced an option that controls splitting the footer BB which is disabled by default.
So, no functional change at the moment.

[dendibakh]

Ping.

[dendibakh]

@bjope, please review.

[bjope]

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

348

Are we sure that this iteration is safe now when we may end up adding new blocks while iterating.
I think the iterator still will be valid in the sense that it still points to a BB (or end) and that we can continue iterating. I also assume that it will be deterministic, right?

The inserted blocks may or may not end up being analysed here (afaict) depending on if their predecessors in BB order already have been analysed or not). But since the inserted blocks never will be diamond heads that isn't a problem right now, since this loop only care about diamond heads.

Maybe it is worth adding a comment that mergeStores can insert new blocks, and that it isn't sure if they will be included in the iteration or not?

[bjope]

LGTM (apart from the comment about maybe saying something about the iteration order).

@lebedev.ri had comments about the test case earlier, so you could wait another day or two before landing this, to see if someone else got a different opinion (I don't know that much about this, but the code seem to do what you are aiming for afaict).

[dendibakh]

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

348

I took the idea from SimplifyCFG which does similar splitting. In SimplifyCFG we have the same implementation, i.e. we iterate over BBs in a function:
https://github.com/llvm/llvm-project/blob/master/llvm/lib/Transforms/Scalar/SimplifyCFGPass.cpp#L162
simplifyCFG() internally does the splitting.

Anyway, thanks for a good point. I will add a comment.

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

llvm/lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

 //
 // The algorithm iteratively hoists two loads to the same address out of a
 // diamond (hammock) and merges them into a single load in the header. Similar
 // it sinks and merges two stores to the tail block (footer). The algorithm
 // iterates over the instructions of one side of the diamond and attempts to
 // find a matching load/store on the other side. It hoists / sinks when it
 // thinks it safe to do so.  This optimization helps with eg. hiding load
 // latencies, triggering if-conversion, and reducing static code size.
 //
 // NOTE: This code no longer performs load hoisting, it is subsumed by GVNHoist.
 //
 //===----------------------------------------------------------------------===//
 //
 //
 // Example:
 // Diamond shaped code before merge:
 //

 private:
   BasicBlock *getDiamondTail(BasicBlock *BB);
   bool isDiamondHead(BasicBlock *BB);
   // Routines for sinking stores
   StoreInst *canSinkFromBlock(BasicBlock *BB, StoreInst *SI);
   PHINode *getPHIOperand(BasicBlock *BB, StoreInst *S0, StoreInst *S1);
   bool isStoreSinkBarrierInRange(const Instruction &Start,
                                  const Instruction &End, MemoryLocation Loc);
-  bool sinkStore(BasicBlock *BB, StoreInst *SinkCand, StoreInst *ElseInst);
+  bool canSinkStoresAndGEPs(StoreInst *S0, StoreInst *S1) const;
+  void sinkStoresAndGEPs(BasicBlock *BB, StoreInst *SinkCand,
+                         StoreInst *ElseInst);
   bool mergeStores(BasicBlock *BB);
 };
 } // end anonymous namespace
 
 ///
 /// Return tail block of a diamond.
 ///
 BasicBlock *MergedLoadStoreMotion::getDiamondTail(BasicBlock *BB) {
auto *NewPN = PHINode::Create(Opd1->getType(), 2, Opd2->getName() + ".sink",
                                 &BB->front());
   NewPN->applyMergedLocation(S0->getDebugLoc(), S1->getDebugLoc());
   NewPN->addIncoming(Opd1, S0->getParent());
   NewPN->addIncoming(Opd2, S1->getParent());
   return NewPN;
 }
 
 ///
+/// Check if 2 stores can be sunk together with corresponding GEPs
+///
+bool MergedLoadStoreMotion::canSinkStoresAndGEPs(StoreInst *S0,
+                                                 StoreInst *S1) const {
+  auto *A0 = dyn_cast<Instruction>(S0->getPointerOperand());
+  auto *A1 = dyn_cast<Instruction>(S1->getPointerOperand());
+  return A0 && A1 && A0->isIdenticalTo(A1) && A0->hasOneUse() &&
+         (A0->getParent() == S0->getParent()) && A1->hasOneUse() &&
+         (A1->getParent() == S1->getParent()) && isa<GetElementPtrInst>(A0);
+}
+
+///
 /// Merge two stores to same address and sink into \p BB
 ///
 /// Also sinks GEP instruction computing the store address
 ///
-bool MergedLoadStoreMotion::sinkStore(BasicBlock *BB, StoreInst *S0,
+void MergedLoadStoreMotion::sinkStoresAndGEPs(BasicBlock *BB, StoreInst *S0,
                                               StoreInst *S1) {
   // Only one definition?
   auto *A0 = dyn_cast<Instruction>(S0->getPointerOperand());
   auto *A1 = dyn_cast<Instruction>(S1->getPointerOperand());
-  if (A0 && A1 && A0->isIdenticalTo(A1) && A0->hasOneUse() &&
-      (A0->getParent() == S0->getParent()) && A1->hasOneUse() &&
-      (A1->getParent() == S1->getParent()) && isa<GetElementPtrInst>(A0)) {
-    LLVM_DEBUG(dbgs() << "Sink Instruction into BB \n"; BB->dump();
-               dbgs() << "Instruction Left\n"; S0->dump(); dbgs() << "\n";
-               dbgs() << "Instruction Right\n"; S1->dump(); dbgs() << "\n");
-    // Hoist the instruction.
-    BasicBlock::iterator InsertPt = BB->getFirstInsertionPt();
-    // Intersect optional metadata.
-    S0->andIRFlags(S1);
-    S0->dropUnknownNonDebugMetadata();
 
-    // Create the new store to be inserted at the join point.
-    StoreInst *SNew = cast<StoreInst>(S0->clone());
-    Instruction *ANew = A0->clone();
-    SNew->insertBefore(&*InsertPt);
-    ANew->insertBefore(SNew);
 
-    assert(S0->getParent() == A0->getParent());
-    assert(S1->getParent() == A1->getParent());
 
-    // New PHI operand? Use it.
-    if (PHINode *NewPN = getPHIOperand(BB, S0, S1))
-      SNew->setOperand(0, NewPN);
-    S0->eraseFromParent();
-    S1->eraseFromParent();
-    A0->replaceAllUsesWith(ANew);
-    A0->eraseFromParent();
-    A1->replaceAllUsesWith(ANew);
-    A1->eraseFromParent();
-    return true;
-  }
-  return false;
 }
 
 ///
 /// True when two stores are equivalent and can sink into the footer
 ///
-/// Starting from a diamond tail block, iterate over the instructions in one
-/// predecessor block and try to match a store in the second predecessor.
+/// Starting from a diamond head block, iterate over the instructions in one
+/// successor block and try to match a store in the second successor.
 ///
-bool MergedLoadStoreMotion::mergeStores(BasicBlock *T) {
+bool MergedLoadStoreMotion::mergeStores(BasicBlock *HeadBB) {
 
   bool MergedStores = false;
-  assert(T && "Footer of a diamond cannot be empty");
-
-  pred_iterator PI = pred_begin(T), E = pred_end(T);
-  assert(PI != E);
-  BasicBlock *Pred0 = *PI;
-  ++PI;
-  BasicBlock *Pred1 = *PI;
-  ++PI;
+  BasicBlock *TailBB = getDiamondTail(HeadBB);
+  BasicBlock *SinkBB = TailBB;
+  assert(SinkBB && "Footer of a diamond cannot be empty");
+
+  succ_iterator SI = succ_begin(HeadBB);
+  assert(SI != succ_end(HeadBB));
+  BasicBlock *Pred0 = *SI;
+  ++SI;
+  assert(SI != succ_end(HeadBB));
+  BasicBlock *Pred1 = *SI;
   // tail block  of a diamond/hammock?
   if (Pred0 == Pred1)
     return false; // No.
-  if (PI != E)
-    return false; // No. More than 2 predecessors.
-
-  // #Instructions in Succ1 for Compile Time Control
+  // #Instructions in Pred1 for Compile Time Control
   auto InstsNoDbg = Pred1->instructionsWithoutDebug();
   int Size1 = std::distance(InstsNoDbg.begin(), InstsNoDbg.end());
   int NStores = 0;
 
   for (BasicBlock::reverse_iterator RBI = Pred0->rbegin(), RBE = Pred0->rend();
        RBI != RBE;) {
 
     Instruction *I = &*RBI;
     ++RBI;
 
     // Don't sink non-simple (atomic, volatile) stores.
     auto *S0 = dyn_cast<StoreInst>(I);
     if (!S0 || !S0->isSimple())
       continue;
 
     ++NStores;
     if (NStores * Size1 >= MagicCompileTimeControl)
       break;
     if (StoreInst *S1 = canSinkFromBlock(Pred1, S0)) {
-      bool Res = sinkStore(T, S0, S1);
-      MergedStores |= Res;
+      if (!canSinkStoresAndGEPs(S0, S1))
         // Don't attempt to sink below stores that had to stick around
         // But after removal of a store and some of its feeding
         // instruction search again from the beginning since the iterator
         // is likely stale at this point.
-      if (!Res)
         break;
+
+      if (SinkBB == TailBB && TailBB->hasNPredecessorsOrMore(3)) {
+        // We have more than 2 predecessors. Insert a new block
+        // postdominating 2 predecessors we're going to sink from.
+        SinkBB = SplitBlockPredecessors(TailBB, {Pred0, Pred1}, ".sink.split");
+        if (!SinkBB)
+          break;
+      }
+
+      MergedStores = true;
+      sinkStoresAndGEPs(SinkBB, S0, S1);
       RBI = Pred0->rbegin();
       RBE = Pred0->rend();
       LLVM_DEBUG(dbgs() << "Search again\n"; Instruction *I = &*RBI; I->dump());
     }
   }
   return MergedStores;
 }
 
 bool MergedLoadStoreMotion::run(Function &F, AliasAnalysis &AA) {
   this->AA = &AA;
 
   bool Changed = false;
   LLVM_DEBUG(dbgs() << "Instruction Merger\n");
 
   // Merge unconditional branches, allowing PRE to catch more
   // optimization opportunities.
   for (Function::iterator FI = F.begin(), FE = F.end(); FI != FE;) {
     BasicBlock *BB = &*FI++;
 
     // Hoist equivalent loads and sink stores
     // outside diamonds when possible
     if (isDiamondHead(BB)) {
-      Changed |= mergeStores(getDiamondTail(BB));
+      Changed |= mergeStores(BB);
     }
   }
   return Changed;
 }
 
 namespace {
 class MergedLoadStoreMotionLegacyPass : public FunctionPass {
 public:
bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
     MergedLoadStoreMotion Impl;
     return Impl.run(F, getAnalysis<AAResultsWrapperPass>().getAAResults());
   }
 
 private:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesCFG();
     AU.addRequired<AAResultsWrapperPass>();
     AU.addPreserved<GlobalsAAWrapperPass>();
   }
 };
 
 char MergedLoadStoreMotionLegacyPass::ID = 0;
 } // anonymous namespace
 

 PreservedAnalyses
 MergedLoadStoreMotionPass::run(Function &F, FunctionAnalysisManager &AM) {
   MergedLoadStoreMotion Impl;
   auto &AA = AM.getResult<AAManager>(F);
   if (!Impl.run(F, AA))
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
-  PA.preserveSet<CFGAnalyses>();
   PA.preserve<GlobalsAA>();
   return PA;
 }

llvm/test/Transforms/InstMerge/st_sink_split_bb.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; Test to make sure that a new block is inserted if we
+; have more than 2 predecessors for the block we're going to sink to.
+; RUN: opt -basicaa -memdep -mldst-motion -S < %s | FileCheck %s
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+; Function Attrs: nounwind uwtable
+define dso_local void @st_sink_split_bb(i32* nocapture %arg, i32* nocapture %arg1, i1 zeroext %arg2, i1 zeroext %arg3) local_unnamed_addr {
+; CHECK-LABEL: @st_sink_split_bb(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    br i1 [[ARG2:%.*]], label [[BB4:%.*]], label [[BB5:%.*]]
+; CHECK:       bb4:
+; CHECK-NEXT:    store i32 1, i32* [[ARG:%.*]], align 4
+; CHECK-NEXT:    br label [[BB9:%.*]]
+; CHECK:       bb5:
+; CHECK-NEXT:    br i1 [[ARG3:%.*]], label [[BB6:%.*]], label [[BB7:%.*]]
+; CHECK:       bb6:
+; CHECK-NEXT:    store i32 2, i32* [[ARG]], align 4
+; CHECK-NEXT:    br label [[BB9_SINK_SPLIT:%.*]]
+; CHECK:       bb7:
+; CHECK-NEXT:    store i32 3, i32* [[ARG]], align 4
+; CHECK-NEXT:    br label [[BB9_SINK_SPLIT]]
+; CHECK:       bb9.sink.split:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i32, i32* [[ARG1:%.*]], i64 1
+; CHECK-NEXT:    store i32 3, i32* [[TMP0]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, i32* [[ARG1]], i64 2
+; CHECK-NEXT:    store i32 3, i32* [[TMP1]], align 4
+; CHECK-NEXT:    br label [[BB9]]
+; CHECK:       bb9:
+; CHECK-NEXT:    ret void
+;
+bb:
+  br i1 %arg2, label %bb4, label %bb5
+
+bb4:                                              ; preds = %bb
+  store i32 1, i32* %arg, align 4
+  br label %bb9
+
+bb5:                                              ; preds = %bb
+  br i1 %arg3, label %bb6, label %bb7
+
+bb6:                                              ; preds = %bb5
+  store i32 2, i32* %arg, align 4
+  %tmp1 = getelementptr inbounds i32, i32* %arg1, i64 1
+  store i32 3, i32* %tmp1, align 4
+  %tmp = getelementptr inbounds i32, i32* %arg1, i64 2
+  store i32 3, i32* %tmp, align 4
+  br label %bb9
+
+bb7:                                              ; preds = %bb5
+  store i32 3, i32* %arg, align 4
+  %tmp2 = getelementptr inbounds i32, i32* %arg1, i64 1
+  store i32 3, i32* %tmp2, align 4
+  %tmp8 = getelementptr inbounds i32, i32* %arg1, i64 2
+  store i32 3, i32* %tmp8, align 4
+  br label %bb9
+
+
+bb9:                                              ; preds = %bb7, %bb6, %bb4
+  ret void
+}