diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp index 3f0d000ae6985..83e5664aafa2f 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp @@ -1443,13 +1443,93 @@ void VPlanTransforms::addActiveLaneMask( HeaderMask->replaceAllUsesWith(LaneMask); } +/// Try to convert \p CurRecipe to a corresponding EVL-based recipe. Returns +/// nullptr if no EVL-based recipe could be created. +/// \p HeaderMask Header Mask. +/// \p CurRecipe Recipe to be transform. +/// \p TypeInfo VPlan-based type analysis. +/// \p AllOneMask The vector mask parameter of vector-predication intrinsics. +/// \p EVL The explicit vector length parameter of vector-predication +/// intrinsics. +static VPRecipeBase *createEVLRecipe(VPValue *HeaderMask, + VPRecipeBase &CurRecipe, + VPTypeAnalysis &TypeInfo, + VPValue &AllOneMask, VPValue &EVL) { + using namespace llvm::VPlanPatternMatch; + auto GetNewMask = [&](VPValue *OrigMask) -> VPValue * { + assert(OrigMask && "Unmasked recipe when folding tail"); + return HeaderMask == OrigMask ? nullptr : OrigMask; + }; + + return TypeSwitch(&CurRecipe) + .Case([&](VPWidenLoadRecipe *L) { + VPValue *NewMask = GetNewMask(L->getMask()); + return new VPWidenLoadEVLRecipe(*L, EVL, NewMask); + }) + .Case([&](VPWidenStoreRecipe *S) { + VPValue *NewMask = GetNewMask(S->getMask()); + return new VPWidenStoreEVLRecipe(*S, EVL, NewMask); + }) + .Case([&](VPWidenRecipe *W) -> VPRecipeBase * { + unsigned Opcode = W->getOpcode(); + if (!Instruction::isBinaryOp(Opcode) && !Instruction::isUnaryOp(Opcode)) + return nullptr; + return new VPWidenEVLRecipe(*W, EVL); + }) + .Case([&](VPReductionRecipe *Red) { + VPValue *NewMask = GetNewMask(Red->getCondOp()); + return new VPReductionEVLRecipe(*Red, EVL, NewMask); + }) + .Case( + [&](auto *CR) -> VPRecipeBase * { + Intrinsic::ID VPID; + if (auto *CallR = dyn_cast(CR)) + VPID = + VPIntrinsic::getForIntrinsic(CallR->getVectorIntrinsicID()); + else if (auto *CastR = dyn_cast(CR)) + VPID = VPIntrinsic::getForOpcode(CastR->getOpcode()); + assert(VPID != Intrinsic::not_intrinsic && "Expected VP intrinsic"); + assert(VPIntrinsic::getMaskParamPos(VPID) && + VPIntrinsic::getVectorLengthParamPos(VPID) && + "Expected VP intrinsic"); + + SmallVector Ops(CR->operands()); + Ops.push_back(&AllOneMask); + Ops.push_back(&EVL); + return new VPWidenIntrinsicRecipe( + VPID, Ops, TypeInfo.inferScalarType(CR), CR->getDebugLoc()); + }) + .Case([&](VPWidenSelectRecipe *Sel) { + SmallVector Ops(Sel->operands()); + Ops.push_back(&EVL); + return new VPWidenIntrinsicRecipe(Intrinsic::vp_select, Ops, + TypeInfo.inferScalarType(Sel), + Sel->getDebugLoc()); + }) + .Case([&](VPInstruction *VPI) -> VPRecipeBase * { + VPValue *LHS, *RHS; + // Transform select with a header mask condition + // select(header_mask, LHS, RHS) + // into vector predication merge. + // vp.merge(all-true, LHS, RHS, EVL) + if (!match(VPI, m_Select(m_Specific(HeaderMask), m_VPValue(LHS), + m_VPValue(RHS)))) + return nullptr; + // Use all true as the condition because this transformation is + // limited to selects whose condition is a header mask. + return new VPWidenIntrinsicRecipe( + Intrinsic::vp_merge, {&AllOneMask, LHS, RHS, &EVL}, + TypeInfo.inferScalarType(LHS), VPI->getDebugLoc()); + }) + .Default([&](VPRecipeBase *R) { return nullptr; }); +} + /// Replace recipes with their EVL variants. static void transformRecipestoEVLRecipes(VPlan &Plan, VPValue &EVL) { - using namespace llvm::VPlanPatternMatch; Type *CanonicalIVType = Plan.getCanonicalIV()->getScalarType(); VPTypeAnalysis TypeInfo(CanonicalIVType); LLVMContext &Ctx = CanonicalIVType->getContext(); - SmallVector HeaderMasks = collectAllHeaderMasks(Plan); + VPValue *AllOneMask = Plan.getOrAddLiveIn(ConstantInt::getTrue(Ctx)); for (VPUser *U : Plan.getVF().users()) { if (auto *R = dyn_cast(U)) @@ -1461,110 +1541,22 @@ static void transformRecipestoEVLRecipes(VPlan &Plan, VPValue &EVL) { for (VPValue *HeaderMask : collectAllHeaderMasks(Plan)) { for (VPUser *U : collectUsersRecursively(HeaderMask)) { auto *CurRecipe = cast(U); - auto GetNewMask = [&](VPValue *OrigMask) -> VPValue * { - assert(OrigMask && "Unmasked recipe when folding tail"); - return HeaderMask == OrigMask ? nullptr : OrigMask; - }; - - VPRecipeBase *NewRecipe = - TypeSwitch(CurRecipe) - .Case([&](VPWidenLoadRecipe *L) { - VPValue *NewMask = GetNewMask(L->getMask()); - return new VPWidenLoadEVLRecipe(*L, EVL, NewMask); - }) - .Case([&](VPWidenStoreRecipe *S) { - VPValue *NewMask = GetNewMask(S->getMask()); - return new VPWidenStoreEVLRecipe(*S, EVL, NewMask); - }) - .Case([&](VPWidenRecipe *W) -> VPRecipeBase * { - unsigned Opcode = W->getOpcode(); - if (!Instruction::isBinaryOp(Opcode) && - !Instruction::isUnaryOp(Opcode)) - return nullptr; - return new VPWidenEVLRecipe(*W, EVL); - }) - .Case([&](VPReductionRecipe *Red) { - VPValue *NewMask = GetNewMask(Red->getCondOp()); - return new VPReductionEVLRecipe(*Red, EVL, NewMask); - }) - .Case( - [&](VPWidenIntrinsicRecipe *CallR) -> VPRecipeBase * { - Intrinsic::ID VPID = VPIntrinsic::getForIntrinsic( - CallR->getVectorIntrinsicID()); - assert(VPID != Intrinsic::not_intrinsic && - "Expected vp.casts Instrinsic"); - assert(VPIntrinsic::getMaskParamPos(VPID) && - VPIntrinsic::getVectorLengthParamPos(VPID) && - "Expected VP intrinsic"); - - SmallVector Ops(CallR->operands()); - VPValue *Mask = - Plan.getOrAddLiveIn(ConstantInt::getTrue(Ctx)); - Ops.push_back(Mask); - Ops.push_back(&EVL); - return new VPWidenIntrinsicRecipe( - VPID, Ops, TypeInfo.inferScalarType(CallR), - CallR->getDebugLoc()); - }) - .Case( - [&](VPWidenCastRecipe *CastR) -> VPRecipeBase * { - Intrinsic::ID VPID = - VPIntrinsic::getForOpcode(CastR->getOpcode()); - assert(VPID != Intrinsic::not_intrinsic && - "Expected vp.casts Instrinsic"); - - SmallVector Ops(CastR->operands()); - assert(VPIntrinsic::getMaskParamPos(VPID) && - VPIntrinsic::getVectorLengthParamPos(VPID) && - "Expected VP intrinsic"); - VPValue *Mask = - Plan.getOrAddLiveIn(ConstantInt::getTrue(Ctx)); - Ops.push_back(Mask); - Ops.push_back(&EVL); - return new VPWidenIntrinsicRecipe( - VPID, Ops, TypeInfo.inferScalarType(CastR), - CastR->getDebugLoc()); - }) - .Case([&](VPWidenSelectRecipe *Sel) { - SmallVector Ops(Sel->operands()); - Ops.push_back(&EVL); - return new VPWidenIntrinsicRecipe(Intrinsic::vp_select, Ops, - TypeInfo.inferScalarType(Sel), - Sel->getDebugLoc()); - }) - .Case([&](VPInstruction *VPI) -> VPRecipeBase * { - VPValue *LHS, *RHS; - // Transform select with a header mask condition - // select(header_mask, LHS, RHS) - // into vector predication merge. - // vp.merge(all-true, LHS, RHS, EVL) - if (!match(VPI, m_Select(m_Specific(HeaderMask), m_VPValue(LHS), - m_VPValue(RHS)))) - return nullptr; - // Use all true as the condition because this transformation is - // limited to selects whose condition is a header mask. - VPValue *AllTrue = - Plan.getOrAddLiveIn(ConstantInt::getTrue(Ctx)); - return new VPWidenIntrinsicRecipe( - Intrinsic::vp_merge, {AllTrue, LHS, RHS, &EVL}, - TypeInfo.inferScalarType(LHS), VPI->getDebugLoc()); - }) - .Default([&](VPRecipeBase *R) { return nullptr; }); - - if (!NewRecipe) + VPRecipeBase *EVLRecipe = + createEVLRecipe(HeaderMask, *CurRecipe, TypeInfo, *AllOneMask, EVL); + if (!EVLRecipe) continue; - [[maybe_unused]] unsigned NumDefVal = NewRecipe->getNumDefinedValues(); + [[maybe_unused]] unsigned NumDefVal = EVLRecipe->getNumDefinedValues(); assert(NumDefVal == CurRecipe->getNumDefinedValues() && "New recipe must define the same number of values as the " "original."); assert( NumDefVal <= 1 && "Only supports recipes with a single definition or without users."); - NewRecipe->insertBefore(CurRecipe); - if (isa(NewRecipe)) { + EVLRecipe->insertBefore(CurRecipe); + if (isa(EVLRecipe)) { VPValue *CurVPV = CurRecipe->getVPSingleValue(); - CurVPV->replaceAllUsesWith(NewRecipe->getVPSingleValue()); + CurVPV->replaceAllUsesWith(EVLRecipe->getVPSingleValue()); } // Defer erasing recipes till the end so that we don't invalidate the // VPTypeAnalysis cache.