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 UnifyDivergentExitNodes 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 "llvm/ADT/ArrayRef.h" 24 #include "llvm/ADT/SmallPtrSet.h" 25 #include "llvm/ADT/SmallVector.h" 26 #include "llvm/ADT/StringRef.h" 27 #include "llvm/Analysis/LegacyDivergenceAnalysis.h" 28 #include "llvm/Analysis/PostDominators.h" 29 #include "llvm/Analysis/TargetTransformInfo.h" 30 #include "llvm/IR/BasicBlock.h" 31 #include "llvm/IR/CFG.h" 32 #include "llvm/IR/Constants.h" 33 #include "llvm/IR/Function.h" 34 #include "llvm/IR/InstrTypes.h" 35 #include "llvm/IR/Instructions.h" 36 #include "llvm/IR/Intrinsics.h" 37 #include "llvm/IR/Type.h" 38 #include "llvm/InitializePasses.h" 39 #include "llvm/Pass.h" 40 #include "llvm/Support/Casting.h" 41 #include "llvm/Transforms/Scalar.h" 42 #include "llvm/Transforms/Utils.h" 43 #include "llvm/Transforms/Utils/Local.h" 44 45 using namespace llvm; 46 47 #define DEBUG_TYPE "amdgpu-unify-divergent-exit-nodes" 48 49 namespace { 50 51 class AMDGPUUnifyDivergentExitNodes : public FunctionPass { 52 public: 53 static char ID; // Pass identification, replacement for typeid 54 55 AMDGPUUnifyDivergentExitNodes() : FunctionPass(ID) { 56 initializeAMDGPUUnifyDivergentExitNodesPass(*PassRegistry::getPassRegistry()); 57 } 58 59 // We can preserve non-critical-edgeness when we unify function exit nodes 60 void getAnalysisUsage(AnalysisUsage &AU) const override; 61 bool runOnFunction(Function &F) override; 62 }; 63 64 } // end anonymous namespace 65 66 char AMDGPUUnifyDivergentExitNodes::ID = 0; 67 68 char &llvm::AMDGPUUnifyDivergentExitNodesID = AMDGPUUnifyDivergentExitNodes::ID; 69 70 INITIALIZE_PASS_BEGIN(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE, 71 "Unify divergent function exit nodes", false, false) 72 INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass) 73 INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis) 74 INITIALIZE_PASS_END(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE, 75 "Unify divergent function exit nodes", false, false) 76 77 void AMDGPUUnifyDivergentExitNodes::getAnalysisUsage(AnalysisUsage &AU) const{ 78 // TODO: Preserve dominator tree. 79 AU.addRequired<PostDominatorTreeWrapperPass>(); 80 81 AU.addRequired<LegacyDivergenceAnalysis>(); 82 83 // No divergent values are changed, only blocks and branch edges. 84 AU.addPreserved<LegacyDivergenceAnalysis>(); 85 86 // We preserve the non-critical-edgeness property 87 AU.addPreservedID(BreakCriticalEdgesID); 88 89 // This is a cluster of orthogonal Transforms 90 AU.addPreservedID(LowerSwitchID); 91 FunctionPass::getAnalysisUsage(AU); 92 93 AU.addRequired<TargetTransformInfoWrapperPass>(); 94 } 95 96 /// \returns true if \p BB is reachable through only uniform branches. 97 /// XXX - Is there a more efficient way to find this? 98 static bool isUniformlyReached(const LegacyDivergenceAnalysis &DA, 99 BasicBlock &BB) { 100 SmallVector<BasicBlock *, 8> Stack; 101 SmallPtrSet<BasicBlock *, 8> Visited; 102 103 for (BasicBlock *Pred : predecessors(&BB)) 104 Stack.push_back(Pred); 105 106 while (!Stack.empty()) { 107 BasicBlock *Top = Stack.pop_back_val(); 108 if (!DA.isUniform(Top->getTerminator())) 109 return false; 110 111 for (BasicBlock *Pred : predecessors(Top)) { 112 if (Visited.insert(Pred).second) 113 Stack.push_back(Pred); 114 } 115 } 116 117 return true; 118 } 119 120 static BasicBlock *unifyReturnBlockSet(Function &F, 121 ArrayRef<BasicBlock *> ReturningBlocks, 122 const TargetTransformInfo &TTI, 123 StringRef Name) { 124 // Otherwise, we need to insert a new basic block into the function, add a PHI 125 // nodes (if the function returns values), and convert all of the return 126 // instructions into unconditional branches. 127 BasicBlock *NewRetBlock = BasicBlock::Create(F.getContext(), Name, &F); 128 129 PHINode *PN = nullptr; 130 if (F.getReturnType()->isVoidTy()) { 131 ReturnInst::Create(F.getContext(), nullptr, NewRetBlock); 132 } else { 133 // If the function doesn't return void... add a PHI node to the block... 134 PN = PHINode::Create(F.getReturnType(), ReturningBlocks.size(), 135 "UnifiedRetVal"); 136 NewRetBlock->getInstList().push_back(PN); 137 ReturnInst::Create(F.getContext(), PN, NewRetBlock); 138 } 139 140 // Loop over all of the blocks, replacing the return instruction with an 141 // unconditional branch. 142 for (BasicBlock *BB : ReturningBlocks) { 143 // Add an incoming element to the PHI node for every return instruction that 144 // is merging into this new block... 145 if (PN) 146 PN->addIncoming(BB->getTerminator()->getOperand(0), BB); 147 148 // Remove and delete the return inst. 149 BB->getTerminator()->eraseFromParent(); 150 BranchInst::Create(NewRetBlock, BB); 151 } 152 153 for (BasicBlock *BB : ReturningBlocks) { 154 // Cleanup possible branch to unconditional branch to the return. 155 simplifyCFG(BB, TTI, {2}); 156 } 157 158 return NewRetBlock; 159 } 160 161 bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) { 162 auto &PDT = getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree(); 163 if (PDT.getRoots().size() <= 1) 164 return false; 165 166 LegacyDivergenceAnalysis &DA = getAnalysis<LegacyDivergenceAnalysis>(); 167 168 // Loop over all of the blocks in a function, tracking all of the blocks that 169 // return. 170 SmallVector<BasicBlock *, 4> ReturningBlocks; 171 SmallVector<BasicBlock *, 4> UnreachableBlocks; 172 173 // Dummy return block for infinite loop. 174 BasicBlock *DummyReturnBB = nullptr; 175 176 for (BasicBlock *BB : PDT.getRoots()) { 177 if (isa<ReturnInst>(BB->getTerminator())) { 178 if (!isUniformlyReached(DA, *BB)) 179 ReturningBlocks.push_back(BB); 180 } else if (isa<UnreachableInst>(BB->getTerminator())) { 181 if (!isUniformlyReached(DA, *BB)) 182 UnreachableBlocks.push_back(BB); 183 } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) { 184 185 ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext()); 186 if (DummyReturnBB == nullptr) { 187 DummyReturnBB = BasicBlock::Create(F.getContext(), 188 "DummyReturnBlock", &F); 189 Type *RetTy = F.getReturnType(); 190 Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy); 191 ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB); 192 ReturningBlocks.push_back(DummyReturnBB); 193 } 194 195 if (BI->isUnconditional()) { 196 BasicBlock *LoopHeaderBB = BI->getSuccessor(0); 197 BI->eraseFromParent(); // Delete the unconditional branch. 198 // Add a new conditional branch with a dummy edge to the return block. 199 BranchInst::Create(LoopHeaderBB, DummyReturnBB, BoolTrue, BB); 200 } else { // Conditional branch. 201 // Create a new transition block to hold the conditional branch. 202 BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock"); 203 204 // Create a branch that will always branch to the transition block and 205 // references DummyReturnBB. 206 BB->getTerminator()->eraseFromParent(); 207 BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB); 208 } 209 } 210 } 211 212 if (!UnreachableBlocks.empty()) { 213 BasicBlock *UnreachableBlock = nullptr; 214 215 if (UnreachableBlocks.size() == 1) { 216 UnreachableBlock = UnreachableBlocks.front(); 217 } else { 218 UnreachableBlock = BasicBlock::Create(F.getContext(), 219 "UnifiedUnreachableBlock", &F); 220 new UnreachableInst(F.getContext(), UnreachableBlock); 221 222 for (BasicBlock *BB : UnreachableBlocks) { 223 // Remove and delete the unreachable inst. 224 BB->getTerminator()->eraseFromParent(); 225 BranchInst::Create(UnreachableBlock, BB); 226 } 227 } 228 229 if (!ReturningBlocks.empty()) { 230 // Don't create a new unreachable inst if we have a return. The 231 // structurizer/annotator can't handle the multiple exits 232 233 Type *RetTy = F.getReturnType(); 234 Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy); 235 // Remove and delete the unreachable inst. 236 UnreachableBlock->getTerminator()->eraseFromParent(); 237 238 Function *UnreachableIntrin = 239 Intrinsic::getDeclaration(F.getParent(), Intrinsic::amdgcn_unreachable); 240 241 // Insert a call to an intrinsic tracking that this is an unreachable 242 // point, in case we want to kill the active lanes or something later. 243 CallInst::Create(UnreachableIntrin, {}, "", UnreachableBlock); 244 245 // Don't create a scalar trap. We would only want to trap if this code was 246 // really reached, but a scalar trap would happen even if no lanes 247 // actually reached here. 248 ReturnInst::Create(F.getContext(), RetVal, UnreachableBlock); 249 ReturningBlocks.push_back(UnreachableBlock); 250 } 251 } 252 253 // Now handle return blocks. 254 if (ReturningBlocks.empty()) 255 return false; // No blocks return 256 257 if (ReturningBlocks.size() == 1) 258 return false; // Already has a single return block 259 260 const TargetTransformInfo &TTI 261 = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F); 262 263 unifyReturnBlockSet(F, ReturningBlocks, TTI, "UnifiedReturnBlock"); 264 return true; 265 } 266