diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp index 25557f1d5d651..fe5652bd85bda 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp @@ -4268,6 +4268,34 @@ static bool canNarrowLoad(VPWidenRecipe *WideMember0, unsigned OpIdx, return false; } +static bool canNarrowOps(ArrayRef Ops) { + SmallVector Ops0; + auto *WideMember0 = dyn_cast(Ops[0]); + if (!WideMember0) + return false; + + for (const auto &[_, V] : enumerate(Ops)) { + auto *R = dyn_cast(V); + if (!R || R->getOpcode() != WideMember0->getOpcode() || + R->getNumOperands() > 2) + return false; + } + + for (unsigned Idx = 0; Idx != WideMember0->getNumOperands(); ++Idx) { + SmallVector OpsI; + for (VPValue *Op : Ops) + OpsI.push_back(Op->getDefiningRecipe()->getOperand(Idx)); + + if (any_of(enumerate(OpsI), [WideMember0, Idx](const auto &P) { + const auto &[OpIdx, OpV] = P; + return !canNarrowLoad(WideMember0, Idx, OpV, OpIdx); + })) + return false; + } + + return true; +} + /// Returns true if \p IR is a full interleave group with factor and number of /// members both equal to \p VF. The interleave group must also access the full /// vector width \p VectorRegWidth. @@ -4441,22 +4469,8 @@ void VPlanTransforms::narrowInterleaveGroups(VPlan &Plan, ElementCount VF, // Check if all values feeding InterleaveR are matching wide recipes, which // operands that can be narrowed. - auto *WideMember0 = - dyn_cast_or_null(InterleaveR->getStoredValues()[0]); - if (!WideMember0) + if (!canNarrowOps(InterleaveR->getStoredValues())) return; - for (const auto &[I, V] : enumerate(InterleaveR->getStoredValues())) { - auto *R = dyn_cast_or_null(V); - if (!R || R->getOpcode() != WideMember0->getOpcode() || - R->getNumOperands() > 2) - return; - if (any_of(enumerate(R->operands()), - [WideMember0, Idx = I](const auto &P) { - const auto &[OpIdx, OpV] = P; - return !canNarrowLoad(WideMember0, OpIdx, OpV, Idx); - })) - return; - } StoreGroups.push_back(InterleaveR); }