[VPlan] Simplify HCFG construction of region blocks (NFC).

[fhahn]

Summary

Update the logic to update the successors and predecessors of region
blocks directly. This adds special handling for header and latch blocks
in place, and removes the separate loop to fix up the region blocks.

Helps to simplify D158333.

[Ayal]

Title should indicate this is an HCFG simplification patch.

Worth indicating that this refactoring assists D158333.

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

53

Mapping BB's to their VPBB's except for header BB's that are mapped to the enclosing VPRegion may be confusing (exacerbated by reusing VPBB or using VPB for VPBlockBase), and should be documented.

Would it be better to keep BB2VPBB as is, mapping BB's to their VPBasicBlocks, and augment it with a getRegionOfHeader() function that checks if a VPBB is the entry of its enclosing region and if so returns the region, otherwise returns the given VPBasicBlock?

85

Would two methods be clearer, one taking a VPBasicBlock and the other a VPRegionBlock?

90

Should avoid connecting VPBBPred-VPBB that are latch-header, handled above?

92

Two cases:

1. VPBBPred's parent == VPBB's grandparent: VPBBPred-VPBB are preheader-header, handled above?
2. VPBBPred's grandparent == VPBB's parent: VPBBPred-VPBB are latch-exit, handled here?

92

Can be simplified - trying to select one of two expected predecessors?

135

Independent cleanup?

159

nit: deserves brackets.

184

This check if BB is a header should probably be generalized (e.g., into isHeader(BasicBlock*)) and used below.

326

Why/Is the above change needed? HeaderVPBB sounds like a VPBasicBlock, rather than its region.

352

nit: "during the RPO traversal" - Recipes will be created for each successor when it is reached later during this RPO traversal.

353–354

nit: assert VPBB and SuccVPBB have same parent.

nit: set predecessors (checking if header?) and early continue;

353–354

Worth dealing here with four cases explicitly: BB is header or not, BB is latch or not?

Setting predecessors: easier to check if BB is header of its loop, and if so hook the region of its VPBB (setting its entry) to a single predecessor - the VPBB of preheader - having a distinct Region? Otherwise setVPBBPredsFromBB.

Setting successors: easier to check if BB is latch of its loop, and if so hook the region of its VPBB (setting its exit) to a single successor - the VPBB of exit - having a distinct Region? Otherwise set VPBB's successor(s).

354–381

nit: consider asserting at the outset that NumSuccs is 1 or 2, to save indentation.

358

SuccVPBB >> ExitVPBB, or OutOfRegionSuccVPBB.

381–382

Better to check if BB is a header before calling setVPBBPredsFromBB?

391

An empty VPBB was created, and set as successor, but not hooked up to its predecessor, which is done here. Setting predecessors separately from setting successors (rather than together as bidirectional edges), presumably to keep predecessors in order.
Instructions of loop exit weren't visited, nor are they visited here.

394

nit: independent cleanup?

[fhahn]

Address latest comments and update description, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

53

Undone changes here, getOrCreateVPB checks if the created VPBB is the entry for a region. I still retained getOrCreateVPB, because it is convenient to get the block to set the successors, e.g. when connecting the pre-header to the inner loop header.

85

Moved out, thanks!

90

Yes, the check here is to catch setting the predecessor for the exit block, as it needs to connect to the predecessor region. Added a comment.

92

In the else case, VPBB Is the exit block of PredVPBB's region, which isn't handled elsewhere (checking elsewhere would be more complicated I think.

92

Retained loop to add an assert that we only have. single predecessors.

135

This was needed after changing the map, undone now.

159

added, thanks!

184

This code is not needed in the latest version

326

We need to connect the pre-header to the top region and set the name of the header block (to avoid unnecessary test changes). Update to clarify that.

353–354

Added assert to setOneSuccessor and fixed an issue uncovered where the parent for regions wasn't set properly.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

353–354

Updated to handle header and latch separately while setting successors. To simplify the latter, I updated the code to collect the successors in a vector, as this is used to check for the latch.

358

adjusted, thanks!

381–382

Updated and moved before setting successors. Checking if it's a header can be done by checking if the VPBB is the entry to its parent region

391

Yes, predecessors are only set once the block has been visited in RPO for that reason, so it needs to be done explicitly here.

394

Was only needed after changing the map, undone, thanks!

[fhahn]

Undo uneeded change

[fhahn]

ping :)

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

97

This assumes/asserts BB is not a header block, with a preheader PredVPBB having VPBB's region as a single successor.

This assumes successors have already been set, before setting predecessors.

nit: assert first; add error message.

106

Comment needs updating.

109

Better to distinguish between latch and preheader by checking if enclosing regions are the same, as in setVPBBPredsFromBB(), rather than relying on having already set the successor(s) of preheader but not yet of the latch? The latch should have no successors even after processing it.

168

CurrentLoop >> LoopOfBB
ParentR >> RegionOfVPBB

170

May be clearer to do something like

bool RegionExists = Loop2Region.contains(CurrentLoop);
if (RegionExists)
  ParentR = Loop2Region[CurrentLoop];
else {
  ParentR = new VPRegionBlock(...);
  ParentR->setEntry(VPBB); // Or feed it in the constructor.
  Loop2Region.insert({CurrentLoop, ParentR});
}

178

nit: worth separating the two if's with an empty line.

182

Should VPBB's parent be set to ParentR also for header blocks which create the region when visited?

330

Simpler to do something like

VPBlockBase *HeaderVPBB = getOrCreateVPBB(TheLoop->getHeader());
HeaderVPBB->setName("vector.body");
ThePreheaderVPBB->setOneSuccessor(HeaderVPBB->getParent());

?

345

Check if BB is header of its loop according to LI, and assert that Region's entry is VPBB, rather than the other way around? I.e., have isHeaderBB() rather than isHeaderVPBB().

353–354

Proposed above to check for latch using parental regions than number of successors.

363

nit: count >> contains

366

"one of its predecessors is the region" >> "one of its successors is the header, which is set to its region by getOrCreateVPB(), in constrast to getOrCreateVPBB()".

This is confusing.
Can we keep using getOrCreateVPBB() only, checking if BB is pre/header/latch using LI, and processing its (predecessors and) successors accordingly?

385

else {

Already asserted above that numSuccs is 1 or 2.

386

Make sure first successor corresponds to True?

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

319

buildPlain/HierarchicalCFG() is called following VPlan::createInitialVPlan(), which creates VPBB's (and will create more, plus Region, in D158333) directly rather than via getOrCreateVPBB(); so need to revisit them and update BB2VPBB, Loop2Region, setParent. Worth a comment.

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

97

Reordered assert, added message, thanks!

106

Updated, thanks!

109

Adjusted to check the parent region, thanks!

168

Renamed, thanks!

170

adjusted, thanks!

178

Moved the code to the earlier if

182

That should be handled by setEntry IIUC

319

Added a comment, thanks!

330

Simplified, thanks!

345

Flipped, thanks!

363

Adjusted, thanks!

366

Updated to us LI to check for latch and use early continue.

385

Adjusted, thanks!

386

Added assert, thanks!

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

143

isHeaderVPBB used only in assert, define under !NDEBUG?
Suggest below to reuse it outside of asserts.

158–182

nit: BlockIt is not needed below, can simply check

// Retrieve VPBB if exists.
if (BB2VPBB.contains(BB))
  return BB2VPBB[BB];

185

nit: can early exit here via

if (!LoopOfBB)
  return VPBB;

and define RegionOfVPBB inside then/else when setting it.

188

nit: assert BB-is-header xor RegionExists?

194

VPBB->setParent(RegionOfVPBB);
here too?

195

When called with BB being the header of TheLoop, its parentLoop may be null or not; in any case, will Loop2Region[TheLoop->getParentLoop()] surely return null?
Can set the parent of RegionOfVPBB only if LoopOfBB != TheLoop.

197

343

nit: can handle the simpler non-header case first.

396

Trying to sketch how this could be done using only getOrCreateVPBB(), i.e., w/o getOrCreateVPB(), ordered according to simplicity:

unsigned NumSuccs = succ_size(BB);
if (NumSuccs == 1) {
  auto *Successor = getOrCreateVPBB(BB->getSingleSuccessor());
  VPBB->setOneSuccessor(isHeaderVPBB(Successor) ? Successor->getParent()
                                                : Successor);
  continue;
}

assert(NumSuccs == 2 && "block must have one or two successors");
// Look up the branch condition to get the corresponding VPValue
// representing the condition bit in VPlan (which may be in another VPBB).
assert(IRDef2VPValue.contains(BI->getCondition()) &&
       "Missing condition bit in IRDef2VPValue!");
VPBasicBlock *Successor0 = getOrCreateVPBB(BI->getSuccessor(0));
VPBasicBlock *Successor1 = getOrCreateVPBB(BI->getSuccessor(1));
if (!isLatchBB(BB, LoopForBB)) {
  VPBB->setTwoSuccessors(Successor0, Successor1);
  continue;
}
// For a latch we need to set the successor of the region rather than that
// of VPBB and it should be set to the exit, i.e., non-header successor.
Region->setOneSuccessor(isHeaderVPBB(Successor0) ? Successor1 : Successor0);
Region->setExiting(VPBB);

401

Should the parents of Region and ExitVPBB already be set, presumably to the same parent?
(As will be asserted when setting one to be the successor of the other.)

[Ayal]

Title should indicate this is an HCFG simplification patch.

[fhahn]

Address latest comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

158–182

Yes, but that would require 2 lookups I think?

185

Adjusted, thanks!

188

added, thanks!

194

this will be done by setEntry

195

It should be guaranteed to be null for TheLoop->getParentLoop(), as no entries for it will be added.

197

adjusted, thanks!

343

reordered, thanks!

396

Adjusted, thanks!

401

Yep, parent of the region is already set, removed, thanks!

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

86

Perhaps the following would more clearly update the current code by handling the exit block case separately, than during traversal over all predecessors:

if (auto *LatchBB = getLatchOfExit(BB)) {
  auto *PredRegion = getOrCreateVPBB(LatchBB)->getParent();
  assert(VPBB == cast<VPBasicBlock>(PredRegion->getSingleSuccessor()) &&
         "successor must already be set for PredRegion; it must have VPBB "
         "as single successor");
  VPBB->setPredecessors({PredRegion});
  return;
}
// Collect VPBB predecessors.
SmallVector<VPBlockBase *, 2> VPBBPreds;
for (BasicBlock *Pred : predecessors(BB))
  VPBBPreds.push_back(getOrCreateVPBB(Pred));
VPBB->setPredecessors(VPBBPreds);

where a lambda or function would check isExitBB(), analogous to isHeaderBB(), possibly along with providing the Latch:

BasicBlock* getLatchOfExit(BasicBlock *BB) {
  auto *SinglePred = BB->getSinglePredecessor();
  if (!SinglePred || LI->getLoopFor(SinglePred) == LI->getLoopFor(BB))
    return nullptr;
  // can assert SinglePred is the latch of its loop.
  return SinglePred;
}

?

107

Worth augmenting the comment with a note that this loop not only searches for the preheader (for which it could early-break once found or employ find_if()) but also (gets or) creates VPBB's for (all) its latch(es).

This loop over predecessors has a trip count of 2; would unrolling it be clearer?

assert(pred_size(BB) == 2 && "A preheader and single latch are expected");
auto PredBBI = pred_begin(BB);
auto *PredVPBB0 = getOrCreateVPBB(*PredBBI);
auto *PredVPBB1 = getOrCreateVPBB(*++PredBBI);
int Pred0InRegion = PredVPBB0->getParent() == Region;
assert(Pred0InRegion ^ PredVPBB1->getParent() == Region &&
       "Expecting latch to be in region but not preheader");
Region->setPredecessors({Pred0InRegion ? PredVPBB1 : PredVPBB0});

158–182

Yes, but such 2 lookups should be folded by the compiler rather than the programmer...
We do use contains()/[] below for Loop2Region.
Would this work better:

// Retrieve VPBB if exists.
if (auto BlockIt = BB2VPBB.find(BB) != BB2VPBB.end())
  return BlockIt->second;

?

192–194

396

Note that checking if BB == LoopForBB->getLoopLatch() implies a single latch, compared to checking if LoopForBB->isLoopLatch(BB). This is fine, and independent of handling loops with multiple exiting blocks.

[fhahn]

Address latest comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

86

update, thanks!

107

Now that we liberally use loop info, we can just use that here as well to identify the loop predecessor.

158–182

Yes, but such 2 lookups should be folded by the compiler rather than the programmer...

Agreed, but I think in this case the current compilers won't be able to do so :(

Would this work better:

I don't think this will work, as BlockIt would be a bool I think (result of the compare).

Probably better to use lookup, which returns nullptr if the key doesn't exist.

192–194

I think the current coding standard recommends using branches for both if and else when either requires them (https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements). Left the code as is for now.

396

Ack!

[Ayal]

This looks good to me, thanks for accommodating! Adding a couple of last minor comments.

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

90

This comment should be replaced by an assert, or dropped.

Note (somewhere) that exit blocks are assumed to have a single predecessor.

107

Using LoopOfBB->getLoopPredecessor() is indeed much simpler!
Note that the latch VPBB will now be created (potentially) later when visiting its predecessor rather than here when visiting its successor (the header).

158–182

Yes, lookup looks fine!
The default value constructed in case of absence would be a nullptr, given that value type is VPBasicBlock*.
(A parenthesis might make BlockIt the desired iterator.)

189

Use lookup here as well instead of contains and []?
The default value constructed in case of absence would be a nullptr, given that value type is VPRegionBlock*.

196

Worth a comment, something like:
// If BB's loop is nested inside another loop within VPlan's scope, the header of that enclosing loop was already visited and its region constructed and recorded in Loop2Region. That region is now set as the parent of VPBB's region. Otherwise it is set to null.

343–345

Very well, let's be consistent.

[fhahn]

Thanks!

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

90

Right, adjusted, thanks!

189

Updated, thanks!

196

added, thanks!

llvm/lib/Transforms/Vectorize/VPlan.h

VPBlockBase *getSingleHierarchicalPredecessor() {
     return getEnclosingBlockWithPredecessors()->getSinglePredecessor();
   }
 
   /// Set a given VPBlockBase \p Successor as the single successor of this
   /// VPBlockBase. This VPBlockBase is not added as predecessor of \p Successor.
   /// This VPBlockBase must have no successors.
   void setOneSuccessor(VPBlockBase *Successor) {
     assert(Successors.empty() && "Setting one successor when others exist.");
+    assert(Successor->getParent() == getParent() &&
+           "connected blocks must have the same parent");
     appendSuccessor(Successor);
   }
 
   /// Set two given VPBlockBases \p IfTrue and \p IfFalse to be the two
   /// successors of this VPBlockBase. This VPBlockBase is not added as
   /// predecessor of \p IfTrue or \p IfFalse. This VPBlockBase must have no
   /// successors.
   void setTwoSuccessors(VPBlockBase *IfTrue, VPBlockBase *IfFalse) {

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

private:
   // Builder of the VPlan instruction-level representation.
   VPBuilder VPIRBuilder;
 
   // NOTE: The following maps are intentionally destroyed after the plain CFG
   // construction because subsequent VPlan-to-VPlan transformation may
   // invalidate them.
   // Map incoming BasicBlocks to their newly-created VPBasicBlocks.
   DenseMap<BasicBlock *, VPBasicBlock *> BB2VPBB;
   // Map incoming Value definitions to their newly-created VPValues.
   DenseMap<Value *, VPValue *> IRDef2VPValue;
 
   // Hold phi node's that need to be fixed once the plain CFG has been built.
   SmallVector<PHINode *, 8> PhisToFix;
 
   /// Maps loops in the original IR to their corresponding region.
   DenseMap<Loop *, VPRegionBlock *> Loop2Region;
 
   // Utility functions.
   void setVPBBPredsFromBB(VPBasicBlock *VPBB, BasicBlock *BB);
+  void setRegionPredsFromBB(VPRegionBlock *VPBB, BasicBlock *BB);
   void fixPhiNodes();
+  VPBlockBase *getOrCreateVPB(BasicBlock *BB);
   VPBasicBlock *getOrCreateVPBB(BasicBlock *BB);
 #ifndef NDEBUG
   bool isExternalDef(Value *Val);
 #endif
   VPValue *getOrCreateVPOperand(Value *IRVal);
   void createVPInstructionsForVPBB(VPBasicBlock *VPBB, BasicBlock *BB);
 
 public:
   PlainCFGBuilder(Loop *Lp, LoopInfo *LI, VPlan &P)
       : TheLoop(Lp), LI(LI), Plan(P) {}
 
   /// Build plain CFG for TheLoop  and connects it to Plan's entry.
   void buildPlainCFG();
 };
 } // anonymous namespace
 
 // Set predecessors of \p VPBB in the same order as they are in \p BB. \p VPBB
 // must have no predecessors.
 void PlainCFGBuilder::setVPBBPredsFromBB(VPBasicBlock *VPBB, BasicBlock *BB) {
   SmallVector<VPBlockBase *, 8> VPBBPreds;
   // Collect VPBB predecessors.
-  for (BasicBlock *Pred : predecessors(BB))
-    VPBBPreds.push_back(getOrCreateVPBB(Pred));
+  for (BasicBlock *Pred : predecessors(BB)) {
+    auto *PredVPBB = getOrCreateVPBB(Pred);
+    auto *PredRegion = PredVPBB->getParent();
+    if (PredRegion == VPBB->getParent())
+      VPBBPreds.push_back(PredVPBB);
+    else {
+      // VPBB is the exit block of PredVPBB's region. Its predecessor is that
+      // region.
+      VPBBPreds.push_back(PredRegion);
+      assert(VPBB == cast<VPBasicBlock>(PredRegion->getSingleSuccessor()));
+    }
+  }
 
   VPBB->setPredecessors(VPBBPreds);
 }
 
+void PlainCFGBuilder::setRegionPredsFromBB(VPRegionBlock *Region,
+                                           BasicBlock *BB) {
+  // If VPBB is a region block, we are dealing with a loop header block. Connect
+  // the region to the loop preheader, identified by checking the number of
+  // successors (the latch won't have any successors yet, as it has not been
+  // processed at this time.
+  VPBasicBlock *PreheaderVPBB = nullptr;
+  for (BasicBlock *Pred : predecessors(BB)) {
+    auto *PredVPBB = getOrCreateVPBB(Pred);
+    if (PredVPBB->getNumSuccessors() > 0) {
+      assert(!PreheaderVPBB && "more than one preheader?");
+      PreheaderVPBB = PredVPBB;
+    }
+  }
+
+  Region->setPredecessors({PreheaderVPBB});
+}
+
 // Add operands to VPInstructions representing phi nodes from the input IR.
 void PlainCFGBuilder::fixPhiNodes() {
   for (auto *Phi : PhisToFix) {
     assert(IRDef2VPValue.count(Phi) && "Missing VPInstruction for PHINode.");
     VPValue *VPVal = IRDef2VPValue[Phi];
     assert(isa<VPWidenPHIRecipe>(VPVal) &&
            "Expected WidenPHIRecipe for phi node.");
     auto *VPPhi = cast<VPWidenPHIRecipe>(VPVal);
     assert(VPPhi->getNumOperands() == 0 &&
            "Expected VPInstruction with no operands.");
 
     for (unsigned I = 0; I != Phi->getNumOperands(); ++I)
       VPPhi->addIncoming(getOrCreateVPOperand(Phi->getIncomingValue(I)),
                          BB2VPBB[Phi->getIncomingBlock(I)]);
   }
 }
 
+static bool isHeaderVPBB(VPBasicBlock *VPBB) {
+  return VPBB->getParent() && VPBB->getParent()->getEntry() == VPBB;
+}
+
+VPBlockBase *PlainCFGBuilder::getOrCreateVPB(BasicBlock *BB) {
+  VPBasicBlock *VPBB = getOrCreateVPBB(BB);
+  if (isHeaderVPBB(VPBB))
+    return VPBB->getParent();
+  return VPBB;
+}
+
 // Create a new empty VPBasicBlock for an incoming BasicBlock in the region
 // corresponding to the containing loop  or retrieve an existing one if it was
 // already created. If no region exists yet for the loop containing \p BB, a new
 // one is created.
 VPBasicBlock *PlainCFGBuilder::getOrCreateVPBB(BasicBlock *BB) {
   auto BlockIt = BB2VPBB.find(BB);
-  if (BlockIt != BB2VPBB.end())
+  if (BlockIt != BB2VPBB.end()) {
     // Retrieve existing VPBB.
     return BlockIt->second;
+  }
+
+  // Create new VPBB.
+  LLVM_DEBUG(dbgs() << "Creating VPBasicBlock for " << BB->getName() << "\n");
+  VPBasicBlock *VPBB = new VPBasicBlock(BB->getName());
+  BB2VPBB[BB] = VPBB;
 
   // Get or create a region for the loop containing BB.
   Loop *CurrentLoop = LI->getLoopFor(BB);
   VPRegionBlock *ParentR = nullptr;
   if (CurrentLoop) {
     auto Iter = Loop2Region.insert({CurrentLoop, nullptr});
     if (Iter.second)
       Iter.first->second = new VPRegionBlock(
           CurrentLoop->getHeader()->getName().str(), false /*isReplicator*/);
     ParentR = Iter.first->second;
-  }
-
-  // Create new VPBB.
-  LLVM_DEBUG(dbgs() << "Creating VPBasicBlock for " << BB->getName() << "\n");
-  VPBasicBlock *VPBB = new VPBasicBlock(BB->getName());
-  BB2VPBB[BB] = VPBB;
+    // Set the region's parent to the region of the parent loop unless it is
+    // TheLoop.
+    if (CurrentLoop != TheLoop)
+      ParentR->setParent(Loop2Region[CurrentLoop->getParentLoop()]);
+    if (Iter.second)
+      ParentR->setEntry(VPBB);
+    else
       VPBB->setParent(ParentR);
+  }
   return VPBB;
 }
 
 #ifndef NDEBUG
 // Return true if \p Val is considered an external definition. An external
 // definition is either:
 // 1. A Value that is not an Instruction. This will be refined in the future.
 // 2. An Instruction that is outside of the CFG snippet represented in VPlan,
 // i.e., is not part of: a) the loop nest, b) outermost loop PH and, c)
 // outermost loop exits.
 bool PlainCFGBuilder::isExternalDef(Value *Val) {
   // All the Values that are not Instructions are considered external
   // definitions for now.
   Instruction *Inst = dyn_cast<Instruction>(Val);
   if (!Inst)
     return true;
 
   BasicBlock *InstParent = Inst->getParent();
   assert(InstParent && "Expected instruction parent.");
 
   // Check whether Instruction definition is in loop PH.
   BasicBlock *PH = TheLoop->getLoopPreheader();
void PlainCFGBuilder::buildPlainCFG() {
   // based on predecessors traversal.
 
   // Loop PH needs to be explicitly visited since it's not taken into account by
   // LoopBlocksDFS.
   BasicBlock *ThePreheaderBB = TheLoop->getLoopPreheader();
   assert((ThePreheaderBB->getTerminator()->getNumSuccessors() == 1) &&
          "Unexpected loop preheader");
   VPBasicBlock *ThePreheaderVPBB = Plan.getEntry();
   BB2VPBB[ThePreheaderBB] = ThePreheaderVPBB;
   ThePreheaderVPBB->setName("vector.ph");
   for (auto &I : *ThePreheaderBB) {
     if (I.getType()->isVoidTy())
       continue;
     IRDef2VPValue[&I] = Plan.getVPValueOrAddLiveIn(&I);
   }
-  // Create empty VPBB for Loop H so that we can link PH->H.
-  VPBlockBase *HeaderVPBB = getOrCreateVPBB(TheLoop->getHeader());
-  HeaderVPBB->setName("vector.body");
-  ThePreheaderVPBB->setOneSuccessor(HeaderVPBB);
+  // Create region (and header block) for the outer loop, so that we can link
+  // PH->Region.
+  auto *TopRegion = cast<VPRegionBlock>(getOrCreateVPB(TheLoop->getHeader()));
+  ThePreheaderVPBB->setOneSuccessor(TopRegion);
+  TopRegion->getEntry()->setName("vector.body");
 
   LoopBlocksRPO RPO(TheLoop);
   RPO.perform(LI);
 
   for (BasicBlock *BB : RPO) {
     // Create or retrieve the VPBasicBlock for this BB and create its
     // VPInstructions.
     VPBasicBlock *VPBB = getOrCreateVPBB(BB);
+    VPRegionBlock *Region = VPBB->getParent();
     createVPInstructionsForVPBB(VPBB, BB);
 
+    // Set VPBB predecessors in the same order as they are in the incoming BB.
+    if (isHeaderVPBB(VPBB)) {
+      // BB is a loop header, set the predecessor for the region.
+      assert(LI->getLoopFor(BB)->getHeader() == BB);
+      setRegionPredsFromBB(Region, BB);
+    } else
+      setVPBBPredsFromBB(VPBB, BB);
+
     // Set VPBB successors. We create empty VPBBs for successors if they don't
     // exist already. Recipes will be created when the successor is visited
     // during the RPO traversal.
     Instruction *TI = BB->getTerminator();
-    assert(TI && "Terminator expected.");
-    unsigned NumSuccs = TI->getNumSuccessors();
-
-    if (NumSuccs == 1) {
-      VPBasicBlock *SuccVPBB = getOrCreateVPBB(TI->getSuccessor(0));
-      assert(SuccVPBB && "VPBB Successor not found.");
-      VPBB->setOneSuccessor(SuccVPBB);
-    } else if (NumSuccs == 2) {
-      VPBasicBlock *SuccVPBB0 = getOrCreateVPBB(TI->getSuccessor(0));
-      assert(SuccVPBB0 && "Successor 0 not found.");
-      VPBasicBlock *SuccVPBB1 = getOrCreateVPBB(TI->getSuccessor(1));
-      assert(SuccVPBB1 && "Successor 1 not found.");
-
-      // Get VPBB's condition bit.
-      assert(isa<BranchInst>(TI) && "Unsupported terminator!");
-      // Look up the branch condition to get the corresponding VPValue
-      // representing the condition bit in VPlan (which may be in another VPBB).
-      assert(IRDef2VPValue.count(cast<BranchInst>(TI)->getCondition()) &&
+    assert(isa<BranchInst>(TI) && "BranchInst terminator expected.");
+    SmallVector<VPBlockBase *> Successors;
+    for (BasicBlock *Succ : successors(BB))
+      Successors.push_back(getOrCreateVPB(Succ));
+    assert((Successors.size() == 1 || Successors.size() == 2) &&
+           "block must have one or two successors");
+    assert(Successors.size() == 1 ||
+           IRDef2VPValue.count(cast<BranchInst>(TI)->getCondition()) &&
                "Missing condition bit in IRDef2VPValue!");
 
-      // Link successors.
-      VPBB->setTwoSuccessors(SuccVPBB0, SuccVPBB1);
-    } else
-      llvm_unreachable("Number of successors not supported.");
+    // If VPBB is the loop latch i.e. one of its predecessor is the region
+    // itself, update the region's exiting block and the successor of the
+    // region.
+    VPBlockBase *VPB = VPBB;
+    if (any_of(Successors,
+               [Region](VPBlockBase *Succ) { return Succ == Region; })) {
+      assert(Successors.size() == 2 && "latch must have 2 successors");
+      VPBlockBase *ExitVPBB =
+          Successors[0] == Region ? Successors[1] : Successors[0];
+      Region->setExiting(VPBB);
+      Region->setParent(ExitVPBB->getParent());
+      Successors = {ExitVPBB};
+      // For the latch, we need to set the successor for the region.
+      VPB = Region;
+    }
 
-    // Set VPBB predecessors in the same order as they are in the incoming BB.
-    setVPBBPredsFromBB(VPBB, BB);
+    // Link successors.
+    if (Successors.size() == 1)
+      VPB->setOneSuccessor(Successors[0]);
+    else if (Successors.size() == 2)
+      VPB->setTwoSuccessors(Successors[0], Successors[1]);
   }
 
   // 2. Process outermost loop exit. We created an empty VPBB for the loop
   // single exit BB during the RPO traversal of the loop body but Instructions
   // weren't visited because it's not part of the the loop.
   BasicBlock *LoopExitBB = TheLoop->getUniqueExitBlock();
   assert(LoopExitBB && "Loops with multiple exits are not supported.");
-  VPBasicBlock *LoopExitVPBB = BB2VPBB[LoopExitBB];
+  VPBasicBlock *LoopExitVPBB = getOrCreateVPBB(LoopExitBB);
   // Loop exit was already set as successor of the loop exiting BB.
   // We only set its predecessor VPBB now.
   setVPBBPredsFromBB(LoopExitVPBB, LoopExitBB);
 
-  // 3. Fix up region blocks for loops. For each loop,
-  //   * use the header block as entry to the corresponding region,
-  //   * use the latch block as exit of the corresponding region,
-  //   * set the region as successor of the loop pre-header, and
-  //   * set the exit block as successor to the region.
-  SmallVector<Loop *> LoopWorkList;
-  LoopWorkList.push_back(TheLoop);
-  while (!LoopWorkList.empty()) {
-    Loop *L = LoopWorkList.pop_back_val();
-    BasicBlock *Header = L->getHeader();
-    BasicBlock *Exiting = L->getLoopLatch();
-    assert(Exiting == L->getExitingBlock() &&
-           "Latch must be the only exiting block");
-    VPRegionBlock *Region = Loop2Region[L];
-    VPBasicBlock *HeaderVPBB = getOrCreateVPBB(Header);
-    VPBasicBlock *ExitingVPBB = getOrCreateVPBB(Exiting);
-
-    // Disconnect backedge and pre-header from header.
-    VPBasicBlock *PreheaderVPBB = getOrCreateVPBB(L->getLoopPreheader());
-    VPBlockUtils::disconnectBlocks(PreheaderVPBB, HeaderVPBB);
-    VPBlockUtils::disconnectBlocks(ExitingVPBB, HeaderVPBB);
-
-    Region->setParent(PreheaderVPBB->getParent());
-    Region->setEntry(HeaderVPBB);
-    VPBlockUtils::connectBlocks(PreheaderVPBB, Region);
-
-    // Disconnect exit block from exiting (=latch) block, set exiting block and
-    // connect region to exit block.
-    VPBasicBlock *ExitVPBB = getOrCreateVPBB(L->getExitBlock());
-    VPBlockUtils::disconnectBlocks(ExitingVPBB, ExitVPBB);
-    Region->setExiting(ExitingVPBB);
-    VPBlockUtils::connectBlocks(Region, ExitVPBB);
-
-    // Queue sub-loops for processing.
-    LoopWorkList.append(L->begin(), L->end());
-  }
-  // 4. The whole CFG has been built at this point so all the input Values must
+  // 3. The whole CFG has been built at this point so all the input Values must
   // have a VPlan couterpart. Fix VPlan phi nodes by adding their corresponding
   // VPlan operands.
   fixPhiNodes();
 }
 
 void VPlanHCFGBuilder::buildPlainCFG() {
   PlainCFGBuilder PCFGBuilder(TheLoop, LI, Plan);
   PCFGBuilder.buildPlainCFG();
 }