Skip to content

Commit 1395cd0

Browse files
committed
[VPlan] Support multi-exit loops in HCFG builder.
Update HCFG construction to support multi-exit loops. If there is no unique exit block, map the middle block of the initial plan to the exit block from the latch. This further unifies HCFG construction and prepares for use to also build an initial VPlan (VPlan0) for inner loops. Effectively NFC as this isn't used on the default code path yet.
1 parent 2131115 commit 1395cd0

File tree

2 files changed

+155
-14
lines changed

2 files changed

+155
-14
lines changed

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

Lines changed: 46 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -350,10 +350,25 @@ void PlainCFGBuilder::buildPlainCFG() {
350350
// new vector preheader); here we're interested in setting BB2VPBB to the
351351
// latter.
352352
BB2VPBB[ThePreheaderBB] = VectorPreheaderVPBB;
353-
BasicBlock *LoopExitBB = TheLoop->getUniqueExitBlock();
354353
Loop2Region[LI->getLoopFor(TheLoop->getHeader())] = TheRegion;
355-
assert(LoopExitBB && "Loops with multiple exits are not supported.");
356-
BB2VPBB[LoopExitBB] = cast<VPBasicBlock>(TheRegion->getSingleSuccessor());
354+
BasicBlock *ExitBB = TheLoop->getUniqueExitBlock();
355+
if (!ExitBB) {
356+
// If there is no unique exit block, we must exit via the latch. This exit
357+
// is mapped to the middle block in the input plan.
358+
BasicBlock *Latch = TheLoop->getLoopLatch();
359+
auto *Br = cast<BranchInst>(Latch->getTerminator());
360+
if (TheLoop->contains(Br->getSuccessor(0))) {
361+
assert(!TheLoop->contains(Br->getSuccessor(1)) &&
362+
"latch must exit the loop");
363+
ExitBB = Br->getSuccessor(1);
364+
} else {
365+
assert(!TheLoop->contains(Br->getSuccessor(0)) &&
366+
"latch must exit the loop");
367+
ExitBB = Br->getSuccessor(0);
368+
}
369+
}
370+
assert(ExitBB && "Must have a unique exit block or also exit via the latch.");
371+
BB2VPBB[ExitBB] = cast<VPBasicBlock>(TheRegion->getSingleSuccessor());
357372

358373
// The existing vector region's entry and exiting VPBBs correspond to the loop
359374
// header and latch.
@@ -423,21 +438,38 @@ void PlainCFGBuilder::buildPlainCFG() {
423438
// representing the condition bit in VPlan (which may be in another VPBB).
424439
assert(IRDef2VPValue.contains(BI->getCondition()) &&
425440
"Missing condition bit in IRDef2VPValue!");
426-
VPBasicBlock *Successor0 = getOrCreateVPBB(BI->getSuccessor(0));
427-
VPBasicBlock *Successor1 = getOrCreateVPBB(BI->getSuccessor(1));
428-
if (!LoopForBB || BB != LoopForBB->getLoopLatch()) {
429-
VPBB->setTwoSuccessors(Successor0, Successor1);
430-
continue;
431-
}
432-
// For a latch we need to set the successor of the region rather than that
433-
// of VPBB and it should be set to the exit, i.e., non-header successor,
434-
// except for the top region, whose successor was set when creating VPlan's
435-
// skeleton.
436-
if (TheRegion != Region) {
441+
442+
BasicBlock *IRSucc0 = BI->getSuccessor(0);
443+
BasicBlock *IRSucc1 = BI->getSuccessor(1);
444+
VPBasicBlock *Successor0 = getOrCreateVPBB(IRSucc0);
445+
VPBasicBlock *Successor1 = getOrCreateVPBB(IRSucc1);
446+
if (BB == LoopForBB->getLoopLatch()) {
447+
// For a latch we need to set the successor of the region rather than that
448+
// of VPBB and it should be set to the exit, i.e., non-header successor,
449+
// except for the top region, whose successor was set when creating
450+
// VPlan's skeleton.
451+
assert(TheRegion != Region &&
452+
"Latch of the top region should have been handled earlier");
437453
Region->setOneSuccessor(isHeaderVPBB(Successor0) ? Successor1
438454
: Successor0);
439455
Region->setExiting(VPBB);
456+
continue;
440457
}
458+
459+
// Don't connect any blocks outside the current loop except the latch for
460+
// now. The latch is handled above.
461+
if (LoopForBB) {
462+
if (!LoopForBB->contains(IRSucc0)) {
463+
VPBB->setOneSuccessor(Successor1);
464+
continue;
465+
}
466+
if (!LoopForBB->contains(IRSucc1)) {
467+
VPBB->setOneSuccessor(Successor0);
468+
continue;
469+
}
470+
}
471+
472+
VPBB->setTwoSuccessors(Successor0, Successor1);
441473
}
442474

443475
// 2. The whole CFG has been built at this point so all the input Values must

llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,5 +234,114 @@ TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
234234
EXPECT_EQ(VecBB->end(), Iter);
235235
}
236236

237+
TEST_F(VPlanHCFGTest, testBuildHCFGInnerLoopMultiExit) {
238+
const char *ModuleString =
239+
"define void @f(ptr %A, i64 %N) {\n"
240+
"entry:\n"
241+
" br label %loop.header\n"
242+
"loop.header:\n"
243+
" %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.latch ]\n"
244+
" %arr.idx = getelementptr inbounds i32, ptr %A, i64 %iv\n"
245+
" %l1 = load i32, ptr %arr.idx, align 4\n"
246+
" %c = icmp eq i32 %l1, 0\n"
247+
" br i1 %c, label %exit.1, label %loop.latch\n"
248+
"loop.latch:\n"
249+
" %res = add i32 %l1, 10\n"
250+
" store i32 %res, ptr %arr.idx, align 4\n"
251+
" %iv.next = add i64 %iv, 1\n"
252+
" %exitcond = icmp ne i64 %iv.next, %N\n"
253+
" br i1 %exitcond, label %loop.header, label %exit.2\n"
254+
"exit.1:\n"
255+
" ret void\n"
256+
"exit.2:\n"
257+
" ret void\n"
258+
"}\n";
259+
260+
Module &M = parseModule(ModuleString);
261+
262+
Function *F = M.getFunction("f");
263+
BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor();
264+
auto Plan = buildHCFG(LoopHeader);
265+
266+
#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
267+
// Add an external value to check we do not print the list of external values,
268+
// as this is not required with the new printing.
269+
Plan->getOrAddLiveIn(&*F->arg_begin());
270+
std::string FullDump;
271+
raw_string_ostream OS(FullDump);
272+
Plan->printDOT(OS);
273+
const char *ExpectedStr = R"(digraph VPlan {
274+
graph [labelloc=t, fontsize=30; label="Vectorization Plan\n for UF\>=1\nLive-in vp\<%0\> = vector-trip-count\nLive-in ir\<%N\> = original trip-count\n"]
275+
node [shape=rect, fontname=Courier, fontsize=30]
276+
edge [fontname=Courier, fontsize=30]
277+
compound=true
278+
N0 [label =
279+
"ir-bb\<entry\>:\l" +
280+
"Successor(s): vector.ph\l"
281+
]
282+
N0 -> N1 [ label=""]
283+
N1 [label =
284+
"vector.ph:\l" +
285+
"Successor(s): vector loop\l"
286+
]
287+
N1 -> N2 [ label="" lhead=cluster_N3]
288+
subgraph cluster_N3 {
289+
fontname=Courier
290+
label="\<x1\> vector loop"
291+
N2 [label =
292+
"vector.body:\l" +
293+
" WIDEN-PHI ir\<%iv\> = phi ir\<0\>, ir\<%iv.next\>\l" +
294+
" EMIT ir\<%arr.idx\> = getelementptr ir\<%A\>, ir\<%iv\>\l" +
295+
" EMIT ir\<%l1\> = load ir\<%arr.idx\>\l" +
296+
" EMIT ir\<%c\> = icmp ir\<%l1\>, ir\<0\>\l" +
297+
"Successor(s): loop.latch\l"
298+
]
299+
N2 -> N4 [ label=""]
300+
N4 [label =
301+
"loop.latch:\l" +
302+
" EMIT ir\<%res\> = add ir\<%l1\>, ir\<10\>\l" +
303+
" EMIT store ir\<%res\>, ir\<%arr.idx\>\l" +
304+
" EMIT ir\<%iv.next\> = add ir\<%iv\>, ir\<1\>\l" +
305+
" EMIT ir\<%exitcond\> = icmp ir\<%iv.next\>, ir\<%N\>\l" +
306+
"Successor(s): vector.latch\l"
307+
]
308+
N4 -> N5 [ label=""]
309+
N5 [label =
310+
"vector.latch:\l" +
311+
"No successors\l"
312+
]
313+
}
314+
N5 -> N6 [ label="" ltail=cluster_N3]
315+
N6 [label =
316+
"middle.block:\l" +
317+
" EMIT vp\<%cmp.n\> = icmp eq ir\<%N\>, vp\<%0\>\l" +
318+
" EMIT branch-on-cond vp\<%cmp.n\>\l" +
319+
"Successor(s): ir-bb\<exit.2\>, scalar.ph\l"
320+
]
321+
N6 -> N7 [ label="T"]
322+
N6 -> N8 [ label="F"]
323+
N7 [label =
324+
"ir-bb\<exit.2\>:\l" +
325+
"No successors\l"
326+
]
327+
N8 [label =
328+
"scalar.ph:\l" +
329+
"Successor(s): ir-bb\<loop.header\>\l"
330+
]
331+
N8 -> N9 [ label=""]
332+
N9 [label =
333+
"ir-bb\<loop.header\>:\l" +
334+
" IR %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.latch ]\l" +
335+
" IR %arr.idx = getelementptr inbounds i32, ptr %A, i64 %iv\l" +
336+
" IR %l1 = load i32, ptr %arr.idx, align 4\l" +
337+
" IR %c = icmp eq i32 %l1, 0\l" +
338+
"No successors\l"
339+
]
340+
}
341+
)";
342+
EXPECT_EQ(ExpectedStr, FullDump);
343+
#endif
344+
}
345+
237346
} // namespace
238347
} // namespace llvm

0 commit comments

Comments
 (0)