1 //===-- AMDGPULowerModuleLDSPass.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 // This pass eliminates LDS uses from non-kernel functions.
10 //
11 // The strategy is to create a new struct with a field for each LDS variable
12 // and allocate that struct at the same address for every kernel. Uses of the
13 // original LDS variables are then replaced with compile time offsets from that
14 // known address. AMDGPUMachineFunction allocates the LDS global.
15 //
16 // Local variables with constant annotation or non-undef initializer are passed
17 // through unchanged for simplication or error diagnostics in later passes.
18 //
19 // To reduce the memory overhead variables that are only used by kernels are
20 // excluded from this transform. The analysis to determine whether a variable
21 // is only used by a kernel is cheap and conservative so this may allocate
22 // a variable in every kernel when it was not strictly necessary to do so.
23 //
24 // A possible future refinement is to specialise the structure per-kernel, so
25 // that fields can be elided based on more expensive analysis.
26 //
27 //===----------------------------------------------------------------------===//
28
29 #include "AMDGPU.h"
30 #include "Utils/AMDGPUBaseInfo.h"
31 #include "Utils/AMDGPULDSUtils.h"
32 #include "llvm/ADT/STLExtras.h"
33 #include "llvm/IR/Constants.h"
34 #include "llvm/IR/DerivedTypes.h"
35 #include "llvm/IR/IRBuilder.h"
36 #include "llvm/IR/InlineAsm.h"
37 #include "llvm/IR/Instructions.h"
38 #include "llvm/InitializePasses.h"
39 #include "llvm/Pass.h"
40 #include "llvm/Support/Debug.h"
41 #include "llvm/Transforms/Utils/ModuleUtils.h"
42 #include <algorithm>
43 #include <vector>
44
45 #define DEBUG_TYPE "amdgpu-lower-module-lds"
46
47 using namespace llvm;
48
49 namespace {
50
51 class AMDGPULowerModuleLDS : public ModulePass {
52
removeFromUsedList(Module & M,StringRef Name,SmallPtrSetImpl<Constant * > & ToRemove)53 static void removeFromUsedList(Module &M, StringRef Name,
54 SmallPtrSetImpl<Constant *> &ToRemove) {
55 GlobalVariable *GV = M.getNamedGlobal(Name);
56 if (!GV || ToRemove.empty()) {
57 return;
58 }
59
60 SmallVector<Constant *, 16> Init;
61 auto *CA = cast<ConstantArray>(GV->getInitializer());
62 for (auto &Op : CA->operands()) {
63 // ModuleUtils::appendToUsed only inserts Constants
64 Constant *C = cast<Constant>(Op);
65 if (!ToRemove.contains(C->stripPointerCasts())) {
66 Init.push_back(C);
67 }
68 }
69
70 if (Init.size() == CA->getNumOperands()) {
71 return; // none to remove
72 }
73
74 GV->eraseFromParent();
75
76 if (!Init.empty()) {
77 ArrayType *ATy =
78 ArrayType::get(Type::getInt8PtrTy(M.getContext()), Init.size());
79 GV =
80 new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage,
81 ConstantArray::get(ATy, Init), Name);
82 GV->setSection("llvm.metadata");
83 }
84 }
85
86 static void
removeFromUsedLists(Module & M,const std::vector<GlobalVariable * > & LocalVars)87 removeFromUsedLists(Module &M,
88 const std::vector<GlobalVariable *> &LocalVars) {
89 SmallPtrSet<Constant *, 32> LocalVarsSet;
90 for (size_t I = 0; I < LocalVars.size(); I++) {
91 if (Constant *C = dyn_cast<Constant>(LocalVars[I]->stripPointerCasts())) {
92 LocalVarsSet.insert(C);
93 }
94 }
95 removeFromUsedList(M, "llvm.used", LocalVarsSet);
96 removeFromUsedList(M, "llvm.compiler.used", LocalVarsSet);
97 }
98
markUsedByKernel(IRBuilder<> & Builder,Function * Func,GlobalVariable * SGV)99 static void markUsedByKernel(IRBuilder<> &Builder, Function *Func,
100 GlobalVariable *SGV) {
101 // The llvm.amdgcn.module.lds instance is implicitly used by all kernels
102 // that might call a function which accesses a field within it. This is
103 // presently approximated to 'all kernels' if there are any such functions
104 // in the module. This implicit use is reified as an explicit use here so
105 // that later passes, specifically PromoteAlloca, account for the required
106 // memory without any knowledge of this transform.
107
108 // An operand bundle on llvm.donothing works because the call instruction
109 // survives until after the last pass that needs to account for LDS. It is
110 // better than inline asm as the latter survives until the end of codegen. A
111 // totally robust solution would be a function with the same semantics as
112 // llvm.donothing that takes a pointer to the instance and is lowered to a
113 // no-op after LDS is allocated, but that is not presently necessary.
114
115 LLVMContext &Ctx = Func->getContext();
116
117 Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI());
118
119 FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {});
120
121 Function *Decl =
122 Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {});
123
124 Value *UseInstance[1] = {Builder.CreateInBoundsGEP(
125 SGV->getValueType(), SGV, ConstantInt::get(Type::getInt32Ty(Ctx), 0))};
126
127 Builder.CreateCall(FTy, Decl, {},
128 {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)},
129 "");
130 }
131
132 public:
133 static char ID;
134
AMDGPULowerModuleLDS()135 AMDGPULowerModuleLDS() : ModulePass(ID) {
136 initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry());
137 }
138
runOnModule(Module & M)139 bool runOnModule(Module &M) override {
140 LLVMContext &Ctx = M.getContext();
141 const DataLayout &DL = M.getDataLayout();
142 SmallPtrSet<GlobalValue *, 32> UsedList = AMDGPU::getUsedList(M);
143
144 // Find variables to move into new struct instance
145 std::vector<GlobalVariable *> FoundLocalVars =
146 AMDGPU::findVariablesToLower(M, UsedList);
147
148 if (FoundLocalVars.empty()) {
149 // No variables to rewrite, no changes made.
150 return false;
151 }
152
153 // Sort by alignment, descending, to minimise padding.
154 // On ties, sort by size, descending, then by name, lexicographical.
155 llvm::stable_sort(
156 FoundLocalVars,
157 [&](const GlobalVariable *LHS, const GlobalVariable *RHS) -> bool {
158 Align ALHS = AMDGPU::getAlign(DL, LHS);
159 Align ARHS = AMDGPU::getAlign(DL, RHS);
160 if (ALHS != ARHS) {
161 return ALHS > ARHS;
162 }
163
164 TypeSize SLHS = DL.getTypeAllocSize(LHS->getValueType());
165 TypeSize SRHS = DL.getTypeAllocSize(RHS->getValueType());
166 if (SLHS != SRHS) {
167 return SLHS > SRHS;
168 }
169
170 // By variable name on tie for predictable order in test cases.
171 return LHS->getName() < RHS->getName();
172 });
173
174 std::vector<GlobalVariable *> LocalVars;
175 LocalVars.reserve(FoundLocalVars.size()); // will be at least this large
176 {
177 // This usually won't need to insert any padding, perhaps avoid the alloc
178 uint64_t CurrentOffset = 0;
179 for (size_t I = 0; I < FoundLocalVars.size(); I++) {
180 GlobalVariable *FGV = FoundLocalVars[I];
181 Align DataAlign = AMDGPU::getAlign(DL, FGV);
182
183 uint64_t DataAlignV = DataAlign.value();
184 if (uint64_t Rem = CurrentOffset % DataAlignV) {
185 uint64_t Padding = DataAlignV - Rem;
186
187 // Append an array of padding bytes to meet alignment requested
188 // Note (o + (a - (o % a)) ) % a == 0
189 // (offset + Padding ) % align == 0
190
191 Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding);
192 LocalVars.push_back(new GlobalVariable(
193 M, ATy, false, GlobalValue::InternalLinkage, UndefValue::get(ATy),
194 "", nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
195 false));
196 CurrentOffset += Padding;
197 }
198
199 LocalVars.push_back(FGV);
200 CurrentOffset += DL.getTypeAllocSize(FGV->getValueType());
201 }
202 }
203
204 std::vector<Type *> LocalVarTypes;
205 LocalVarTypes.reserve(LocalVars.size());
206 std::transform(
207 LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes),
208 [](const GlobalVariable *V) -> Type * { return V->getValueType(); });
209
210 StructType *LDSTy = StructType::create(
211 Ctx, LocalVarTypes, llvm::StringRef("llvm.amdgcn.module.lds.t"));
212
213 Align MaxAlign =
214 AMDGPU::getAlign(DL, LocalVars[0]); // was sorted on alignment
215
216 GlobalVariable *SGV = new GlobalVariable(
217 M, LDSTy, false, GlobalValue::InternalLinkage, UndefValue::get(LDSTy),
218 "llvm.amdgcn.module.lds", nullptr, GlobalValue::NotThreadLocal,
219 AMDGPUAS::LOCAL_ADDRESS, false);
220 SGV->setAlignment(MaxAlign);
221 appendToCompilerUsed(
222 M, {static_cast<GlobalValue *>(
223 ConstantExpr::getPointerBitCastOrAddrSpaceCast(
224 cast<Constant>(SGV), Type::getInt8PtrTy(Ctx)))});
225
226 // The verifier rejects used lists containing an inttoptr of a constant
227 // so remove the variables from these lists before replaceAllUsesWith
228 removeFromUsedLists(M, LocalVars);
229
230 // Replace uses of ith variable with a constantexpr to the ith field of the
231 // instance that will be allocated by AMDGPUMachineFunction
232 Type *I32 = Type::getInt32Ty(Ctx);
233 for (size_t I = 0; I < LocalVars.size(); I++) {
234 GlobalVariable *GV = LocalVars[I];
235 Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)};
236 GV->replaceAllUsesWith(
237 ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx));
238 GV->eraseFromParent();
239 }
240
241 // Mark kernels with asm that reads the address of the allocated structure
242 // This is not necessary for lowering. This lets other passes, specifically
243 // PromoteAlloca, accurately calculate how much LDS will be used by the
244 // kernel after lowering.
245 {
246 IRBuilder<> Builder(Ctx);
247 SmallPtrSet<Function *, 32> Kernels;
248 for (auto &I : M.functions()) {
249 Function *Func = &I;
250 if (AMDGPU::isKernelCC(Func) && !Kernels.contains(Func)) {
251 markUsedByKernel(Builder, Func, SGV);
252 Kernels.insert(Func);
253 }
254 }
255 }
256 return true;
257 }
258 };
259
260 } // namespace
261 char AMDGPULowerModuleLDS::ID = 0;
262
263 char &llvm::AMDGPULowerModuleLDSID = AMDGPULowerModuleLDS::ID;
264
265 INITIALIZE_PASS(AMDGPULowerModuleLDS, DEBUG_TYPE,
266 "Lower uses of LDS variables from non-kernel functions", false,
267 false)
268
createAMDGPULowerModuleLDSPass()269 ModulePass *llvm::createAMDGPULowerModuleLDSPass() {
270 return new AMDGPULowerModuleLDS();
271 }
272
run(Module & M,ModuleAnalysisManager &)273 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M,
274 ModuleAnalysisManager &) {
275 return AMDGPULowerModuleLDS().runOnModule(M) ? PreservedAnalyses::none()
276 : PreservedAnalyses::all();
277 }
278