1//===-- SPIRVStructurizer.cpp ----------------------*- C++ -*-===//
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//===----------------------------------------------------------------------===//
10
11#include "Analysis/SPIRVConvergenceRegionAnalysis.h"
12#include "SPIRV.h"
13#include "SPIRVSubtarget.h"
14#include "SPIRVUtils.h"
15#include "llvm/ADT/DenseMap.h"
16#include "llvm/ADT/STLExtras.h"
17#include "llvm/ADT/SmallPtrSet.h"
18#include "llvm/Analysis/LoopInfo.h"
19#include "llvm/IR/CFG.h"
20#include "llvm/IR/Dominators.h"
21#include "llvm/IR/IRBuilder.h"
22#include "llvm/IR/IntrinsicInst.h"
23#include "llvm/IR/Intrinsics.h"
24#include "llvm/IR/IntrinsicsSPIRV.h"
25#include "llvm/InitializePasses.h"
26#include "llvm/Transforms/Utils.h"
27#include "llvm/Transforms/Utils/Cloning.h"
28#include "llvm/Transforms/Utils/LoopSimplify.h"
29#include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
30#include <optional>
31#include <stack>
32
33using namespace llvm;
34using namespace SPIRV;
35
36using BlockSet = SmallPtrSet<BasicBlock *, 0>;
37using Edge = std::pair<BasicBlock *, BasicBlock *>;
38
39// Returns the exact convergence region in the tree defined by `Node` for which
40// `BB` is the header, nullptr otherwise.
41static const ConvergenceRegion *
42getRegionForHeader(const ConvergenceRegion *Node, BasicBlock *BB) {
43 if (Node->Entry == BB)
44 return Node;
45
46 for (auto *Child : Node->Children) {
47 const auto *CR = getRegionForHeader(Node: Child, BB);
48 if (CR != nullptr)
49 return CR;
50 }
51 return nullptr;
52}
53
54// Returns the single BasicBlock exiting the convergence region `CR`,
55// nullptr if no such exit exists.
56static BasicBlock *getExitFor(const ConvergenceRegion *CR) {
57 SmallPtrSet<BasicBlock *, 0> ExitTargets;
58 for (BasicBlock *Exit : CR->Exits) {
59 for (BasicBlock *Successor : successors(BB: Exit)) {
60 if (CR->Blocks.count(Ptr: Successor) == 0)
61 ExitTargets.insert(Ptr: Successor);
62 }
63 }
64
65 assert(ExitTargets.size() <= 1);
66 if (ExitTargets.size() == 0)
67 return nullptr;
68
69 return *ExitTargets.begin();
70}
71
72// Returns the merge block designated by I if I is a merge instruction, nullptr
73// otherwise.
74static BasicBlock *getDesignatedMergeBlock(Instruction *I) {
75 IntrinsicInst *II = dyn_cast_or_null<IntrinsicInst>(Val: I);
76 if (II == nullptr)
77 return nullptr;
78
79 if (II->getIntrinsicID() != Intrinsic::spv_loop_merge &&
80 II->getIntrinsicID() != Intrinsic::spv_selection_merge)
81 return nullptr;
82
83 BlockAddress *BA = cast<BlockAddress>(Val: II->getOperand(i_nocapture: 0));
84 return BA->getBasicBlock();
85}
86
87// Returns the continue block designated by I if I is an OpLoopMerge, nullptr
88// otherwise.
89static BasicBlock *getDesignatedContinueBlock(Instruction *I) {
90 IntrinsicInst *II = dyn_cast_or_null<IntrinsicInst>(Val: I);
91 if (II == nullptr)
92 return nullptr;
93
94 if (II->getIntrinsicID() != Intrinsic::spv_loop_merge)
95 return nullptr;
96
97 BlockAddress *BA = cast<BlockAddress>(Val: II->getOperand(i_nocapture: 1));
98 return BA->getBasicBlock();
99}
100
101// Returns true if Header has one merge instruction which designated Merge as
102// merge block.
103static bool isDefinedAsSelectionMergeBy(BasicBlock &Header, BasicBlock &Merge) {
104 for (auto &I : Header) {
105 BasicBlock *MB = getDesignatedMergeBlock(I: &I);
106 if (MB == &Merge)
107 return true;
108 }
109 return false;
110}
111
112// Returns true if the BB has one OpLoopMerge instruction.
113static bool hasLoopMergeInstruction(BasicBlock &BB) {
114 for (auto &I : BB)
115 if (getDesignatedContinueBlock(I: &I))
116 return true;
117 return false;
118}
119
120// Returns true is I is an OpSelectionMerge or OpLoopMerge instruction, false
121// otherwise.
122static bool isMergeInstruction(Instruction *I) {
123 return getDesignatedMergeBlock(I) != nullptr;
124}
125
126// Return all the merge instructions contained in BB.
127// Note: the SPIR-V spec doesn't allow a single BB to contain more than 1 merge
128// instruction, but this can happen while we structurize the CFG.
129static std::vector<Instruction *> getMergeInstructions(BasicBlock &BB) {
130 std::vector<Instruction *> Output;
131 for (Instruction &I : BB)
132 if (isMergeInstruction(I: &I))
133 Output.push_back(x: &I);
134 return Output;
135}
136
137// Bundles the header/merge/continue block sets for a function, computed in a
138// single scan since they all classify the same instructions. Callers only
139// needing a subset of them still share the single underlying scan.
140struct HeaderMergeContinueBlocks {
141 // Blocks in F having at least one OpLoopMerge or OpSelectionMerge
142 // instruction.
143 SmallPtrSet<BasicBlock *, 2> Header;
144 // Blocks in F referenced by at least 1 OpSelectionMerge/OpLoopMerge
145 // instruction.
146 SmallPtrSet<BasicBlock *, 2> Merge;
147 // Blocks in F referenced as continue target by at least 1 OpLoopMerge
148 // instruction.
149 SmallPtrSet<BasicBlock *, 2> Continue;
150
151 HeaderMergeContinueBlocks(Function &F) {
152 for (BasicBlock &BB : F) {
153 for (Instruction &I : BB) {
154 if (BasicBlock *MB = getDesignatedMergeBlock(I: &I)) {
155 Header.insert(Ptr: &BB);
156 Merge.insert(Ptr: MB);
157 }
158 if (BasicBlock *CB = getDesignatedContinueBlock(I: &I))
159 Continue.insert(Ptr: CB);
160 }
161 }
162 }
163};
164
165// Do a preorder traversal of the CFG starting from the BB |Start|.
166// Calls |op| on each basic block encountered during the traversal. Returning
167// false prunes that block's successors without stopping other pending paths.
168static void visit(BasicBlock &Start, std::function<bool(BasicBlock *)> op) {
169 std::stack<BasicBlock *> ToVisit;
170 SmallPtrSet<BasicBlock *, 8> Seen;
171
172 ToVisit.push(x: &Start);
173 Seen.insert(Ptr: ToVisit.top());
174 while (ToVisit.size() != 0) {
175 BasicBlock *BB = ToVisit.top();
176 ToVisit.pop();
177
178 if (!op(BB))
179 continue;
180
181 for (auto Succ : successors(BB)) {
182 if (Seen.contains(Ptr: Succ))
183 continue;
184 ToVisit.push(x: Succ);
185 Seen.insert(Ptr: Succ);
186 }
187 }
188}
189
190// Replaces the conditional and unconditional branch targets of |BB| by
191// |NewTarget| if the target was |OldTarget|. This function also makes sure the
192// associated merge instruction gets updated accordingly.
193static void replaceIfBranchTargets(BasicBlock *BB, BasicBlock *OldTarget,
194 BasicBlock *NewTarget) {
195 auto *BI = cast<CondBrInst>(Val: BB->getTerminator());
196
197 // 1. Replace all matching successors.
198 for (size_t i = 0; i < BI->getNumSuccessors(); i++) {
199 if (BI->getSuccessor(i) == OldTarget)
200 BI->setSuccessor(idx: i, NewSucc: NewTarget);
201 }
202
203 // Branch had 2 successors, maybe now both are the same?
204 if (BI->getSuccessor(i: 0) != BI->getSuccessor(i: 1))
205 return;
206
207 // Note: we may end up here because the original IR had such branches.
208 // This means Target is not necessarily equal to NewTarget.
209 IRBuilder<> Builder(BB);
210 Builder.SetInsertPoint(BI);
211 Builder.CreateBr(Dest: BI->getSuccessor(i: 0));
212 BI->eraseFromParent();
213
214 // The branch was the only instruction, nothing else to do.
215 if (BB->size() == 1)
216 return;
217
218 // Otherwise, we need to check: was there an OpSelectionMerge before this
219 // branch? If we removed the OpBranchConditional, we must also remove the
220 // OpSelectionMerge. This is not valid for OpLoopMerge:
221 IntrinsicInst *II =
222 dyn_cast<IntrinsicInst>(Val: BB->getTerminator()->getPrevNode());
223 if (!II || II->getIntrinsicID() != Intrinsic::spv_selection_merge)
224 return;
225
226 Constant *C = cast<Constant>(Val: II->getOperand(i_nocapture: 0));
227 II->eraseFromParent();
228 if (!C->isConstantUsed())
229 C->destroyConstant();
230}
231
232// Replaces the target of branch instruction in |BB| with |NewTarget| if it
233// was |OldTarget|. This function also fixes the associated merge instruction.
234// Note: this function does not simplify branching instructions, it only updates
235// targets. See also: simplifyBranches.
236static void replaceBranchTargets(BasicBlock *BB, BasicBlock *OldTarget,
237 BasicBlock *NewTarget) {
238 auto *T = BB->getTerminator();
239 if (isa<ReturnInst>(Val: T))
240 return;
241 if (auto *BI = dyn_cast<UncondBrInst>(Val: T)) {
242 if (BI->getSuccessor() == OldTarget)
243 BI->setSuccessor(NewTarget);
244 return;
245 }
246
247 if (isa<CondBrInst>(Val: T))
248 return replaceIfBranchTargets(BB, OldTarget, NewTarget);
249
250 if (auto *SI = dyn_cast<SwitchInst>(Val: T)) {
251 for (size_t i = 0; i < SI->getNumSuccessors(); i++) {
252 if (SI->getSuccessor(idx: i) == OldTarget)
253 SI->setSuccessor(idx: i, NewSucc: NewTarget);
254 }
255 return;
256 }
257
258 assert(false && "Unhandled terminator type.");
259}
260
261namespace {
262// Given a reducible CFG, produces a structurized CFG in the SPIR-V sense,
263// adding merge instructions when required.
264class SPIRVStructurizerImpl {
265 LoopInfo &LI;
266 ConvergenceRegionInfo &RegionInfo;
267
268 struct DivergentConstruct;
269 // Represents a list of condition/loops/switch constructs.
270 // See SPIR-V 2.11.2. Structured Control-flow Constructs for the list of
271 // constructs.
272 using ConstructList = std::vector<std::unique_ptr<DivergentConstruct>>;
273
274 // Represents a divergent construct in the SPIR-V sense.
275 // Such constructs are represented by a header (entry), a merge block (exit),
276 // and possibly a continue block (back-edge). A construct can contain other
277 // constructs, but their boundaries do not cross.
278 struct DivergentConstruct {
279 BasicBlock *Header = nullptr;
280 BasicBlock *Merge = nullptr;
281 BasicBlock *Continue = nullptr;
282
283 DivergentConstruct *Parent = nullptr;
284 ConstructList Children;
285 };
286
287 // An helper class to clean the construct boundaries.
288 // It is used to gather the list of blocks that should belong to each
289 // divergent construct, and possibly modify CFG edges when exits would cross
290 // the boundary of multiple constructs.
291 struct Splitter {
292 Function &F;
293 DomTreeBuilder::BBDomTree DT;
294 DomTreeBuilder::BBPostDomTree PDT;
295 std::optional<PartialOrderingVisitor> POV;
296
297 Splitter(Function &F) : F(F) { invalidate(); }
298
299 void invalidate() {
300 PDT.recalculate(Func&: F);
301 POV.emplace(args&: F);
302 }
303
304 const DomTreeBuilder::BBDomTree &getDT() const {
305 return POV->getDominatorTree();
306 }
307
308 // Returns the list of blocks that belong to a SPIR-V loop construct,
309 // including the continue construct.
310 std::vector<BasicBlock *> getLoopConstructBlocks(BasicBlock *Header,
311 BasicBlock *Merge) {
312 const DomTreeBuilder::BBDomTree &DT = getDT();
313 assert(DT.dominates(Header, Merge));
314 std::vector<BasicBlock *> Output;
315 visit(Start&: *Header, op: [&](BasicBlock *BB) {
316 if (BB == Merge)
317 return false;
318 if (DT.dominates(A: Merge, B: BB) || !DT.dominates(A: Header, B: BB))
319 return false;
320 Output.push_back(x: BB);
321 return true;
322 });
323 return Output;
324 }
325
326 // Returns the list of blocks that belong to a SPIR-V selection construct.
327 std::vector<BasicBlock *>
328 getSelectionConstructBlocks(DivergentConstruct *Node) {
329 const DomTreeBuilder::BBDomTree &DT = getDT();
330 assert(DT.dominates(Node->Header, Node->Merge));
331 BlockSet OutsideBlocks;
332 OutsideBlocks.insert(Ptr: Node->Merge);
333
334 for (DivergentConstruct *It = Node->Parent; It != nullptr;
335 It = It->Parent) {
336 OutsideBlocks.insert(Ptr: It->Merge);
337 if (It->Continue)
338 OutsideBlocks.insert(Ptr: It->Continue);
339 }
340
341 std::vector<BasicBlock *> Output;
342 visit(Start&: *Node->Header, op: [&](BasicBlock *BB) {
343 if (OutsideBlocks.count(Ptr: BB) != 0)
344 return false;
345 if (DT.dominates(A: Node->Merge, B: BB) || !DT.dominates(A: Node->Header, B: BB))
346 return false;
347 Output.push_back(x: BB);
348 return true;
349 });
350 return Output;
351 }
352
353 // Splits the given edges by recreating proxy nodes so that the destination
354 // has unique incoming edges from this region.
355 //
356 // clang-format off
357 //
358 // In SPIR-V, constructs must have a single exit/merge.
359 // Given nodes A and B in the construct, a node C outside, and the following edges.
360 // A -> C
361 // B -> C
362 //
363 // In such cases, we must create a new exit node D, that belong to the construct to make is viable:
364 // A -> D -> C
365 // B -> D -> C
366 //
367 // This is fine (assuming C has no PHI nodes), but requires handling the merge instruction here.
368 // By adding a proxy node, we create a regular divergent shape which can easily be regularized later on.
369 // A -> D -> D1 -> C
370 // B -> D -> D2 -> C
371 //
372 // A, B, D belongs to the construct. D is the exit. D1 and D2 are empty.
373 //
374 // clang-format on
375 std::vector<Edge>
376 createAliasBlocksForComplexEdges(std::vector<Edge> Edges) {
377 SmallPtrSet<BasicBlock *, 0> Seen;
378 std::vector<Edge> Output;
379 Output.reserve(n: Edges.size());
380
381 for (auto &[Src, Dst] : Edges) {
382 auto [Iterator, Inserted] = Seen.insert(Ptr: Src);
383 if (!Inserted) {
384 // Src already a source node. Cannot have 2 edges from A to B.
385 // Creating alias source block.
386 BasicBlock *NewSrc = BasicBlock::Create(
387 Context&: F.getContext(), Name: Src->getName() + ".new.src", Parent: &F);
388 replaceBranchTargets(BB: Src, OldTarget: Dst, NewTarget: NewSrc);
389 IRBuilder<> Builder(NewSrc);
390 Builder.CreateBr(Dest: Dst);
391 Src = NewSrc;
392 }
393
394 Output.emplace_back(args&: Src, args&: Dst);
395 }
396
397 return Output;
398 }
399
400 // Given a construct defined by |Header|, and a list of exiting edges
401 // |Edges|, creates a new single exit node, fixing up those edges.
402 BasicBlock *createSingleExitNode(BasicBlock *Header,
403 std::vector<Edge> &Edges) {
404
405 std::vector<Edge> FixedEdges = createAliasBlocksForComplexEdges(Edges);
406
407 std::vector<BasicBlock *> Dsts;
408 DenseMap<BasicBlock *, ConstantInt *> DstToIndex;
409 auto NewExit = BasicBlock::Create(Context&: F.getContext(),
410 Name: Header->getName() + ".new.exit", Parent: &F);
411 IRBuilder<> ExitBuilder(NewExit);
412 for (auto &[Src, Dst] : FixedEdges) {
413 if (DstToIndex.count(Val: Dst) != 0)
414 continue;
415 DstToIndex.try_emplace(Key: Dst, Args: ExitBuilder.getInt32(C: DstToIndex.size()));
416 Dsts.push_back(x: Dst);
417 }
418
419 if (Dsts.size() == 1) {
420 for (auto &[Src, Dst] : FixedEdges) {
421 replaceBranchTargets(BB: Src, OldTarget: Dst, NewTarget: NewExit);
422 }
423 ExitBuilder.CreateBr(Dest: Dsts[0]);
424 return NewExit;
425 }
426
427 AllocaInst *Variable = createVariable(F, Type: ExitBuilder.getInt32Ty());
428 for (auto &[Src, Dst] : FixedEdges) {
429 IRBuilder<> B2(Src);
430 B2.SetInsertPoint(Src->getFirstInsertionPt());
431 B2.CreateStore(Val: DstToIndex[Dst], Ptr: Variable);
432 replaceBranchTargets(BB: Src, OldTarget: Dst, NewTarget: NewExit);
433 }
434
435 Value *Load = ExitBuilder.CreateLoad(Ty: ExitBuilder.getInt32Ty(), Ptr: Variable);
436
437 // If we can avoid an OpSwitch, generate an OpBranch. Reason is some
438 // OpBranch are allowed to exist without a new OpSelectionMerge if one of
439 // the branch is the parent's merge node, while OpSwitches are not.
440 if (Dsts.size() == 2) {
441 Value *Condition =
442 ExitBuilder.CreateCmp(Pred: CmpInst::ICMP_EQ, LHS: DstToIndex[Dsts[0]], RHS: Load);
443 ExitBuilder.CreateCondBr(Cond: Condition, True: Dsts[0], False: Dsts[1]);
444 return NewExit;
445 }
446
447 SwitchInst *Sw = ExitBuilder.CreateSwitch(V: Load, Dest: Dsts[0], NumCases: Dsts.size() - 1);
448 for (BasicBlock *BB : drop_begin(RangeOrContainer&: Dsts))
449 Sw->addCase(OnVal: DstToIndex[BB], Dest: BB);
450 return NewExit;
451 }
452 };
453
454 // Creates a new basic block in F with a single OpUnreachable instruction.
455 BasicBlock *CreateUnreachable(Function &F) {
456 BasicBlock *BB = BasicBlock::Create(Context&: F.getContext(), Name: "unreachable", Parent: &F);
457 IRBuilder<> Builder(BB);
458 Builder.CreateUnreachable();
459 return BB;
460 }
461
462 // Add OpLoopMerge instruction on cycles.
463 bool addMergeForLoops(Function &F) {
464 auto *TopLevelRegion = RegionInfo.getTopLevelRegion();
465
466 bool Modified = false;
467 for (auto &BB : F) {
468 // Not a loop header. Ignoring for now.
469 if (!LI.isLoopHeader(BB: &BB))
470 continue;
471 auto *L = LI.getLoopFor(BB: &BB);
472
473 // This loop header is not the entrance of a convergence region. Ignoring
474 // this block.
475 auto *CR = getRegionForHeader(Node: TopLevelRegion, BB: &BB);
476 if (CR == nullptr)
477 continue;
478
479 IRBuilder<> Builder(&BB);
480
481 auto *Merge = getExitFor(CR);
482 // We are indeed in a loop, but there are no exits (infinite loop).
483 // This could be caused by a bad shader, but also could be an artifact
484 // from an earlier optimization. It is not always clear if structurally
485 // reachable means runtime reachable, so we cannot error-out. What we must
486 // do however is to make is legal on the SPIR-V point of view, hence
487 // adding an unreachable merge block.
488 if (Merge == nullptr) {
489 UncondBrInst *Br = cast<UncondBrInst>(Val: BB.getTerminator());
490 Merge = CreateUnreachable(F);
491 Builder.SetInsertPoint(Br);
492 Builder.CreateCondBr(Cond: Builder.getFalse(), True: Merge, False: Br->getSuccessor(i: 0));
493 Br->eraseFromParent();
494 }
495
496 auto *Continue = L->getLoopLatch();
497
498 Builder.SetInsertPoint(BB.getTerminator());
499 auto MergeAddress = BlockAddress::get(F: Merge->getParent(), BB: Merge);
500 auto ContinueAddress = BlockAddress::get(F: Continue->getParent(), BB: Continue);
501 SmallVector<Value *, 2> Args = {MergeAddress, ContinueAddress};
502 SmallVector<unsigned, 1> LoopControlImms =
503 getSpirvLoopControlOperandsFromLoopMetadata(L);
504 for (unsigned Imm : LoopControlImms)
505 Args.emplace_back(Args: ConstantInt::get(Ty: Builder.getInt32Ty(), V: Imm));
506 Builder.CreateIntrinsic(ID: Intrinsic::spv_loop_merge, Args: {Args});
507 Modified = true;
508 }
509
510 return Modified;
511 }
512
513 // Adds an OpSelectionMerge to the immediate dominator or each node with an
514 // in-degree of 2 or more which is not already the merge target of an
515 // OpLoopMerge/OpSelectionMerge.
516 bool addMergeForNodesWithMultiplePredecessors(Function &F) {
517 DomTreeBuilder::BBDomTree DT;
518 DT.recalculate(Func&: F);
519
520 bool Modified = false;
521 for (auto &BB : F) {
522 if (pred_size(BB: &BB) <= 1)
523 continue;
524
525 if (hasLoopMergeInstruction(BB) && pred_size(BB: &BB) <= 2)
526 continue;
527
528 assert(DT.getNode(&BB)->getIDom());
529 BasicBlock *Header = DT.getNode(BB: &BB)->getIDom()->getBlock();
530
531 if (isDefinedAsSelectionMergeBy(Header&: *Header, Merge&: BB))
532 continue;
533
534 IRBuilder<> Builder(Header);
535 Builder.SetInsertPoint(Header->getTerminator());
536
537 auto MergeAddress = BlockAddress::get(F: BB.getParent(), BB: &BB);
538 createOpSelectMerge(Builder: &Builder, MergeAddress);
539
540 Modified = true;
541 }
542
543 return Modified;
544 }
545
546 // When a block has multiple OpSelectionMerge/OpLoopMerge instructions, sorts
547 // them to put the "largest" first. A merge instruction is defined as larger
548 // than another when its target merge block post-dominates the other target's
549 // merge block. (This ordering should match the nesting ordering of the source
550 // HLSL).
551 bool sortSelectionMerge(PartialOrderingVisitor &Visitor, BasicBlock &Block) {
552 std::vector<Instruction *> MergeInstructions;
553 for (Instruction &I : Block)
554 if (isMergeInstruction(I: &I))
555 MergeInstructions.push_back(x: &I);
556
557 if (MergeInstructions.size() <= 1)
558 return false;
559
560 Instruction *InsertionPoint = *MergeInstructions.begin();
561
562 llvm::sort(C&: MergeInstructions,
563 Comp: [&Visitor](Instruction *Left, Instruction *Right) {
564 if (Left == Right)
565 return false;
566 BasicBlock *RightMerge = getDesignatedMergeBlock(I: Right);
567 BasicBlock *LeftMerge = getDesignatedMergeBlock(I: Left);
568 return !Visitor.compare(LHS: RightMerge, RHS: LeftMerge);
569 });
570
571 for (Instruction *I : MergeInstructions) {
572 I->moveBefore(InsertPos: InsertionPoint->getIterator());
573 InsertionPoint = I;
574 }
575
576 return true;
577 }
578
579 // Sorts selection merge headers in |F|.
580 // A is sorted before B if the merge block designated by B is an ancestor of
581 // the one designated by A.
582 bool sortSelectionMergeHeaders(Function &F) {
583 bool Modified = false;
584 PartialOrderingVisitor Visitor(F);
585 for (BasicBlock &BB : F) {
586 Modified |= sortSelectionMerge(Visitor, Block&: BB);
587 }
588 return Modified;
589 }
590
591 // Split basic blocks containing multiple OpLoopMerge/OpSelectionMerge
592 // instructions so each basic block contains only a single merge instruction.
593 bool splitBlocksWithMultipleHeaders(Function &F) {
594 std::stack<BasicBlock *> Work;
595 for (auto &BB : F) {
596 std::vector<Instruction *> MergeInstructions = getMergeInstructions(BB);
597 if (MergeInstructions.size() <= 1)
598 continue;
599 Work.push(x: &BB);
600 }
601
602 const bool Modified = Work.size() > 0;
603 while (Work.size() > 0) {
604 BasicBlock *Header = Work.top();
605 Work.pop();
606
607 std::vector<Instruction *> MergeInstructions =
608 getMergeInstructions(BB&: *Header);
609 for (unsigned i = 1; i < MergeInstructions.size(); i++) {
610 BasicBlock *NewBlock =
611 Header->splitBasicBlock(I: MergeInstructions[i], BBName: "new.header");
612
613 if (getDesignatedContinueBlock(I: MergeInstructions[0]) == nullptr) {
614 BasicBlock *Unreachable = CreateUnreachable(F);
615
616 Instruction *Term = Header->getTerminator();
617 IRBuilder<> Builder(Header);
618 Builder.SetInsertPoint(Term);
619 Builder.CreateCondBr(Cond: Builder.getTrue(), True: NewBlock, False: Unreachable);
620 Term->eraseFromParent();
621 }
622
623 Header = NewBlock;
624 }
625 }
626
627 return Modified;
628 }
629
630 // Adds an OpSelectionMerge to each block with an out-degree >= 2 which
631 // doesn't already have an OpSelectionMerge.
632 bool addMergeForDivergentBlocks(Function &F) {
633 DomTreeBuilder::BBPostDomTree PDT;
634 PDT.recalculate(Func&: F);
635 bool Modified = false;
636
637 HeaderMergeContinueBlocks Blocks(F);
638 auto &MergeBlocks = Blocks.Merge;
639 auto &ContinueBlocks = Blocks.Continue;
640
641 for (auto &BB : F) {
642 if (getMergeInstructions(BB).size() != 0)
643 continue;
644
645 std::vector<BasicBlock *> Candidates;
646 for (BasicBlock *Successor : successors(BB: &BB)) {
647 if (MergeBlocks.contains(Ptr: Successor))
648 continue;
649 if (ContinueBlocks.contains(Ptr: Successor))
650 continue;
651 Candidates.push_back(x: Successor);
652 }
653
654 if (Candidates.size() <= 1)
655 continue;
656
657 Modified = true;
658 BasicBlock *Merge = Candidates[0];
659
660 auto MergeAddress = BlockAddress::get(F: Merge->getParent(), BB: Merge);
661 IRBuilder<> Builder(&BB);
662 Builder.SetInsertPoint(BB.getTerminator());
663 createOpSelectMerge(Builder: &Builder, MergeAddress);
664 }
665
666 return Modified;
667 }
668
669 // Gather all the exit nodes for the construct header by |Header| and
670 // containing the blocks |Construct|.
671 std::vector<Edge> getExitsFrom(const BlockSet &Construct,
672 BasicBlock &Header) {
673 std::vector<Edge> Output;
674 visit(Start&: Header, op: [&](BasicBlock *Item) {
675 if (Construct.count(Ptr: Item) == 0)
676 return false;
677
678 for (BasicBlock *Successor : successors(BB: Item)) {
679 if (Construct.count(Ptr: Successor) == 0)
680 Output.emplace_back(args&: Item, args&: Successor);
681 }
682 return true;
683 });
684
685 return Output;
686 }
687
688 // Build a divergent construct tree searching from |BB|.
689 // If |Parent| is not null, this tree is attached to the parent's tree.
690 void constructDivergentConstruct(BlockSet &Visited, Splitter &S,
691 BasicBlock *BB, DivergentConstruct *Parent) {
692 if (Visited.count(Ptr: BB) != 0)
693 return;
694 Visited.insert(Ptr: BB);
695
696 auto MIS = getMergeInstructions(BB&: *BB);
697 if (MIS.size() == 0) {
698 for (BasicBlock *Successor : successors(BB))
699 constructDivergentConstruct(Visited, S, BB: Successor, Parent);
700 return;
701 }
702
703 assert(MIS.size() == 1);
704 Instruction *MI = MIS[0];
705
706 BasicBlock *Merge = getDesignatedMergeBlock(I: MI);
707 BasicBlock *Continue = getDesignatedContinueBlock(I: MI);
708
709 auto Output = std::make_unique<DivergentConstruct>();
710 Output->Header = BB;
711 Output->Merge = Merge;
712 Output->Continue = Continue;
713 Output->Parent = Parent;
714
715 constructDivergentConstruct(Visited, S, BB: Merge, Parent);
716 if (Continue)
717 constructDivergentConstruct(Visited, S, BB: Continue, Parent: Output.get());
718
719 for (BasicBlock *Successor : successors(BB))
720 constructDivergentConstruct(Visited, S, BB: Successor, Parent: Output.get());
721
722 if (Parent)
723 Parent->Children.emplace_back(args: std::move(Output));
724 }
725
726 // Returns the blocks belonging to the divergent construct |Node|.
727 BlockSet getConstructBlocks(Splitter &S, DivergentConstruct *Node) {
728 assert(Node->Header && Node->Merge);
729
730 if (Node->Continue) {
731 auto LoopBlocks = S.getLoopConstructBlocks(Header: Node->Header, Merge: Node->Merge);
732 return BlockSet(LoopBlocks.begin(), LoopBlocks.end());
733 }
734
735 auto SelectionBlocks = S.getSelectionConstructBlocks(Node);
736 return BlockSet(SelectionBlocks.begin(), SelectionBlocks.end());
737 }
738
739 // Fixup the construct |Node| to respect a set of rules defined by the SPIR-V
740 // spec.
741 bool fixupConstruct(Splitter &S, DivergentConstruct *Node) {
742 bool Modified = false;
743 for (auto &Child : Node->Children)
744 Modified |= fixupConstruct(S, Node: Child.get());
745
746 // This construct is the root construct. Does not represent any real
747 // construct, just a way to access the first level of the forest.
748 if (Node->Parent == nullptr)
749 return Modified;
750
751 // This node's parent is the root. Meaning this is a top-level construct.
752 // There can be multiple exists, but all are guaranteed to exit at most 1
753 // construct since we are at first level.
754 if (Node->Parent->Header == nullptr)
755 return Modified;
756
757 // Health check for the structure.
758 assert(Node->Header && Node->Merge);
759 assert(Node->Parent->Header && Node->Parent->Merge);
760
761 BlockSet ConstructBlocks = getConstructBlocks(S, Node);
762 auto Edges = getExitsFrom(Construct: ConstructBlocks, Header&: *Node->Header);
763
764 // No edges exiting the construct.
765 if (Edges.size() < 1)
766 return Modified;
767
768 bool HasBadEdge = Node->Merge == Node->Parent->Merge ||
769 Node->Merge == Node->Parent->Continue;
770 // BasicBlock *Target = Edges[0].second;
771 for (auto &[Src, Dst] : Edges) {
772 // - Breaking from a selection construct: S is a selection construct, S is
773 // the innermost structured
774 // control-flow construct containing A, and B is the merge block for S
775 // - Breaking from the innermost loop: S is the innermost loop construct
776 // containing A,
777 // and B is the merge block for S
778 if (Node->Merge == Dst)
779 continue;
780
781 // Entering the innermost loop’s continue construct: S is the innermost
782 // loop construct containing A, and B is the continue target for S
783 if (Node->Continue == Dst)
784 continue;
785
786 // TODO: what about cases branching to another case in the switch? Seems
787 // to work, but need to double check.
788 HasBadEdge = true;
789 }
790
791 if (!HasBadEdge)
792 return Modified;
793
794 // Create a single exit node gathering all exit edges.
795 BasicBlock *NewExit = S.createSingleExitNode(Header: Node->Header, Edges);
796
797 // Fixup this construct's merge node to point to the new exit.
798 // Note: this algorithm fixes inner-most divergence construct first. So
799 // recursive structures sharing a single merge node are fixed from the
800 // inside toward the outside.
801 auto MergeInstructions = getMergeInstructions(BB&: *Node->Header);
802 assert(MergeInstructions.size() == 1);
803 Instruction *I = MergeInstructions[0];
804 BlockAddress *BA = cast<BlockAddress>(Val: I->getOperand(i: 0));
805 if (BA->getBasicBlock() == Node->Merge) {
806 auto MergeAddress = BlockAddress::get(F: NewExit->getParent(), BB: NewExit);
807 I->setOperand(i: 0, Val: MergeAddress);
808 }
809
810 // Clean up of the possible dangling BockAddr operands to prevent MIR
811 // comments about "address of removed block taken".
812 if (!BA->isConstantUsed())
813 BA->destroyConstant();
814
815 Node->Merge = NewExit;
816 // Regenerate the dom trees.
817 S.invalidate();
818 return true;
819 }
820
821 bool splitCriticalEdges(Function &F) {
822 Splitter S(F);
823
824 DivergentConstruct Root;
825 BlockSet Visited;
826 constructDivergentConstruct(Visited, S, BB: &*F.begin(), Parent: &Root);
827 return fixupConstruct(S, Node: &Root);
828 }
829
830 // Simplify branches when possible:
831 // - if the 2 sides of a conditional branch are the same, transforms it to an
832 // unconditional branch.
833 // - if a switch has only 2 distinct successors, converts it to a conditional
834 // branch.
835 bool simplifyBranches(Function &F) {
836 bool Modified = false;
837
838 for (BasicBlock &BB : F) {
839 SwitchInst *SI = dyn_cast<SwitchInst>(Val: BB.getTerminator());
840 if (!SI)
841 continue;
842 if (SI->getNumCases() > 1)
843 continue;
844
845 Modified = true;
846 IRBuilder<> Builder(&BB);
847 Builder.SetInsertPoint(SI);
848
849 if (SI->getNumCases() == 0) {
850 Builder.CreateBr(Dest: SI->getDefaultDest());
851 } else {
852 Value *Condition =
853 Builder.CreateCmp(Pred: CmpInst::ICMP_EQ, LHS: SI->getCondition(),
854 RHS: SI->case_begin()->getCaseValue());
855 Builder.CreateCondBr(Cond: Condition, True: SI->case_begin()->getCaseSuccessor(),
856 False: SI->getDefaultDest());
857 }
858 SI->eraseFromParent();
859 }
860
861 return Modified;
862 }
863
864 // Makes sure every case target in |F| is unique. If 2 cases branch to the
865 // same basic block, one of the targets is updated so it jumps to a new basic
866 // block ending with a single unconditional branch to the original target.
867 bool splitSwitchCases(Function &F) {
868 bool Modified = false;
869
870 for (BasicBlock &BB : F) {
871 SwitchInst *SI = dyn_cast<SwitchInst>(Val: BB.getTerminator());
872 if (!SI)
873 continue;
874
875 BlockSet Seen;
876 Seen.insert(Ptr: SI->getDefaultDest());
877
878 auto It = SI->case_begin();
879 while (It != SI->case_end()) {
880 BasicBlock *Target = It->getCaseSuccessor();
881
882 // Don't Split. Just remove cases branching to the default destination
883 // to prevent spurious extra successors thus preserving single-exit
884 // convergence regions (i.e. if a merged exit is default & a case).
885 if (Target == SI->getDefaultDest()) {
886 Modified = true;
887 It = SI->removeCase(I: It);
888 continue;
889 }
890
891 if (Seen.count(Ptr: Target) == 0) {
892 Seen.insert(Ptr: Target);
893 ++It;
894 continue;
895 }
896
897 Modified = true;
898 BasicBlock *NewTarget =
899 BasicBlock::Create(Context&: F.getContext(), Name: "new.sw.case", Parent: &F);
900 IRBuilder<> Builder(NewTarget);
901 Builder.CreateBr(Dest: Target);
902 SI->addCase(OnVal: It->getCaseValue(), Dest: NewTarget);
903 It = SI->removeCase(I: It);
904 }
905 }
906
907 return Modified;
908 }
909
910 // Removes blocks not contributing to any structured CFG. This assumes there
911 // is no PHI nodes.
912 bool removeUselessBlocks(Function &F) {
913 std::vector<BasicBlock *> ToRemove;
914
915 HeaderMergeContinueBlocks Blocks(F);
916 auto &MergeBlocks = Blocks.Merge;
917 auto &ContinueBlocks = Blocks.Continue;
918
919 for (BasicBlock &BB : F) {
920 if (BB.size() != 1)
921 continue;
922
923 if (isa<ReturnInst>(Val: BB.getTerminator()))
924 continue;
925
926 if (MergeBlocks.count(Ptr: &BB) != 0 || ContinueBlocks.count(Ptr: &BB) != 0)
927 continue;
928
929 if (BB.getUniqueSuccessor() == nullptr)
930 continue;
931
932 BasicBlock *Successor = BB.getUniqueSuccessor();
933 std::vector<BasicBlock *> Predecessors(predecessors(BB: &BB).begin(),
934 predecessors(BB: &BB).end());
935 for (BasicBlock *Predecessor : Predecessors)
936 replaceBranchTargets(BB: Predecessor, OldTarget: &BB, NewTarget: Successor);
937 ToRemove.push_back(x: &BB);
938 }
939
940 for (BasicBlock *BB : ToRemove)
941 BB->eraseFromParent();
942
943 return ToRemove.size() != 0;
944 }
945
946 bool addHeaderToRemainingDivergentDAG(Function &F) {
947 bool Modified = false;
948
949 HeaderMergeContinueBlocks Blocks(F);
950 auto &MergeBlocks = Blocks.Merge;
951 auto &ContinueBlocks = Blocks.Continue;
952 auto &HeaderBlocks = Blocks.Header;
953
954 DomTreeBuilder::BBDomTree DT;
955 DomTreeBuilder::BBPostDomTree PDT;
956 PDT.recalculate(Func&: F);
957 DT.recalculate(Func&: F);
958
959 for (BasicBlock &BB : F) {
960 if (HeaderBlocks.count(Ptr: &BB) != 0)
961 continue;
962 if (succ_size(BB: &BB) < 2)
963 continue;
964
965 size_t CandidateEdges = 0;
966 for (BasicBlock *Successor : successors(BB: &BB)) {
967 if (MergeBlocks.count(Ptr: Successor) != 0 ||
968 ContinueBlocks.count(Ptr: Successor) != 0)
969 continue;
970 if (HeaderBlocks.count(Ptr: Successor) != 0)
971 continue;
972 CandidateEdges += 1;
973 }
974
975 if (CandidateEdges <= 1)
976 continue;
977
978 BasicBlock *Header = &BB;
979 BasicBlock *Merge = PDT.getNode(BB: &BB)->getIDom()->getBlock();
980
981 bool HasBadBlock = false;
982 visit(Start&: *Header, op: [&](const BasicBlock *Node) {
983 if (DT.dominates(A: Header, B: Node))
984 return false;
985 if (PDT.dominates(A: Merge, B: Node))
986 return false;
987 if (Node == Header || Node == Merge)
988 return true;
989
990 HasBadBlock |= MergeBlocks.count(Ptr: Node) != 0 ||
991 ContinueBlocks.count(Ptr: Node) != 0 ||
992 HeaderBlocks.count(Ptr: Node) != 0;
993 return !HasBadBlock;
994 });
995
996 if (HasBadBlock)
997 continue;
998
999 Modified = true;
1000
1001 if (Merge == nullptr) {
1002 Merge = *successors(BB: Header).begin();
1003 IRBuilder<> Builder(Header);
1004 Builder.SetInsertPoint(Header->getTerminator());
1005
1006 auto MergeAddress = BlockAddress::get(F: Merge->getParent(), BB: Merge);
1007 createOpSelectMerge(Builder: &Builder, MergeAddress);
1008 continue;
1009 }
1010
1011 Instruction *SplitInstruction = Merge->getTerminator();
1012 if (isMergeInstruction(I: SplitInstruction->getPrevNode()))
1013 SplitInstruction = SplitInstruction->getPrevNode();
1014 BasicBlock *NewMerge =
1015 Merge->splitBasicBlockBefore(I: SplitInstruction, BBName: "new.merge");
1016
1017 IRBuilder<> Builder(Header);
1018 Builder.SetInsertPoint(Header->getTerminator());
1019
1020 auto MergeAddress = BlockAddress::get(F: NewMerge->getParent(), BB: NewMerge);
1021 createOpSelectMerge(Builder: &Builder, MergeAddress);
1022 }
1023
1024 return Modified;
1025 }
1026
1027public:
1028 SPIRVStructurizerImpl(LoopInfo &LI, ConvergenceRegionInfo &RegionInfo)
1029 : LI(LI), RegionInfo(RegionInfo) {}
1030
1031 bool run(Function &F) {
1032 bool Modified = false;
1033
1034 // In LLVM, Switches are allowed to have several cases branching to the same
1035 // basic block. This is allowed in SPIR-V, but can make structurizing SPIR-V
1036 // harder, so first remove edge cases.
1037 Modified |= splitSwitchCases(F);
1038
1039 // LLVM allows conditional branches to have both side jumping to the same
1040 // block. It also allows switched to have a single default, or just one
1041 // case. Cleaning this up now.
1042 Modified |= simplifyBranches(F);
1043
1044 // At this state, we should have a reducible CFG with cycles.
1045 // STEP 1: Adding OpLoopMerge instructions to loop headers.
1046 Modified |= addMergeForLoops(F);
1047
1048 // STEP 2: adding OpSelectionMerge to each node with an in-degree >= 2.
1049 Modified |= addMergeForNodesWithMultiplePredecessors(F);
1050
1051 // STEP 3:
1052 // Sort selection merge, the largest construct goes first.
1053 // This simplifies the next step.
1054 Modified |= sortSelectionMergeHeaders(F);
1055
1056 // STEP 4: As this stage, we can have a single basic block with multiple
1057 // OpLoopMerge/OpSelectionMerge instructions. Splitting this block so each
1058 // BB has a single merge instruction.
1059 Modified |= splitBlocksWithMultipleHeaders(F);
1060
1061 // STEP 5: In the previous steps, we added merge blocks the loops and
1062 // natural merge blocks (in-degree >= 2). What remains are conditions with
1063 // an exiting branch (return, unreachable). In such case, we must start from
1064 // the header, and add headers to divergent construct with no headers.
1065 Modified |= addMergeForDivergentBlocks(F);
1066
1067 // STEP 6: At this stage, we have several divergent construct defines by a
1068 // header and a merge block. But their boundaries have no constraints: a
1069 // construct exit could be outside of the parents' construct exit. Such
1070 // edges are called critical edges. What we need is to split those edges
1071 // into several parts. Each part exiting the parent's construct by its merge
1072 // block.
1073 Modified |= splitCriticalEdges(F);
1074
1075 // STEP 7: The previous steps possibly created a lot of "proxy" blocks.
1076 // Blocks with a single unconditional branch, used to create a valid
1077 // divergent construct tree. Some nodes are still requires (e.g: nodes
1078 // allowing a valid exit through the parent's merge block). But some are
1079 // left-overs of past transformations, and could cause actual validation
1080 // issues. E.g: the SPIR-V spec allows a construct to break to the parents
1081 // loop construct without an OpSelectionMerge, but this requires a straight
1082 // jump. If a proxy block lies between the conditional branch and the
1083 // parent's merge, the CFG is not valid.
1084 Modified |= removeUselessBlocks(F);
1085
1086 // STEP 8: Final fix-up steps: our tree boundaries are correct, but some
1087 // blocks are branching with no header. Those are often simple conditional
1088 // branches with 1 or 2 returning edges. Adding a header for those.
1089 Modified |= addHeaderToRemainingDivergentDAG(F);
1090
1091 // STEP 9: sort basic blocks to match both the LLVM & SPIR-V requirements.
1092 Modified |= sortBlocks(F);
1093
1094 return Modified;
1095 }
1096
1097 void createOpSelectMerge(IRBuilder<> *Builder, BlockAddress *MergeAddress) {
1098 Instruction *BBTerminatorInst = Builder->GetInsertBlock()->getTerminator();
1099
1100 MDNode *MDNode = BBTerminatorInst->getMetadata(Kind: "hlsl.controlflow.hint");
1101
1102 ConstantInt *BranchHint = ConstantInt::get(Ty: Builder->getInt32Ty(), V: 0);
1103
1104 if (MDNode) {
1105 assert(MDNode->getNumOperands() == 2 &&
1106 "invalid metadata hlsl.controlflow.hint");
1107 BranchHint = mdconst::extract<ConstantInt>(MD: MDNode->getOperand(I: 1));
1108 }
1109
1110 SmallVector<Value *, 2> Args = {MergeAddress, BranchHint};
1111
1112 Builder->CreateIntrinsic(ID: Intrinsic::spv_selection_merge,
1113 OverloadTypes: {MergeAddress->getType()}, Args);
1114 }
1115};
1116
1117class SPIRVStructurizer : public FunctionPass {
1118public:
1119 static char ID;
1120
1121 SPIRVStructurizer() : FunctionPass(ID) {}
1122
1123 bool runOnFunction(Function &F) override {
1124 LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
1125 ConvergenceRegionInfo &RegionInfo =
1126 getAnalysis<SPIRVConvergenceRegionAnalysisWrapperPass>()
1127 .getRegionInfo();
1128 return SPIRVStructurizerImpl(LI, RegionInfo).run(F);
1129 }
1130
1131 void getAnalysisUsage(AnalysisUsage &AU) const override {
1132 AU.addRequired<LoopInfoWrapperPass>();
1133 AU.addRequired<SPIRVConvergenceRegionAnalysisWrapperPass>();
1134
1135 AU.addPreserved<SPIRVConvergenceRegionAnalysisWrapperPass>();
1136 FunctionPass::getAnalysisUsage(AU);
1137 }
1138};
1139} // anonymous namespace
1140
1141char SPIRVStructurizer::ID = 0;
1142
1143INITIALIZE_PASS_BEGIN(SPIRVStructurizer, "spirv-structurizer",
1144 "structurize SPIRV", false, false)
1145INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
1146INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
1147INITIALIZE_PASS_DEPENDENCY(SPIRVConvergenceRegionAnalysisWrapperPass)
1148
1149INITIALIZE_PASS_END(SPIRVStructurizer, "spirv-structurizer",
1150 "structurize SPIRV", false, false)
1151
1152FunctionPass *llvm::createSPIRVStructurizerPass() {
1153 return new SPIRVStructurizer();
1154}
1155
1156PreservedAnalyses SPIRVStructurizerPass::run(Function &F,
1157 FunctionAnalysisManager &AM) {
1158 LoopInfo &LI = AM.getResult<LoopAnalysis>(IR&: F);
1159 ConvergenceRegionInfo &RegionInfo =
1160 AM.getResult<SPIRVConvergenceRegionAnalysis>(IR&: F);
1161 return SPIRVStructurizerImpl(LI, RegionInfo).run(F)
1162 ? PreservedAnalyses::none()
1163 : PreservedAnalyses::all();
1164}
1165