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-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// \file
10/// This file implements predication for VPlans.
11///
12//===----------------------------------------------------------------------===//
13
14#include "VPRecipeBuilder.h"
15#include "VPlan.h"
16#include "VPlanCFG.h"
17#include "VPlanDominatorTree.h"
18#include "VPlanPatternMatch.h"
19#include "VPlanTransforms.h"
20#include "VPlanUtils.h"
21#include "llvm/ADT/PostOrderIterator.h"
22
23using namespace llvm;
24using namespace VPlanPatternMatch;
25
26namespace {
27class VPPredicator {
28 /// Builder to construct recipes to compute masks.
29 VPBuilder Builder;
30
31 /// Dominator tree for the VPlan.
32 VPDominatorTree VPDT;
33
34 /// Post-dominator tree for the VPlan.
35 VPPostDominatorTree VPPDT;
36
37 /// Post-dominator frontier for the VPlan.
38 VPPostDominanceFrontier VPPDF;
39
40 /// When we if-convert we need to create edge masks. We have to cache values
41 /// so that we don't end up with exponential recursion/IR.
42 using EdgeMaskCacheTy =
43 DenseMap<std::pair<const VPBasicBlock *, const VPBasicBlock *>,
44 VPValue *>;
45 using BlockMaskCacheTy = DenseMap<const VPBasicBlock *, VPValue *>;
46 EdgeMaskCacheTy EdgeMaskCache;
47
48 BlockMaskCacheTy BlockMaskCache;
49
50 /// Create an edge mask for every destination of cases and/or default.
51 void createSwitchEdgeMasks(const VPInstruction *SI);
52
53 /// Computes and return the predicate of the edge between \p Src and \p Dst,
54 /// possibly inserting new recipes at \p Dst (using Builder's insertion point)
55 VPValue *createEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst);
56
57 /// Record \p Mask as the *entry* mask of \p VPBB, which is expected to not
58 /// already have a mask.
59 void setBlockInMask(const VPBasicBlock *VPBB, VPValue *Mask) {
60 // TODO: Include the masks as operands in the predicated VPlan directly to
61 // avoid keeping the map of masks beyond the predication transform.
62 assert(!getBlockInMask(VPBB) && "Mask already set");
63 BlockMaskCache[VPBB] = Mask;
64 }
65
66 /// Record \p Mask as the mask of the edge from \p Src to \p Dst. The edge is
67 /// expected to not have a mask already.
68 VPValue *setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst,
69 VPValue *Mask) {
70 assert(Src != Dst && "Src and Dst must be different");
71 assert(!getEdgeMask(Src, Dst) && "Mask already set");
72 return EdgeMaskCache[{Src, Dst}] = Mask;
73 }
74
75 /// Returns where to insert new masks in \p VPBB.
76 VPBasicBlock::iterator getMaskInsertPoint(VPBasicBlock *VPBB) {
77 if (VPValue *Mask = getBlockInMask(VPBB))
78 if (VPRecipeBase *MaskR = Mask->getDefiningRecipe())
79 if (MaskR->getParent() == VPBB) // In-mask may be the IDom's.
80 return std::next(x: MaskR->getIterator());
81 return VPBB->getFirstNonPhi();
82 }
83
84 using EdgeTy = std::pair<const VPBasicBlock *, const VPBasicBlock *>;
85
86 /// Compute the set of edges that are "furthest up" in the CFG for each
87 /// incoming value of \p Phi.
88 MapVector<EdgeTy, VPValue *> computeBlendEdges(VPPhi *Phi);
89
90 /// Given a set of \p Edges that each can reach \p VPBB, return the OR of all
91 /// edges, or an equivalent block in-mask.
92 VPValue *createBlendMaskForEdges(ArrayRef<EdgeTy> Edges, VPBasicBlock *VPBB);
93
94public:
95 VPPredicator(VPlan &Plan) : VPDT(Plan), VPPDT(Plan), VPPDF(VPPDT) {}
96
97 /// Returns the *entry* mask for \p VPBB.
98 VPValue *getBlockInMask(const VPBasicBlock *VPBB) const {
99 return BlockMaskCache.lookup(Val: VPBB);
100 }
101
102 /// Returns the precomputed predicate of the edge from \p Src to \p Dst.
103 VPValue *getEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst) const {
104 return EdgeMaskCache.lookup(Val: {Src, Dst});
105 }
106
107 /// Compute the predicate of \p VPBB.
108 void createBlockInMask(VPBasicBlock *VPBB);
109
110 /// Convert phi recipes in \p VPBB to VPBlendRecipes.
111 void convertPhisToBlends(VPBasicBlock *VPBB);
112};
113} // namespace
114
115VPValue *VPPredicator::createEdgeMask(const VPBasicBlock *Src,
116 const VPBasicBlock *Dst) {
117 assert(is_contained(Dst->getPredecessors(), Src) && "Invalid edge");
118
119 // Look for cached value.
120 VPValue *EdgeMask = getEdgeMask(Src, Dst);
121 if (EdgeMask)
122 return EdgeMask;
123
124 VPValue *SrcMask = getBlockInMask(VPBB: Src);
125
126 // If there's a single successor, there's no terminator recipe.
127 if (Src->getNumSuccessors() == 1)
128 return setEdgeMask(Src, Dst, Mask: SrcMask);
129
130 auto *Term = cast<VPInstruction>(Val: Src->getTerminator());
131 if (Term->getOpcode() == Instruction::Switch) {
132 createSwitchEdgeMasks(SI: Term);
133 return getEdgeMask(Src, Dst);
134 }
135
136 assert(Term->getOpcode() == VPInstruction::BranchOnCond &&
137 "Unsupported terminator");
138 if (Src->getSuccessors()[0] == Src->getSuccessors()[1])
139 return setEdgeMask(Src, Dst, Mask: SrcMask);
140
141 EdgeMask = Term->getOperand(N: 0);
142 assert(EdgeMask && "No Edge Mask found for condition");
143
144 if (Src->getSuccessors()[0] != Dst)
145 EdgeMask = Builder.createNot(Operand: EdgeMask, DL: Term->getDebugLoc());
146
147 if (SrcMask) { // Otherwise block in-mask is all-one, no need to AND.
148 // The bitwise 'And' of SrcMask and EdgeMask introduces new UB if SrcMask
149 // is false and EdgeMask is poison. Avoid that by using 'LogicalAnd'
150 // instead which generates 'select i1 SrcMask, i1 EdgeMask, i1 false'.
151 EdgeMask = Builder.createLogicalAnd(LHS: SrcMask, RHS: EdgeMask, DL: Term->getDebugLoc());
152 }
153
154 return setEdgeMask(Src, Dst, Mask: EdgeMask);
155}
156
157void VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {
158 // Start inserting after the block's phis, which be replaced by blends later.
159 Builder.setInsertPoint(TheBB: VPBB, IP: VPBB->getFirstNonPhi());
160
161 // Reuse the mask of the immediate dominator if the VPBB post-dominates the
162 // immediate dominator.
163 auto *IDom = VPDT.getNode(BB: VPBB)->getIDom();
164 assert(IDom && "Block in loop must have immediate dominator");
165 auto *IDomBB = cast<VPBasicBlock>(Val: IDom->getBlock());
166 if (VPPDT.properlyDominates(A: VPBB, B: IDomBB)) {
167 setBlockInMask(VPBB, Mask: getBlockInMask(VPBB: IDomBB));
168 return;
169 }
170 // All-one mask is modelled as no-mask following the convention for masked
171 // load/store/gather/scatter. Initialize BlockMask to no-mask.
172 VPValue *BlockMask = nullptr;
173 // This is the block mask. We OR all unique incoming edges.
174 for (auto *Predecessor : SetVector<VPBlockBase *>(
175 VPBB->getPredecessors().begin(), VPBB->getPredecessors().end())) {
176 VPValue *EdgeMask = createEdgeMask(Src: cast<VPBasicBlock>(Val: Predecessor), Dst: VPBB);
177 if (!EdgeMask) { // Mask of predecessor is all-one so mask of block is
178 // too.
179 setBlockInMask(VPBB, Mask: EdgeMask);
180 return;
181 }
182
183 if (!BlockMask) { // BlockMask has its initial nullptr value.
184 BlockMask = EdgeMask;
185 continue;
186 }
187
188 BlockMask = Builder.createOr(LHS: BlockMask, RHS: EdgeMask, DL: {});
189 }
190
191 setBlockInMask(VPBB, Mask: BlockMask);
192}
193
194void VPPredicator::createSwitchEdgeMasks(const VPInstruction *SI) {
195 const VPBasicBlock *Src = SI->getParent();
196
197 // Create masks where SI is a switch. We create masks for all edges from SI's
198 // parent block at the same time. This is more efficient, as we can create and
199 // collect compares for all cases once.
200 VPValue *Cond = SI->getOperand(N: 0);
201 VPBasicBlock *DefaultDst = cast<VPBasicBlock>(Val: Src->getSuccessors()[0]);
202 MapVector<VPBasicBlock *, SmallVector<VPValue *>> Dst2Compares;
203 for (const auto &[Idx, Succ] : enumerate(First: drop_begin(RangeOrContainer: Src->getSuccessors()))) {
204 VPBasicBlock *Dst = cast<VPBasicBlock>(Val: Succ);
205 assert(!getEdgeMask(Src, Dst) && "Edge masks already created");
206 // Cases whose destination is the same as default are redundant and can
207 // be ignored - they will get there anyhow.
208 if (Dst == DefaultDst)
209 continue;
210 auto &Compares = Dst2Compares[Dst];
211 VPValue *V = SI->getOperand(N: Idx + 1);
212 Compares.push_back(Elt: Builder.createICmp(Pred: CmpInst::ICMP_EQ, A: Cond, B: V));
213 }
214
215 // We need to handle 2 separate cases below for all entries in Dst2Compares,
216 // which excludes destinations matching the default destination.
217 VPValue *SrcMask = getBlockInMask(VPBB: Src);
218 VPValue *DefaultMask = nullptr;
219 for (const auto &[Dst, Conds] : Dst2Compares) {
220 // 1. Dst is not the default destination. Dst is reached if any of the
221 // cases with destination == Dst are taken. Join the conditions for each
222 // case whose destination == Dst using an OR.
223 VPValue *Mask = Conds[0];
224 for (VPValue *V : drop_begin(RangeOrContainer: Conds))
225 Mask = Builder.createOr(LHS: Mask, RHS: V);
226 if (SrcMask)
227 Mask = Builder.createLogicalAnd(LHS: SrcMask, RHS: Mask);
228 setEdgeMask(Src, Dst, Mask);
229
230 // 2. Create the mask for the default destination, which is reached if
231 // none of the cases with destination != default destination are taken.
232 // Join the conditions for each case where the destination is != Dst using
233 // an OR and negate it.
234 DefaultMask = DefaultMask ? Builder.createOr(LHS: DefaultMask, RHS: Mask) : Mask;
235 }
236
237 if (DefaultMask) {
238 DefaultMask = Builder.createNot(Operand: DefaultMask);
239 if (SrcMask)
240 DefaultMask = Builder.createLogicalAnd(LHS: SrcMask, RHS: DefaultMask);
241 } else {
242 // There are no destinations other than the default destination, so this is
243 // an unconditional branch.
244 DefaultMask = SrcMask;
245 }
246 setEdgeMask(Src, Dst: DefaultDst, Mask: DefaultMask);
247}
248
249// Start by keeping track of what edges lead to which value. Then see if any
250// node has the same value for all outgoing edges. If so then propagate that
251// value up to every node it postdominates. E.g:
252//
253// Entry Edges = {C->ɸ : %x, D->ɸ : %x, F->ɸ : %y}
254// / \ [C,D,F all outgoing edges equal: go up postdom frontier]
255// A B ~> {A->C : %x, A->D : %x, Entry->B : %y}
256// / \ |\ [A all outgoing edges equal: go up postdom frontier]
257// C D | E ~> {Entry->A : %x, Entry->B : %y}
258// \ \ |/
259// \ | F
260// \ | /
261// ɸ = phi [%x, C], [%x, D], [%y, F]
262MapVector<VPPredicator::EdgeTy, VPValue *>
263VPPredicator::computeBlendEdges(VPPhi *Phi) {
264 MapVector<EdgeTy, VPValue *> Edges;
265
266 // Mark the given edge as providing the value \p V.
267 auto AddEdge = [&Edges](const VPBlockBase *From, const VPBlockBase *To,
268 VPValue *V) {
269 EdgeTy Edge = {cast<VPBasicBlock>(Val: From), cast<VPBasicBlock>(Val: To)};
270 assert((!Edges.contains(Edge) || Edges.lookup(Edge) == V) &&
271 "Clobbering an edge?");
272 Edges[Edge] = V;
273 };
274
275 for (auto [InVal, InVPBB] : Phi->incoming_values_and_blocks())
276 AddEdge(InVPBB, Phi->getParent(), InVal);
277
278 SetVector<const VPBlockBase *> Worklist(from_range, Phi->incoming_blocks());
279 while (!Worklist.empty()) {
280 auto *VPBB = cast<VPBasicBlock>(Val: Worklist.pop_back_val());
281
282 // Check that all outgoing edges from VPBB have the same value.
283 SmallVector<EdgeTy> OutEdges;
284 for (const VPBlockBase *Succ : VPBB->getSuccessors())
285 OutEdges.emplace_back(Args&: VPBB, Args: cast<VPBasicBlock>(Val: Succ));
286 auto OutVals =
287 map_range(C&: OutEdges, F: [&Edges](EdgeTy E) { return Edges.lookup(Key: E); });
288 VPValue *Common = *OutVals.begin();
289 if (!Common || !all_equal(Range&: OutVals))
290 continue;
291
292 // They have the same value: we can move the edges up.
293 for (EdgeTy Edge : OutEdges)
294 Edges.erase(Key: Edge);
295
296 // Iterate up through the post dominance frontier.
297 assert(VPPDF.find(VPBB) != VPPDF.end() &&
298 "VPBB must have a post-dominance frontier entry");
299 for (const VPBlockBase *Frontier : VPPDF.find(B: VPBB)->second) {
300 for (const VPBlockBase *FrontierSucc : Frontier->getSuccessors())
301 if (VPPDT.dominates(A: VPBB, B: FrontierSucc))
302 AddEdge(Frontier, FrontierSucc, Common);
303 Worklist.insert(X: cast<VPBasicBlock>(Val: Frontier));
304 }
305 }
306
307 return Edges;
308}
309
310VPValue *VPPredicator::createBlendMaskForEdges(ArrayRef<EdgeTy> Edges,
311 VPBasicBlock *VPBB) {
312 // If the nearest common postdominator to all of Edges destinations isn't VPBB
313 // then we can use its block in-mask. E.g:
314 //
315 // A ... B
316 // \ \ /
317 // \ C
318 // \ /
319 // ... D ...
320 // \ | /
321 // VPBB
322 //
323 // If the edges are A->D and B->C, PostDom will be D. We can reuse Ds block
324 // in-mask.
325 const VPBasicBlock *PostDom = Edges[0].second;
326 for (auto [_, DstVPBB] : drop_begin(RangeOrContainer&: Edges))
327 PostDom =
328 cast<VPBasicBlock>(Val: VPPDT.findNearestCommonDominator(A: PostDom, B: DstVPBB));
329 assert(VPPDT.dominates(VPBB, PostDom) && "VPBB doesn't postdominate edges");
330 if (PostDom != VPBB)
331 return getBlockInMask(VPBB: PostDom);
332
333 // Otherwise, compute the disjunction of edges.
334 VPValue *Mask = nullptr;
335 for (auto [Src, ConstDst] : Edges) {
336 auto *Dst = const_cast<VPBasicBlock *>(ConstDst);
337 VPValue *EdgeMask;
338 {
339 VPBuilder::InsertPointGuard Guard(Builder);
340 Builder.setInsertPoint(TheBB: Dst, IP: getMaskInsertPoint(VPBB: Dst));
341 EdgeMask = createEdgeMask(Src, Dst);
342 }
343 Mask = Mask ? Builder.createOr(LHS: Mask, RHS: EdgeMask) : EdgeMask;
344 }
345 return Mask;
346}
347
348void VPPredicator::convertPhisToBlends(VPBasicBlock *VPBB) {
349 Builder.setInsertPoint(TheBB: VPBB, IP: getMaskInsertPoint(VPBB));
350
351 SmallVector<VPPhi *> Phis;
352 for (VPRecipeBase &R : VPBB->phis())
353 Phis.push_back(Elt: cast<VPPhi>(Val: &R));
354 for (VPPhi *PhiR : Phis) {
355 // The non-header Phi is converted into a Blend recipe below,
356 // so we don't have to worry about the insertion order and we can just use
357 // the builder. At this point we generate the predication tree. There may
358 // be duplications since this is a simple recursive scan, but future
359 // optimizations will clean it up.
360
361 auto NotPoison = make_filter_range(Range: PhiR->incoming_values(), Pred: [](VPValue *V) {
362 return !match(V, P: m_Poison());
363 });
364 if (all_equal(Range&: NotPoison)) {
365 PhiR->replaceAllUsesWith(New: NotPoison.empty() ? PhiR->getIncomingValue(Idx: 0)
366 : *NotPoison.begin());
367 PhiR->eraseFromParent();
368 continue;
369 }
370
371 MapVector<VPValue *, SmallVector<EdgeTy>> InValEdgesMap;
372 for (auto [Edge, Val] : computeBlendEdges(Phi: PhiR))
373 InValEdgesMap[Val].push_back(Elt: Edge);
374 auto InValEdges = InValEdgesMap.takeVector();
375
376 // Sort the incoming value order to match PhiR as much as possible.
377 llvm::stable_sort(Range&: InValEdges, C: [&PhiR](auto &L, auto &R) {
378 auto InVs = PhiR->incoming_values();
379 return std::distance(InVs.begin(), find(InVs, L.first)) <
380 std::distance(InVs.begin(), find(InVs, R.first));
381 });
382
383 SmallVector<VPValue *, 2> OperandsWithMask;
384 for (const auto &[InVPV, Edges] : InValEdges) {
385 OperandsWithMask.push_back(Elt: InVPV);
386 OperandsWithMask.push_back(Elt: createBlendMaskForEdges(Edges, VPBB));
387 }
388 PHINode *IRPhi = cast_or_null<PHINode>(Val: PhiR->getUnderlyingValue());
389 auto *Blend =
390 new VPBlendRecipe(IRPhi, OperandsWithMask, *PhiR, PhiR->getDebugLoc());
391 Builder.insert(R: Blend);
392 PhiR->replaceAllUsesWith(New: Blend);
393 PhiR->eraseFromParent();
394 }
395}
396
397void VPlanTransforms::introduceMasksAndLinearize(VPlan &Plan) {
398 // Nested loop regions (outer-loop vectorization) are not supported yet.
399 if (Plan.isOuterLoop())
400 return;
401 VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
402 // Scan the body of the loop in a topological order to visit each basic block
403 // after having visited its predecessor basic blocks.
404 VPBasicBlock *Header = LoopRegion->getEntryBasicBlock();
405 ReversePostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>> RPOT(
406 Header);
407 VPPredicator Predicator(Plan);
408 for (VPBlockBase *VPB : RPOT) {
409 // Non-outer regions with VPBBs only are supported at the moment.
410 auto *VPBB = cast<VPBasicBlock>(Val: VPB);
411 // Introduce the mask for VPBB, which may introduce needed edge masks, and
412 // convert all phi recipes of VPBB to blend recipes unless VPBB is the
413 // header.
414 if (VPBB != Header)
415 Predicator.createBlockInMask(VPBB);
416
417 VPValue *BlockMask = Predicator.getBlockInMask(VPBB);
418 if (!BlockMask)
419 continue;
420
421 // Mask all VPInstructions in the block.
422 for (VPRecipeBase &R : *VPBB) {
423 if (auto *VPI = dyn_cast<VPInstruction>(Val: &R))
424 VPI->addMask(Mask: BlockMask);
425 }
426 }
427
428 for (VPBlockBase *VPBB : reverse(C&: RPOT))
429 if (VPBB != Header)
430 Predicator.convertPhisToBlends(VPBB: cast<VPBasicBlock>(Val: VPBB));
431
432 // Linearize the blocks of the loop into one serial chain.
433 VPBlockBase *PrevVPBB = nullptr;
434 for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(Range&: RPOT)) {
435 auto Successors = to_vector(Range&: VPBB->getSuccessors());
436 if (Successors.size() > 1)
437 VPBB->getTerminator()->eraseFromParent();
438
439 // Flatten the CFG in the loop. To do so, first disconnect VPBB from its
440 // successors. Then connect VPBB to the previously visited VPBB.
441 for (auto *Succ : Successors)
442 VPBlockUtils::disconnectBlocks(From: VPBB, To: Succ);
443 if (PrevVPBB)
444 VPBlockUtils::connectBlocks(From: PrevVPBB, To: VPBB);
445
446 PrevVPBB = VPBB;
447 }
448}
449