[llvm-branch-commits] [llvm] [VPlan] Reassociate header mask in the predicator during mask creation (PR #219061)

via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Fri Aug 28 08:37:50 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-backend-risc-v

Author: Andrei Elovikov (eas)

<details>
<summary>Changes</summary>

Now that the header mask is a `VPRegionValue` it's easier to just emit masks in the predicator in the re-associated form. It seems to be slightly more effective than later reassociation as well.

I've originally implemented it on top of partial linearization PR to address some of its regression, but then decided that I can just move it earlier in the stack.

AI-assisted.

---
Full diff: https://github.com/llvm/llvm-project/pull/219061.diff


5 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp (+38-6) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (-29) 
- (modified) llvm/test/Transforms/LoopVectorize/RISCV/low-trip-count.ll (+9-10) 
- (modified) llvm/test/Transforms/LoopVectorize/RISCV/tail-folding-complex-mask.ll (+3-1) 
- (modified) llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll (-2) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp b/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
index c9e1eae12eb99..bf70900837b5a 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
@@ -56,6 +56,12 @@ class VPPredicator {
   /// possibly inserting new recipes at \p Dst (using Builder's insertion point)
   VPValue *createEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst);
 
+  /// Create a logical-and, keeping the header mask as the outermost operand.
+  VPValue *createMaskAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL);
+
+  /// Create a logical-or, factoring out a common header mask if present.
+  VPValue *createMaskOr(VPValue *LHS, VPValue *RHS, DebugLoc DL);
+
   /// Record \p Mask as the *entry* mask of \p VPBB, which is expected to not
   /// already have a mask.
   void setBlockInMask(const VPBasicBlock *VPBB, VPValue *Mask) {
@@ -118,6 +124,32 @@ class VPPredicator {
 };
 } // namespace
 
+VPValue *VPPredicator::createMaskAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL) {
+  VPValue *HeaderMask = Plan.getVectorLoopRegion()->getHeaderMask();
+  VPValue *Remainder = nullptr;
+  if (!HeaderMask || !match(LHS, m_RemoveMask(HeaderMask, Remainder)))
+    return Builder.createLogicalAnd(LHS, RHS, DL);
+
+  if (!Remainder)
+    return Builder.createLogicalAnd(HeaderMask, RHS, DL);
+  return Builder.createLogicalAnd(
+      HeaderMask, Builder.createLogicalAnd(Remainder, RHS, DL), DL);
+}
+
+VPValue *VPPredicator::createMaskOr(VPValue *LHS, VPValue *RHS, DebugLoc DL) {
+  VPValue *HeaderMask = Plan.getVectorLoopRegion()->getHeaderMask();
+  VPValue *LHSRemainder = nullptr;
+  VPValue *RHSRemainder = nullptr;
+  if (!HeaderMask || !match(LHS, m_RemoveMask(HeaderMask, LHSRemainder)) ||
+      !match(RHS, m_RemoveMask(HeaderMask, RHSRemainder)))
+    return Builder.createOr(LHS, RHS, DL);
+
+  if (!LHSRemainder || !RHSRemainder)
+    return HeaderMask;
+  return Builder.createLogicalAnd(
+      HeaderMask, Builder.createOr(LHSRemainder, RHSRemainder, DL), DL);
+}
+
 VPValue *VPPredicator::createEdgeMask(const VPBasicBlock *Src,
                                       const VPBasicBlock *Dst) {
   assert(is_contained(Dst->getPredecessors(), Src) && "Invalid edge");
@@ -154,7 +186,7 @@ VPValue *VPPredicator::createEdgeMask(const VPBasicBlock *Src,
     // The bitwise 'And' of SrcMask and EdgeMask introduces new UB if SrcMask
     // is false and EdgeMask is poison. Avoid that by using 'LogicalAnd'
     // instead which generates 'select i1 SrcMask, i1 EdgeMask, i1 false'.
-    EdgeMask = Builder.createLogicalAnd(SrcMask, EdgeMask, Term->getDebugLoc());
+    EdgeMask = createMaskAnd(SrcMask, EdgeMask, Term->getDebugLoc());
   }
 
   return setEdgeMask(Src, Dst, EdgeMask);
@@ -191,7 +223,7 @@ void VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {
       continue;
     }
 
-    BlockMask = Builder.createOr(BlockMask, EdgeMask, {});
+    BlockMask = createMaskOr(BlockMask, EdgeMask, {});
   }
 
   setBlockInMask(VPBB, BlockMask);
@@ -230,11 +262,11 @@ void VPPredicator::createSwitchEdgeMasks(const VPInstruction *SI) {
     for (VPValue *V : drop_begin(Conds))
       Mask = Builder.createOr(Mask, V);
     if (SrcMask)
-      Mask = Builder.createLogicalAnd(SrcMask, Mask);
+      Mask = createMaskAnd(SrcMask, Mask, {});
     setEdgeMask(Src, Dst, Mask);
 
     // 2. Create the mask for the default destination, which is reached if
-    // none of the cases with destination != default destination are taken.
+    // none of the cases with destination != Dst are taken.
     // Join the conditions for each case where the destination is != Dst using
     // an OR and negate it.
     DefaultMask = DefaultMask ? Builder.createOr(DefaultMask, Mask) : Mask;
@@ -243,7 +275,7 @@ void VPPredicator::createSwitchEdgeMasks(const VPInstruction *SI) {
   if (DefaultMask) {
     DefaultMask = Builder.createNot(DefaultMask);
     if (SrcMask)
-      DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
+      DefaultMask = createMaskAnd(SrcMask, DefaultMask, {});
   } else {
     // There are no destinations other than the default destination, so this is
     // an unconditional branch.
@@ -346,7 +378,7 @@ VPValue *VPPredicator::createBlendMaskForEdges(ArrayRef<EdgeTy> Edges,
       Builder.setInsertPoint(Dst, getMaskInsertPoint(Dst));
       EdgeMask = createEdgeMask(Src, Dst);
     }
-    Mask = Mask ? Builder.createOr(Mask, EdgeMask) : EdgeMask;
+    Mask = Mask ? createMaskOr(Mask, EdgeMask, {}) : EdgeMask;
   }
   return Mask;
 }
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index dd85eeaa8592c..7499a02a5b897 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1661,33 +1661,6 @@ void VPlanTransforms::simplifyReverses(VPlan &Plan) {
         R.getVPSingleValue()->replaceAllUsesWith(X);
 }
 
-/// Reassociate (headermask && x) && y -> headermask && (x && y) to allow the
-/// header mask to be simplified further when tail folding, e.g. in
-/// optimizeEVLMasks.
-static void reassociateHeaderMask(VPlan &Plan) {
-  VPValue *HeaderMask = Plan.getVectorLoopRegion()->getHeaderMask();
-  if (!HeaderMask)
-    return;
-
-  SmallVector<VPUser *> Worklist;
-  for (VPUser *U : HeaderMask->users())
-    if (match(U, m_LogicalAnd(m_Specific(HeaderMask), m_VPValue())))
-      append_range(Worklist, cast<VPSingleDefRecipe>(U)->users());
-
-  while (!Worklist.empty()) {
-    auto *R = dyn_cast<VPSingleDefRecipe>(Worklist.pop_back_val());
-    VPValue *X, *Y;
-    if (!R || !match(R, m_LogicalAnd(
-                            m_LogicalAnd(m_Specific(HeaderMask), m_VPValue(X)),
-                            m_VPValue(Y))))
-      continue;
-    append_range(Worklist, R->users());
-    VPBuilder Builder(R);
-    R->replaceAllUsesWith(
-        Builder.createLogicalAnd(HeaderMask, Builder.createLogicalAnd(X, Y)));
-  }
-}
-
 static std::optional<Instruction::BinaryOps>
 getUnmaskedDivRemOpcode(Intrinsic::ID ID) {
   switch (ID) {
@@ -2567,14 +2540,12 @@ bool VPlanTransforms::removeBranchOnConst(VPlan &Plan, bool OnlyLatches) {
 void VPlanTransforms::optimize(VPlan &Plan) {
   RUN_VPLAN_PASS(removeRedundantInductionCasts, Plan);
 
-  RUN_VPLAN_PASS(reassociateHeaderMask, Plan);
   RUN_VPLAN_PASS(simplifyRecipes, Plan);
   RUN_VPLAN_PASS(removeDeadRecipes, Plan);
   RUN_VPLAN_PASS(simplifyBlends, Plan);
   RUN_VPLAN_PASS(legalizeAndOptimizeInductions, Plan);
   RUN_VPLAN_PASS(narrowToSingleScalarRecipes, Plan);
   RUN_VPLAN_PASS(removeRedundantExpandSCEVRecipes, Plan);
-  RUN_VPLAN_PASS(reassociateHeaderMask, Plan);
   RUN_VPLAN_PASS(simplifyRecipes, Plan);
   RUN_VPLAN_PASS(removeBranchOnConst, Plan, /*OnlyLatches=*/false);
   RUN_VPLAN_PASS(simplifyReverses, Plan);
diff --git a/llvm/test/Transforms/LoopVectorize/RISCV/low-trip-count.ll b/llvm/test/Transforms/LoopVectorize/RISCV/low-trip-count.ll
index c4e3f51eff4d6..822e67b1b4394 100644
--- a/llvm/test/Transforms/LoopVectorize/RISCV/low-trip-count.ll
+++ b/llvm/test/Transforms/LoopVectorize/RISCV/low-trip-count.ll
@@ -280,18 +280,17 @@ define void @const_tc_with_predicated_store(i1 %c1, i1 %c2, i1 %c3, ptr %dst) #1
 ; CHECK-NEXT:    [[BROADCAST_SPLAT2:%.*]] = shufflevector <vscale x 4 x i1> [[BROADCAST_SPLATINSERT1]], <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer
 ; CHECK-NEXT:    [[BROADCAST_SPLATINSERT3:%.*]] = insertelement <vscale x 4 x i1> poison, i1 [[C1:%.*]], i64 0
 ; CHECK-NEXT:    [[BROADCAST_SPLAT4:%.*]] = shufflevector <vscale x 4 x i1> [[BROADCAST_SPLATINSERT3]], <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP12:%.*]] = xor <vscale x 4 x i1> [[BROADCAST_SPLAT4]], splat (i1 true)
-; CHECK-NEXT:    [[TMP1:%.*]] = xor <vscale x 4 x i1> [[BROADCAST_SPLAT2]], splat (i1 true)
-; CHECK-NEXT:    [[TMP13:%.*]] = select <vscale x 4 x i1> [[TMP12]], <vscale x 4 x i1> [[TMP1]], <vscale x 4 x i1> zeroinitializer
-; CHECK-NEXT:    [[TMP2:%.*]] = or <vscale x 4 x i1> [[TMP13]], [[BROADCAST_SPLAT4]]
-; CHECK-NEXT:    [[PREDPHI:%.*]] = select i1 [[C1]], <vscale x 4 x float> splat (float 1.000000e+00), <vscale x 4 x float> zeroinitializer
 ; CHECK-NEXT:    [[BROADCAST_SPLATINSERT4:%.*]] = insertelement <vscale x 4 x i1> poison, i1 [[C3:%.*]], i64 0
-; CHECK-NEXT:    [[BROADCAST_SPLAT:%.*]] = shufflevector <vscale x 4 x i1> [[BROADCAST_SPLATINSERT4]], <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[BROADCAST_SPLAT5:%.*]] = shufflevector <vscale x 4 x i1> [[BROADCAST_SPLATINSERT4]], <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP12:%.*]] = xor <vscale x 4 x i1> [[BROADCAST_SPLAT5]], splat (i1 true)
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <vscale x 4 x i1> [[BROADCAST_SPLAT4]], splat (i1 true)
+; CHECK-NEXT:    [[TMP13:%.*]] = select <vscale x 4 x i1> [[TMP12]], <vscale x 4 x i1> [[TMP1]], <vscale x 4 x i1> zeroinitializer
+; CHECK-NEXT:    [[TMP3:%.*]] = or <vscale x 4 x i1> [[TMP13]], [[BROADCAST_SPLAT5]]
+; CHECK-NEXT:    [[PREDPHI:%.*]] = select i1 [[C3]], <vscale x 4 x float> splat (float 1.000000e+00), <vscale x 4 x float> zeroinitializer
+; CHECK-NEXT:    [[TMP10:%.*]] = select <vscale x 4 x i1> [[TMP3]], <vscale x 4 x i1> [[BROADCAST_SPLAT2]], <vscale x 4 x i1> zeroinitializer
+; CHECK-NEXT:    [[PREDPHI5:%.*]] = select <vscale x 4 x i1> [[TMP10]], <vscale x 4 x float> [[PREDPHI]], <vscale x 4 x float> splat (float 2.000000e+00)
 ; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
 ; CHECK:       vector.body:
-; CHECK-NEXT:    [[TMP6:%.*]] = call <vscale x 4 x i1> @llvm.vp.merge.nxv4i1(<vscale x 4 x i1> splat (i1 true), <vscale x 4 x i1> [[TMP2]], <vscale x 4 x i1> zeroinitializer, i32 57)
-; CHECK-NEXT:    [[TMP10:%.*]] = select <vscale x 4 x i1> [[TMP6]], <vscale x 4 x i1> [[BROADCAST_SPLAT]], <vscale x 4 x i1> zeroinitializer
-; CHECK-NEXT:    [[PREDPHI5:%.*]] = select <vscale x 4 x i1> [[TMP10]], <vscale x 4 x float> [[PREDPHI]], <vscale x 4 x float> splat (float 2.000000e+00)
 ; CHECK-NEXT:    call void @llvm.vp.store.nxv4f32.p0(<vscale x 4 x float> [[PREDPHI5]], ptr align 4 [[DST:%.*]], <vscale x 4 x i1> splat (i1 true), i32 57)
 ; CHECK-NEXT:    br label [[MIDDLE_BLOCK:%.*]]
 ; CHECK:       middle.block:
@@ -360,7 +359,7 @@ define i8 @mul_non_pow_2_low_trip_count(ptr noalias %a) {
 ; CHECK-NEXT:    [[MUL]] = mul i8 [[TMP5]], [[RDX]]
 ; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV]], 1
 ; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i64 [[IV_NEXT]], 10
-; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    [[MUL_LCSSA:%.*]] = phi i8 [ [[MUL]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    ret i8 [[MUL_LCSSA]]
diff --git a/llvm/test/Transforms/LoopVectorize/RISCV/tail-folding-complex-mask.ll b/llvm/test/Transforms/LoopVectorize/RISCV/tail-folding-complex-mask.ll
index 1ffe82b366a38..0792fc784b211 100644
--- a/llvm/test/Transforms/LoopVectorize/RISCV/tail-folding-complex-mask.ll
+++ b/llvm/test/Transforms/LoopVectorize/RISCV/tail-folding-complex-mask.ll
@@ -21,6 +21,8 @@ define void @test(i64 %n, ptr noalias %src0, ptr noalias %src1, ptr noalias %src
 ; IF-EVL-NEXT:    [[TMP1:%.*]] = or <vscale x 4 x i1> [[BROADCAST_SPLAT]], [[BROADCAST_SPLAT2]]
 ; IF-EVL-NEXT:    [[TMP3:%.*]] = select <vscale x 4 x i1> [[TMP2]], <vscale x 4 x i1> [[TMP1]], <vscale x 4 x i1> zeroinitializer
 ; IF-EVL-NEXT:    [[TMP4:%.*]] = or <vscale x 4 x i1> [[BROADCAST_SPLAT]], [[TMP3]]
+; IF-EVL-NEXT:    [[TMP5:%.*]] = xor <vscale x 4 x i1> [[TMP1]], splat (i1 true)
+; IF-EVL-NEXT:    [[TMP6:%.*]] = select <vscale x 4 x i1> [[TMP2]], <vscale x 4 x i1> [[TMP5]], <vscale x 4 x i1> zeroinitializer
 ; IF-EVL-NEXT:    [[BROADCAST_SPLATINSERT3:%.*]] = insertelement <vscale x 4 x i1> poison, i1 [[C3]], i64 0
 ; IF-EVL-NEXT:    [[BROADCAST_SPLAT4:%.*]] = shufflevector <vscale x 4 x i1> [[BROADCAST_SPLATINSERT3]], <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer
 ; IF-EVL-NEXT:    br label %[[VECTOR_BODY:.*]]
@@ -34,7 +36,7 @@ define void @test(i64 %n, ptr noalias %src0, ptr noalias %src1, ptr noalias %src
 ; IF-EVL-NEXT:    [[TMP11:%.*]] = getelementptr i32, ptr [[SRC1]], i64 [[EVL_BASED_IV]]
 ; IF-EVL-NEXT:    [[VP_OP_LOAD7:%.*]] = call <vscale x 4 x i32> @llvm.vp.load.nxv4i32.p0(ptr align 4 [[TMP11]], <vscale x 4 x i1> [[TMP4]], i32 [[TMP7]])
 ; IF-EVL-NEXT:    [[TMP12:%.*]] = add <vscale x 4 x i32> [[VP_OP_LOAD7]], [[PREDPHI]]
-; IF-EVL-NEXT:    [[PREDPHI8:%.*]] = call <vscale x 4 x i32> @llvm.vp.merge.nxv4i32(<vscale x 4 x i1> [[TMP4]], <vscale x 4 x i32> [[TMP12]], <vscale x 4 x i32> zeroinitializer, i32 [[TMP7]])
+; IF-EVL-NEXT:    [[PREDPHI8:%.*]] = select <vscale x 4 x i1> [[TMP6]], <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> [[TMP12]]
 ; IF-EVL-NEXT:    [[TMP18:%.*]] = getelementptr i32, ptr [[SRC2]], i64 [[EVL_BASED_IV]]
 ; IF-EVL-NEXT:    [[WIDE_MASKED_LOAD:%.*]] = call <vscale x 4 x i32> @llvm.vp.load.nxv4i32.p0(ptr align 4 [[TMP18]], <vscale x 4 x i1> [[BROADCAST_SPLAT4]], i32 [[TMP7]])
 ; IF-EVL-NEXT:    [[TMP19:%.*]] = add <vscale x 4 x i32> [[WIDE_MASKED_LOAD]], [[PREDPHI8]]
diff --git a/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll b/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
index c2bd046efae87..21e67e691a795 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
@@ -48,14 +48,12 @@
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::truncateToMinimalBitwidths
 ; CHECK-BEFORE: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::optimize
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] removeRedundantInductionCasts
-; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] reassociateHeaderMask
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] simplifyRecipes
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] removeDeadRecipes
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] simplifyBlends
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] legalizeAndOptimizeInductions
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] narrowToSingleScalarRecipes
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] removeRedundantExpandSCEVRecipes
-; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] reassociateHeaderMask at 2
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] simplifyRecipes at 2
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] removeBranchOnConst
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] simplifyReverses

``````````

</details>


https://github.com/llvm/llvm-project/pull/219061


More information about the llvm-branch-commits mailing list