Ensure SplitEdge to return the new block between the two given blocks

bangtianliu · Whitney Tsang · commit cf638d793c48 · 2020-12-15T23:32:29.000Z
This PR implements the function splitBasicBlockBefore to address an issue that occurred during SplitEdge(BB, Succ, ...), inside splitBlockBefore. The issue occurs in SplitEdge when the Succ has a single predecessor and the edge between the BB and Succ is not critical. This produces the result ‘BB->Succ->New’. The new function splitBasicBlockBefore was added to splitBlockBefore to handle the issue and now produces the correct result ‘BB->New->Succ’. Below is an example of splitting the block bb1 at its first instruction. /// Original IR bb0: br bb1 bb1: %0 = mul i32 1, 2 br bb2 bb2: /// IR after splitEdge(bb0, bb1) using splitBasicBlock bb0: br bb1 bb1: br bb1.split bb1.split: %0 = mul i32 1, 2 br bb2 bb2: /// IR after splitEdge(bb0, bb1) using splitBasicBlockBefore bb0: br bb1.split bb1.split br bb1 bb1: %0 = mul i32 1, 2 br bb2 bb2: Differential Revision: https://reviews.llvm.org/D92200
diff --git a/llvm/include/llvm/IR/BasicBlock.h b/llvm/include/llvm/IR/BasicBlock.h
@@ -396,22 +396,49 @@ class BasicBlock final : public Value, // Basic blocks are data objects also
 
   /// Split the basic block into two basic blocks at the specified instruction.
   ///
-  /// Note that all instructions BEFORE the specified iterator stay as part of
-  /// the original basic block, an unconditional branch is added to the original
-  /// BB, and the rest of the instructions in the BB are moved to the new BB,
-  /// including the old terminator.  The newly formed BasicBlock is returned.
-  /// This function invalidates the specified iterator.
+  /// If \p Before is true, splitBasicBlockBefore handles the
+  /// block splitting. Otherwise, execution proceeds as described below.
+  ///
+  /// Note that all instructions BEFORE the specified iterator
+  /// stay as part of the original basic block, an unconditional branch is added
+  /// to the original BB, and the rest of the instructions in the BB are moved
+  /// to the new BB, including the old terminator.  The newly formed basic block
+  /// is returned. This function invalidates the specified iterator.
   ///
   /// Note that this only works on well formed basic blocks (must have a
-  /// terminator), and 'I' must not be the end of instruction list (which would
-  /// cause a degenerate basic block to be formed, having a terminator inside of
-  /// the basic block).
+  /// terminator), and \p 'I' must not be the end of instruction list (which
+  /// would cause a degenerate basic block to be formed, having a terminator
+  /// inside of the basic block).
   ///
   /// Also note that this doesn't preserve any passes. To split blocks while
   /// keeping loop information consistent, use the SplitBlock utility function.
-  BasicBlock *splitBasicBlock(iterator I, const Twine &BBName = "");
-  BasicBlock *splitBasicBlock(Instruction *I, const Twine &BBName = "") {
-    return splitBasicBlock(I->getIterator(), BBName);
+  BasicBlock *splitBasicBlock(iterator I, const Twine &BBName = "",
+                              bool Before = false);
+  BasicBlock *splitBasicBlock(Instruction *I, const Twine &BBName = "",
+                              bool Before = false) {
+    return splitBasicBlock(I->getIterator(), BBName, Before);
+  }
+
+  /// Split the basic block into two basic blocks at the specified instruction
+  /// and insert the new basic blocks as the predecessor of the current block.
+  ///
+  /// This function ensures all instructions AFTER and including the specified
+  /// iterator \p I are part of the original basic block. All Instructions
+  /// BEFORE the iterator \p I are moved to the new BB and an unconditional
+  /// branch is added to the new BB. The new basic block is returned.
+  ///
+  /// Note that this only works on well formed basic blocks (must have a
+  /// terminator), and \p 'I' must not be the end of instruction list (which
+  /// would cause a degenerate basic block to be formed, having a terminator
+  /// inside of the basic block).  \p 'I' cannot be a iterator for a PHINode
+  /// with multiple incoming blocks.
+  ///
+  /// Also note that this doesn't preserve any passes. To split blocks while
+  /// keeping loop information consistent, use the SplitBlockBefore utility
+  /// function.
+  BasicBlock *splitBasicBlockBefore(iterator I, const Twine &BBName = "");
+  BasicBlock *splitBasicBlockBefore(Instruction *I, const Twine &BBName = "") {
+    return splitBasicBlockBefore(I->getIterator(), BBName);
   }
 
   /// Returns true if there are any uses of this basic block other than
diff --git a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
@@ -244,19 +244,33 @@ unsigned SplitAllCriticalEdges(Function &F,
                                const CriticalEdgeSplittingOptions &Options =
                                    CriticalEdgeSplittingOptions());
 
-/// Split the edge connecting specified block.
+/// Split the edge connecting the specified blocks, and return the newly created
+/// basic block between \p From and \p To.
 BasicBlock *SplitEdge(BasicBlock *From, BasicBlock *To,
                       DominatorTree *DT = nullptr, LoopInfo *LI = nullptr,
                       MemorySSAUpdater *MSSAU = nullptr);
 
-/// Split the specified block at the specified instruction - everything before
-/// SplitPt stays in Old and everything starting with SplitPt moves to a new
-/// block. The two blocks are joined by an unconditional branch and the loop
-/// info is updated.
+/// Split the specified block at the specified instruction.
+///
+/// If \p Before is true, splitBlockBefore handles the block
+/// splitting. Otherwise, execution proceeds as described below.
+///
+/// Everything before \p SplitPt stays in \p Old and everything starting with \p
+/// SplitPt moves to a new block. The two blocks are joined by an unconditional
+/// branch. The new block with name \p BBName is returned.
 BasicBlock *SplitBlock(BasicBlock *Old, Instruction *SplitPt,
                        DominatorTree *DT = nullptr, LoopInfo *LI = nullptr,
                        MemorySSAUpdater *MSSAU = nullptr,
-                       const Twine &BBName = "");
+                       const Twine &BBName = "", bool Before = false);
+
+/// Split the specified block at the specified instruction \p SplitPt.
+/// All instructions before \p SplitPt are moved to a new block and all
+/// instructions after \p SplitPt stay in the old block. The new block and the
+/// old block are joined by inserting an unconditional branch to the end of the
+/// new block. The new block with name \p BBName is returned.
+BasicBlock *splitBlockBefore(BasicBlock *Old, Instruction *SplitPt,
+                             DominatorTree *DT, LoopInfo *LI,
+                             MemorySSAUpdater *MSSAU, const Twine &BBName = "");
 
 /// This method introduces at least one new basic block into the function and
 /// moves some of the predecessors of BB to be predecessors of the new block.
diff --git a/llvm/lib/IR/BasicBlock.cpp b/llvm/lib/IR/BasicBlock.cpp
@@ -372,7 +372,11 @@ bool BasicBlock::isLegalToHoistInto() const {
   return !Term->isExceptionalTerminator();
 }
 
-BasicBlock *BasicBlock::splitBasicBlock(iterator I, const Twine &BBName) {
+BasicBlock *BasicBlock::splitBasicBlock(iterator I, const Twine &BBName,
+                                        bool Before) {
+  if (Before)
+    return splitBasicBlockBefore(I, BBName);
+
   assert(getTerminator() && "Can't use splitBasicBlock on degenerate BB!");
   assert(I != InstList.end() &&
          "Trying to get me to create degenerate basic block!");
@@ -399,6 +403,40 @@ BasicBlock *BasicBlock::splitBasicBlock(iterator I, const Twine &BBName) {
   return New;
 }
 
+BasicBlock *BasicBlock::splitBasicBlockBefore(iterator I, const Twine &BBName) {
+  assert(getTerminator() &&
+         "Can't use splitBasicBlockBefore on degenerate BB!");
+  assert(I != InstList.end() &&
+         "Trying to get me to create degenerate basic block!");
+
+  assert((!isa<PHINode>(*I) || getSinglePredecessor()) &&
+         "cannot split on multi incoming phis");
+
+  BasicBlock *New = BasicBlock::Create(getContext(), BBName, getParent(), this);
+  // Save DebugLoc of split point before invalidating iterator.
+  DebugLoc Loc = I->getDebugLoc();
+  // Move all of the specified instructions from the original basic block into
+  // the new basic block.
+  New->getInstList().splice(New->end(), this->getInstList(), begin(), I);
+
+  // Loop through all of the predecessors of the 'this' block (which will be the
+  // predecessors of the New block), replace the specified successor 'this'
+  // block to point at the New block and update any PHI nodes in 'this' block.
+  // If there were PHI nodes in 'this' block, the PHI nodes are updated
+  // to reflect that the incoming branches will be from the New block and not
+  // from predecessors of the 'this' block.
+  for (BasicBlock *Pred : predecessors(this)) {
+    Instruction *TI = Pred->getTerminator();
+    TI->replaceSuccessorWith(this, New);
+    this->replacePhiUsesWith(Pred, New);
+  }
+  // Add a branch instruction from  "New" to "this" Block.
+  BranchInst *BI = BranchInst::Create(this, New);
+  BI->setDebugLoc(Loc);
+
+  return New;
+}
+
 void BasicBlock::replacePhiUsesWith(BasicBlock *Old, BasicBlock *New) {
   // N.B. This might not be a complete BasicBlock, so don't assume
   // that it ends with a non-phi instruction.
diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -510,7 +510,7 @@ BasicBlock *llvm::SplitEdge(BasicBlock *BB, BasicBlock *Succ, DominatorTree *DT,
     // block.
     assert(SP == BB && "CFG broken");
     SP = nullptr;
-    return SplitBlock(Succ, &Succ->front(), DT, LI, MSSAU);
+    return SplitBlock(Succ, &Succ->front(), DT, LI, MSSAU, "", /*Before=*/true);
   }
 
   // Otherwise, if BB has a single successor, split it at the bottom of the
@@ -537,7 +537,10 @@ llvm::SplitAllCriticalEdges(Function &F,
 
 BasicBlock *llvm::SplitBlock(BasicBlock *Old, Instruction *SplitPt,
                              DominatorTree *DT, LoopInfo *LI,
-                             MemorySSAUpdater *MSSAU, const Twine &BBName) {
+                             MemorySSAUpdater *MSSAU, const Twine &BBName,
+                             bool Before) {
+  if (Before)
+    return splitBlockBefore(Old, SplitPt, DT, LI, MSSAU, BBName);
   BasicBlock::iterator SplitIt = SplitPt->getIterator();
   while (isa<PHINode>(SplitIt) || SplitIt->isEHPad())
     ++SplitIt;
@@ -569,6 +572,51 @@ BasicBlock *llvm::SplitBlock(BasicBlock *Old, Instruction *SplitPt,
   return New;
 }
 
+BasicBlock *llvm::splitBlockBefore(BasicBlock *Old, Instruction *SplitPt,
+                                   DominatorTree *DT, LoopInfo *LI,
+                                   MemorySSAUpdater *MSSAU,
+                                   const Twine &BBName) {
+
+  BasicBlock::iterator SplitIt = SplitPt->getIterator();
+  while (isa<PHINode>(SplitIt) || SplitIt->isEHPad())
+    ++SplitIt;
+  std::string Name = BBName.str();
+  BasicBlock *New = Old->splitBasicBlock(
+      SplitIt, Name.empty() ? Old->getName() + ".split" : Name,
+      /* Before=*/true);
+
+  // The new block lives in whichever loop the old one did. This preserves
+  // LCSSA as well, because we force the split point to be after any PHI nodes.
+  if (LI)
+    if (Loop *L = LI->getLoopFor(Old))
+      L->addBasicBlockToLoop(New, *LI);
+
+  if (DT) {
+    DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
+    SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
+    // New dominates Old. The predecessor nodes of the Old node dominate
+    // New node.
+    DTUpdates.push_back({DominatorTree::Insert, New, Old});
+    for (BasicBlock *Pred : predecessors(New))
+      if (DT->getNode(Pred)) {
+        DTUpdates.push_back({DominatorTree::Insert, Pred, New});
+        DTUpdates.push_back({DominatorTree::Delete, Pred, Old});
+      }
+
+    DTU.applyUpdates(DTUpdates);
+    DTU.flush();
+
+    // Move MemoryAccesses still tracked in Old, but part of New now.
+    // Update accesses in successor blocks accordingly.
+    if (MSSAU) {
+      MSSAU->applyUpdates(DTUpdates, *DT);
+      if (VerifyMemorySSA)
+        MSSAU->getMemorySSA()->verifyMemorySSA();
+    }
+  }
+  return New;
+}
+
 /// Update DominatorTree, LoopInfo, and LCCSA analysis information.
 static void UpdateAnalysisInformation(BasicBlock *OldBB, BasicBlock *NewBB,
                                       ArrayRef<BasicBlock *> Preds,
diff --git a/llvm/test/CodeGen/AMDGPU/call-constexpr.ll b/llvm/test/CodeGen/AMDGPU/call-constexpr.ll
@@ -100,7 +100,7 @@ define amdgpu_kernel void @test_bitcast_use_workitem_id_x() #0 {
 ; OPT-LABEL: @test_invoke(
 ; OPT: %1 = bitcast float 2.000000e+00 to i32
 ; OPT: %val = invoke i32 @ident_i32(i32 %1)
-; OPT-NEXT: to label %continue unwind label %broken
+; OPT-NEXT: to label %continue.split unwind label %broken
 ; OPT-LABEL: continue.split:
 ; OPT: bitcast i32 %val to float
 @_ZTIi = external global i8*
diff --git a/llvm/test/Transforms/LoopUnswitch/2011-11-18-SimpleSwitch.ll b/llvm/test/Transforms/LoopUnswitch/2011-11-18-SimpleSwitch.ll
@@ -36,7 +36,7 @@
 ; CHECK:      loop_begin.us1:                                   ; preds = %loop_begin.backedge.us5, %.split.split.us
 ; CHECK-NEXT:   %var_val.us2 = load i32, i32* %var
 ; CHECK-NEXT:   switch i32 2, label %default.us-lcssa.us-lcssa.us [
-; CHECK-NEXT:     i32 1, label %inc.us4
+; CHECK-NEXT:     i32 1, label %inc.split.us
 ; CHECK-NEXT:     i32 2, label %dec.us3
 ; CHECK-NEXT:   ]
 
@@ -50,15 +50,15 @@
 ; CHECK:      loop_begin:                                       ; preds = %loop_begin.backedge, %.split.split
 ; CHECK-NEXT:   %var_val = load i32, i32* %var
 ; CHECK-NEXT:   switch i32 %c, label %default.us-lcssa.us-lcssa [
-; CHECK-NEXT:     i32 1, label %inc
-; CHECK-NEXT:     i32 2, label %dec
+; CHECK-NEXT:     i32 1, label %inc.split
+; CHECK-NEXT:     i32 2, label %dec.split
 ; CHECK-NEXT:   ]
 
-; CHECK:      inc:                                              ; preds = %loop_begin
-; CHECK-NEXT:   br i1 true, label %us-unreachable.us-lcssa, label %inc.split
+; CHECK:      inc.split:                                        ; preds = %loop_begin
+; CHECK-NEXT:   br i1 true, label %us-unreachable.us-lcssa, label %inc
 
-; CHECK:      dec:                                              ; preds = %loop_begin
-; CHECK-NEXT:   br i1 true, label %us-unreachable6, label %dec.split
+; CHECK:      dec.split:                                        ; preds = %loop_begin
+; CHECK-NEXT:   br i1 true, label %us-unreachable6, label %dec
 
 define i32 @test(i32* %var) {
   %mem = alloca i32
diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
