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 191 // If there's only one exit, we don't need to do anything. 192 if (PDT.root_size() <= 1) 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 for (BasicBlock *BB : PDT.roots()) { 210 if (isa<ReturnInst>(BB->getTerminator())) { 211 if (!isUniformlyReached(DA, *BB)) 212 ReturningBlocks.push_back(BB); 213 } else if (isa<UnreachableInst>(BB->getTerminator())) { 214 if (!isUniformlyReached(DA, *BB)) 215 UnreachableBlocks.push_back(BB); 216 } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) { 217 218 ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext()); 219 if (DummyReturnBB == nullptr) { 220 DummyReturnBB = BasicBlock::Create(F.getContext(), 221 "DummyReturnBlock", &F); 222 Type *RetTy = F.getReturnType(); 223 Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy); 224 ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB); 225 ReturningBlocks.push_back(DummyReturnBB); 226 } 227 228 if (BI->isUnconditional()) { 229 BasicBlock *LoopHeaderBB = BI->getSuccessor(0); 230 BI->eraseFromParent(); // Delete the unconditional branch. 231 // Add a new conditional branch with a dummy edge to the return block. 232 BranchInst::Create(LoopHeaderBB, DummyReturnBB, BoolTrue, BB); 233 Updates.push_back({DominatorTree::Insert, BB, DummyReturnBB}); 234 } else { // Conditional branch. 235 SmallVector<BasicBlock *, 2> Successors(succ_begin(BB), succ_end(BB)); 236 237 // Create a new transition block to hold the conditional branch. 238 BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock"); 239 240 Updates.reserve(Updates.size() + 2 * Successors.size() + 2); 241 242 // 'Successors' become successors of TransitionBB instead of BB, 243 // and TransitionBB becomes a single successor of BB. 244 Updates.push_back({DominatorTree::Insert, BB, TransitionBB}); 245 for (BasicBlock *Successor : Successors) { 246 Updates.push_back({DominatorTree::Insert, TransitionBB, Successor}); 247 Updates.push_back({DominatorTree::Delete, BB, Successor}); 248 } 249 250 // Create a branch that will always branch to the transition block and 251 // references DummyReturnBB. 252 BB->getTerminator()->eraseFromParent(); 253 BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB); 254 Updates.push_back({DominatorTree::Insert, BB, DummyReturnBB}); 255 } 256 Changed = true; 257 } 258 } 259 260 if (!UnreachableBlocks.empty()) { 261 BasicBlock *UnreachableBlock = nullptr; 262 263 if (UnreachableBlocks.size() == 1) { 264 UnreachableBlock = UnreachableBlocks.front(); 265 } else { 266 UnreachableBlock = BasicBlock::Create(F.getContext(), 267 "UnifiedUnreachableBlock", &F); 268 new UnreachableInst(F.getContext(), UnreachableBlock); 269 270 Updates.reserve(Updates.size() + UnreachableBlocks.size()); 271 for (BasicBlock *BB : UnreachableBlocks) { 272 // Remove and delete the unreachable inst. 273 BB->getTerminator()->eraseFromParent(); 274 BranchInst::Create(UnreachableBlock, BB); 275 Updates.push_back({DominatorTree::Insert, BB, UnreachableBlock}); 276 } 277 Changed = true; 278 } 279 280 if (!ReturningBlocks.empty()) { 281 // Don't create a new unreachable inst if we have a return. The 282 // structurizer/annotator can't handle the multiple exits 283 284 Type *RetTy = F.getReturnType(); 285 Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy); 286 // Remove and delete the unreachable inst. 287 UnreachableBlock->getTerminator()->eraseFromParent(); 288 289 Function *UnreachableIntrin = 290 Intrinsic::getDeclaration(F.getParent(), Intrinsic::amdgcn_unreachable); 291 292 // Insert a call to an intrinsic tracking that this is an unreachable 293 // point, in case we want to kill the active lanes or something later. 294 CallInst::Create(UnreachableIntrin, {}, "", UnreachableBlock); 295 296 // Don't create a scalar trap. We would only want to trap if this code was 297 // really reached, but a scalar trap would happen even if no lanes 298 // actually reached here. 299 ReturnInst::Create(F.getContext(), RetVal, UnreachableBlock); 300 ReturningBlocks.push_back(UnreachableBlock); 301 Changed = true; 302 } 303 } 304 305 // FIXME: add PDT here once simplifycfg is ready. 306 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager); 307 if (RequireAndPreserveDomTree) 308 DTU.applyUpdates(Updates); 309 Updates.clear(); 310 311 // Now handle return blocks. 312 if (ReturningBlocks.empty()) 313 return Changed; // No blocks return 314 315 if (ReturningBlocks.size() == 1) 316 return Changed; // Already has a single return block 317 318 unifyReturnBlockSet(F, DTU, ReturningBlocks, "UnifiedReturnBlock"); 319 return true; 320 } 321