1//===-- SPIRVMergeRegionExitTargets.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// Merge the multiple exit targets of a convergence region into a single block.
10// Each exit target will be assigned a constant value, and a phi node + switch
11// will allow the new exit target to re-route to the correct basic block.
12//
13//===----------------------------------------------------------------------===//
14
15#include "Analysis/SPIRVConvergenceRegionAnalysis.h"
16#include "SPIRV.h"
17#include "SPIRVSubtarget.h"
18#include "SPIRVUtils.h"
19#include "llvm/ADT/DenseMap.h"
20#include "llvm/ADT/SmallPtrSet.h"
21#include "llvm/Analysis/LoopInfo.h"
22#include "llvm/IR/Dominators.h"
23#include "llvm/IR/IRBuilder.h"
24#include "llvm/IR/Intrinsics.h"
25#include "llvm/InitializePasses.h"
26#include "llvm/Transforms/Utils/Cloning.h"
27#include "llvm/Transforms/Utils/LoopSimplify.h"
28#include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
29
30using namespace llvm;
31
32namespace {
33
34// Run the pass on the given convergence region, ignoring the sub-regions.
35// Returns true if the CFG changed, false otherwise.
36static bool runOnConvergenceRegionNoRecurse(LoopInfo &LI,
37 SPIRV::ConvergenceRegion *CR) {
38 // Gather all the exit targets for this region.
39 SmallPtrSet<BasicBlock *, 4> ExitTargets;
40 for (BasicBlock *Exit : CR->Exits) {
41 for (BasicBlock *Target : successors(BB: Exit)) {
42 if (CR->Blocks.count(Ptr: Target) == 0)
43 ExitTargets.insert(Ptr: Target);
44 }
45 }
46
47 // If we have zero or one exit target, nothing do to.
48 if (ExitTargets.size() <= 1)
49 return false;
50
51 // Create the new single exit target.
52 auto F = CR->Entry->getParent();
53 auto NewExitTarget = BasicBlock::Create(Context&: F->getContext(), Name: "new.exit", Parent: F);
54 IRBuilder<> Builder(NewExitTarget);
55
56 AllocaInst *Variable = createVariable(F&: *F, Type: Builder.getInt32Ty());
57
58 // CodeGen output needs to be stable. Using the set as-is would order
59 // the targets differently depending on the allocation pattern.
60 // Sorting per basic-block ordering in the function.
61 std::vector<BasicBlock *> SortedExitTargets;
62 std::vector<BasicBlock *> SortedExits;
63 for (BasicBlock &BB : *F) {
64 if (ExitTargets.count(Ptr: &BB) != 0)
65 SortedExitTargets.push_back(x: &BB);
66 if (CR->Exits.count(Ptr: &BB) != 0)
67 SortedExits.push_back(x: &BB);
68 }
69
70 // Creating one constant per distinct exit target. This will be route to the
71 // correct target.
72 DenseMap<BasicBlock *, ConstantInt *> TargetToValue;
73 for (BasicBlock *Target : SortedExitTargets)
74 TargetToValue.insert(
75 KV: std::make_pair(x&: Target, y: Builder.getInt32(C: TargetToValue.size())));
76
77 // Creating one variable per exit node, set to the constant matching the
78 // targeted external block.
79 std::vector<std::pair<BasicBlock *, Value *>> ExitToVariable;
80 for (auto Exit : SortedExits) {
81 llvm::Value *Value = createExitVariable(BB: Exit, TargetToValue);
82 IRBuilder<> B2(Exit);
83 B2.SetInsertPoint(Exit->getFirstInsertionPt());
84 B2.CreateStore(Val: Value, Ptr: Variable);
85 ExitToVariable.emplace_back(args: std::make_pair(x&: Exit, y&: Value));
86 }
87
88 llvm::Value *Load = Builder.CreateLoad(Ty: Builder.getInt32Ty(), Ptr: Variable);
89
90 // Creating the switch to jump to the correct exit target.
91 llvm::SwitchInst *Sw = Builder.CreateSwitch(V: Load, Dest: SortedExitTargets[0],
92 NumCases: SortedExitTargets.size() - 1);
93 for (size_t i = 1; i < SortedExitTargets.size(); i++) {
94 BasicBlock *BB = SortedExitTargets[i];
95 Sw->addCase(OnVal: TargetToValue[BB], Dest: BB);
96 }
97
98 // Fix exit branches to redirect to the new exit.
99 for (auto Exit : CR->Exits) {
100 Instruction *T = Exit->getTerminator();
101 for (auto I = succ_begin(I: T), E = succ_end(I: T); I != E; ++I)
102 if (ExitTargets.contains(Ptr: *I))
103 I.getUse()->set(NewExitTarget);
104 }
105
106 CR = CR->Parent;
107 while (CR) {
108 CR->Blocks.insert(Ptr: NewExitTarget);
109 CR = CR->Parent;
110 }
111
112 return true;
113}
114
115/// Run the pass on the given convergence region and sub-regions (DFS).
116/// Returns true if a region/sub-region was modified, false otherwise.
117/// This returns as soon as one region/sub-region has been modified.
118static bool runOnConvergenceRegion(LoopInfo &LI, SPIRV::ConvergenceRegion *CR) {
119 for (auto *Child : CR->Children)
120 if (runOnConvergenceRegion(LI, CR: Child))
121 return true;
122
123 return runOnConvergenceRegionNoRecurse(LI, CR);
124}
125
126#if !NDEBUG
127/// Validates each edge exiting the region has the same destination basic
128/// block.
129static void validateRegionExits(const SPIRV::ConvergenceRegion *CR) {
130 for (auto *Child : CR->Children)
131 validateRegionExits(Child);
132
133 SmallPtrSet<BasicBlock *, 0> ExitTargets;
134 for (auto *Exit : CR->Exits) {
135 for (auto *BB : successors(Exit)) {
136 if (CR->Blocks.count(BB) == 0)
137 ExitTargets.insert(BB);
138 }
139 }
140
141 assert(ExitTargets.size() <= 1);
142}
143#endif
144
145static bool runImpl(Function &F, LoopInfo &LI,
146 SPIRV::ConvergenceRegionInfo &RegionInfo) {
147 auto *TopLevelRegion = RegionInfo.getWritableTopLevelRegion();
148
149 // FIXME: very inefficient method: each time a region is modified, we bubble
150 // back up, and recompute the whole convergence region tree. Once the
151 // algorithm is completed and test coverage good enough, rewrite this pass
152 // to be efficient instead of simple.
153 bool Modified = false;
154 while (runOnConvergenceRegion(LI, CR: TopLevelRegion)) {
155 Modified = true;
156 }
157
158#if !defined(NDEBUG) || defined(EXPENSIVE_CHECKS)
159 validateRegionExits(TopLevelRegion);
160#endif
161 return Modified;
162}
163
164class SPIRVMergeRegionExitTargetsLegacy : public FunctionPass {
165public:
166 static char ID;
167
168 SPIRVMergeRegionExitTargetsLegacy() : FunctionPass(ID) {}
169
170 bool runOnFunction(Function &F) override {
171 LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
172 auto &RegionInfo = getAnalysis<SPIRVConvergenceRegionAnalysisWrapperPass>()
173 .getRegionInfo();
174 return runImpl(F, LI, RegionInfo);
175 }
176
177 void getAnalysisUsage(AnalysisUsage &AU) const override {
178 AU.addRequired<LoopInfoWrapperPass>();
179 AU.addRequired<SPIRVConvergenceRegionAnalysisWrapperPass>();
180
181 AU.addPreserved<SPIRVConvergenceRegionAnalysisWrapperPass>();
182 FunctionPass::getAnalysisUsage(AU);
183 }
184};
185} // namespace
186
187PreservedAnalyses
188SPIRVMergeRegionExitTargetsPass::run(Function &F, FunctionAnalysisManager &AM) {
189 auto &LI = AM.getResult<LoopAnalysis>(IR&: F);
190 auto &RegionInfo = AM.getResult<SPIRVConvergenceRegionAnalysis>(IR&: F);
191 return runImpl(F, LI, RegionInfo) ? PreservedAnalyses::none()
192 : PreservedAnalyses::all();
193}
194
195char SPIRVMergeRegionExitTargetsLegacy::ID = 0;
196
197INITIALIZE_PASS_BEGIN(SPIRVMergeRegionExitTargetsLegacy,
198 "split-region-exit-blocks",
199 "SPIRV split region exit blocks", false, false)
200INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
201INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
202INITIALIZE_PASS_DEPENDENCY(SPIRVConvergenceRegionAnalysisWrapperPass)
203
204INITIALIZE_PASS_END(SPIRVMergeRegionExitTargetsLegacy,
205 "split-region-exit-blocks",
206 "SPIRV split region exit blocks", false, false)
207
208FunctionPass *llvm::createSPIRVMergeRegionExitTargetsPass() {
209 return new SPIRVMergeRegionExitTargetsLegacy();
210}
211