!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp

@@ -386,51 +386,60 @@ std::unique_ptr<VPlan> VPlanTransforms::buildPlainCFG(
 /// Checks if \p HeaderVPB is a loop header block in the plain CFG; that is, it
 /// has exactly 2 predecessors (preheader and latch), where the block
 /// dominates the latch and the preheader dominates the block. If it is a
-/// header block return true, making sure the preheader appears first and
-/// the latch second. Otherwise return false.
-static bool canonicalHeader(VPBlockBase *HeaderVPB,
-                            const VPDominatorTree &VPDT) {
+/// header block return true and canonicalize the predecessors of the header
+/// (making sure the preheader appears first and the latch second) and the
+/// successors of the latch (making sure the loop exit comes first). Otherwise
+/// return false.
+static bool canonicalHeaderAndLatch(VPBlockBase *HeaderVPB,
+                                    const VPDominatorTree &VPDT) {
   ArrayRef<VPBlockBase *> Preds = HeaderVPB->getPredecessors();
   if (Preds.size() != 2)
     return false;
 
   auto *PreheaderVPBB = Preds[0];
   auto *LatchVPBB = Preds[1];
-  if (VPDT.dominates(PreheaderVPBB, HeaderVPB) &&
-      VPDT.dominates(HeaderVPB, LatchVPBB))
-    return true;
-
-  std::swap(PreheaderVPBB, LatchVPBB);
+  if (!VPDT.dominates(PreheaderVPBB, HeaderVPB) ||
+      !VPDT.dominates(HeaderVPB, LatchVPBB)) {
+    std::swap(PreheaderVPBB, LatchVPBB);
 
-  if (VPDT.dominates(PreheaderVPBB, HeaderVPB) &&
-      VPDT.dominates(HeaderVPB, LatchVPBB)) {
-    // Canonicalize predecessors of header so that preheader is first and latch
-    // second.
-    HeaderVPB->swapPredecessors();
-    for (VPRecipeBase &R : cast<VPBasicBlock>(HeaderVPB)->phis())
-      R.swapOperands();
-    return true;
+    if (!VPDT.dominates(PreheaderVPBB, HeaderVPB) ||
+        !VPDT.dominates(HeaderVPB, LatchVPBB)) {
+      return false;
+    } else {
+      // Canonicalize predecessors of header so that preheader is first and
+      // latch second.
+      HeaderVPB->swapPredecessors();
+      for (VPRecipeBase &R : cast<VPBasicBlock>(HeaderVPB)->phis())
+        R.swapOperands();
+    }
   }
 
-  return false;
-}
-
-/// Create a new VPRegionBlock for the loop starting at \p HeaderVPB.
-static void createLoopRegion(VPlan &Plan, VPBlockBase *HeaderVPB) {
-  auto *PreheaderVPBB = HeaderVPB->getPredecessors()[0];
-  auto *LatchVPBB = HeaderVPB->getPredecessors()[1];
-
-  // We are canonicalizing the successors of the latch when introducing the
-  // region. We will exit the region of the latch condition is true; invert the
+  // The two successors of conditional branch match the condition, with the
+  // first successor corresponding to true and the second to false. We
+  // canonicalize the successors of the latch when introducing the region, such
+  // that the latch exits the region when its condition is true; invert the
   // original condition if the original CFG branches to the header on true.
   if (!LatchVPBB->getSingleSuccessor() &&
       LatchVPBB->getSuccessors()[0] == HeaderVPB) {
+    assert(LatchVPBB->getNumSuccessors() == 2 && "Must have 2 successors");
     auto *Term = cast<VPBasicBlock>(LatchVPBB)->getTerminator();
+    assert(cast<VPInstruction>(Term)->getOpcode() ==
+               VPInstruction::BranchOnCond &&
+           "terminator must be a BranchOnCond");
     auto *Not = new VPInstruction(VPInstruction::Not, {Term->getOperand(0)});
     Not->insertBefore(Term);
     Term->setOperand(0, Not);
+    LatchVPBB->swapSuccessors();
   }
 
+  return true;
+}
+
+/// Create a new VPRegionBlock for the loop starting at \p HeaderVPB.
+static void createLoopRegion(VPlan &Plan, VPBlockBase *HeaderVPB) {
+  auto *PreheaderVPBB = HeaderVPB->getPredecessors()[0];
+  auto *LatchVPBB = HeaderVPB->getPredecessors()[1];
+
   VPBlockUtils::disconnectBlocks(PreheaderVPBB, HeaderVPB);
   VPBlockUtils::disconnectBlocks(LatchVPBB, HeaderVPB);
   VPBlockBase *Succ = LatchVPBB->getSingleSuccessor();
@@ -458,7 +467,7 @@ void VPlanTransforms::createLoopRegions(VPlan &Plan, Type *InductionTy,
   VPDominatorTree VPDT;
   VPDT.recalculate(Plan);
   for (VPBlockBase *HeaderVPB : vp_depth_first_shallow(Plan.getEntry()))
-    if (canonicalHeader(HeaderVPB, VPDT))
+    if (canonicalHeaderAndLatch(HeaderVPB, VPDT))
       createLoopRegion(Plan, HeaderVPB);
 
   VPRegionBlock *TopRegion = Plan.getVectorLoopRegion();