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