335 lines · cpp
1//===-- VPlanPredicator.cpp - VPlan predicator ----------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8///9/// \file10/// This file implements predication for VPlans.11///12//===----------------------------------------------------------------------===//13 14#include "VPRecipeBuilder.h"15#include "VPlan.h"16#include "VPlanCFG.h"17#include "VPlanPatternMatch.h"18#include "VPlanTransforms.h"19#include "VPlanUtils.h"20#include "llvm/ADT/PostOrderIterator.h"21 22using namespace llvm;23using namespace VPlanPatternMatch;24 25namespace {26class VPPredicator {27 /// Builder to construct recipes to compute masks.28 VPBuilder Builder;29 30 /// When we if-convert we need to create edge masks. We have to cache values31 /// so that we don't end up with exponential recursion/IR.32 using EdgeMaskCacheTy =33 DenseMap<std::pair<const VPBasicBlock *, const VPBasicBlock *>,34 VPValue *>;35 using BlockMaskCacheTy = DenseMap<VPBasicBlock *, VPValue *>;36 EdgeMaskCacheTy EdgeMaskCache;37 38 BlockMaskCacheTy BlockMaskCache;39 40 /// Create an edge mask for every destination of cases and/or default.41 void createSwitchEdgeMasks(VPInstruction *SI);42 43 /// Computes and return the predicate of the edge between \p Src and \p Dst,44 /// possibly inserting new recipes at \p Dst (using Builder's insertion point)45 VPValue *createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst);46 47 /// Record \p Mask as the *entry* mask of \p VPBB, which is expected to not48 /// already have a mask.49 void setBlockInMask(VPBasicBlock *VPBB, VPValue *Mask) {50 // TODO: Include the masks as operands in the predicated VPlan directly to51 // avoid keeping the map of masks beyond the predication transform.52 assert(!getBlockInMask(VPBB) && "Mask already set");53 BlockMaskCache[VPBB] = Mask;54 }55 56 /// Record \p Mask as the mask of the edge from \p Src to \p Dst. The edge is57 /// expected to not have a mask already.58 VPValue *setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst,59 VPValue *Mask) {60 assert(Src != Dst && "Src and Dst must be different");61 assert(!getEdgeMask(Src, Dst) && "Mask already set");62 return EdgeMaskCache[{Src, Dst}] = Mask;63 }64 65public:66 /// Returns the *entry* mask for \p VPBB.67 VPValue *getBlockInMask(VPBasicBlock *VPBB) const {68 return BlockMaskCache.lookup(VPBB);69 }70 71 /// Returns the precomputed predicate of the edge from \p Src to \p Dst.72 VPValue *getEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst) const {73 return EdgeMaskCache.lookup({Src, Dst});74 }75 76 /// Compute and return the mask for the vector loop header block.77 void createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail);78 79 /// Compute and return the predicate of \p VPBB, assuming that the header80 /// block of the loop is set to True, or to the loop mask when tail folding.81 VPValue *createBlockInMask(VPBasicBlock *VPBB);82 83 /// Convert phi recipes in \p VPBB to VPBlendRecipes.84 void convertPhisToBlends(VPBasicBlock *VPBB);85 86 const BlockMaskCacheTy getBlockMaskCache() const { return BlockMaskCache; }87};88} // namespace89 90VPValue *VPPredicator::createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst) {91 assert(is_contained(Dst->getPredecessors(), Src) && "Invalid edge");92 93 // Look for cached value.94 VPValue *EdgeMask = getEdgeMask(Src, Dst);95 if (EdgeMask)96 return EdgeMask;97 98 VPValue *SrcMask = getBlockInMask(Src);99 100 // If there's a single successor, there's no terminator recipe.101 if (Src->getNumSuccessors() == 1)102 return setEdgeMask(Src, Dst, SrcMask);103 104 auto *Term = cast<VPInstruction>(Src->getTerminator());105 if (Term->getOpcode() == Instruction::Switch) {106 createSwitchEdgeMasks(Term);107 return getEdgeMask(Src, Dst);108 }109 110 assert(Term->getOpcode() == VPInstruction::BranchOnCond &&111 "Unsupported terminator");112 if (Src->getSuccessors()[0] == Src->getSuccessors()[1])113 return setEdgeMask(Src, Dst, SrcMask);114 115 EdgeMask = Term->getOperand(0);116 assert(EdgeMask && "No Edge Mask found for condition");117 118 if (Src->getSuccessors()[0] != Dst)119 EdgeMask = Builder.createNot(EdgeMask, Term->getDebugLoc());120 121 if (SrcMask) { // Otherwise block in-mask is all-one, no need to AND.122 // The bitwise 'And' of SrcMask and EdgeMask introduces new UB if SrcMask123 // is false and EdgeMask is poison. Avoid that by using 'LogicalAnd'124 // instead which generates 'select i1 SrcMask, i1 EdgeMask, i1 false'.125 EdgeMask = Builder.createLogicalAnd(SrcMask, EdgeMask, Term->getDebugLoc());126 }127 128 return setEdgeMask(Src, Dst, EdgeMask);129}130 131VPValue *VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {132 // Start inserting after the block's phis, which be replaced by blends later.133 Builder.setInsertPoint(VPBB, VPBB->getFirstNonPhi());134 // All-one mask is modelled as no-mask following the convention for masked135 // load/store/gather/scatter. Initialize BlockMask to no-mask.136 VPValue *BlockMask = nullptr;137 // This is the block mask. We OR all unique incoming edges.138 for (auto *Predecessor : SetVector<VPBlockBase *>(139 VPBB->getPredecessors().begin(), VPBB->getPredecessors().end())) {140 VPValue *EdgeMask = createEdgeMask(cast<VPBasicBlock>(Predecessor), VPBB);141 if (!EdgeMask) { // Mask of predecessor is all-one so mask of block is142 // too.143 setBlockInMask(VPBB, EdgeMask);144 return EdgeMask;145 }146 147 if (!BlockMask) { // BlockMask has its initial nullptr value.148 BlockMask = EdgeMask;149 continue;150 }151 152 BlockMask = Builder.createOr(BlockMask, EdgeMask, {});153 }154 155 setBlockInMask(VPBB, BlockMask);156 return BlockMask;157}158 159void VPPredicator::createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail) {160 if (!FoldTail) {161 setBlockInMask(HeaderVPBB, nullptr);162 return;163 }164 165 // Introduce the early-exit compare IV <= BTC to form header block mask.166 // This is used instead of IV < TC because TC may wrap, unlike BTC. Start by167 // constructing the desired canonical IV in the header block as its first168 // non-phi instructions.169 170 auto &Plan = *HeaderVPBB->getPlan();171 auto *IV =172 new VPWidenCanonicalIVRecipe(HeaderVPBB->getParent()->getCanonicalIV());173 Builder.setInsertPoint(HeaderVPBB, HeaderVPBB->getFirstNonPhi());174 Builder.insert(IV);175 176 VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();177 VPValue *BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);178 setBlockInMask(HeaderVPBB, BlockMask);179}180 181void VPPredicator::createSwitchEdgeMasks(VPInstruction *SI) {182 VPBasicBlock *Src = SI->getParent();183 184 // Create masks where SI is a switch. We create masks for all edges from SI's185 // parent block at the same time. This is more efficient, as we can create and186 // collect compares for all cases once.187 VPValue *Cond = SI->getOperand(0);188 VPBasicBlock *DefaultDst = cast<VPBasicBlock>(Src->getSuccessors()[0]);189 MapVector<VPBasicBlock *, SmallVector<VPValue *>> Dst2Compares;190 for (const auto &[Idx, Succ] : enumerate(drop_begin(Src->getSuccessors()))) {191 VPBasicBlock *Dst = cast<VPBasicBlock>(Succ);192 assert(!getEdgeMask(Src, Dst) && "Edge masks already created");193 // Cases whose destination is the same as default are redundant and can194 // be ignored - they will get there anyhow.195 if (Dst == DefaultDst)196 continue;197 auto &Compares = Dst2Compares[Dst];198 VPValue *V = SI->getOperand(Idx + 1);199 Compares.push_back(Builder.createICmp(CmpInst::ICMP_EQ, Cond, V));200 }201 202 // We need to handle 2 separate cases below for all entries in Dst2Compares,203 // which excludes destinations matching the default destination.204 VPValue *SrcMask = getBlockInMask(Src);205 VPValue *DefaultMask = nullptr;206 for (const auto &[Dst, Conds] : Dst2Compares) {207 // 1. Dst is not the default destination. Dst is reached if any of the208 // cases with destination == Dst are taken. Join the conditions for each209 // case whose destination == Dst using an OR.210 VPValue *Mask = Conds[0];211 for (VPValue *V : drop_begin(Conds))212 Mask = Builder.createOr(Mask, V);213 if (SrcMask)214 Mask = Builder.createLogicalAnd(SrcMask, Mask);215 setEdgeMask(Src, Dst, Mask);216 217 // 2. Create the mask for the default destination, which is reached if218 // none of the cases with destination != default destination are taken.219 // Join the conditions for each case where the destination is != Dst using220 // an OR and negate it.221 DefaultMask = DefaultMask ? Builder.createOr(DefaultMask, Mask) : Mask;222 }223 224 if (DefaultMask) {225 DefaultMask = Builder.createNot(DefaultMask);226 if (SrcMask)227 DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);228 }229 setEdgeMask(Src, DefaultDst, DefaultMask);230}231 232void VPPredicator::convertPhisToBlends(VPBasicBlock *VPBB) {233 SmallVector<VPPhi *> Phis;234 for (VPRecipeBase &R : VPBB->phis())235 Phis.push_back(cast<VPPhi>(&R));236 for (VPPhi *PhiR : Phis) {237 // The non-header Phi is converted into a Blend recipe below,238 // so we don't have to worry about the insertion order and we can just use239 // the builder. At this point we generate the predication tree. There may240 // be duplications since this is a simple recursive scan, but future241 // optimizations will clean it up.242 243 SmallVector<VPValue *, 2> OperandsWithMask;244 for (const auto &[InVPV, InVPBB] : PhiR->incoming_values_and_blocks()) {245 OperandsWithMask.push_back(InVPV);246 VPValue *EdgeMask = getEdgeMask(InVPBB, VPBB);247 if (!EdgeMask) {248 assert(all_equal(PhiR->incoming_values()) &&249 "Distinct incoming values with one having a full mask");250 break;251 }252 253 OperandsWithMask.push_back(EdgeMask);254 }255 PHINode *IRPhi = cast_or_null<PHINode>(PhiR->getUnderlyingValue());256 auto *Blend =257 new VPBlendRecipe(IRPhi, OperandsWithMask, PhiR->getDebugLoc());258 Builder.insert(Blend);259 PhiR->replaceAllUsesWith(Blend);260 PhiR->eraseFromParent();261 }262}263 264DenseMap<VPBasicBlock *, VPValue *>265VPlanTransforms::introduceMasksAndLinearize(VPlan &Plan, bool FoldTail) {266 VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();267 // Scan the body of the loop in a topological order to visit each basic block268 // after having visited its predecessor basic blocks.269 VPBasicBlock *Header = LoopRegion->getEntryBasicBlock();270 ReversePostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>> RPOT(271 Header);272 VPPredicator Predicator;273 for (VPBlockBase *VPB : RPOT) {274 // Non-outer regions with VPBBs only are supported at the moment.275 auto *VPBB = cast<VPBasicBlock>(VPB);276 // Introduce the mask for VPBB, which may introduce needed edge masks, and277 // convert all phi recipes of VPBB to blend recipes unless VPBB is the278 // header.279 if (VPBB == Header) {280 Predicator.createHeaderMask(Header, FoldTail);281 continue;282 }283 284 Predicator.createBlockInMask(VPBB);285 Predicator.convertPhisToBlends(VPBB);286 }287 288 // Linearize the blocks of the loop into one serial chain.289 VPBlockBase *PrevVPBB = nullptr;290 for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {291 auto Successors = to_vector(VPBB->getSuccessors());292 if (Successors.size() > 1)293 VPBB->getTerminator()->eraseFromParent();294 295 // Flatten the CFG in the loop. To do so, first disconnect VPBB from its296 // successors. Then connect VPBB to the previously visited VPBB.297 for (auto *Succ : Successors)298 VPBlockUtils::disconnectBlocks(VPBB, Succ);299 if (PrevVPBB)300 VPBlockUtils::connectBlocks(PrevVPBB, VPBB);301 302 PrevVPBB = VPBB;303 }304 305 // If we folded the tail and introduced a header mask, any extract of the306 // last element must be updated to extract from the last active lane of the307 // header mask instead (i.e., the lane corresponding to the last active308 // iteration).309 if (FoldTail) {310 assert(Plan.getExitBlocks().size() == 1 &&311 "only a single-exit block is supported currently");312 VPBasicBlock *EB = Plan.getExitBlocks().front();313 assert(EB->getSinglePredecessor() == Plan.getMiddleBlock() &&314 "the exit block must have middle block as single predecessor");315 316 VPBuilder B(Plan.getMiddleBlock()->getTerminator());317 for (auto &P : EB->phis()) {318 auto *ExitIRI = cast<VPIRPhi>(&P);319 VPValue *Inc = ExitIRI->getIncomingValue(0);320 VPValue *Op;321 if (!match(Inc, m_ExtractLastElement(m_VPValue(Op))))322 continue;323 324 // Compute the index of the last active lane.325 VPValue *HeaderMask = Predicator.getBlockInMask(Header);326 VPValue *LastActiveLane =327 B.createNaryOp(VPInstruction::LastActiveLane, HeaderMask);328 auto *Ext =329 B.createNaryOp(VPInstruction::ExtractLane, {LastActiveLane, Op});330 Inc->replaceAllUsesWith(Ext);331 }332 }333 return Predicator.getBlockMaskCache();334}335