[llvm] 1395cd0 - [VPlan] Support multi-exit loops in HCFG builder.

Florian Hahn via llvm-commits llvm-commits at lists.llvm.org
Sat Jan 25 13:55:36 PST 2025


Author: Florian Hahn
Date: 2025-01-25T21:55:15Z
New Revision: 1395cd015f2edf26f8c2567870183d63f4fdd753

URL: https://github.com/llvm/llvm-project/commit/1395cd015f2edf26f8c2567870183d63f4fdd753
DIFF: https://github.com/llvm/llvm-project/commit/1395cd015f2edf26f8c2567870183d63f4fdd753.diff

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

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp
    llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp b/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp
index 32723e5db9c457..5a2e5d7cfee48d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp
+++ b/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

diff  --git a/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp
index dcdaf008e10fe4..d787a6c9771942 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp
+++ b/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


        


More information about the llvm-commits mailing list