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

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

@@ -350,10 +350,25 @@ void PlainCFGBuilder::buildPlainCFG() {
   // new vector preheader); here we're interested in setting BB2VPBB to the
   // latter.
   BB2VPBB[ThePreheaderBB] = VectorPreheaderVPBB;
-  BasicBlock *LoopExitBB = TheLoop->getUniqueExitBlock();
   Loop2Region[LI->getLoopFor(TheLoop->getHeader())] = TheRegion;
-  assert(LoopExitBB && "Loops with multiple exits are not supported.");
-  BB2VPBB[LoopExitBB] = cast<VPBasicBlock>(TheRegion->getSingleSuccessor());
+  BasicBlock *ExitBB = TheLoop->getUniqueExitBlock();
+  if (!ExitBB) {
+    // If there is no unique exit block, we must exit via the latch. This exit
+    // is mapped to the middle block in the input plan.
+    BasicBlock *Latch = TheLoop->getLoopLatch();
+    auto *Br = cast<BranchInst>(Latch->getTerminator());
+    if (TheLoop->contains(Br->getSuccessor(0))) {
+      assert(!TheLoop->contains(Br->getSuccessor(1)) &&
+             "latch must exit the loop");
+      ExitBB = Br->getSuccessor(1);
+    } else {
+      assert(!TheLoop->contains(Br->getSuccessor(0)) &&
+             "latch must exit the loop");
+      ExitBB = Br->getSuccessor(0);
+    }
+  }
+  assert(ExitBB && "Must have a unique exit block or also exit via the latch.");
+  BB2VPBB[ExitBB] = cast<VPBasicBlock>(TheRegion->getSingleSuccessor());
 
   // The existing vector region's entry and exiting VPBBs correspond to the loop
   // header and latch.
@@ -423,21 +438,38 @@ void PlainCFGBuilder::buildPlainCFG() {
     // 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 (!LoopForBB || BB != LoopForBB->getLoopLatch()) {
-      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,
-    // except for the top region, whose successor was set when creating VPlan's
-    // skeleton.
-    if (TheRegion != Region) {
+
+    BasicBlock *IRSucc0 = BI->getSuccessor(0);
+    BasicBlock *IRSucc1 = BI->getSuccessor(1);
+    VPBasicBlock *Successor0 = getOrCreateVPBB(IRSucc0);
+    VPBasicBlock *Successor1 = getOrCreateVPBB(IRSucc1);
+    if (BB == LoopForBB->getLoopLatch()) {
+      // 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,
+      // except for the top region, whose successor was set when creating
+      // VPlan's skeleton.
+      assert(TheRegion != Region &&
+             "Latch of the top region should have been handled earlier");
       Region->setOneSuccessor(isHeaderVPBB(Successor0) ? Successor1
                                                        : Successor0);
       Region->setExiting(VPBB);
+      continue;
     }
+
+    // Don't connect any blocks outside the current loop except the latch for
+    // now. The latch is handled above.
+    if (LoopForBB) {
+      if (!LoopForBB->contains(IRSucc0)) {
+        VPBB->setOneSuccessor(Successor1);
+        continue;
+      }
+      if (!LoopForBB->contains(IRSucc1)) {
+        VPBB->setOneSuccessor(Successor0);
+        continue;
+      }
+    }
+
+    VPBB->setTwoSuccessors(Successor0, Successor1);
   }
 
   // 2. The whole CFG has been built at this point so all the input Values must

llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

@@ -234,5 +234,114 @@ TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
   EXPECT_EQ(VecBB->end(), Iter);
 }
 
+TEST_F(VPlanHCFGTest, testBuildHCFGInnerLoopMultiExit) {
+  const char *ModuleString =
+      "define void @f(ptr %A, i64 %N) {\n"
+      "entry:\n"
+      "  br label %loop.header\n"
+      "loop.header:\n"
+      "  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.latch ]\n"
+      "  %arr.idx = getelementptr inbounds i32, ptr %A, i64 %iv\n"
+      "  %l1 = load i32, ptr %arr.idx, align 4\n"
+      "  %c = icmp eq i32 %l1, 0\n"
+      "  br i1 %c, label %exit.1, label %loop.latch\n"
+      "loop.latch:\n"
+      "  %res = add i32 %l1, 10\n"
+      "  store i32 %res, ptr %arr.idx, align 4\n"
+      "  %iv.next = add i64 %iv, 1\n"
+      "  %exitcond = icmp ne i64 %iv.next, %N\n"
+      "  br i1 %exitcond, label %loop.header, label %exit.2\n"
+      "exit.1:\n"
+      "  ret void\n"
+      "exit.2:\n"
+      "  ret void\n"
+      "}\n";
+
+  Module &M = parseModule(ModuleString);
+
+  Function *F = M.getFunction("f");
+  BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor();
+  auto Plan = buildHCFG(LoopHeader);
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  // Add an external value to check we do not print the list of external values,
+  // as this is not required with the new printing.
+  Plan->getOrAddLiveIn(&*F->arg_begin());
+  std::string FullDump;
+  raw_string_ostream OS(FullDump);
+  Plan->printDOT(OS);
+  const char *ExpectedStr = R"(digraph VPlan {
+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"]
+node [shape=rect, fontname=Courier, fontsize=30]
+edge [fontname=Courier, fontsize=30]
+compound=true
+  N0 [label =
+    "ir-bb\<entry\>:\l" +
+    "Successor(s): vector.ph\l"
+  ]
+  N0 -> N1 [ label=""]
+  N1 [label =
+    "vector.ph:\l" +
+    "Successor(s): vector loop\l"
+  ]
+  N1 -> N2 [ label="" lhead=cluster_N3]
+  subgraph cluster_N3 {
+    fontname=Courier
+    label="\<x1\> vector loop"
+    N2 [label =
+      "vector.body:\l" +
+      "  WIDEN-PHI ir\<%iv\> = phi ir\<0\>, ir\<%iv.next\>\l" +
+      "  EMIT ir\<%arr.idx\> = getelementptr ir\<%A\>, ir\<%iv\>\l" +
+      "  EMIT ir\<%l1\> = load ir\<%arr.idx\>\l" +
+      "  EMIT ir\<%c\> = icmp ir\<%l1\>, ir\<0\>\l" +
+      "Successor(s): loop.latch\l"
+    ]
+    N2 -> N4 [ label=""]
+    N4 [label =
+      "loop.latch:\l" +
+      "  EMIT ir\<%res\> = add ir\<%l1\>, ir\<10\>\l" +
+      "  EMIT store ir\<%res\>, ir\<%arr.idx\>\l" +
+      "  EMIT ir\<%iv.next\> = add ir\<%iv\>, ir\<1\>\l" +
+      "  EMIT ir\<%exitcond\> = icmp ir\<%iv.next\>, ir\<%N\>\l" +
+      "Successor(s): vector.latch\l"
+    ]
+    N4 -> N5 [ label=""]
+    N5 [label =
+      "vector.latch:\l" +
+      "No successors\l"
+    ]
+  }
+  N5 -> N6 [ label="" ltail=cluster_N3]
+  N6 [label =
+    "middle.block:\l" +
+    "  EMIT vp\<%cmp.n\> = icmp eq ir\<%N\>, vp\<%0\>\l" +
+    "  EMIT branch-on-cond vp\<%cmp.n\>\l" +
+    "Successor(s): ir-bb\<exit.2\>, scalar.ph\l"
+  ]
+  N6 -> N7 [ label="T"]
+  N6 -> N8 [ label="F"]
+  N7 [label =
+    "ir-bb\<exit.2\>:\l" +
+    "No successors\l"
+  ]
+  N8 [label =
+    "scalar.ph:\l" +
+    "Successor(s): ir-bb\<loop.header\>\l"
+  ]
+  N8 -> N9 [ label=""]
+  N9 [label =
+    "ir-bb\<loop.header\>:\l" +
+    "  IR   %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.latch ]\l" +
+    "  IR   %arr.idx = getelementptr inbounds i32, ptr %A, i64 %iv\l" +
+    "  IR   %l1 = load i32, ptr %arr.idx, align 4\l" +
+    "  IR   %c = icmp eq i32 %l1, 0\l" +
+    "No successors\l"
+  ]
+}
+)";
+  EXPECT_EQ(ExpectedStr, FullDump);
+#endif
+}
+
 } // namespace
 } // namespace llvm