[VPlan] Update vector latch terminator edge to exit block after execution.

[fhahn]

Summary

Instead of setting the successor to the exit using CFG.ExitBB, set it to
nullptr initially. The successor to the exit block is later set either
through createEmptyBasicBlock or after VPlan execution (because at the
moment, no block is created by VPlan for the exit block, the existing
one is reused).

This also enables BranchOnCond to be used as terminator for the exiting
block of the topmost vector region.

Depends on D126618.

[fhahn]

Move code to set branch to ExitBB (aka middle.block) to VPBasicBlock::execute.

[Ayal]

llvm/lib/Transforms/Vectorize/VPlan.cpp

265–266

Set here only fallthru successor 0 to NewBB, or only backedge successor 1 to Header?

Intent is to set the backward successor of a conditional (loop) branch when the branch is created, but to set forward successor(s) later when they are created/visited, for all conditional branches? All such successors will not be able to reuse previous BB so must call createEmptyBasicBlock() except for reusing ExitBB.

Note that LoopRegion in this case refers to a nested loop rather than the outer/top/vector loop - createEmptyBasicBlock() call below excludes final Exit VPBB.

290

PredVPBB >> PredVPBlock or ExitingVPBlock? (Expecting BB to mean BasicBlock rather than BlockBase)

292

Add error message?

777–782

(nit: can set IfFalse to null for both above, then setSuccessor(1) to Header if parent isLatch()), as we're setting Successor(0)...

798–799

(nit: is this still needed, or subject to future cleanup?)

802

Here there's only a single destination/successor 0: "destination(s) later when they are" >> "destination later when it is".

Setting IfTrue first to Builder.GetInsertBlock() and then replacing it with nullptr to circumvent an assert of CreateCondBr?

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/VPlan.cpp

265–266

Intent is to set the backward successor of a conditional (loop) branch when the branch is created, but to set forward successor(s) later when they are created/visited, for all conditional branches?

Exactly! I had another look and the extra code here to handle the exiting block separately is not needed any longer I think with this patch.

It can be handled in the else case above, if we use getHierarchicalSuccessors instead of getSuccessors. This seems simpler and I updated the patch in that respect.

290

Changed to PredVPB, which I think is used elsewhere.

292

Done, also switched to using getSingleSuccessor

777–782

Sounds great, I'll update D126618 accordingly.

798–799

It's not needed any longer, removed in 4f1c86e3d5ef, thanks!

802

Here there's only a single destination/successor 0: "destination(s) later when they are" >> "destination later when it is".

Updated! I also updated the comment to mention backward dest = header, forward dest = exit/middle block.

Setting IfTrue first to Builder.GetInsertBlock() and then replacing it with nullptr to circumvent an assert of CreateCondBr?

Yes, unfortunately the first destination is required to be valid block.

[Ayal]

Ship it!

llvm/lib/Transforms/Vectorize/VPlan.cpp

260–262

nit: assert and use TermBr here?

Can comment that this scan sets each forward successor here when it is created, and excludes backedges. A backward successor is set when the branch is created (as successor 1).

295

Can comment that Exit block of a loop is always set to be successor 0 of the Exiting block.

802

Setting IfTrue first to Builder.GetInsertBlock() and then replacing it with nullptr to circumvent an assert of CreateCondBr?

Yes, unfortunately the first destination is required to be valid block.

Understood. Might be worth a comment.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

260–262

Done & comment added, thanks!

295

added the comment, thanks!

802

Extended the comment, thanks!

llvm/lib/Transforms/Vectorize/VPlan.cpp

VPBasicBlock::createEmptyBasicBlock(VPTransformState::CFGState &CFG) {
   BasicBlock *PrevBB = CFG.PrevBB;
   BasicBlock *NewBB = BasicBlock::Create(PrevBB->getContext(), getName(),
                                          PrevBB->getParent(), CFG.ExitBB);
   LLVM_DEBUG(dbgs() << "LV: created " << NewBB->getName() << '\n');
 
   // Hook up the new basic block to its predecessors.
   for (VPBlockBase *PredVPBlock : getHierarchicalPredecessors()) {
     VPBasicBlock *PredVPBB = PredVPBlock->getExitingBasicBlock();
-    auto &PredVPSuccessors = PredVPBB->getSuccessors();
+    auto &PredVPSuccessors = PredVPBB->getHierarchicalSuccessors();
     BasicBlock *PredBB = CFG.VPBB2IRBB[PredVPBB];
 
     assert(PredBB && "Predecessor basic-block not found building successor.");
     auto *PredBBTerminator = PredBB->getTerminator();
     LLVM_DEBUG(dbgs() << "LV: draw edge from" << PredBB->getName() << '\n');
 
     auto *TermBr = dyn_cast<BranchInst>(PredBBTerminator);
     if (isa<UnreachableInst>(PredBBTerminator)) {
       assert(PredVPSuccessors.size() == 1 &&
              "Predecessor ending w/o branch must have single successor.");
       DebugLoc DL = PredBBTerminator->getDebugLoc();
       PredBBTerminator->eraseFromParent();
       auto *Br = BranchInst::Create(NewBB, PredBB);
       Br->setDebugLoc(DL);
     } else if (TermBr && !TermBr->isConditional()) {
       TermBr->setSuccessor(0, NewBB);
-    } else if (PredVPSuccessors.size() == 2) {
+    } else {
+      // Set each forward successor here when it is created, excluding
+      // backedges. A backward successor is set when the branch is created.
       unsigned idx = PredVPSuccessors.front() == this ? 0 : 1;
-      assert(!PredBBTerminator->getSuccessor(idx) &&
+      assert(!TermBr->getSuccessor(idx) &&
              "Trying to reset an existing successor block.");
-      PredBBTerminator->setSuccessor(idx, NewBB);
-    } else {
-      // PredVPBB is the exiting block of a loop region. Connect its successor
-      // outside the region.
-      auto *LoopRegion = cast<VPRegionBlock>(PredVPBB->getParent());
-      assert(!LoopRegion->isReplicator() &&
-             "predecessor must be in a loop region");
-      assert(PredVPSuccessors.empty() &&
-             LoopRegion->getExitingBasicBlock() == PredVPBB &&
-             "PredVPBB must be the exiting block of its parent region");
-      assert(this == LoopRegion->getSingleSuccessor() &&
-             "the current block must be the single successor of the region");
-      PredBBTerminator->setSuccessor(0, NewBB);
-      PredBBTerminator->setSuccessor(
-          1, CFG.VPBB2IRBB[LoopRegion->getEntryBasicBlock()]);
+      TermBr->setSuccessor(idx, NewBB);
     }
   }
   return NewBB;
 }
 
 void VPBasicBlock::execute(VPTransformState *State) {
   bool Replica = State->Instance && !State->Instance->isFirstIteration();
   VPBasicBlock *PrevVPBB = State->CFG.PrevVPBB;
   VPBlockBase *SingleHPred = nullptr;
   BasicBlock *NewBB = State->CFG.PrevBB; // Reuse it if possible.
 
   auto IsLoopRegion = [](VPBlockBase *BB) {
     auto *R = dyn_cast<VPRegionBlock>(BB);
     return R && !R->isReplicator();
   };
 
   // 1. Create an IR basic block, or reuse the last one or ExitBB if possible.
   if (getPlan()->getVectorLoopRegion()->getSingleSuccessor() == this) {
     // ExitBB can be re-used for the exit block of the Plan.
     NewBB = State->CFG.ExitBB;
     State->CFG.PrevBB = NewBB;
+
+    // Update the branch instruction in the predecessor to branch to ExitBB.
+    VPBlockBase *PredVPB = getSingleHierarchicalPredecessor();
+    VPBasicBlock *ExitingVPBB = PredVPB->getExitingBasicBlock();
+    assert(PredVPB->getSingleSuccessor() == this &&
+           "predecessor must have the current block as only successor");
+    BasicBlock *ExitingBB = State->CFG.VPBB2IRBB[ExitingVPBB];
+    // The Exit block of a loop is always set to be successor 0 of the Exiting
+    // block.
+    cast<BranchInst>(ExitingBB->getTerminator())->setSuccessor(0, NewBB);
   } else if (PrevVPBB && /* A */
              !((SingleHPred = getSingleHierarchicalPredecessor()) &&
                SingleHPred->getExitingBasicBlock() == PrevVPBB &&
                PrevVPBB->getSingleHierarchicalSuccessor() &&
                (SingleHPred->getParent() == getEnclosingLoopRegion() &&
                 !IsLoopRegion(SingleHPred))) &&         /* B */
              !(Replica && getPredecessors().empty())) { /* C */
     // The last IR basic block is reused, as an optimization, in three cases:
if (Part != 0)
       break;
     Value *Cond = State.get(getOperand(0), VPIteration(Part, 0));
     VPRegionBlock *ParentRegion = getParent()->getParent();
     VPBasicBlock *Header = ParentRegion->getEntryBasicBlock();
 
     // Replace the temporary unreachable terminator with a new conditional
     // branch, hooking it up to backward destination for exiting blocks now and
     // to forward destination(s) later when they are created.
     BranchInst *CondBr =
         Builder.CreateCondBr(Cond, Builder.GetInsertBlock(), nullptr);
 
     if (getParent()->isExiting())
       CondBr->setSuccessor(1, State.CFG.VPBB2IRBB[Header]);
 
     CondBr->setSuccessor(0, nullptr);
     Builder.GetInsertBlock()->getTerminator()->eraseFromParent();
     break;
   }
   case VPInstruction::BranchOnCount: {
     if (Part != 0)
       break;
     // First create the compare.
     Value *IV = State.get(getOperand(0), Part);
     Value *TC = State.get(getOperand(1), Part);
     Value *Cond = Builder.CreateICmpEQ(IV, TC);
 
     // Now create the branch.
     auto *Plan = getParent()->getPlan();
     VPRegionBlock *TopRegion = Plan->getVectorLoopRegion();
     VPBasicBlock *Header = TopRegion->getEntry()->getEntryBasicBlock();
-    // TODO: Once the exit block is modeled in VPlan, use it instead of going
-    // through State.CFG.ExitBB.
-    BasicBlock *Exit = State.CFG.ExitBB;
 
-    Builder.CreateCondBr(Cond, Exit, State.CFG.VPBB2IRBB[Header]);
+    // Replace the temporary unreachable terminator with a new conditional
+    // branch, hooking it up to backward destination (the header) now and to the
+    // forward destination (the exit/middle block) later when it is created.
+    // Note that CreateCondBr expects a valid BB as first argument, so we need
+    // to set it to nullptr later.
+    BranchInst *CondBr = Builder.CreateCondBr(Cond, Builder.GetInsertBlock(),
+                                              State.CFG.VPBB2IRBB[Header]);
+    CondBr->setSuccessor(0, nullptr);
     Builder.GetInsertBlock()->getTerminator()->eraseFromParent();
     break;
   }
   default:
     llvm_unreachable("Unsupported opcode for instruction");
   }
 }