[LV] Introduce conditional vector basic block.

This patch add the transformation that convert flatten control flow with conditional vector basic block. This transformation can help program skip masked operations without any active lane. First, this transformation will collect all masked stores and operands bottom-up. And put these msaked operations into a new vector basic block. Second, this transformation will split original vector loop and insert the new basic block between split blocks. And update the conditional branch in the orignal blocks. E.g. Before: { vector.loop: ... BranchOnCount %IV, %TC Successors middle.block, vector.loop } After: { vector.loop: ... %any.active.mask = any-of(%mask) BranchOnCond %any.active.mask Successors vector.if.bb, vector.loop.split vector.if.bb: ... (Masked operations) Successors vector.loop.split vector.loop.split: ... BranchOnCount %IV, %TC Successors middle.block, vector.loop }

Author: ElvisWang123

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -350,6 +350,10 @@ static cl::opt<bool> PreferPredicatedReductionSelect(
  cl::desc(
  "Prefer predicating a reduction operation over an after loop select."));
 
+static cl::opt<bool> PreferControlFlow(
+ "prefer-control-flow", cl::init(false), cl::Hidden,
+ cl::desc("Generate control flow inside the vector region."));
+
 cl::opt<bool> llvm::EnableVPlanNativePath(
  "enable-vplan-native-path", cl::Hidden,
  cl::desc("Enable VPlan-native vectorization path with "
@@ -4251,6 +4255,10 @@ VectorizationFactor LoopVectorizationPlanner::selectVectorizationFactor() {
  case VPInstruction::ExplicitVectorLength:
  C += VPI->cost(VF, CostCtx);
  break;
+ case VPInstruction::AnyOf:
+ if (!VPI->getUnderlyingValue())
+ C += VPI->cost(VF, CostCtx);
+ break;
  default:
  break;
  }
@@ -8296,6 +8304,8 @@ void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
  VPlanTransforms::runPass(VPlanTransforms::truncateToMinimalBitwidths,
  *Plan, CM.getMinimalBitwidths());
  VPlanTransforms::runPass(VPlanTransforms::optimize, *Plan);
+ if (PreferControlFlow || TTI.preferControlFlow())
+ VPlanTransforms::optimizeConditionalVPBB(*Plan);
  // TODO: try to put it close to addActiveLaneMask().
  if (CM.foldTailWithEVL())
  VPlanTransforms::runPass(VPlanTransforms::addExplicitVectorLength,

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -4680,3 +4680,159 @@ void VPlanTransforms::addExitUsersForFirstOrderRecurrences(VPlan &Plan,
  }
  }
 }
+
+void VPlanTransforms::optimizeConditionalVPBB(VPlan &Plan) {
+ VPDominatorTree VPDT(Plan);
+
+ VPValue *HeaderMask = findHeaderMask(Plan);
+
+ // Get the mask from the store recipes.
+ auto GetMask = [&HeaderMask](VPRecipeBase &R) -> VPValue * {
+ using namespace llvm::VPlanPatternMatch;
+ if (!isa<VPWidenStoreRecipe>(R))
+ return nullptr;
+ VPValue *OrigMask = cast<VPWidenMemoryRecipe>(R).getMask();
+ if (!OrigMask || OrigMask == HeaderMask ||
+ match(OrigMask, m_VPInstruction<VPInstruction::ActiveLaneMask>(
+ m_VPValue(), m_VPValue())))
+ return nullptr;
+
+ return OrigMask;
+ };
+
+ // First, collect all masked stores.
+ SmallVector<std::pair<VPRecipeBase *, VPValue *>> MaskedStores;
+ ReversePostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>> RPOT(
+ Plan.getEntry());
+ for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
+ for (VPRecipeBase &R : *VPBB) {
+ if (VPValue *Mask = GetMask(R))
+ MaskedStores.emplace_back(&R, Mask);
+ }
+ }
+
+ if (MaskedStores.empty())
+ return;
+
+ DenseSet<VPRecipeBase *> Candidates;
+ auto AddOperandsToCandidates = [&Candidates](VPRecipeBase *R) {
+ for (VPValue *Op : R->operands())
+ if (VPRecipeBase *OpR = Op->getDefiningRecipe())
+ Candidates.insert(OpR);
+ };
+
+ SmallVector<SetVector<VPRecipeBase *>> Tries;
+ while (!MaskedStores.empty()) {
+ auto [SR, M] = MaskedStores.pop_back_val();
+ Candidates.clear();
+ AddOperandsToCandidates(SR);
+
+ SetVector<VPRecipeBase *> CurrentTree;
+ CurrentTree.insert(SR);
+
+ VPBasicBlock *MaskBlock =
+ M->hasDefiningRecipe() ? M->getDefiningRecipe()->getParent() : nullptr;
+
+ // Don't move recipes before the mask and PHI recipes.
+ auto End = MaskBlock == SR->getParent()
+ ? M->getDefiningRecipe()->getReverseIterator()
+ : SR->getParent()->getFirstNonPhi()->getReverseIterator();
+ // Also don't move the recipes through any recipe that may have side effects
+ // or write to memory.
+ for (auto It = std::next(SR->getReverseIterator()); It != End; ++It) {
+ if (It->mayHaveSideEffects() || It->mayWriteToMemory()) {
+ End = It;
+ break;
+ }
+ }
+
+ // Greedily add all recipes that are used to compute the stored value to the
+ // tree. All users of the added recipe must dominate the store
+ // recipe.
+ for (VPRecipeBase &R : make_range(SR->getReverseIterator(), End)) {
+ // Recipe is not part of the tree
+ if (!Candidates.contains(&R))
+ continue;
+
+ if (any_of(R.definedValues(), [&SR = SR, &VPDT](VPValue *Def) {
+ for (VPUser *U : Def->users()) {
+ if (auto *UR = dyn_cast<VPRecipeBase>(U)) {
+ if (UR == SR || VPDT.properlyDominates(UR, SR))
+ continue;
+ }
+ return true;
+ }
+ return false;
+ }))
+ continue;
+
+ CurrentTree.insert(&R);
+ AddOperandsToCandidates(&R);
+ }
+ // The previous traversal could have added recipes that are used by
+ // non-added recipes, which need to be removed from the list.
+ SmallDenseSet<VPRecipeBase *, 8> ToRemove;
+ bool Changed;
+ do {
+ Changed = false;
+ for (VPRecipeBase *R : CurrentTree) {
+ if (ToRemove.contains(R))
+ continue;
+ if (any_of(R->definedValues(), [&](VPValue *Def) {
+ for (VPUser *U : Def->users()) {
+ if (auto *UR = dyn_cast<VPRecipeBase>(U))
+ if (!CurrentTree.contains(UR) || ToRemove.contains(UR))
+ return true;
+ }
+ return false;
+ })) {
+ Changed = true;
+ ToRemove.insert(R);
+ }
+ }
+ } while (Changed);
+
+ for (VPRecipeBase *R : ToRemove)
+ CurrentTree.remove(R);
+
+ if (CurrentTree.size() > 1)
+ Tries.push_back(CurrentTree);
+ }
+
+ for (const auto &List : Tries) {
+ VPRecipeBase *SR = List.front();
+ VPValue *M = cast<VPWidenMemoryRecipe>(SR)->getMask();
+ assert(M && "Mask VPValue must exist at this point");
+ auto Recipes = reverse(List.getArrayRef());
+
+ // Split the current basic block at the store recipe point so that
+ // a predicated block can be added in between.
+ VPBasicBlock *ParentBB = SR->getParent();
+ VPBasicBlock *ContBB = ParentBB->splitAt(SR->getIterator());
+
+ // Create VPBB and insert it between ParentBB and ContBB.
+ VPBasicBlock *IfBB = Plan.createVPBasicBlock("vector.if.bb");
+ VPBlockUtils::insertBlockAfter(IfBB, ParentBB);
+ if (ContBB->getNumSuccessors() == 0)
+ ParentBB->getEnclosingLoopRegion()->setExiting(ContBB);
+
+ // Move recipes into the conditional block.
+ for (VPRecipeBase *R : Recipes)
+ R->moveBefore(*IfBB, IfBB->end());
+
+ // Add the condition and branch in the parent block.
+ auto *ActiveLane =
+ new VPInstruction(VPInstruction::AnyOf, {M}, nullptr, "any.of.mask");
+
+ auto *BranchOnCond =
+ new VPInstruction(VPInstruction::BranchOnCond, ActiveLane);
+ ParentBB->appendRecipe(ActiveLane);
+ ParentBB->appendRecipe(BranchOnCond);
+
+ // Set proper predecessors and successors for the conditional block.
+ ParentBB->clearSuccessors();
+ ParentBB->setSuccessors({IfBB, ContBB});
+ ContBB->clearPredecessors();
+ ContBB->setPredecessors({ParentBB, IfBB});
+ }
+}

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -378,6 +378,29 @@ struct VPlanTransforms {
  /// users in the original exit block using the VPIRInstruction wrapping to the
  /// LCSSA phi.
  static void addExitUsersForFirstOrderRecurrences(VPlan &Plan, VFRange &Range);
+
+ /// Try to convert flattened control flow into a conditional vector basic
+ /// block. If there are no active bits in the mask, it will skip all masked
+ /// operations. This transformation will collect all masked operations
+ /// bottom-up from the masked stores and put all masked operations in a new
+ /// vector basic block. The original vector.loop will be split and the newly
+ /// created basic block will be inserted in between.
+ ///
+ ///
+ /// [ ] <-- vector.loop
+ /// ^ | %any.active.mask = any-of(%Mask)
+ /// / | Branch-On-Count %any.active.mask, 0
+ /// / |\
+ /// | (T)| \ (F)
+ /// | | v
+ /// | | [ ] <-- vector.if.bb (masked operations)
+ /// | | |
+ /// | | v
+ /// | +-->[ ] <-- vector.loop.split
+ /// | | |
+ /// +---------+ v
+ /// [ ] <-- middle.block
+ static void optimizeConditionalVPBB(VPlan &Plan);
 };
 
 } // namespace llvm

llvm/test/Transforms/LoopVectorize/RISCV/vplan-conditional-basic-block.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -p loop-vectorize -force-vector-width=4 -S -mtriple=riscv64 -mattr=+v %s | FileCheck %s
+; RUN: opt -p loop-vectorize -force-vector-width=4 -S -mtriple=riscv64 -mattr=+v -prefer-control-flow %s | FileCheck %s
 
 define void @test(i32 %control1, i32 %control2, i32 %target, i32 %reg.4.val, ptr %reg.24.val) {
 ; CHECK-LABEL: define void @test(
@@ -28,20 +28,27 @@ define void @test(i32 %control1, i32 %control2, i32 %target, i32 %reg.4.val, ptr
 ; CHECK-NEXT: [[BROADCAST_SPLAT2:%.*]] = shufflevector <4 x i64> [[BROADCAST_SPLATINSERT1]], <4 x i64> poison, <4 x i32> zeroinitializer
 ; CHECK-NEXT: br label %[[VECTOR_BODY:.*]]
 ; CHECK: [[VECTOR_BODY]]:
-; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY]] ]
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, %[[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], %[[VECTOR_BODY_SPLIT:.*]] ]
 ; CHECK-NEXT: [[TMP2:%.*]] = getelementptr i64, ptr [[REG_24_VAL]], i64 [[INDEX]]
-; CHECK-NEXT: [[TMP4:%.*]] = getelementptr inbounds i64, ptr [[TMP2]], i32 4
+; CHECK-NEXT: [[TMP4:%.*]] = getelementptr i64, ptr [[TMP2]], i32 4
 ; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <4 x i64>, ptr [[TMP2]], align 8
 ; CHECK-NEXT: [[WIDE_LOAD3:%.*]] = load <4 x i64>, ptr [[TMP4]], align 8
 ; CHECK-NEXT: [[TMP5:%.*]] = and <4 x i64> [[WIDE_LOAD]], [[BROADCAST_SPLAT2]]
 ; CHECK-NEXT: [[TMP6:%.*]] = and <4 x i64> [[WIDE_LOAD3]], [[BROADCAST_SPLAT2]]
 ; CHECK-NEXT: [[TMP7:%.*]] = icmp eq <4 x i64> [[TMP5]], [[BROADCAST_SPLAT2]]
 ; CHECK-NEXT: [[TMP8:%.*]] = icmp eq <4 x i64> [[TMP6]], [[BROADCAST_SPLAT2]]
+; CHECK-NEXT: [[TMP12:%.*]] = freeze <4 x i1> [[TMP7]]
+; CHECK-NEXT: [[TMP13:%.*]] = freeze <4 x i1> [[TMP8]]
+; CHECK-NEXT: [[TMP14:%.*]] = or <4 x i1> [[TMP12]], [[TMP13]]
+; CHECK-NEXT: [[TMP11:%.*]] = call i1 @llvm.vector.reduce.or.v4i1(<4 x i1> [[TMP14]])
+; CHECK-NEXT: br i1 [[TMP11]], label %[[VECTOR_IF_BB:.*]], label %[[VECTOR_BODY_SPLIT]]
+; CHECK: [[VECTOR_IF_BB]]:
 ; CHECK-NEXT: [[TMP9:%.*]] = xor <4 x i64> [[WIDE_LOAD]], [[BROADCAST_SPLAT]]
 ; CHECK-NEXT: [[TMP10:%.*]] = xor <4 x i64> [[WIDE_LOAD3]], [[BROADCAST_SPLAT]]
-; CHECK-NEXT: [[TMP12:%.*]] = getelementptr i64, ptr [[TMP2]], i32 4
-; CHECK-NEXT: call void @llvm.masked.store.v4i64.p0(<4 x i64> [[TMP9]], ptr [[TMP2]], i32 8, <4 x i1> [[TMP7]])
-; CHECK-NEXT: call void @llvm.masked.store.v4i64.p0(<4 x i64> [[TMP10]], ptr [[TMP12]], i32 8, <4 x i1> [[TMP8]])
+; CHECK-NEXT: call void @llvm.masked.store.v4i64.p0(<4 x i64> [[TMP9]], ptr align 8 [[TMP2]], <4 x i1> [[TMP7]])
+; CHECK-NEXT: call void @llvm.masked.store.v4i64.p0(<4 x i64> [[TMP10]], ptr align 8 [[TMP4]], <4 x i1> [[TMP8]])
+; CHECK-NEXT: br label %[[VECTOR_BODY_SPLIT]]
+; CHECK: [[VECTOR_BODY_SPLIT]]:
 ; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 8
 ; CHECK-NEXT: [[TMP26:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
 ; CHECK-NEXT: br i1 [[TMP26]], label %[[MIDDLE_BLOCK:.*]], label %[[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]