diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -1386,6 +1386,45 @@ 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, + std::forward_iterator_tag, VPBasicBlock> { + explicit vpbb_iterator_adaptor(const df_iterator &Other) + : iterator_adaptor_base(Other) {} + + VPBasicBlock &operator*() { return *cast(*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 basicblocks() { + VPRegionBlock *TopRegion = dyn_cast(Entry); + assert((!TopRegion || TopRegion->getNumSuccessors() == 0) && + "We need a single top-level region at the moment"); + assert((!TopRegion || isa(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>> + loop_basicblocks() { + VPRegionBlock *TopRegion = dyn_cast(Entry); + std::function 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. diff --git a/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp --- a/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp +++ b/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp @@ -9,6 +9,7 @@ #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 { @@ -46,7 +47,9 @@ EXPECT_EQ(1u, Entry->getNumSuccessors()); EXPECT_EQ(nullptr, Entry->getCondBit()); - VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock(); + // VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock(); + VPBasicBlock *VecBB = cast(Entry->getSingleSuccessor()); + EXPECT_NE(VecBB, Entry); EXPECT_EQ(7u, VecBB->size()); EXPECT_EQ(2u, VecBB->getNumPredecessors()); EXPECT_EQ(2u, VecBB->getNumSuccessors()); @@ -87,6 +90,25 @@ EXPECT_EQ(IndvarAdd, ICmp->getOperand(0)); EXPECT_EQ(VecBB->getCondBit(), ICmp); + VPRegionBlock *TopRegion = cast(Plan->getEntry()); + { + SmallVector 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 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 DeadInstructions; VPlanTransforms::VPInstructionsToVPRecipes(Plan, &Inductions, @@ -127,7 +149,7 @@ EXPECT_EQ(0u, Entry->getNumPredecessors()); EXPECT_EQ(1u, Entry->getNumSuccessors()); - VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock(); + VPBasicBlock *VecBB = cast(Entry->getSingleSuccessor()); EXPECT_EQ(6u, VecBB->size()); EXPECT_EQ(2u, VecBB->getNumPredecessors()); EXPECT_EQ(2u, VecBB->getNumSuccessors());