[llvm] [VPlan] Simplify worklist in reassociateHeaderMask. NFC (PR #181595)

[Luke Lau via llvm-commits]

https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/181595

>From cc367096d851f4dcafa6d8701f15bed5a195acc7 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Mon, 16 Feb 2026 10:55:11 +0800
Subject: [PATCH 1/2] [VPlan] Simplify worklist in reassociateHeaderMask. NFC

Addresses review comments from https://github.com/llvm/llvm-project/pull/180898#pullrequestreview-3791945590. We don't need to recursively collect direct users of the header mask, we can do that as a separate step so that the main worklist loop only handles potentially reassociatable candidates.

Also add back mention of tail folding to comment.
---
 .../Transforms/Vectorize/VPlanTransforms.cpp  | 30 ++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 58ab147226ee2..8ff9f8aabf454 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1607,27 +1607,29 @@ void VPlanTransforms::simplifyRecipes(VPlan &Plan) {
 }
 
 /// Reassociate (headermask && x) && y -> headermask && (x && y) to allow the
-/// header mask to be simplified further, e.g. in optimizeEVLMasks.
+/// header mask to be simplified further when tail folding, e.g. in
+/// optimizeEVLMasks.
 static void reassociateHeaderMask(VPlan &Plan) {
   VPValue *HeaderMask = vputils::findHeaderMask(Plan);
   if (!HeaderMask)
     return;
-  SmallVector<VPUser *> Worklist(HeaderMask->users());
+
+  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());
+
+  VPValue *X, *Y;
   while (!Worklist.empty()) {
     auto *R = dyn_cast<VPSingleDefRecipe>(Worklist.pop_back_val());
-    if (!R)
+    if (!R || !match(R, m_LogicalAnd(
+                            m_LogicalAnd(m_Specific(HeaderMask), m_VPValue(X)),
+                            m_VPValue(Y))))
       continue;
-    VPValue *X, *Y;
-    if (match(R, m_LogicalAnd(m_Specific(HeaderMask), m_VPValue())))
-      Worklist.append(R->user_begin(), R->user_end());
-    else if (match(R, m_LogicalAnd(
-                          m_LogicalAnd(m_Specific(HeaderMask), m_VPValue(X)),
-                          m_VPValue(Y)))) {
-      VPBuilder Builder(R);
-      Worklist.append(R->user_begin(), R->user_end());
-      R->replaceAllUsesWith(
-          Builder.createLogicalAnd(HeaderMask, Builder.createLogicalAnd(X, Y)));
-    }
+    append_range(Worklist, R->users());
+    VPBuilder Builder(R);
+    R->replaceAllUsesWith(
+        Builder.createLogicalAnd(HeaderMask, Builder.createLogicalAnd(X, Y)));
   }
 }
 

>From 058cd23526f041a55fa546f8a551913e7fbcb05c Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Tue, 17 Feb 2026 22:01:58 +0800
Subject: [PATCH 2/2] Move VPValue def, add todo for reassociating branch
 condition

---
 llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp               | 2 +-
 .../LoopVectorize/AArch64/force-target-instruction-cost.ll      | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 8ff9f8aabf454..f7d8a81df450d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1619,9 +1619,9 @@ static void reassociateHeaderMask(VPlan &Plan) {
     if (match(U, m_LogicalAnd(m_Specific(HeaderMask), m_VPValue())))
       append_range(Worklist, cast<VPSingleDefRecipe>(U)->users());
 
-  VPValue *X, *Y;
   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))))
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/force-target-instruction-cost.ll b/llvm/test/Transforms/LoopVectorize/AArch64/force-target-instruction-cost.ll
index 892403415b335..e1eb86e2e3a57 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/force-target-instruction-cost.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/force-target-instruction-cost.ll
@@ -165,6 +165,8 @@ exit:
   ret void
 }
 
+; TODO: The branch condition could be reassociated to be simplified further, see
+; the diff in 3482a9c6cba57b4e605c2b99fb6d90a023439f9b
 define void @test_exit_branch_cost(ptr %dst, ptr noalias %x.ptr, ptr noalias %y.ptr, ptr noalias %dst.1, i1 %c.4, ptr %src, ptr noalias %dst.3, i1 %c.3, ptr noalias %dst.2) {
 ; COMMON-LABEL: define void @test_exit_branch_cost(
 ; COMMON-SAME: ptr [[DST:%.*]], ptr noalias [[X_PTR:%.*]], ptr noalias [[Y_PTR:%.*]], ptr noalias [[DST_1:%.*]], i1 [[C_4:%.*]], ptr [[SRC:%.*]], ptr noalias [[DST_3:%.*]], i1 [[C_3:%.*]], ptr noalias [[DST_2:%.*]]) {