xref: /freebsd/contrib/llvm-project/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp (revision 7ef62cebc2f965b0f640263e179276928885e33d)
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 
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)
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?
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 
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 
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