[LV][EVL] Reimplement method for extracting new mask. nfc

[Mel-Chen]

When 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.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Mel Chen (Mel-Chen)

Changes

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 --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.

[Mel-Chen]

@lukel97 @alexey-bataev @fhahn
Just after I submitted the PR, I suddenly thought of another approach.
What if we simply replace the header mask with an all-true mask after all EVL recipes are converted? Wouldn’t that be simpler? :D

[lukel97]

@lukel97 @alexey-bataev @fhahn

Just after I submitted the PR, I suddenly thought of another approach.

What if we simply replace the header mask with an all-true mask after all EVL recipes are converted? Wouldn’t that be simpler? :D

Oh, and then let simplifyRecipes remove the all ones mask from the logical and? That seems like a good idea to avoid creating new masks

Why do we need to do it after all EVL recipes are converted though. Can we do it as we're converting?

[lukel97]

Btw, is this PR NFC? I think this is good to have but it looks to be more powerful than the existing logic. Does it remove any more header masks in llvm-test-suite/SPEC?

[Mel-Chen]

Oh, and then let simplifyRecipes remove the all ones mask from the logical and? That seems like a good idea to avoid creating new masks

Yes.

Why do we need to do it after all EVL recipes are converted though. Can we do it as we're converting?

I’m not sure whether there could be cases where the same mask is shared. If so, we might need to generate a new mask for each user. But in any case, I think simply replacing the header mask with an all-true mask at the end would be better and simpler — we wouldn’t even need GetNewMask. This approach is currently blocked by #154072
, but once that patch lands, we should be able to directly replace the header mask with an all-true mask.

Btw, is this PR NFC? I think this is good to have but it looks to be more powerful than the existing logic. Does it remove any more header masks in llvm-test-suite/SPEC?

Unfortunately, I haven’t found any cases, which suggests that in practice we might not encounter such a complex situation.

[fhahn]

@lukel97 @alexey-bataev @fhahn Just after I submitted the PR, I suddenly thought of another approach. What if we simply replace the header mask with an all-true mask after all EVL recipes are converted? Wouldn’t that be simpler? :D

Yes that would be better, if it is doable.

But I guess we need to make sure all users of the mask got converted to EVL-based recipes?