|
| 1 | +//===-- AMDGPUCloneModuleLDSPass.cpp ------------------------------*- C++ -*-=// |
| 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 | +// The purpose of this pass is to ensure that the combined module contains |
| 10 | +// as many LDS global variables as there are kernels that (indirectly) access |
| 11 | +// them. As LDS variables behave like C++ static variables, it is important that |
| 12 | +// each partition contains a unique copy of the variable on a per kernel basis. |
| 13 | +// This representation also prepares the combined module to eliminate |
| 14 | +// cross-module false dependencies of LDS variables. This pass runs prior to the |
| 15 | +// AMDGPULowerModuleLDS pass in the fullLTO pipeline and is used to improve |
| 16 | +// the functionality of --lto-partitions. |
| 17 | +// |
| 18 | +// This pass operates as follows: |
| 19 | +// 1. Firstly, traverse the call graph from each kernel to determine the number |
| 20 | +// of kernels calling each device function. |
| 21 | +// 2. For each LDS global variable GV, determine the function F that defines it. |
| 22 | +// Collect it's caller functions. Clone F and GV, and finally insert a |
| 23 | +// call/invoke instruction in each caller function. |
| 24 | +// |
| 25 | +//===----------------------------------------------------------------------===// |
| 26 | + |
| 27 | +#include "AMDGPU.h" |
| 28 | +#include "llvm/ADT/DepthFirstIterator.h" |
| 29 | +#include "llvm/ADT/Twine.h" |
| 30 | +#include "llvm/Analysis/CallGraph.h" |
| 31 | +#include "llvm/IR/InstrTypes.h" |
| 32 | +#include "llvm/IR/Instructions.h" |
| 33 | +#include "llvm/Passes/PassBuilder.h" |
| 34 | +#include "llvm/Support/ScopedPrinter.h" |
| 35 | +#include "llvm/Transforms/Utils/Cloning.h" |
| 36 | + |
| 37 | +using namespace llvm; |
| 38 | + |
| 39 | +#define DEBUG_TYPE "amdgpu-clone-module-lds" |
| 40 | + |
| 41 | +static cl::opt<unsigned int> MaxCountForClonedFunctions( |
| 42 | + "clone-lds-functions-max-count", cl::init(16), cl::Hidden, |
| 43 | + cl::desc("Specify a limit to the number of clones of a function")); |
| 44 | + |
| 45 | +/// Return the function that defines \p GV |
| 46 | +/// \param GV The global variable in question |
| 47 | +/// \return The function defining \p GV |
| 48 | +static Function *getFunctionDefiningGV(GlobalVariable &GV) { |
| 49 | + SmallVector<User *> Worklist(GV.users()); |
| 50 | + while (!Worklist.empty()) { |
| 51 | + User *U = Worklist.pop_back_val(); |
| 52 | + if (auto *Inst = dyn_cast<Instruction>(U)) |
| 53 | + return Inst->getFunction(); |
| 54 | + if (auto *Op = dyn_cast<Operator>(U)) |
| 55 | + append_range(Worklist, Op->users()); |
| 56 | + } |
| 57 | + return nullptr; |
| 58 | +}; |
| 59 | + |
| 60 | +PreservedAnalyses AMDGPUCloneModuleLDSPass::run(Module &M, |
| 61 | + ModuleAnalysisManager &AM) { |
| 62 | + if (MaxCountForClonedFunctions.getValue() == 1) |
| 63 | + return PreservedAnalyses::all(); |
| 64 | + |
| 65 | + bool Changed = false; |
| 66 | + auto &CG = AM.getResult<CallGraphAnalysis>(M); |
| 67 | + |
| 68 | + // For each function in the call graph, determine the number |
| 69 | + // of ancestor-caller kernels. |
| 70 | + DenseMap<Function *, unsigned int> KernelRefsToFuncs; |
| 71 | + for (auto &Fn : M) { |
| 72 | + if (Fn.getCallingConv() != CallingConv::AMDGPU_KERNEL) |
| 73 | + continue; |
| 74 | + for (auto I = df_begin(&CG), E = df_end(&CG); I != E; ++I) { |
| 75 | + if (auto *F = I->getFunction()) |
| 76 | + KernelRefsToFuncs[F]++; |
| 77 | + } |
| 78 | + } |
| 79 | + |
| 80 | + DenseMap<GlobalVariable *, Function *> GVToFnMap; |
| 81 | + for (auto &GV : M.globals()) { |
| 82 | + if (GVToFnMap.contains(&GV) || |
| 83 | + GV.getAddressSpace() != AMDGPUAS::LOCAL_ADDRESS || !GV.hasInitializer()) |
| 84 | + continue; |
| 85 | + |
| 86 | + auto *OldF = getFunctionDefiningGV(GV); |
| 87 | + GVToFnMap.insert({&GV, OldF}); |
| 88 | + LLVM_DEBUG(dbgs() << "Found LDS " << GV.getName() << " used in function " |
| 89 | + << OldF->getName() << '\n'); |
| 90 | + |
| 91 | + // Collect all call instructions to OldF |
| 92 | + SmallVector<Instruction *> InstsCallingOldF; |
| 93 | + for (auto &I : OldF->uses()) { |
| 94 | + if (auto *CI = dyn_cast<CallBase>(I.getUser())) |
| 95 | + InstsCallingOldF.push_back(CI); |
| 96 | + } |
| 97 | + |
| 98 | + // Create as many clones of the function containing LDS global as |
| 99 | + // there are kernels calling the function (including the function |
| 100 | + // already defining the LDS global). Respectively, clone the |
| 101 | + // LDS global and the call instructions to the function. |
| 102 | + LLVM_DEBUG(dbgs() << "\tFunction is referenced by " |
| 103 | + << KernelRefsToFuncs[OldF] << " kernels.\n"); |
| 104 | + for (unsigned int ID = 0; |
| 105 | + ID + 1 < std::min(KernelRefsToFuncs[OldF], |
| 106 | + MaxCountForClonedFunctions.getValue()); |
| 107 | + ++ID) { |
| 108 | + // Clone LDS global variable |
| 109 | + auto *NewGV = new GlobalVariable( |
| 110 | + M, GV.getValueType(), GV.isConstant(), GlobalValue::InternalLinkage, |
| 111 | + PoisonValue::get(GV.getValueType()), |
| 112 | + GV.getName() + ".clone." + Twine(ID), &GV, |
| 113 | + GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, false); |
| 114 | + NewGV->copyAttributesFrom(&GV); |
| 115 | + NewGV->copyMetadata(&GV, 0); |
| 116 | + NewGV->setComdat(GV.getComdat()); |
| 117 | + LLVM_DEBUG(dbgs() << "Inserting LDS clone with name " << NewGV->getName() |
| 118 | + << '\n'); |
| 119 | + |
| 120 | + // Clone function |
| 121 | + ValueToValueMapTy VMap; |
| 122 | + VMap[&GV] = NewGV; |
| 123 | + auto *NewF = CloneFunction(OldF, VMap); |
| 124 | + NewF->setName(OldF->getName() + ".clone." + Twine(ID)); |
| 125 | + LLVM_DEBUG(dbgs() << "Inserting function clone with name " |
| 126 | + << NewF->getName() << '\n'); |
| 127 | + |
| 128 | + // Create a new CallInst to call the cloned function |
| 129 | + for (auto *Inst : InstsCallingOldF) { |
| 130 | + Instruction *I = Inst->clone(); |
| 131 | + I->setName(Inst->getName() + ".clone." + Twine(ID)); |
| 132 | + if (auto *CI = dyn_cast<CallBase>(I)) |
| 133 | + CI->setCalledOperand(NewF); |
| 134 | + I->insertAfter(Inst); |
| 135 | + LLVM_DEBUG(dbgs() << "Inserting inst: " << *I << '\n'); |
| 136 | + } |
| 137 | + Changed = true; |
| 138 | + } |
| 139 | + } |
| 140 | + return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all(); |
| 141 | +} |
0 commit comments