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<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. 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<VPBasicBlock>(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<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, @@ -127,7 +149,7 @@ 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());