1 //===- AMDGPUUnifyDivergentExitNodes.cpp ----------------------------------===//
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 // This is a variant of the UnifyFunctionExitNodes pass. Rather than ensuring
10 // there is at most one ret and one unreachable instruction, it ensures there is
11 // at most one divergent exiting block.
12 //
13 // StructurizeCFG can't deal with multi-exit regions formed by branches to
14 // multiple return nodes. It is not desirable to structurize regions with
15 // uniform branches, so unifying those to the same return block as divergent
16 // branches inhibits use of scalar branching. It still can't deal with the case
17 // where one branch goes to return, and one unreachable. Replace unreachable in
18 // this case with a return.
19 //
20 //===----------------------------------------------------------------------===//
21
22 #include "AMDGPU.h"
23 #include "SIDefines.h"
24 #include "llvm/ADT/ArrayRef.h"
25 #include "llvm/ADT/SmallPtrSet.h"
26 #include "llvm/ADT/SmallVector.h"
27 #include "llvm/ADT/StringRef.h"
28 #include "llvm/Analysis/DomTreeUpdater.h"
29 #include "llvm/Analysis/LegacyDivergenceAnalysis.h"
30 #include "llvm/Analysis/PostDominators.h"
31 #include "llvm/Analysis/TargetTransformInfo.h"
32 #include "llvm/IR/BasicBlock.h"
33 #include "llvm/IR/CFG.h"
34 #include "llvm/IR/Constants.h"
35 #include "llvm/IR/Dominators.h"
36 #include "llvm/IR/Function.h"
37 #include "llvm/IR/IRBuilder.h"
38 #include "llvm/IR/InstrTypes.h"
39 #include "llvm/IR/Instructions.h"
40 #include "llvm/IR/Intrinsics.h"
41 #include "llvm/IR/IntrinsicsAMDGPU.h"
42 #include "llvm/IR/Type.h"
43 #include "llvm/InitializePasses.h"
44 #include "llvm/Pass.h"
45 #include "llvm/Support/Casting.h"
46 #include "llvm/Transforms/Scalar.h"
47 #include "llvm/Transforms/Utils.h"
48 #include "llvm/Transforms/Utils/Local.h"
49
50 using namespace llvm;
51
52 #define DEBUG_TYPE "amdgpu-unify-divergent-exit-nodes"
53
54 namespace {
55
56 class AMDGPUUnifyDivergentExitNodes : public FunctionPass {
57 private:
58 const TargetTransformInfo *TTI = nullptr;
59
60 public:
61 static char ID; // Pass identification, replacement for typeid
62
AMDGPUUnifyDivergentExitNodes()63 AMDGPUUnifyDivergentExitNodes() : FunctionPass(ID) {
64 initializeAMDGPUUnifyDivergentExitNodesPass(*PassRegistry::getPassRegistry());
65 }
66
67 // We can preserve non-critical-edgeness when we unify function exit nodes
68 void getAnalysisUsage(AnalysisUsage &AU) const override;
69 BasicBlock *unifyReturnBlockSet(Function &F, DomTreeUpdater &DTU,
70 ArrayRef<BasicBlock *> ReturningBlocks,
71 StringRef Name);
72 bool runOnFunction(Function &F) override;
73 };
74
75 } // end anonymous namespace
76
77 char AMDGPUUnifyDivergentExitNodes::ID = 0;
78
79 char &llvm::AMDGPUUnifyDivergentExitNodesID = AMDGPUUnifyDivergentExitNodes::ID;
80
81 INITIALIZE_PASS_BEGIN(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,
82 "Unify divergent function exit nodes", false, false)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)83 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
84 INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
85 INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
86 INITIALIZE_PASS_END(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,
87 "Unify divergent function exit nodes", false, false)
88
89 void AMDGPUUnifyDivergentExitNodes::getAnalysisUsage(AnalysisUsage &AU) const{
90 if (RequireAndPreserveDomTree)
91 AU.addRequired<DominatorTreeWrapperPass>();
92
93 AU.addRequired<PostDominatorTreeWrapperPass>();
94
95 AU.addRequired<LegacyDivergenceAnalysis>();
96
97 if (RequireAndPreserveDomTree) {
98 AU.addPreserved<DominatorTreeWrapperPass>();
99 // FIXME: preserve PostDominatorTreeWrapperPass
100 }
101
102 // No divergent values are changed, only blocks and branch edges.
103 AU.addPreserved<LegacyDivergenceAnalysis>();
104
105 // We preserve the non-critical-edgeness property
106 AU.addPreservedID(BreakCriticalEdgesID);
107
108 // This is a cluster of orthogonal Transforms
109 AU.addPreservedID(LowerSwitchID);
110 FunctionPass::getAnalysisUsage(AU);
111
112 AU.addRequired<TargetTransformInfoWrapperPass>();
113 }
114
115 /// \returns true if \p BB is reachable through only uniform branches.
116 /// XXX - Is there a more efficient way to find this?
isUniformlyReached(const LegacyDivergenceAnalysis & DA,BasicBlock & BB)117 static bool isUniformlyReached(const LegacyDivergenceAnalysis &DA,
118 BasicBlock &BB) {
119 SmallVector<BasicBlock *, 8> Stack(predecessors(&BB));
120 SmallPtrSet<BasicBlock *, 8> Visited;
121
122 while (!Stack.empty()) {
123 BasicBlock *Top = Stack.pop_back_val();
124 if (!DA.isUniform(Top->getTerminator()))
125 return false;
126
127 for (BasicBlock *Pred : predecessors(Top)) {
128 if (Visited.insert(Pred).second)
129 Stack.push_back(Pred);
130 }
131 }
132
133 return true;
134 }
135
unifyReturnBlockSet(Function & F,DomTreeUpdater & DTU,ArrayRef<BasicBlock * > ReturningBlocks,StringRef Name)136 BasicBlock *AMDGPUUnifyDivergentExitNodes::unifyReturnBlockSet(
137 Function &F, DomTreeUpdater &DTU, ArrayRef<BasicBlock *> ReturningBlocks,
138 StringRef Name) {
139 // Otherwise, we need to insert a new basic block into the function, add a PHI
140 // nodes (if the function returns values), and convert all of the return
141 // instructions into unconditional branches.
142 BasicBlock *NewRetBlock = BasicBlock::Create(F.getContext(), Name, &F);
143 IRBuilder<> B(NewRetBlock);
144
145 PHINode *PN = nullptr;
146 if (F.getReturnType()->isVoidTy()) {
147 B.CreateRetVoid();
148 } else {
149 // If the function doesn't return void... add a PHI node to the block...
150 PN = B.CreatePHI(F.getReturnType(), ReturningBlocks.size(),
151 "UnifiedRetVal");
152 B.CreateRet(PN);
153 }
154
155 // Loop over all of the blocks, replacing the return instruction with an
156 // unconditional branch.
157 std::vector<DominatorTree::UpdateType> Updates;
158 Updates.reserve(ReturningBlocks.size());
159 for (BasicBlock *BB : ReturningBlocks) {
160 // Add an incoming element to the PHI node for every return instruction that
161 // is merging into this new block...
162 if (PN)
163 PN->addIncoming(BB->getTerminator()->getOperand(0), BB);
164
165 // Remove and delete the return inst.
166 BB->getTerminator()->eraseFromParent();
167 BranchInst::Create(NewRetBlock, BB);
168 Updates.push_back({DominatorTree::Insert, BB, NewRetBlock});
169 }
170
171 if (RequireAndPreserveDomTree)
172 DTU.applyUpdates(Updates);
173 Updates.clear();
174
175 for (BasicBlock *BB : ReturningBlocks) {
176 // Cleanup possible branch to unconditional branch to the return.
177 simplifyCFG(BB, *TTI, RequireAndPreserveDomTree ? &DTU : nullptr,
178 SimplifyCFGOptions().bonusInstThreshold(2));
179 }
180
181 return NewRetBlock;
182 }
183
runOnFunction(Function & F)184 bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
185 DominatorTree *DT = nullptr;
186 if (RequireAndPreserveDomTree)
187 DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
188
189 auto &PDT = getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
190 if (PDT.root_size() == 0 ||
191 (PDT.root_size() == 1 &&
192 !isa<BranchInst>(PDT.getRoot()->getTerminator())))
193 return false;
194
195 LegacyDivergenceAnalysis &DA = getAnalysis<LegacyDivergenceAnalysis>();
196 TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
197
198 // Loop over all of the blocks in a function, tracking all of the blocks that
199 // return.
200 SmallVector<BasicBlock *, 4> ReturningBlocks;
201 SmallVector<BasicBlock *, 4> UnreachableBlocks;
202
203 // Dummy return block for infinite loop.
204 BasicBlock *DummyReturnBB = nullptr;
205
206 bool Changed = false;
207 std::vector<DominatorTree::UpdateType> Updates;
208
209 // TODO: For now we unify all exit blocks, even though they are uniformly
210 // reachable, if there are any exits not uniformly reached. This is to
211 // workaround the limitation of structurizer, which can not handle multiple
212 // function exits. After structurizer is able to handle multiple function
213 // exits, we should only unify UnreachableBlocks that are not uniformly
214 // reachable.
215 bool HasDivergentExitBlock = llvm::any_of(
216 PDT.roots(), [&](auto BB) { return !isUniformlyReached(DA, *BB); });
217
218 for (BasicBlock *BB : PDT.roots()) {
219 if (isa<ReturnInst>(BB->getTerminator())) {
220 if (HasDivergentExitBlock)
221 ReturningBlocks.push_back(BB);
222 } else if (isa<UnreachableInst>(BB->getTerminator())) {
223 if (HasDivergentExitBlock)
224 UnreachableBlocks.push_back(BB);
225 } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) {
226
227 ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext());
228 if (DummyReturnBB == nullptr) {
229 DummyReturnBB = BasicBlock::Create(F.getContext(),
230 "DummyReturnBlock", &F);
231 Type *RetTy = F.getReturnType();
232 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(RetTy);
233 ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);
234 ReturningBlocks.push_back(DummyReturnBB);
235 }
236
237 if (BI->isUnconditional()) {
238 BasicBlock *LoopHeaderBB = BI->getSuccessor(0);
239 BI->eraseFromParent(); // Delete the unconditional branch.
240 // Add a new conditional branch with a dummy edge to the return block.
241 BranchInst::Create(LoopHeaderBB, DummyReturnBB, BoolTrue, BB);
242 Updates.push_back({DominatorTree::Insert, BB, DummyReturnBB});
243 } else { // Conditional branch.
244 SmallVector<BasicBlock *, 2> Successors(successors(BB));
245
246 // Create a new transition block to hold the conditional branch.
247 BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock");
248
249 Updates.reserve(Updates.size() + 2 * Successors.size() + 2);
250
251 // 'Successors' become successors of TransitionBB instead of BB,
252 // and TransitionBB becomes a single successor of BB.
253 Updates.push_back({DominatorTree::Insert, BB, TransitionBB});
254 for (BasicBlock *Successor : Successors) {
255 Updates.push_back({DominatorTree::Insert, TransitionBB, Successor});
256 Updates.push_back({DominatorTree::Delete, BB, Successor});
257 }
258
259 // Create a branch that will always branch to the transition block and
260 // references DummyReturnBB.
261 BB->getTerminator()->eraseFromParent();
262 BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB);
263 Updates.push_back({DominatorTree::Insert, BB, DummyReturnBB});
264 }
265 Changed = true;
266 }
267 }
268
269 if (!UnreachableBlocks.empty()) {
270 BasicBlock *UnreachableBlock = nullptr;
271
272 if (UnreachableBlocks.size() == 1) {
273 UnreachableBlock = UnreachableBlocks.front();
274 } else {
275 UnreachableBlock = BasicBlock::Create(F.getContext(),
276 "UnifiedUnreachableBlock", &F);
277 new UnreachableInst(F.getContext(), UnreachableBlock);
278
279 Updates.reserve(Updates.size() + UnreachableBlocks.size());
280 for (BasicBlock *BB : UnreachableBlocks) {
281 // Remove and delete the unreachable inst.
282 BB->getTerminator()->eraseFromParent();
283 BranchInst::Create(UnreachableBlock, BB);
284 Updates.push_back({DominatorTree::Insert, BB, UnreachableBlock});
285 }
286 Changed = true;
287 }
288
289 if (!ReturningBlocks.empty()) {
290 // Don't create a new unreachable inst if we have a return. The
291 // structurizer/annotator can't handle the multiple exits
292
293 Type *RetTy = F.getReturnType();
294 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(RetTy);
295 // Remove and delete the unreachable inst.
296 UnreachableBlock->getTerminator()->eraseFromParent();
297
298 Function *UnreachableIntrin =
299 Intrinsic::getDeclaration(F.getParent(), Intrinsic::amdgcn_unreachable);
300
301 // Insert a call to an intrinsic tracking that this is an unreachable
302 // point, in case we want to kill the active lanes or something later.
303 CallInst::Create(UnreachableIntrin, {}, "", UnreachableBlock);
304
305 // Don't create a scalar trap. We would only want to trap if this code was
306 // really reached, but a scalar trap would happen even if no lanes
307 // actually reached here.
308 ReturnInst::Create(F.getContext(), RetVal, UnreachableBlock);
309 ReturningBlocks.push_back(UnreachableBlock);
310 Changed = true;
311 }
312 }
313
314 // FIXME: add PDT here once simplifycfg is ready.
315 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
316 if (RequireAndPreserveDomTree)
317 DTU.applyUpdates(Updates);
318 Updates.clear();
319
320 // Now handle return blocks.
321 if (ReturningBlocks.empty())
322 return Changed; // No blocks return
323
324 if (ReturningBlocks.size() == 1)
325 return Changed; // Already has a single return block
326
327 unifyReturnBlockSet(F, DTU, ReturningBlocks, "UnifiedReturnBlock");
328 return true;
329 }
330