lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp

   1 //===- AMDGPUUnifyDivergentExitNodes.cpp ----------------------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This is a variant of the UnifyDivergentExitNodes pass. Rather than ensuring
  11 // there is at most one ret and one unreachable instruction, it ensures there is
  12 // at most one divergent exiting block.
  13 //
  14 // StructurizeCFG can't deal with multi-exit regions formed by branches to
  15 // multiple return nodes. It is not desirable to structurize regions with
  16 // uniform branches, so unifying those to the same return block as divergent
  17 // branches inhibits use of scalar branching. It still can't deal with the case
  18 // where one branch goes to return, and one unreachable. Replace unreachable in
  19 // this case with a return.
  20 //
  21 //===----------------------------------------------------------------------===//
  22
  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/LegacyDivergenceAnalysis.h"
  29 #include "llvm/Analysis/PostDominators.h"
  30 #include "llvm/Analysis/TargetTransformInfo.h"
  31 #include "llvm/Transforms/Utils/Local.h"
  32 #include "llvm/IR/BasicBlock.h"
  33 #include "llvm/IR/CFG.h"
  34 #include "llvm/IR/Constants.h"
  35 #include "llvm/IR/Function.h"
  36 #include "llvm/IR/InstrTypes.h"
  37 #include "llvm/IR/Instructions.h"
  38 #include "llvm/IR/Intrinsics.h"
  39 #include "llvm/IR/Type.h"
  40 #include "llvm/Pass.h"
  41 #include "llvm/Support/Casting.h"
  42 #include "llvm/Transforms/Scalar.h"
  43 #include "llvm/Transforms/Utils.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 = BasicBlock::Create(F.getContext(),
 203                                                       "TransitionBlock", &F);
 204
 205         // Move BI from BB to the new transition block.
 206         BI->removeFromParent();
 207         TransitionBB->getInstList().push_back(BI);
 208
 209         // Create a branch that will always branch to the transition block.
 210         BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB);
 211       }
 212     }
 213   }
 214
 215   if (!UnreachableBlocks.empty()) {
 216     BasicBlock *UnreachableBlock = nullptr;
 217
 218     if (UnreachableBlocks.size() == 1) {
 219       UnreachableBlock = UnreachableBlocks.front();
 220     } else {
 221       UnreachableBlock = BasicBlock::Create(F.getContext(),
 222                                             "UnifiedUnreachableBlock", &F);
 223       new UnreachableInst(F.getContext(), UnreachableBlock);
 224
 225       for (BasicBlock *BB : UnreachableBlocks) {
 226         // Remove and delete the unreachable inst.
 227         BB->getTerminator()->eraseFromParent();
 228         BranchInst::Create(UnreachableBlock, BB);
 229       }
 230     }
 231
 232     if (!ReturningBlocks.empty()) {
 233       // Don't create a new unreachable inst if we have a return. The
 234       // structurizer/annotator can't handle the multiple exits
 235
 236       Type *RetTy = F.getReturnType();
 237       Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy);
 238       // Remove and delete the unreachable inst.
 239       UnreachableBlock->getTerminator()->eraseFromParent();
 240
 241       Function *UnreachableIntrin =
 242         Intrinsic::getDeclaration(F.getParent(), Intrinsic::amdgcn_unreachable);
 243
 244       // Insert a call to an intrinsic tracking that this is an unreachable
 245       // point, in case we want to kill the active lanes or something later.
 246       CallInst::Create(UnreachableIntrin, {}, "", UnreachableBlock);
 247
 248       // Don't create a scalar trap. We would only want to trap if this code was
 249       // really reached, but a scalar trap would happen even if no lanes
 250       // actually reached here.
 251       ReturnInst::Create(F.getContext(), RetVal, UnreachableBlock);
 252       ReturningBlocks.push_back(UnreachableBlock);
 253     }
 254   }
 255
 256   // Now handle return blocks.
 257   if (ReturningBlocks.empty())
 258     return false; // No blocks return
 259
 260   if (ReturningBlocks.size() == 1)
 261     return false; // Already has a single return block
 262
 263   const TargetTransformInfo &TTI
 264     = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 265
 266   unifyReturnBlockSet(F, ReturningBlocks, TTI, "UnifiedReturnBlock");
 267   return true;
 268 }