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