[VPlan] Add VPBlockUtils::basicBlocksOnly helper.

[fhahn]

Summary

This patch adds basicBlocksOnly and constBasicBlocksOnly helpers which
take an iterator range over VPBlockBase * and const VPBlockBase *
respectively, and return an iterator range that only includes
VPBasicBlocks.

Unfortunately I could not think of a way to have a single version of
basicBlocksOnly that can take arbitrary iterators and also handle
filtering/mapping of both VPBlockBase * and const VPBlockBase *.
Unfortunately, filter_range does not work with iterators over pointers,
so we need to map it to a reference and back. Mapping back to a pointer
is optional, but more convenient when using the iterator range to
populate containers.

[a.elovikov]

VPlan part looks good to me. I'm not familiar with implementation details of the filtered_iterator to discuss the to_reference->filter->to_pointer approach though. Regardless of that, it's an improvement in interfaces, so LGTM unless someone has a way to make filtering work with pointer directly.

[fhahn]

In D101093#2709842, @a.elovikov wrote:

VPlan part looks good to me. I'm not familiar with implementation details of the filtered_iterator to discuss the to_reference->filter->to_pointer approach though. Regardless of that, it's an improvement in interfaces, so LGTM unless someone has a way to make filtering work with pointer directly.

Thanks for taking a look! If it turns out there's a better way to implement the interface I'm move than happy to adjust this, even after it goes in :)

llvm/lib/Transforms/Vectorize/VPlan.h

static unsigned countSuccessorsNoBE(VPBlockBase *PredBlock,
                                       VPLoopInfo *VPLI) {
     unsigned Count = 0;
     for (VPBlockBase *SuccBlock : PredBlock->getSuccessors()) {
       if (!VPBlockUtils::isBackEdge(PredBlock, SuccBlock, VPLI))
         Count++;
     }
     return Count;
   }
+
+  /// Return an iterator range over \p Iter which only includes \p BlockTy
+  /// blocks. The accesses are casted to \p BlockTy.
+  template <typename BlockTy, typename T> static auto blocksOnly(T Iter) {
+    // We need to first create an iterator range over VPBlockBase & instead of
+    // VPBlockBase * for filter_range to work properly.
+    auto Mapped = map_range(
+        Iter, [](VPBlockBase *Block) -> VPBlockBase & { return *Block; });
+    auto Filter = make_filter_range(
+        Mapped, [](VPBlockBase &Block) { return isa<BlockTy>(&Block); });
+    return map_range(Filter, [](VPBlockBase &Block) -> BlockTy * {
+      return cast<BlockTy>(&Block);
+    });
+  }
 };
 
 class VPInterleavedAccessInfo {
   DenseMap<VPInstruction *, InterleaveGroup<VPInstruction> *>
       InterleaveGroupMap;
 
   /// Type for mapping of instruction based interleave groups to VPInstruction
   /// interleave groups

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

TEST(VPBasicBlockTest, TraversingIteratorTest) {
     EXPECT_EQ(R1BB1, FromIterator[1]);
     EXPECT_EQ(R1BB2, FromIterator[2]);
     EXPECT_EQ(R1BB4, FromIterator[3]);
     EXPECT_EQ(R2, FromIterator[4]);
     EXPECT_EQ(R2BB1, FromIterator[5]);
     EXPECT_EQ(R2BB2, FromIterator[6]);
     EXPECT_EQ(R1BB3, FromIterator[7]);
 
+    // const VPBasicBlocks only.
+    FromIterator.clear();
+    copy(VPBlockUtils::blocksOnly<const VPBasicBlock>(depth_first(Start)),
+         std::back_inserter(FromIterator));
+    EXPECT_EQ(6u, FromIterator.size());
+    EXPECT_EQ(R1BB1, FromIterator[0]);
+    EXPECT_EQ(R1BB2, FromIterator[1]);
+    EXPECT_EQ(R1BB4, FromIterator[2]);
+    EXPECT_EQ(R2BB1, FromIterator[3]);
+    EXPECT_EQ(R2BB2, FromIterator[4]);
+    EXPECT_EQ(R1BB3, FromIterator[5]);
+
+    // VPRegionBlocks only.
+    SmallVector<VPRegionBlock *> FromIteratorVPRegion(
+        VPBlockUtils::blocksOnly<VPRegionBlock>(depth_first(Start)));
+    EXPECT_EQ(2u, FromIteratorVPRegion.size());
+    EXPECT_EQ(R1, FromIteratorVPRegion[0]);
+    EXPECT_EQ(R2, FromIteratorVPRegion[1]);
+
     // Post-order.
     FromIterator.clear();
     copy(post_order(Start), std::back_inserter(FromIterator));
     EXPECT_EQ(8u, FromIterator.size());
     EXPECT_EQ(R2BB2, FromIterator[0]);
     EXPECT_EQ(R2BB1, FromIterator[1]);
     EXPECT_EQ(R2, FromIterator[2]);
     EXPECT_EQ(R1BB4, FromIterator[3]);
TEST(VPBasicBlockTest, TraversingIteratorTest) {
     EXPECT_EQ(VPBB1, FromIterator[0]);
     EXPECT_EQ(R1, FromIterator[1]);
     EXPECT_EQ(R2, FromIterator[2]);
     EXPECT_EQ(R2BB1, FromIterator[3]);
     EXPECT_EQ(R3, FromIterator[4]);
     EXPECT_EQ(R3BB1, FromIterator[5]);
     EXPECT_EQ(VPBB2, FromIterator[6]);
 
+    SmallVector<VPBlockBase *> FromIteratorVPBB;
+    copy(VPBlockUtils::blocksOnly<VPBasicBlock>(depth_first(Start)),
+         std::back_inserter(FromIteratorVPBB));
+    EXPECT_EQ(VPBB1, FromIteratorVPBB[0]);
+    EXPECT_EQ(R2BB1, FromIteratorVPBB[1]);
+    EXPECT_EQ(R3BB1, FromIteratorVPBB[2]);
+    EXPECT_EQ(VPBB2, FromIteratorVPBB[3]);
+
     // Post-order.
     FromIterator.clear();
     copy(post_order(Start), std::back_inserter(FromIterator));
     EXPECT_EQ(7u, FromIterator.size());
     EXPECT_EQ(VPBB2, FromIterator[0]);
     EXPECT_EQ(R3BB1, FromIterator[1]);
     EXPECT_EQ(R3, FromIterator[2]);
     EXPECT_EQ(R2BB1, FromIterator[3]);
     EXPECT_EQ(R2, FromIterator[4]);
     EXPECT_EQ(R1, FromIterator[5]);
     EXPECT_EQ(VPBB1, FromIterator[6]);
 
+    // Post-order, const VPRegionBlocks only.
+    SmallVector<const VPRegionBlock *> FromIteratorVPRegion(
+        VPBlockUtils::blocksOnly<const VPRegionBlock>(post_order(Start)));
+    EXPECT_EQ(3u, FromIteratorVPRegion.size());
+    EXPECT_EQ(R3, FromIteratorVPRegion[0]);
+    EXPECT_EQ(R2, FromIteratorVPRegion[1]);
+    EXPECT_EQ(R1, FromIteratorVPRegion[2]);
+
+    // Post-order, VPBasicBlocks only.
+    FromIterator.clear();
+    copy(VPBlockUtils::blocksOnly<VPBasicBlock>(post_order(Start)),
+         std::back_inserter(FromIterator));
+    EXPECT_EQ(FromIterator.size(), 4u);
+    EXPECT_EQ(VPBB2, FromIterator[0]);
+    EXPECT_EQ(R3BB1, FromIterator[1]);
+    EXPECT_EQ(R2BB1, FromIterator[2]);
+    EXPECT_EQ(VPBB1, FromIterator[3]);
+
     // Use Plan to properly clean up created blocks.
     VPlan Plan;
     Plan.setEntry(VPBB1);
   }
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 TEST(VPBasicBlockTest, print) {