194 lines · cpp
1//===- AMDGPURewriteUndefForPHI.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// This file implements the idea to rewrite undef incoming operand for certain9// PHIs in structurized CFG. This pass only works on IR that has gone through10// StructurizedCFG pass, and this pass has some additional limitation that make11// it can only run after SIAnnotateControlFlow.12//13// To achieve optimal code generation for AMDGPU, we assume that uniformity14// analysis reports the PHI in join block of divergent branch as uniform if15// it has one unique uniform value plus additional undefined/poisoned incoming16// value. That is to say the later compiler pipeline will ensure such PHI always17// return uniform value and ensure it work correctly. Let's take a look at two18// typical patterns in structured CFG that need to be taken care: (In both19// patterns, block %if terminate with divergent branch.)20//21// Pattern A: Block with undefined incoming value dominates defined predecessor22// %if23// | \24// | %then25// | /26// %endif: %phi = phi [%undef, %if], [%uniform, %then]27//28// Pattern B: Block with defined incoming value dominates undefined predecessor29// %if30// | \31// | %then32// | /33// %endif: %phi = phi [%uniform, %if], [%undef, %then]34//35// For pattern A, by reporting %phi as uniform, the later pipeline need to make36// sure it be handled correctly. The backend usually allocates a scalar register37// and if any thread in a wave takes %then path, the scalar register will get38// the %uniform value.39//40// For pattern B, we will replace the undef operand with the other defined value41// in this pass. So the scalar register allocated for such PHI will get correct42// liveness. Without this transformation, the scalar register may be overwritten43// in the %then block.44//45// Limitation note:46// If the join block of divergent threads is a loop header, the pass cannot47// handle it correctly right now. For below case, the undef in %phi should also48// be rewritten. Currently we depend on SIAnnotateControlFlow to split %header49// block to get a separate join block, then we can rewrite the undef correctly.50// %if51// | \52// | %then53// | /54// -> %header: %phi = phi [%uniform, %if], [%undef, %then], [%uniform2, %header]55// | |56// \---57 58#include "AMDGPU.h"59#include "llvm/Analysis/UniformityAnalysis.h"60#include "llvm/IR/BasicBlock.h"61#include "llvm/IR/Constants.h"62#include "llvm/IR/Dominators.h"63#include "llvm/IR/Instructions.h"64#include "llvm/InitializePasses.h"65 66using namespace llvm;67 68#define DEBUG_TYPE "amdgpu-rewrite-undef-for-phi"69 70namespace {71 72class AMDGPURewriteUndefForPHILegacy : public FunctionPass {73public:74 static char ID;75 AMDGPURewriteUndefForPHILegacy() : FunctionPass(ID) {}76 bool runOnFunction(Function &F) override;77 StringRef getPassName() const override {78 return "AMDGPU Rewrite Undef for PHI";79 }80 81 void getAnalysisUsage(AnalysisUsage &AU) const override {82 AU.addRequired<UniformityInfoWrapperPass>();83 AU.addRequired<DominatorTreeWrapperPass>();84 85 AU.addPreserved<DominatorTreeWrapperPass>();86 AU.setPreservesCFG();87 }88};89 90} // end anonymous namespace91char AMDGPURewriteUndefForPHILegacy::ID = 0;92 93INITIALIZE_PASS_BEGIN(AMDGPURewriteUndefForPHILegacy, DEBUG_TYPE,94 "Rewrite undef for PHI", false, false)95INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)96INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)97INITIALIZE_PASS_END(AMDGPURewriteUndefForPHILegacy, DEBUG_TYPE,98 "Rewrite undef for PHI", false, false)99 100bool rewritePHIs(Function &F, UniformityInfo &UA, DominatorTree *DT) {101 bool Changed = false;102 SmallVector<PHINode *> ToBeDeleted;103 for (auto &BB : F) {104 for (auto &PHI : BB.phis()) {105 if (UA.isDivergent(&PHI))106 continue;107 108 // The unique incoming value except undef/poison for the PHI node.109 Value *UniqueDefinedIncoming = nullptr;110 // The divergent block with defined incoming value that dominates all111 // other block with the same incoming value.112 BasicBlock *DominateBB = nullptr;113 // Predecessors with undefined incoming value (excluding loop backedge).114 SmallVector<BasicBlock *> Undefs;115 116 for (unsigned i = 0; i < PHI.getNumIncomingValues(); i++) {117 Value *Incoming = PHI.getIncomingValue(i);118 BasicBlock *IncomingBB = PHI.getIncomingBlock(i);119 120 if (Incoming == &PHI)121 continue;122 123 if (isa<UndefValue>(Incoming)) {124 // Undef from loop backedge will not be replaced.125 if (!DT->dominates(&BB, IncomingBB))126 Undefs.push_back(IncomingBB);127 continue;128 }129 130 if (!UniqueDefinedIncoming) {131 UniqueDefinedIncoming = Incoming;132 DominateBB = IncomingBB;133 } else if (Incoming == UniqueDefinedIncoming) {134 // Update DominateBB if necessary.135 if (DT->dominates(IncomingBB, DominateBB))136 DominateBB = IncomingBB;137 } else {138 UniqueDefinedIncoming = nullptr;139 break;140 }141 }142 // We only need to replace the undef for the PHI which is merging143 // defined/undefined values from divergent threads.144 // TODO: We should still be able to replace undef value if the unique145 // value is a Constant.146 if (!UniqueDefinedIncoming || Undefs.empty() ||147 !UA.isDivergent(DominateBB->getTerminator()))148 continue;149 150 // We only replace the undef when DominateBB truly dominates all the151 // other predecessors with undefined incoming value. Make sure DominateBB152 // dominates BB so that UniqueDefinedIncoming is available in BB and153 // afterwards.154 if (DT->dominates(DominateBB, &BB) && all_of(Undefs, [&](BasicBlock *UD) {155 return DT->dominates(DominateBB, UD);156 })) {157 PHI.replaceAllUsesWith(UniqueDefinedIncoming);158 ToBeDeleted.push_back(&PHI);159 Changed = true;160 }161 }162 }163 164 for (auto *PHI : ToBeDeleted)165 PHI->eraseFromParent();166 167 return Changed;168}169 170bool AMDGPURewriteUndefForPHILegacy::runOnFunction(Function &F) {171 UniformityInfo &UA =172 getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();173 DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();174 return rewritePHIs(F, UA, DT);175}176 177PreservedAnalyses178AMDGPURewriteUndefForPHIPass::run(Function &F, FunctionAnalysisManager &AM) {179 UniformityInfo &UA = AM.getResult<UniformityInfoAnalysis>(F);180 DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);181 bool Changed = rewritePHIs(F, UA, DT);182 if (Changed) {183 PreservedAnalyses PA;184 PA.preserveSet<CFGAnalyses>();185 return PA;186 }187 188 return PreservedAnalyses::all();189}190 191FunctionPass *llvm::createAMDGPURewriteUndefForPHILegacyPass() {192 return new AMDGPURewriteUndefForPHILegacy();193}194