llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.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-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
 191   // If there's only one exit, we don't need to do anything.
 192   if (PDT.root_size() <= 1)
 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   for (BasicBlock *BB : PDT.roots()) {
 210     if (isa<ReturnInst>(BB->getTerminator())) {
 211       if (!isUniformlyReached(DA, *BB))
 212         ReturningBlocks.push_back(BB);
 213     } else if (isa<UnreachableInst>(BB->getTerminator())) {
 214       if (!isUniformlyReached(DA, *BB))
 215         UnreachableBlocks.push_back(BB);
 216     } else if (BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator())) {
 217
 218       ConstantInt *BoolTrue = ConstantInt::getTrue(F.getContext());
 219       if (DummyReturnBB == nullptr) {
 220         DummyReturnBB = BasicBlock::Create(F.getContext(),
 221                                            "DummyReturnBlock", &F);
 222         Type *RetTy = F.getReturnType();
 223         Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy);
 224         ReturnInst::Create(F.getContext(), RetVal, DummyReturnBB);
 225         ReturningBlocks.push_back(DummyReturnBB);
 226       }
 227
 228       if (BI->isUnconditional()) {
 229         BasicBlock *LoopHeaderBB = BI->getSuccessor(0);
 230         BI->eraseFromParent(); // Delete the unconditional branch.
 231         // Add a new conditional branch with a dummy edge to the return block.
 232         BranchInst::Create(LoopHeaderBB, DummyReturnBB, BoolTrue, BB);
 233         Updates.push_back({DominatorTree::Insert, BB, DummyReturnBB});
 234       } else { // Conditional branch.
 235         SmallVector<BasicBlock *, 2> Successors(succ_begin(BB), succ_end(BB));
 236
 237         // Create a new transition block to hold the conditional branch.
 238         BasicBlock *TransitionBB = BB->splitBasicBlock(BI, "TransitionBlock");
 239
 240         Updates.reserve(Updates.size() + 2 * Successors.size() + 2);
 241
 242         // 'Successors' become successors of TransitionBB instead of BB,
 243         // and TransitionBB becomes a single successor of BB.
 244         Updates.push_back({DominatorTree::Insert, BB, TransitionBB});
 245         for (BasicBlock *Successor : Successors) {
 246           Updates.push_back({DominatorTree::Insert, TransitionBB, Successor});
 247           Updates.push_back({DominatorTree::Delete, BB, Successor});
 248         }
 249
 250         // Create a branch that will always branch to the transition block and
 251         // references DummyReturnBB.
 252         BB->getTerminator()->eraseFromParent();
 253         BranchInst::Create(TransitionBB, DummyReturnBB, BoolTrue, BB);
 254         Updates.push_back({DominatorTree::Insert, BB, DummyReturnBB});
 255       }
 256       Changed = true;
 257     }
 258   }
 259
 260   if (!UnreachableBlocks.empty()) {
 261     BasicBlock *UnreachableBlock = nullptr;
 262
 263     if (UnreachableBlocks.size() == 1) {
 264       UnreachableBlock = UnreachableBlocks.front();
 265     } else {
 266       UnreachableBlock = BasicBlock::Create(F.getContext(),
 267                                             "UnifiedUnreachableBlock", &F);
 268       new UnreachableInst(F.getContext(), UnreachableBlock);
 269
 270       Updates.reserve(Updates.size() + UnreachableBlocks.size());
 271       for (BasicBlock *BB : UnreachableBlocks) {
 272         // Remove and delete the unreachable inst.
 273         BB->getTerminator()->eraseFromParent();
 274         BranchInst::Create(UnreachableBlock, BB);
 275         Updates.push_back({DominatorTree::Insert, BB, UnreachableBlock});
 276       }
 277       Changed = true;
 278     }
 279
 280     if (!ReturningBlocks.empty()) {
 281       // Don't create a new unreachable inst if we have a return. The
 282       // structurizer/annotator can't handle the multiple exits
 283
 284       Type *RetTy = F.getReturnType();
 285       Value *RetVal = RetTy->isVoidTy() ? nullptr : UndefValue::get(RetTy);
 286       // Remove and delete the unreachable inst.
 287       UnreachableBlock->getTerminator()->eraseFromParent();
 288
 289       Function *UnreachableIntrin =
 290         Intrinsic::getDeclaration(F.getParent(), Intrinsic::amdgcn_unreachable);
 291
 292       // Insert a call to an intrinsic tracking that this is an unreachable
 293       // point, in case we want to kill the active lanes or something later.
 294       CallInst::Create(UnreachableIntrin, {}, "", UnreachableBlock);
 295
 296       // Don't create a scalar trap. We would only want to trap if this code was
 297       // really reached, but a scalar trap would happen even if no lanes
 298       // actually reached here.
 299       ReturnInst::Create(F.getContext(), RetVal, UnreachableBlock);
 300       ReturningBlocks.push_back(UnreachableBlock);
 301       Changed = true;
 302     }
 303   }
 304
 305   // FIXME: add PDT here once simplifycfg is ready.
 306   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
 307   if (RequireAndPreserveDomTree)
 308     DTU.applyUpdates(Updates);
 309   Updates.clear();
 310
 311   // Now handle return blocks.
 312   if (ReturningBlocks.empty())
 313     return Changed; // No blocks return
 314
 315   if (ReturningBlocks.size() == 1)
 316     return Changed; // Already has a single return block
 317
 318   unifyReturnBlockSet(F, DTU, ReturningBlocks, "UnifiedReturnBlock");
 319   return true;
 320 }