[llvm] [LV][EVL] Reimplement method for extracting new mask. nfc (PR #156827)

[Mel Chen via llvm-commits]

https://github.com/Mel-Chen created https://github.com/llvm/llvm-project/pull/156827

During tail folding by EVL, the functionality of the header mask is replaced by EVL. Therefore, during EVL lowering, the header mask should be excluded from the mask AND-tree, and a new mask should be reconstructed.

>From bf8402892400dfb0f0d8e2132ece8145ea643c94 Mon Sep 17 00:00:00 2001
From: Mel Chen <mel.chen at sifive.com>
Date: Wed, 3 Sep 2025 03:23:03 -0700
Subject: [PATCH 1/2] evl, Rrefine get GetNewMask

---
 .../Transforms/Vectorize/VPlanTransforms.cpp  | 44 +++++++++++++++++--
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index c588453091fcc..c25dcc364171d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2403,10 +2403,46 @@ static VPRecipeBase *optimizeMaskToEVL(VPValue *HeaderMask,
   auto GetNewMask = [&](VPValue *OrigMask) -> VPValue * {
     assert(OrigMask && "Unmasked recipe when folding tail");
     // HeaderMask will be handled using EVL.
-    VPValue *Mask;
-    if (match(OrigMask, m_LogicalAnd(m_Specific(HeaderMask), m_VPValue(Mask))))
-      return Mask;
-    return HeaderMask == OrigMask ? nullptr : OrigMask;
+    bool FoundHeaderMask = false;
+    SmallVector<VPValue *, 2> Operands, Worklist;
+    Worklist.push_back(OrigMask);
+    while (!Worklist.empty()) {
+      VPValue *Cur = Worklist.pop_back_val();
+      // Replace header mask with all-true mask.
+      if (Cur == HeaderMask) {
+        Operands.push_back(&AllOneMask);
+        FoundHeaderMask = true;
+        continue;
+      }
+
+      VPValue *Op1, *Op2;
+      if (match(Cur, m_LogicalAnd(m_VPValue(Op1), m_VPValue(Op2)))) {
+        Worklist.append({Op1, Op2});
+        continue;
+      }
+      // Stop to query if the value is not in the and-tree.
+      Operands.push_back(Cur);
+    }
+
+    // Return original mask if there is no header mask in the and-tree.
+    // FIXME: Should be assertion here?
+    if (!FoundHeaderMask)
+      return OrigMask;
+    // The new mask is all-true if the header mask is in the and-tree and there
+    // is only one operand we need to do logical-and.
+    if (Operands.size() == 1)
+      return nullptr;
+    // Otherwise, create the and-tree for new mask.
+    VPBuilder Builder(cast<VPInstruction>(OrigMask));
+    VPValue *NewMask = Operands.pop_back_val();
+    while (!Operands.empty()) {
+      VPValue *Op = Operands.pop_back_val();
+      // Skip the operand if it is all-true.
+      if (match(Op, m_True()))
+        continue;
+      NewMask = Builder.createLogicalAnd(NewMask, Op);
+    }
+    return NewMask;
   };
 
   /// Adjust any end pointers so that they point to the end of EVL lanes not VF.

>From 13fdece6be87b9e86760da3a82480eb77ae503f3 Mon Sep 17 00:00:00 2001
From: Mel Chen <mel.chen at sifive.com>
Date: Thu, 4 Sep 2025 00:53:02 -0700
Subject: [PATCH 2/2] refine code

---
 llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index c25dcc364171d..388dc013acfc4 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2408,9 +2408,8 @@ static VPRecipeBase *optimizeMaskToEVL(VPValue *HeaderMask,
     Worklist.push_back(OrigMask);
     while (!Worklist.empty()) {
       VPValue *Cur = Worklist.pop_back_val();
-      // Replace header mask with all-true mask.
+      // Skip the header mask, since it will be replaced by an all-true mask.
       if (Cur == HeaderMask) {
-        Operands.push_back(&AllOneMask);
         FoundHeaderMask = true;
         continue;
       }
@@ -2420,7 +2419,7 @@ static VPRecipeBase *optimizeMaskToEVL(VPValue *HeaderMask,
         Worklist.append({Op1, Op2});
         continue;
       }
-      // Stop to query if the value is not in the and-tree.
+      // Stop to query if the value is leaf in the and-tree.
       Operands.push_back(Cur);
     }
 
@@ -2428,18 +2427,16 @@ static VPRecipeBase *optimizeMaskToEVL(VPValue *HeaderMask,
     // FIXME: Should be assertion here?
     if (!FoundHeaderMask)
       return OrigMask;
-    // The new mask is all-true if the header mask is in the and-tree and there
-    // is only one operand we need to do logical-and.
-    if (Operands.size() == 1)
+    // The new mask is all-true if there are no values to logical-and with,
+    // other than the header mask itself.
+    if (Operands.empty())
       return nullptr;
     // Otherwise, create the and-tree for new mask.
+    // TODO: Add a mask cache to avoid generating duplicate masks.
     VPBuilder Builder(cast<VPInstruction>(OrigMask));
     VPValue *NewMask = Operands.pop_back_val();
     while (!Operands.empty()) {
       VPValue *Op = Operands.pop_back_val();
-      // Skip the operand if it is all-true.
-      if (match(Op, m_True()))
-        continue;
       NewMask = Builder.createLogicalAnd(NewMask, Op);
     }
     return NewMask;