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