[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.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

370–371

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

370–371

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.

375

adjusted, thanks!

397

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

400

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

[fhahn]

Undo uneeded change

[fhahn]

ping :)

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

113

Comment needs updating.

116

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.

116

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.

165–166

CurrentLoop >> LoopOfBB
ParentR >> RegionOfVPBB

165–166

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});
}

170

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

174

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

348

Simpler to do something like

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

?

368

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().

370–371

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

374

else {

Already asserted above that numSuccs is 1 or 2.

375

Make sure first successor corresponds to True?

379

nit: count >> contains

382

"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?

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

333

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

113

Updated, thanks!

116

Reordered assert, added message, thanks!

116

Adjusted to check the parent region, thanks!

165–166

Renamed, thanks!

165–166

adjusted, thanks!

170

Moved the code to the earlier if

174

That should be handled by setEntry IIUC

333

Added a comment, thanks!

348

Simplified, thanks!

368

Flipped, thanks!

374

Adjusted, thanks!

375

Added assert, thanks!

379

Adjusted, thanks!

382

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

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

159

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

165–166

nit: BlockIt is not needed below, can simply check

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

175

nit: can early exit here via

if (!LoopOfBB)
  return VPBB;

and define RegionOfVPBB inside then/else when setting it.

178

nit: assert BB-is-header xor RegionExists?

184

VPBB->setParent(RegionOfVPBB);
here too?

185

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.

187

357

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

385

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);

390

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

165–166

Yes, but that would require 2 lookups I think?

175

Adjusted, thanks!

178

added, thanks!

184

this will be done by setEntry

185

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

187

adjusted, thanks!

357

reordered, thanks!

385

Adjusted, thanks!

390

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

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

85–103

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;
}

?

114

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});

165–166

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;

?

182–184

385

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

85–103

update, thanks!

114

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

165–166

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.

182–184

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.

385

Ack!

[Ayal]

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

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

89

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

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

114

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).

165–166

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.)

179

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*.

186

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.

357–359

Very well, let's be consistent.

[fhahn]

Thanks!

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

89

Right, adjusted, thanks!

179

Updated, thanks!

186

added, thanks!