diff --git a/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h b/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h index 8f9ce7a74b58b..35cba04b66cc7 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h +++ b/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h @@ -176,6 +176,33 @@ inline int_pred_ty m_ZeroInt() { /// For vectors, this includes constants with undefined elements. inline int_pred_ty m_One() { return int_pred_ty(); } +struct bind_const_int { + uint64_t &Res; + + bind_const_int(uint64_t &Res) : Res(Res) {} + + bool match(VPValue *VPV) const { + if (!VPV->isLiveIn()) + return false; + Value *V = VPV->getLiveInIRValue(); + if (!V) + return false; + assert(!V->getType()->isVectorTy() && "Unexpected vector live-in"); + const auto *CI = dyn_cast(V); + if (!CI) + return false; + if (auto C = CI->getValue().tryZExtValue()) { + Res = *C; + return true; + } + return false; + } +}; + +/// Match a plain integer constant no wider than 64-bits, capturing it if we +/// match. +inline bind_const_int m_ConstantInt(uint64_t &C) { return C; } + /// Matching combinators template struct match_combine_or { LTy L; diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp index 9996b0167edcb..ea3df95c8e5be 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp @@ -1592,11 +1592,12 @@ static bool tryToReplaceALMWithWideALM(VPlan &Plan, ElementCount VF, m_ActiveLaneMask(m_VPValue(Index), m_VPValue(), m_VPValue())); assert(Index && "Expected index from ActiveLaneMask instruction"); - auto *II = dyn_cast(Index); - if (II && II->getOpcode() == VPInstruction::CanonicalIVIncrementForPart) { - auto Part = cast(II->getOperand(1)->getLiveInIRValue()); - Phis[Part->getZExtValue()] = Phi; - } else + uint64_t Part; + if (match(Index, + m_VPInstruction( + m_VPValue(), m_ConstantInt(Part)))) + Phis[Part] = Phi; + else // Anything other than a CanonicalIVIncrementForPart is part 0 Phis[0] = Phi; }