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