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