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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Mel Chen (Mel-Chen)

<details>
<summary>Changes</summary>

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.

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


1 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (+37-4) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index c588453091fcc..388dc013acfc4 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -2403,10 +2403,43 @@ 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();
+      // Skip the header mask, since it will be replaced by an all-true mask.
+      if (Cur == HeaderMask) {
+        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 leaf 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 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();
+      NewMask = Builder.createLogicalAnd(NewMask, Op);
+    }
+    return NewMask;
   };
 
   /// Adjust any end pointers so that they point to the end of EVL lanes not VF.

``````````

</details>


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