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 "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/PostDominators.h"
30#include "llvm/Analysis/TargetTransformInfo.h"
31#include "llvm/Analysis/UniformityAnalysis.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/BasicBlockUtils.h"
49#include "llvm/Transforms/Utils/Local.h"
50
51using namespace llvm;
52
53#define DEBUG_TYPE "amdgpu-unify-divergent-exit-nodes"
54
55namespace {
56
57class AMDGPUUnifyDivergentExitNodesImpl {
58private:
59 const TargetTransformInfo *TTI = nullptr;
60
61public:
62 AMDGPUUnifyDivergentExitNodesImpl() = delete;
63 AMDGPUUnifyDivergentExitNodesImpl(const TargetTransformInfo *TTI)
64 : TTI(TTI) {}
65
66 // We can preserve non-critical-edgeness when we unify function exit nodes
67 BasicBlock *unifyReturnBlockSet(Function &F, DomTreeUpdater &DTU,
68 ArrayRef<BasicBlock *> ReturningBlocks,
69 StringRef Name);
70 bool run(Function &F, DominatorTree *DT, const PostDominatorTree &PDT,
71 const UniformityInfo &UA);
72};
73
74class AMDGPUUnifyDivergentExitNodes : public FunctionPass {
75public:
76 static char ID;
77 AMDGPUUnifyDivergentExitNodes() : FunctionPass(ID) {}
78 void getAnalysisUsage(AnalysisUsage &AU) const override;
79 bool runOnFunction(Function &F) override;
80};
81} // end anonymous namespace
82
83char AMDGPUUnifyDivergentExitNodes::ID = 0;
84
85char &llvm::AMDGPUUnifyDivergentExitNodesID = AMDGPUUnifyDivergentExitNodes::ID;
86
87INITIALIZE_PASS_BEGIN(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,
88 "Unify divergent function exit nodes", false, false)
89INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
90INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
91INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
92INITIALIZE_PASS_END(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,
93 "Unify divergent function exit nodes", false, false)
94
95void AMDGPUUnifyDivergentExitNodes::getAnalysisUsage(AnalysisUsage &AU) const {
96 if (RequireAndPreserveDomTree)
97 AU.addRequired<DominatorTreeWrapperPass>();
98
99 AU.addRequired<PostDominatorTreeWrapperPass>();
100
101 AU.addRequired<UniformityInfoWrapperPass>();
102
103 if (RequireAndPreserveDomTree) {
104 AU.addPreserved<DominatorTreeWrapperPass>();
105 // FIXME: preserve PostDominatorTreeWrapperPass
106 }
107
108 // We preserve the non-critical-edgeness property
109 AU.addPreservedID(ID&: BreakCriticalEdgesID);
110
111 FunctionPass::getAnalysisUsage(AU);
112
113 AU.addRequired<TargetTransformInfoWrapperPass>();
114}
115
116/// \returns true if \p BB is reachable through only uniform branches.
117/// XXX - Is there a more efficient way to find this?
118static bool isUniformlyReached(const UniformityInfo &UA, BasicBlock &BB) {
119 SmallVector<BasicBlock *, 8> Stack(predecessors(BB: &BB));
120 SmallPtrSet<BasicBlock *, 8> Visited;
121
122 while (!Stack.empty()) {
123 BasicBlock *Top = Stack.pop_back_val();
124 if (!UA.isUniform(I: Top->getTerminator()))
125 return false;
126
127 for (BasicBlock *Pred : predecessors(BB: Top)) {
128 if (Visited.insert(Ptr: Pred).second)
129 Stack.push_back(Elt: Pred);
130 }
131 }
132
133 return true;
134}
135
136BasicBlock *AMDGPUUnifyDivergentExitNodesImpl::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(Context&: F.getContext(), Name, Parent: &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(Ty: F.getReturnType(), NumReservedValues: ReturningBlocks.size(),
151 Name: "UnifiedRetVal");
152 B.CreateRet(V: 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(n: 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(V: BB->getTerminator()->getOperand(i: 0), BB);
164
165 // Remove and delete the return inst.
166 BB->getTerminator()->eraseFromParent();
167 UncondBrInst::Create(IfTrue: NewRetBlock, InsertBefore: BB);
168 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: 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: *TTI, DTU: RequireAndPreserveDomTree ? &DTU : nullptr,
178 Options: SimplifyCFGOptions().bonusInstThreshold(I: 2));
179 }
180
181 return NewRetBlock;
182}
183
184static BasicBlock *
185createDummyReturnBlock(Function &F,
186 SmallVector<BasicBlock *, 4> &ReturningBlocks) {
187 BasicBlock *DummyReturnBB =
188 BasicBlock::Create(Context&: F.getContext(), Name: "DummyReturnBlock", Parent: &F);
189 Type *RetTy = F.getReturnType();
190 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(T: RetTy);
191 ReturnInst::Create(C&: F.getContext(), retVal: RetVal, InsertBefore: DummyReturnBB);
192 ReturningBlocks.push_back(Elt: DummyReturnBB);
193 return DummyReturnBB;
194}
195
196/// Handle conditional branch instructions (-> 2 targets) and callbr
197/// instructions with N targets.
198static void handleNBranch(Function &F, BasicBlock *BB, Instruction *BI,
199 BasicBlock *DummyReturnBB,
200 std::vector<DominatorTree::UpdateType> &Updates) {
201 SmallVector<BasicBlock *, 2> Successors(successors(BB));
202
203 // Create a new transition block to hold the conditional branch.
204 BasicBlock *TransitionBB = BB->splitBasicBlock(I: BI, BBName: "TransitionBlock");
205
206 Updates.reserve(n: Updates.size() + 2 * Successors.size() + 2);
207
208 // 'Successors' become successors of TransitionBB instead of BB,
209 // and TransitionBB becomes a single successor of BB.
210 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: TransitionBB);
211 for (BasicBlock *Successor : Successors) {
212 Updates.emplace_back(args: DominatorTree::Insert, args&: TransitionBB, args&: Successor);
213 Updates.emplace_back(args: DominatorTree::Delete, args&: BB, args&: Successor);
214 }
215
216 // Create a branch that will always branch to the transition block and
217 // references DummyReturnBB.
218 BB->getTerminator()->eraseFromParent();
219 CondBrInst::Create(Cond: ConstantInt::getTrue(Context&: F.getContext()), IfTrue: TransitionBB,
220 IfFalse: DummyReturnBB, InsertBefore: BB);
221 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: DummyReturnBB);
222}
223
224bool AMDGPUUnifyDivergentExitNodesImpl::run(Function &F, DominatorTree *DT,
225 const PostDominatorTree &PDT,
226 const UniformityInfo &UA) {
227 if (PDT.root_size() == 0 ||
228 (PDT.root_size() == 1 && !isa<UncondBrInst, CondBrInst, CallBrInst>(
229 Val: PDT.getRoot()->getTerminator())))
230 return false;
231
232 // Loop over all of the blocks in a function, tracking all of the blocks that
233 // return.
234 SmallVector<BasicBlock *, 4> ReturningBlocks;
235 SmallVector<BasicBlock *, 4> UnreachableBlocks;
236
237 // Dummy return block for infinite loop.
238 BasicBlock *DummyReturnBB = nullptr;
239
240 bool Changed = false;
241 std::vector<DominatorTree::UpdateType> Updates;
242
243 // TODO: For now we unify all exit blocks, even though they are uniformly
244 // reachable, if there are any exits not uniformly reached. This is to
245 // workaround the limitation of structurizer, which can not handle multiple
246 // function exits. After structurizer is able to handle multiple function
247 // exits, we should only unify UnreachableBlocks that are not uniformly
248 // reachable.
249 bool HasDivergentExitBlock = llvm::any_of(
250 Range: PDT.roots(), P: [&](auto BB) { return !isUniformlyReached(UA, *BB); });
251
252 for (BasicBlock *BB : PDT.roots()) {
253 Instruction *Term = BB->getTerminator();
254 if (auto *RI = dyn_cast<ReturnInst>(Val: Term)) {
255 auto *CI = dyn_cast_or_null<CallInst>(Val: RI->getPrevNode());
256 if (CI && CI->isMustTailCall())
257 continue;
258 if (HasDivergentExitBlock)
259 ReturningBlocks.push_back(Elt: BB);
260 } else if (isa<UnreachableInst>(Val: Term)) {
261 if (HasDivergentExitBlock)
262 UnreachableBlocks.push_back(Elt: BB);
263 } else if (UncondBrInst *BI = dyn_cast<UncondBrInst>(Val: Term)) {
264 if (!DummyReturnBB)
265 DummyReturnBB = createDummyReturnBlock(F, ReturningBlocks);
266
267 BasicBlock *LoopHeaderBB = BI->getSuccessor();
268 BI->eraseFromParent(); // Delete the unconditional branch.
269 // Add a new conditional branch with a dummy edge to the return block.
270 CondBrInst::Create(Cond: ConstantInt::getTrue(Context&: F.getContext()), IfTrue: LoopHeaderBB,
271 IfFalse: DummyReturnBB, InsertBefore: BB);
272 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: DummyReturnBB);
273 Changed = true;
274 } else if (isa<CondBrInst, CallBrInst>(Val: Term)) {
275 if (!DummyReturnBB)
276 DummyReturnBB = createDummyReturnBlock(F, ReturningBlocks);
277
278 handleNBranch(F, BB, BI: Term, DummyReturnBB, Updates);
279 Changed = true;
280 } else {
281 llvm_unreachable("unsupported block terminator");
282 }
283 }
284
285 if (!UnreachableBlocks.empty()) {
286 BasicBlock *UnreachableBlock = nullptr;
287
288 if (UnreachableBlocks.size() == 1) {
289 UnreachableBlock = UnreachableBlocks.front();
290 } else {
291 UnreachableBlock = BasicBlock::Create(Context&: F.getContext(),
292 Name: "UnifiedUnreachableBlock", Parent: &F);
293 new UnreachableInst(F.getContext(), UnreachableBlock);
294
295 Updates.reserve(n: Updates.size() + UnreachableBlocks.size());
296 for (BasicBlock *BB : UnreachableBlocks) {
297 // Remove and delete the unreachable inst.
298 BB->getTerminator()->eraseFromParent();
299 UncondBrInst::Create(IfTrue: UnreachableBlock, InsertBefore: BB);
300 Updates.emplace_back(args: DominatorTree::Insert, args&: BB, args&: UnreachableBlock);
301 }
302 Changed = true;
303 }
304
305 if (!ReturningBlocks.empty()) {
306 // Don't create a new unreachable inst if we have a return. The
307 // structurizer/annotator can't handle the multiple exits
308
309 Type *RetTy = F.getReturnType();
310 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(T: RetTy);
311 // Remove and delete the unreachable inst.
312 UnreachableBlock->getTerminator()->eraseFromParent();
313
314 Function *UnreachableIntrin = Intrinsic::getOrInsertDeclaration(
315 M: F.getParent(), id: Intrinsic::amdgcn_unreachable);
316
317 // Insert a call to an intrinsic tracking that this is an unreachable
318 // point, in case we want to kill the active lanes or something later.
319 CallInst::Create(Func: UnreachableIntrin, Args: {}, NameStr: "", InsertBefore: UnreachableBlock);
320
321 // Don't create a scalar trap. We would only want to trap if this code was
322 // really reached, but a scalar trap would happen even if no lanes
323 // actually reached here.
324 ReturnInst::Create(C&: F.getContext(), retVal: RetVal, InsertBefore: UnreachableBlock);
325 ReturningBlocks.push_back(Elt: UnreachableBlock);
326 Changed = true;
327 }
328 }
329
330 // FIXME: add PDT here once simplifycfg is ready.
331 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
332 if (RequireAndPreserveDomTree)
333 DTU.applyUpdates(Updates);
334 Updates.clear();
335
336 // Now handle return blocks.
337 if (ReturningBlocks.empty())
338 return Changed; // No blocks return
339
340 if (ReturningBlocks.size() == 1)
341 return Changed; // Already has a single return block
342
343 unifyReturnBlockSet(F, DTU, ReturningBlocks, Name: "UnifiedReturnBlock");
344 return true;
345}
346
347bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {
348 DominatorTree *DT = nullptr;
349 if (RequireAndPreserveDomTree)
350 DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
351 const auto &PDT =
352 getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
353 const auto &UA = getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
354 const auto *TranformInfo =
355 &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
356 return AMDGPUUnifyDivergentExitNodesImpl(TranformInfo).run(F, DT, PDT, UA);
357}
358
359PreservedAnalyses
360AMDGPUUnifyDivergentExitNodesPass::run(Function &F,
361 FunctionAnalysisManager &AM) {
362 DominatorTree *DT = nullptr;
363 if (RequireAndPreserveDomTree)
364 DT = &AM.getResult<DominatorTreeAnalysis>(IR&: F);
365
366 const auto &PDT = AM.getResult<PostDominatorTreeAnalysis>(IR&: F);
367 const auto &UA = AM.getResult<UniformityInfoAnalysis>(IR&: F);
368 const auto *TransformInfo = &AM.getResult<TargetIRAnalysis>(IR&: F);
369 return AMDGPUUnifyDivergentExitNodesImpl(TransformInfo).run(F, DT, PDT, UA)
370 ? PreservedAnalyses::none()
371 : PreservedAnalyses::all();
372}
373