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

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.