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