376 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This is a variant of the UnifyFunctionExitNodes pass. Rather than ensuring10// there is at most one ret and one unreachable instruction, it ensures there is11// at most one divergent exiting block.12//13// StructurizeCFG can't deal with multi-exit regions formed by branches to14// multiple return nodes. It is not desirable to structurize regions with15// uniform branches, so unifying those to the same return block as divergent16// branches inhibits use of scalar branching. It still can't deal with the case17// where one branch goes to return, and one unreachable. Replace unreachable in18// this case with a return.19//20//===----------------------------------------------------------------------===//21 22#include "AMDGPUUnifyDivergentExitNodes.h"23#include "AMDGPU.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/PostDominators.h"30#include "llvm/Analysis/TargetTransformInfo.h"31#include "llvm/Analysis/UniformityAnalysis.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/BasicBlockUtils.h"49#include "llvm/Transforms/Utils/Local.h"50 51using namespace llvm;52 53#define DEBUG_TYPE "amdgpu-unify-divergent-exit-nodes"54 55namespace {56 57class AMDGPUUnifyDivergentExitNodesImpl {58private:59 const TargetTransformInfo *TTI = nullptr;60 61public:62 AMDGPUUnifyDivergentExitNodesImpl() = delete;63 AMDGPUUnifyDivergentExitNodesImpl(const TargetTransformInfo *TTI)64 : TTI(TTI) {}65 66 // We can preserve non-critical-edgeness when we unify function exit nodes67 BasicBlock *unifyReturnBlockSet(Function &F, DomTreeUpdater &DTU,68 ArrayRef<BasicBlock *> ReturningBlocks,69 StringRef Name);70 bool run(Function &F, DominatorTree *DT, const PostDominatorTree &PDT,71 const UniformityInfo &UA);72};73 74class AMDGPUUnifyDivergentExitNodes : public FunctionPass {75public:76 static char ID;77 AMDGPUUnifyDivergentExitNodes() : FunctionPass(ID) {}78 void getAnalysisUsage(AnalysisUsage &AU) const override;79 bool runOnFunction(Function &F) override;80};81} // end anonymous namespace82 83char AMDGPUUnifyDivergentExitNodes::ID = 0;84 85char &llvm::AMDGPUUnifyDivergentExitNodesID = AMDGPUUnifyDivergentExitNodes::ID;86 87INITIALIZE_PASS_BEGIN(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,88 "Unify divergent function exit nodes", false, false)89INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)90INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)91INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)92INITIALIZE_PASS_END(AMDGPUUnifyDivergentExitNodes, DEBUG_TYPE,93 "Unify divergent function exit nodes", false, false)94 95void AMDGPUUnifyDivergentExitNodes::getAnalysisUsage(AnalysisUsage &AU) const {96 if (RequireAndPreserveDomTree)97 AU.addRequired<DominatorTreeWrapperPass>();98 99 AU.addRequired<PostDominatorTreeWrapperPass>();100 101 AU.addRequired<UniformityInfoWrapperPass>();102 103 if (RequireAndPreserveDomTree) {104 AU.addPreserved<DominatorTreeWrapperPass>();105 // FIXME: preserve PostDominatorTreeWrapperPass106 }107 108 // We preserve the non-critical-edgeness property109 AU.addPreservedID(BreakCriticalEdgesID);110 111 FunctionPass::getAnalysisUsage(AU);112 113 AU.addRequired<TargetTransformInfoWrapperPass>();114}115 116/// \returns true if \p BB is reachable through only uniform branches.117/// XXX - Is there a more efficient way to find this?118static bool isUniformlyReached(const UniformityInfo &UA, 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 (!UA.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 136BasicBlock *AMDGPUUnifyDivergentExitNodesImpl::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 PHI140 // nodes (if the function returns values), and convert all of the return141 // 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 an156 // 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 that161 // 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.emplace_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 184static BasicBlock *185createDummyReturnBlock(Function &F,186 SmallVector<BasicBlock *, 4> &ReturningBlocks) {187 BasicBlock *DummyReturnBB =188 BasicBlock::Create(F.getContext(), "DummyReturnBlock", &F);189 Type *RetTy = F.getReturnType();190 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(RetTy);191 ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);192 ReturningBlocks.push_back(DummyReturnBB);193 return DummyReturnBB;194}195 196/// Handle conditional branch instructions (-> 2 targets) and callbr197/// instructions with N targets.198static void handleNBranch(Function &F, BasicBlock *BB, Instruction *BI,199 BasicBlock *DummyReturnBB,200 std::vector<DominatorTree::UpdateType> &Updates) {201 SmallVector<BasicBlock *, 2> Successors(successors(BB));202 203 // Create a new transition block to hold the conditional branch.204 BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock");205 206 Updates.reserve(Updates.size() + 2 * Successors.size() + 2);207 208 // 'Successors' become successors of TransitionBB instead of BB,209 // and TransitionBB becomes a single successor of BB.210 Updates.emplace_back(DominatorTree::Insert, BB, TransitionBB);211 for (BasicBlock *Successor : Successors) {212 Updates.emplace_back(DominatorTree::Insert, TransitionBB, Successor);213 Updates.emplace_back(DominatorTree::Delete, BB, Successor);214 }215 216 // Create a branch that will always branch to the transition block and217 // references DummyReturnBB.218 BB->getTerminator()->eraseFromParent();219 BranchInst::Create(TransitionBB, DummyReturnBB,220 ConstantInt::getTrue(F.getContext()), BB);221 Updates.emplace_back(DominatorTree::Insert, BB, DummyReturnBB);222}223 224bool AMDGPUUnifyDivergentExitNodesImpl::run(Function &F, DominatorTree *DT,225 const PostDominatorTree &PDT,226 const UniformityInfo &UA) {227 if (PDT.root_size() == 0 ||228 (PDT.root_size() == 1 &&229 !isa<BranchInst, CallBrInst>(PDT.getRoot()->getTerminator())))230 return false;231 232 // Loop over all of the blocks in a function, tracking all of the blocks that233 // return.234 SmallVector<BasicBlock *, 4> ReturningBlocks;235 SmallVector<BasicBlock *, 4> UnreachableBlocks;236 237 // Dummy return block for infinite loop.238 BasicBlock *DummyReturnBB = nullptr;239 240 bool Changed = false;241 std::vector<DominatorTree::UpdateType> Updates;242 243 // TODO: For now we unify all exit blocks, even though they are uniformly244 // reachable, if there are any exits not uniformly reached. This is to245 // workaround the limitation of structurizer, which can not handle multiple246 // function exits. After structurizer is able to handle multiple function247 // exits, we should only unify UnreachableBlocks that are not uniformly248 // reachable.249 bool HasDivergentExitBlock = llvm::any_of(250 PDT.roots(), [&](auto BB) { return !isUniformlyReached(UA, *BB); });251 252 for (BasicBlock *BB : PDT.roots()) {253 if (auto *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {254 auto *CI = dyn_cast_or_null<CallInst>(RI->getPrevNode());255 if (CI && CI->isMustTailCall())256 continue;257 if (HasDivergentExitBlock)258 ReturningBlocks.push_back(BB);259 } else if (isa<UnreachableInst>(BB->getTerminator())) {260 if (HasDivergentExitBlock)261 UnreachableBlocks.push_back(BB);262 } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) {263 if (!DummyReturnBB)264 DummyReturnBB = createDummyReturnBlock(F, ReturningBlocks);265 266 if (BI->isUnconditional()) {267 BasicBlock *LoopHeaderBB = BI->getSuccessor(0);268 BI->eraseFromParent(); // Delete the unconditional branch.269 // Add a new conditional branch with a dummy edge to the return block.270 BranchInst::Create(LoopHeaderBB, DummyReturnBB,271 ConstantInt::getTrue(F.getContext()), BB);272 Updates.emplace_back(DominatorTree::Insert, BB, DummyReturnBB);273 } else {274 handleNBranch(F, BB, BI, DummyReturnBB, Updates);275 }276 Changed = true;277 } else if (CallBrInst *CBI = dyn_cast<CallBrInst>(BB->getTerminator())) {278 if (!DummyReturnBB)279 DummyReturnBB = createDummyReturnBlock(F, ReturningBlocks);280 281 handleNBranch(F, BB, CBI, DummyReturnBB, Updates);282 Changed = true;283 } else {284 llvm_unreachable("unsupported block terminator");285 }286 }287 288 if (!UnreachableBlocks.empty()) {289 BasicBlock *UnreachableBlock = nullptr;290 291 if (UnreachableBlocks.size() == 1) {292 UnreachableBlock = UnreachableBlocks.front();293 } else {294 UnreachableBlock = BasicBlock::Create(F.getContext(),295 "UnifiedUnreachableBlock", &F);296 new UnreachableInst(F.getContext(), UnreachableBlock);297 298 Updates.reserve(Updates.size() + UnreachableBlocks.size());299 for (BasicBlock *BB : UnreachableBlocks) {300 // Remove and delete the unreachable inst.301 BB->getTerminator()->eraseFromParent();302 BranchInst::Create(UnreachableBlock, BB);303 Updates.emplace_back(DominatorTree::Insert, BB, UnreachableBlock);304 }305 Changed = true;306 }307 308 if (!ReturningBlocks.empty()) {309 // Don't create a new unreachable inst if we have a return. The310 // structurizer/annotator can't handle the multiple exits311 312 Type *RetTy = F.getReturnType();313 Value *RetVal = RetTy->isVoidTy() ? nullptr : PoisonValue::get(RetTy);314 // Remove and delete the unreachable inst.315 UnreachableBlock->getTerminator()->eraseFromParent();316 317 Function *UnreachableIntrin = Intrinsic::getOrInsertDeclaration(318 F.getParent(), Intrinsic::amdgcn_unreachable);319 320 // Insert a call to an intrinsic tracking that this is an unreachable321 // point, in case we want to kill the active lanes or something later.322 CallInst::Create(UnreachableIntrin, {}, "", UnreachableBlock);323 324 // Don't create a scalar trap. We would only want to trap if this code was325 // really reached, but a scalar trap would happen even if no lanes326 // actually reached here.327 ReturnInst::Create(F.getContext(), RetVal, UnreachableBlock);328 ReturningBlocks.push_back(UnreachableBlock);329 Changed = true;330 }331 }332 333 // FIXME: add PDT here once simplifycfg is ready.334 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);335 if (RequireAndPreserveDomTree)336 DTU.applyUpdates(Updates);337 Updates.clear();338 339 // Now handle return blocks.340 if (ReturningBlocks.empty())341 return Changed; // No blocks return342 343 if (ReturningBlocks.size() == 1)344 return Changed; // Already has a single return block345 346 unifyReturnBlockSet(F, DTU, ReturningBlocks, "UnifiedReturnBlock");347 return true;348}349 350bool AMDGPUUnifyDivergentExitNodes::runOnFunction(Function &F) {351 DominatorTree *DT = nullptr;352 if (RequireAndPreserveDomTree)353 DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();354 const auto &PDT =355 getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();356 const auto &UA = getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();357 const auto *TranformInfo =358 &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);359 return AMDGPUUnifyDivergentExitNodesImpl(TranformInfo).run(F, DT, PDT, UA);360}361 362PreservedAnalyses363AMDGPUUnifyDivergentExitNodesPass::run(Function &F,364 FunctionAnalysisManager &AM) {365 DominatorTree *DT = nullptr;366 if (RequireAndPreserveDomTree)367 DT = &AM.getResult<DominatorTreeAnalysis>(F);368 369 const auto &PDT = AM.getResult<PostDominatorTreeAnalysis>(F);370 const auto &UA = AM.getResult<UniformityInfoAnalysis>(F);371 const auto *TransformInfo = &AM.getResult<TargetIRAnalysis>(F);372 return AMDGPUUnifyDivergentExitNodesImpl(TransformInfo).run(F, DT, PDT, UA)373 ? PreservedAnalyses::none()374 : PreservedAnalyses::all();375}376