[VPlan] Add basicblocks() and loop_basicblocks iterators.

[fhahn]

Summary

This patch adds 2 convenience iterator ranges, to iterate over all
VPBasicBlocks in a plan. Traversing VPRegions is not supported at the
moment, but we do not create such plans at the moment as far as I know.
They come with an assertion, to implement the additional functionality
once required.

[merge_guards_bot]

Build result: FAILURE -
Log files: console-log.txt, CMakeCache.txt

[gilr]

Traversing VPRegions is not supported at the moment but we do not create such plans at the moment as far as I know

I assume you're referring to the native path? (non-native generates regions for scalarized & predicated instructions)

llvm/lib/Transforms/Vectorize/VPlan.h

1404

basicBlocks()?

1418

loopBasicBlocks()?

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.h

1404

I am not sure. I think for most iterator ranges we use lower case naming, with underscores, but I am not sure why. Personally, I think basicBlocks() looks a bit odd. Same for loopBasicBlocks. But I'd happily change the name, if there's a strong preference.

[fhahn]

In D70734#1764103, @gilr wrote:

Traversing VPRegions is not supported at the moment but we do not create such plans at the moment as far as I know

I assume you're referring to the native path? (non-native generates regions for scalarized & predicated instructions)

Ah right, I missed that code path! For my initial use-case (dropping assumes), the iterator needs to traverse the regions as well then. Let's see if I can get the iterators to co-operate. Do you think it's worth to get the existing iterators in as is? There are a few places that could use them, I think.

[gilr]

In D70734#1765693, @fhahn wrote:

In D70734#1764103, @gilr wrote:

Traversing VPRegions is not supported at the moment but we do not create such plans at the moment as far as I know

I assume you're referring to the native path? (non-native generates regions for scalarized & predicated instructions)

Ah right, I missed that code path! For my initial use-case (dropping assumes), the iterator needs to traverse the regions as well then. Let's see if I can get the iterators to co-operate. Do you think it's worth to get the existing iterators in as is? There are a few places that could use them, I think.

Probably best to bundle it with at least one concrete use, to avoid potential bit rot.

llvm/lib/Transforms/Vectorize/VPlan.h

1404

I think for most iterator ranges we use lower case naming, with underscores

ah, then please keep the current names.

[fhahn]

This requires rebasing and some more work.

[fhahn]

I put up a more general version up in D100175

llvm/lib/Transforms/Vectorize/VPlan.h

VPValue *getVPValue(Value *V) {
     assert(Value2VPValue.count(V) && "Value does not exist in VPlan");
     return Value2VPValue[V];
   }
 
   /// Return the VPLoopInfo analysis for this VPlan.
   VPLoopInfo &getVPLoopInfo() { return VPLInfo; }
   const VPLoopInfo &getVPLoopInfo() const { return VPLInfo; }
 
+  /// An iterator adaptor that converts VPBlockBase * elements to VPBasicBlock
+  /// &. The underlying iterator has to ensure that only VPBasicBlock blocks are
+  /// accessed.
+  struct vpbb_iterator_adaptor
+      : iterator_adaptor_base<vpbb_iterator_adaptor, df_iterator<VPBlockBase *>,
+                              std::forward_iterator_tag, VPBasicBlock> {
+    explicit vpbb_iterator_adaptor(const df_iterator<VPBlockBase *> &Other)
+        : iterator_adaptor_base(Other) {}
+
+    VPBasicBlock &operator*() { return *cast<VPBasicBlock>(*I); }
+  };
+
+  /// Returns an iterator range iterating over all VPBasicBlocks in the Plan in
+  /// depth-first order, including entry (preheader) and exit blocks, as well as
+  /// the loop blocks.
+  iterator_range<vpbb_iterator_adaptor> basicblocks() {
+    VPRegionBlock *TopRegion = dyn_cast<VPRegionBlock>(Entry);
+    assert((!TopRegion || TopRegion->getNumSuccessors() == 0) &&
+           "We need a single top-level region at the moment");
+    assert((!TopRegion || isa<VPBasicBlock>(TopRegion->getEntry())) &&
+           "Nested regions not supported at the moment");
+    VPBlockBase *EntryBB = Entry->getEntryBasicBlock();
+    return make_range(df_begin(EntryBB), df_end(EntryBB));
+  }
+
+  /// Returns an iterator range iterating over all VPBasicBlocks in the loop
+  /// body. This skips the entry (preheader) and exit blocks.
+  iterator_range<filter_iterator<vpbb_iterator_adaptor,
+                                 std::function<bool(VPBlockBase &)>>>
+  loop_basicblocks() {
+    VPRegionBlock *TopRegion = dyn_cast<VPRegionBlock>(Entry);
+    std::function<bool(VPBlockBase &)> Pred = [TopRegion](VPBlockBase &B) {
+      return !TopRegion || &B != TopRegion->getExit();
+    };
+    auto BBRange = basicblocks();
+    return make_filter_range(
+        make_range(std::next(BBRange.begin()), BBRange.end()), Pred);
+  }
+
 private:
   /// Add to the given dominator tree the header block and every new basic block
   /// that was created between it and the latch block, inclusive.
   static void updateDominatorTree(DominatorTree *DT,
                                   BasicBlock *LoopPreHeaderBB,
                                   BasicBlock *LoopLatchBB);
 };
 

llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

 //===- llvm/unittest/Transforms/Vectorize/VPlanHCFGTest.cpp ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "../lib/Transforms/Vectorize/VPlan.h"
 #include "../lib/Transforms/Vectorize/VPlanTransforms.h"
 #include "VPlanTestBase.h"
+#include "llvm/ADT/SmallVector.h"
 #include "gtest/gtest.h"
 
 namespace llvm {
 namespace {
 
 class VPlanHCFGTest : public VPlanTestBase {};
 
 TEST_F(VPlanHCFGTest, testBuildHCFGInnerLoop) {
TEST_F(VPlanHCFGTest, testBuildHCFGInnerLoop) {
   auto Plan = buildHCFG(LoopHeader);
 
   VPBasicBlock *Entry = Plan->getEntry()->getEntryBasicBlock();
   EXPECT_NE(nullptr, Entry->getSingleSuccessor());
   EXPECT_EQ(0u, Entry->getNumPredecessors());
   EXPECT_EQ(1u, Entry->getNumSuccessors());
   EXPECT_EQ(nullptr, Entry->getCondBit());
 
-  VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock();
+  // VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock();
+  VPBasicBlock *VecBB = cast<VPBasicBlock>(Entry->getSingleSuccessor());
+  EXPECT_NE(VecBB, Entry);
   EXPECT_EQ(7u, VecBB->size());
   EXPECT_EQ(2u, VecBB->getNumPredecessors());
   EXPECT_EQ(2u, VecBB->getNumSuccessors());
 
   auto Iter = VecBB->begin();
   VPInstruction *Phi = dyn_cast<VPInstruction>(&*Iter++);
   EXPECT_EQ(Instruction::PHI, Phi->getOpcode());
 
TEST_F(VPlanHCFGTest, testBuildHCFGInnerLoop) {
   EXPECT_EQ(Phi, IndvarAdd->getOperand(0));
 
   VPInstruction *ICmp = dyn_cast<VPInstruction>(&*Iter++);
   EXPECT_EQ(Instruction::ICmp, ICmp->getOpcode());
   EXPECT_EQ(2u, ICmp->getNumOperands());
   EXPECT_EQ(IndvarAdd, ICmp->getOperand(0));
   EXPECT_EQ(VecBB->getCondBit(), ICmp);
 
+  VPRegionBlock *TopRegion = cast<VPRegionBlock>(Plan->getEntry());
+  {
+    SmallVector<VPBasicBlock *, 2> FromIterator;
+    for (VPBasicBlock &VPBB : Plan->basicblocks())
+      FromIterator.push_back(&VPBB);
+    EXPECT_EQ(FromIterator.size(), 3ul);
+    EXPECT_EQ(FromIterator[0], TopRegion->getEntry());
+    EXPECT_EQ(FromIterator[1], VecBB);
+    EXPECT_EQ(FromIterator[2], TopRegion->getExit());
+  }
+
+  {
+    SmallVector<VPBasicBlock *, 2> FromIterator;
+    for (VPBasicBlock &VPBB : Plan->loop_basicblocks())
+      FromIterator.push_back(&VPBB);
+    EXPECT_EQ(FromIterator.size(), 1ul);
+    EXPECT_EQ(FromIterator[0], VecBB);
+  }
+
   LoopVectorizationLegality::InductionList Inductions;
   SmallPtrSet<Instruction *, 1> DeadInstructions;
   VPlanTransforms::VPInstructionsToVPRecipes(Plan, &Inductions,
                                              DeadInstructions);
 }
 
 TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
   const char *ModuleString =
TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
   VPlanTransforms::VPInstructionsToVPRecipes(Plan, &Inductions,
                                              DeadInstructions);
 
   VPBlockBase *Entry = Plan->getEntry()->getEntryBasicBlock();
   EXPECT_NE(nullptr, Entry->getSingleSuccessor());
   EXPECT_EQ(0u, Entry->getNumPredecessors());
   EXPECT_EQ(1u, Entry->getNumSuccessors());
 
-  VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock();
+  VPBasicBlock *VecBB = cast<VPBasicBlock>(Entry->getSingleSuccessor());
   EXPECT_EQ(6u, VecBB->size());
   EXPECT_EQ(2u, VecBB->getNumPredecessors());
   EXPECT_EQ(2u, VecBB->getNumSuccessors());
 
   auto Iter = VecBB->begin();
   auto *Phi = dyn_cast<VPWidenPHIRecipe>(&*Iter++);
   EXPECT_NE(nullptr, Phi);