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 "AMDGPUUnifyDivergentExitNodes.h"
23#include "AMDGPU.h"
24#include "SIDefines.h"
25#include "llvm/ADT/ArrayRef.h"
26#include "llvm/ADT/SmallPtrSet.h"
27#include "llvm/ADT/SmallVector.h"
28#include "llvm/ADT/StringRef.h"
29#include "llvm/Analysis/DomTreeUpdater.h"
30#include "llvm/Analysis/PostDominators.h"
31#include "llvm/Analysis/TargetTransformInfo.h"
32#include "llvm/Analysis/UniformityAnalysis.h"
33#include "llvm/IR/BasicBlock.h"
34#include "llvm/IR/CFG.h"
35#include "llvm/IR/Constants.h"
36#include "llvm/IR/Dominators.h"
37#include "llvm/IR/Function.h"
38#include "llvm/IR/IRBuilder.h"
39#include "llvm/IR/InstrTypes.h"
40#include "llvm/IR/Instructions.h"
41#include "llvm/IR/Intrinsics.h"
42#include "llvm/IR/IntrinsicsAMDGPU.h"
43#include "llvm/IR/Type.h"
44#include "llvm/InitializePasses.h"
45#include "llvm/Pass.h"
46#include "llvm/Support/Casting.h"
47#include "llvm/Transforms/Scalar.h"
48#include "llvm/Transforms/Utils.h"
49#include "llvm/Transforms/Utils/BasicBlockUtils.h"
50#include "llvm/Transforms/Utils/Local.h"
51
52using namespace llvm;
53
54#define DEBUG_TYPE "amdgpu-unify-divergent-exit-nodes"
55
56namespace {
57
58class AMDGPUUnifyDivergentExitNodesImpl {
59private:
60 const TargetTransformInfo *TTI = nullptr;
61
62public:
63 AMDGPUUnifyDivergentExitNodesImpl() = delete;
64 AMDGPUUnifyDivergentExitNodesImpl(const TargetTransformInfo *TTI)
65 : TTI(TTI) {}
66
67 // We can preserve non-critical-edgeness when we unify function exit nodes
68 BasicBlock *unifyReturnBlockSet(Function &F, DomTreeUpdater &DTU,
69 ArrayRef<BasicBlock *> ReturningBlocks,
70 StringRef Name);
71 bool run(Function &F, DominatorTree *DT, const PostDominatorTree &PDT,
72 const UniformityInfo &UA);
73};
74
75class AMDGPUUnifyDivergentExitNodes : public FunctionPass {
76public:
77 static char ID;
78 AMDGPUUnifyDivergentExitNodes() : FunctionPass(ID) {
79 initializeAMDGPUUnifyDivergentExitNodesPass(
80 *PassRegistry::getPassRegistry());
81 }
82 void getAnalysisUsage(AnalysisUsage &AU) const override;
83 bool runOnFunction(Function &F) override;
84};
85} // end anonymous namespace
86
87char AMDGPUUnifyDivergentExitNodes::ID = 0;
88
89char &llvm::AMDGPUUnifyDivergentExitNodesID = AMDGPUUnifyDivergentExitNodes::ID;
90
91INITIALIZE_PASS_BEGIN(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,
92 "Unify divergent function exit nodes", false, false)
93INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
94INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
95INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
96INITIALIZE_PASS_END(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,
97 "Unify divergent function exit nodes", false, false)
98
99void AMDGPUUnifyDivergentExitNodes::getAnalysisUsage(AnalysisUsage &AU) const {
100 if (RequireAndPreserveDomTree)
101 AU.addRequired<DominatorTreeWrapperPass>();
102
103 AU.addRequired<PostDominatorTreeWrapperPass>();
104
105 AU.addRequired<UniformityInfoWrapperPass>();
106
107 if (RequireAndPreserveDomTree) {
108 AU.addPreserved<DominatorTreeWrapperPass>();
109 // FIXME: preserve PostDominatorTreeWrapperPass
110 }
111
112 // We preserve the non-critical-edgeness property
113 AU.addPreservedID(ID&: BreakCriticalEdgesID);
114
115 FunctionPass::getAnalysisUsage(AU);
116
117 AU.addRequired<TargetTransformInfoWrapperPass>();
118}
119
120/// \returns true if \p BB is reachable through only uniform branches.
121/// XXX - Is there a more efficient way to find this?
122static bool isUniformlyReached(const UniformityInfo &UA, BasicBlock &BB) {
123 SmallVector<BasicBlock *, 8> Stack(predecessors(BB: &BB));
124 SmallPtrSet<BasicBlock *, 8> Visited;
125
126 while (!Stack.empty()) {
127 BasicBlock *Top = Stack.pop_back_val();
128 if (!UA.isUniform(I: Top->getTerminator()))
129 return false;
130
131 for (BasicBlock *Pred : predecessors(BB: Top)) {
132 if (Visited.insert(Ptr: Pred).second)
133 Stack.push_back(Elt: Pred);
134 }
135 }
136
137 return true;
138}
139
140BasicBlock *AMDGPUUnifyDivergentExitNodesImpl::unifyReturnBlockSet(
141 Function &F, DomTreeUpdater &DTU, ArrayRef<BasicBlock *> ReturningBlocks,
142 StringRef Name) {
143 // Otherwise, we need to insert a new basic block into the function, add a PHI
144 // nodes (if the function returns values), and convert all of the return
145 // instructions into unconditional branches.
146 BasicBlock *NewRetBlock = BasicBlock::Create(Context&: F.getContext(), Name, Parent: &F);
147 IRBuilder<> B(NewRetBlock);
148
149 PHINode *PN = nullptr;
150 if (F.getReturnType()->isVoidTy()) {
151 B.CreateRetVoid();
152 } else {
153 // If the function doesn't return void... add a PHI node to the block...
154 PN = B.CreatePHI(Ty: F.getReturnType(), NumReservedValues: ReturningBlocks.size(),
155 Name: "UnifiedRetVal");
156 B.CreateRet(V: PN);
157 }
158
159 // Loop over all of the blocks, replacing the return instruction with an
160 // unconditional branch.
161 std::vector<DominatorTree::UpdateType> Updates;
162 Updates.reserve(n: ReturningBlocks.size());
163 for (BasicBlock *BB : ReturningBlocks) {
164 // Add an incoming element to the PHI node for every return instruction that
165 // is merging into this new block...
166 if (PN)
167 PN->addIncoming(V: BB->getTerminator()->getOperand(i: 0), BB);
168
169 // Remove and delete the return inst.
170 BB->getTerminator()->eraseFromParent();
171 BranchInst::Create(IfTrue: NewRetBlock, InsertBefore: BB);
172 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: NewRetBlock);
173 }
174
175 if (RequireAndPreserveDomTree)
176 DTU.applyUpdates(Updates);
177 Updates.clear();
178
179 for (BasicBlock *BB : ReturningBlocks) {
180 // Cleanup possible branch to unconditional branch to the return.
181 simplifyCFG(BB, TTI: *TTI, DTU: RequireAndPreserveDomTree ? &DTU : nullptr,
182 Options: SimplifyCFGOptions().bonusInstThreshold(I: 2));
183 }
184
185 return NewRetBlock;
186}
187
188bool AMDGPUUnifyDivergentExitNodesImpl::run(Function &F, DominatorTree *DT,
189 const PostDominatorTree &PDT,
190 const UniformityInfo &UA) {
191 assert(hasOnlySimpleTerminator(F) && "Unsupported block terminator.");
192
193 if (PDT.root_size() == 0 ||
194 (PDT.root_size() == 1 &&
195 !isa<BranchInst>(Val: PDT.getRoot()->getTerminator())))
196 return false;
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 Range: PDT.roots(), P: [&](auto BB) { return !isUniformlyReached(UA, *BB); });
217
218 for (BasicBlock *BB : PDT.roots()) {
219 if (isa<ReturnInst>(Val: BB->getTerminator())) {
220 if (HasDivergentExitBlock)
221 ReturningBlocks.push_back(Elt: BB);
222 } else if (isa<UnreachableInst>(Val: BB->getTerminator())) {
223 if (HasDivergentExitBlock)
224 UnreachableBlocks.push_back(Elt: BB);
225 } else if (BranchInst *BI = dyn_cast<BranchInst>(Val: BB->getTerminator())) {
226
227 ConstantInt *BoolTrue = ConstantInt::getTrue(Context&: F.getContext());
228 if (DummyReturnBB == nullptr) {
229 DummyReturnBB = BasicBlock::Create(Context&: F.getContext(),
230 Name: "DummyReturnBlock", Parent: &F);
231 Type *RetTy = F.getReturnType();
232 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(T: RetTy);
233 ReturnInst::Create(C&: F.getContext(), retVal: RetVal, InsertBefore: DummyReturnBB);
234 ReturningBlocks.push_back(Elt: DummyReturnBB);
235 }
236
237 if (BI->isUnconditional()) {
238 BasicBlock *LoopHeaderBB = BI->getSuccessor(i: 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(IfTrue: LoopHeaderBB, IfFalse: DummyReturnBB, Cond: BoolTrue, InsertBefore: BB);
242 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: 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(I: BI, BBName: "TransitionBlock");
248
249 Updates.reserve(n: 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.emplace_back(args: DominatorTree::Insert, args&: BB, args&: TransitionBB);
254 for (BasicBlock *Successor : Successors) {
255 Updates.emplace_back(args: DominatorTree::Insert, args&: TransitionBB, args&: Successor);
256 Updates.emplace_back(args: DominatorTree::Delete, args&: BB, args&: 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(IfTrue: TransitionBB, IfFalse: DummyReturnBB, Cond: BoolTrue, InsertBefore: BB);
263 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: 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(Context&: F.getContext(),
276 Name: "UnifiedUnreachableBlock", Parent: &F);
277 new UnreachableInst(F.getContext(), UnreachableBlock);
278
279 Updates.reserve(n: Updates.size() + UnreachableBlocks.size());
280 for (BasicBlock *BB : UnreachableBlocks) {
281 // Remove and delete the unreachable inst.
282 BB->getTerminator()->eraseFromParent();
283 BranchInst::Create(IfTrue: UnreachableBlock, InsertBefore: BB);
284 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: 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(T: RetTy);
295 // Remove and delete the unreachable inst.
296 UnreachableBlock->getTerminator()->eraseFromParent();
297
298 Function *UnreachableIntrin =
299 Intrinsic::getDeclaration(M: F.getParent(), id: 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(Func: UnreachableIntrin, Args: {}, NameStr: "", InsertBefore: 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(C&: F.getContext(), retVal: RetVal, InsertBefore: UnreachableBlock);
309 ReturningBlocks.push_back(Elt: 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, Name: "UnifiedReturnBlock");
328 return true;
329}
330
331bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
332 DominatorTree *DT = nullptr;
333 if (RequireAndPreserveDomTree)
334 DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
335 const auto &PDT =
336 getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
337 const auto &UA = getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
338 const auto *TranformInfo =
339 &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
340 return AMDGPUUnifyDivergentExitNodesImpl(TranformInfo).run(F, DT, PDT, UA);
341}
342
343PreservedAnalyses
344AMDGPUUnifyDivergentExitNodesPass::run(Function &F,
345 FunctionAnalysisManager &AM) {
346 DominatorTree *DT = nullptr;
347 if (RequireAndPreserveDomTree)
348 DT = &AM.getResult<DominatorTreeAnalysis>(IR&: F);
349
350 const auto &PDT = AM.getResult<PostDominatorTreeAnalysis>(IR&: F);
351 const auto &UA = AM.getResult<UniformityInfoAnalysis>(IR&: F);
352 const auto *TransformInfo = &AM.getResult<TargetIRAnalysis>(IR&: F);
353 return AMDGPUUnifyDivergentExitNodesImpl(TransformInfo).run(F, DT, PDT, UA)
354 ? PreservedAnalyses::none()
355 : PreservedAnalyses::all();
356}
357