xref: /llvm-project/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp (revision 8de4db697f2841748a5489d18d9fbcd130ae09bb)
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 
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     for (Constant *C : ToRemove) {
77       C->removeDeadConstantUsers();
78     }
79 
80     if (!Init.empty()) {
81       ArrayType *ATy =
82           ArrayType::get(Type::getInt8PtrTy(M.getContext()), Init.size());
83       GV =
84           new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage,
85                                    ConstantArray::get(ATy, Init), Name);
86       GV->setSection("llvm.metadata");
87     }
88   }
89 
90   static void
91   removeFromUsedLists(Module &M,
92                       const std::vector<GlobalVariable *> &LocalVars) {
93     SmallPtrSet<Constant *, 32> LocalVarsSet;
94     for (size_t I = 0; I < LocalVars.size(); I++) {
95       if (Constant *C = dyn_cast<Constant>(LocalVars[I]->stripPointerCasts())) {
96         LocalVarsSet.insert(C);
97       }
98     }
99     removeFromUsedList(M, "llvm.used", LocalVarsSet);
100     removeFromUsedList(M, "llvm.compiler.used", LocalVarsSet);
101   }
102 
103   static void markUsedByKernel(IRBuilder<> &Builder, Function *Func,
104                                GlobalVariable *SGV) {
105     // The llvm.amdgcn.module.lds instance is implicitly used by all kernels
106     // that might call a function which accesses a field within it. This is
107     // presently approximated to 'all kernels' if there are any such functions
108     // in the module. This implicit use is reified as an explicit use here so
109     // that later passes, specifically PromoteAlloca, account for the required
110     // memory without any knowledge of this transform.
111 
112     // An operand bundle on llvm.donothing works because the call instruction
113     // survives until after the last pass that needs to account for LDS. It is
114     // better than inline asm as the latter survives until the end of codegen. A
115     // totally robust solution would be a function with the same semantics as
116     // llvm.donothing that takes a pointer to the instance and is lowered to a
117     // no-op after LDS is allocated, but that is not presently necessary.
118 
119     LLVMContext &Ctx = Func->getContext();
120 
121     Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI());
122 
123     FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {});
124 
125     Function *Decl =
126         Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {});
127 
128     Value *UseInstance[1] = {Builder.CreateInBoundsGEP(
129         SGV->getValueType(), SGV, ConstantInt::get(Type::getInt32Ty(Ctx), 0))};
130 
131     Builder.CreateCall(FTy, Decl, {},
132                        {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)},
133                        "");
134   }
135 
136 private:
137   SmallPtrSet<GlobalValue *, 32> UsedList;
138 
139 public:
140   static char ID;
141 
142   AMDGPULowerModuleLDS() : ModulePass(ID) {
143     initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry());
144   }
145 
146   bool runOnModule(Module &M) override {
147     UsedList = AMDGPU::getUsedList(M);
148 
149     bool Changed = processUsedLDS(M);
150 
151     for (Function &F : M.functions()) {
152       if (!AMDGPU::isKernelCC(&F))
153         continue;
154       Changed |= processUsedLDS(M, &F);
155     }
156 
157     UsedList.clear();
158     return Changed;
159   }
160 
161 private:
162   bool processUsedLDS(Module &M, Function *F = nullptr) {
163     LLVMContext &Ctx = M.getContext();
164     const DataLayout &DL = M.getDataLayout();
165 
166     // Find variables to move into new struct instance
167     std::vector<GlobalVariable *> FoundLocalVars =
168         AMDGPU::findVariablesToLower(M, UsedList, F);
169 
170     if (FoundLocalVars.empty()) {
171       // No variables to rewrite, no changes made.
172       return false;
173     }
174 
175     // Sort by alignment, descending, to minimise padding.
176     // On ties, sort by size, descending, then by name, lexicographical.
177     llvm::stable_sort(
178         FoundLocalVars,
179         [&](const GlobalVariable *LHS, const GlobalVariable *RHS) -> bool {
180           Align ALHS = AMDGPU::getAlign(DL, LHS);
181           Align ARHS = AMDGPU::getAlign(DL, RHS);
182           if (ALHS != ARHS) {
183             return ALHS > ARHS;
184           }
185 
186           TypeSize SLHS = DL.getTypeAllocSize(LHS->getValueType());
187           TypeSize SRHS = DL.getTypeAllocSize(RHS->getValueType());
188           if (SLHS != SRHS) {
189             return SLHS > SRHS;
190           }
191 
192           // By variable name on tie for predictable order in test cases.
193           return LHS->getName() < RHS->getName();
194         });
195 
196     std::vector<GlobalVariable *> LocalVars;
197     LocalVars.reserve(FoundLocalVars.size()); // will be at least this large
198     {
199       // This usually won't need to insert any padding, perhaps avoid the alloc
200       uint64_t CurrentOffset = 0;
201       for (size_t I = 0; I < FoundLocalVars.size(); I++) {
202         GlobalVariable *FGV = FoundLocalVars[I];
203         Align DataAlign = AMDGPU::getAlign(DL, FGV);
204 
205         uint64_t DataAlignV = DataAlign.value();
206         if (uint64_t Rem = CurrentOffset % DataAlignV) {
207           uint64_t Padding = DataAlignV - Rem;
208 
209           // Append an array of padding bytes to meet alignment requested
210           // Note (o +      (a - (o % a)) ) % a == 0
211           //      (offset + Padding       ) % align == 0
212 
213           Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding);
214           LocalVars.push_back(new GlobalVariable(
215               M, ATy, false, GlobalValue::InternalLinkage, UndefValue::get(ATy),
216               "", nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
217               false));
218           CurrentOffset += Padding;
219         }
220 
221         LocalVars.push_back(FGV);
222         CurrentOffset += DL.getTypeAllocSize(FGV->getValueType());
223       }
224     }
225 
226     std::vector<Type *> LocalVarTypes;
227     LocalVarTypes.reserve(LocalVars.size());
228     std::transform(
229         LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes),
230         [](const GlobalVariable *V) -> Type * { return V->getValueType(); });
231 
232     std::string VarName(
233         F ? (Twine("llvm.amdgcn.kernel.") + F->getName() + ".lds").str()
234           : "llvm.amdgcn.module.lds");
235     StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t");
236 
237     Align MaxAlign =
238         AMDGPU::getAlign(DL, LocalVars[0]); // was sorted on alignment
239 
240     GlobalVariable *SGV = new GlobalVariable(
241         M, LDSTy, false, GlobalValue::InternalLinkage, UndefValue::get(LDSTy),
242         VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
243         false);
244     SGV->setAlignment(MaxAlign);
245     if (!F) {
246       appendToCompilerUsed(
247           M, {static_cast<GlobalValue *>(
248                  ConstantExpr::getPointerBitCastOrAddrSpaceCast(
249                      cast<Constant>(SGV), Type::getInt8PtrTy(Ctx)))});
250     }
251 
252     // The verifier rejects used lists containing an inttoptr of a constant
253     // so remove the variables from these lists before replaceAllUsesWith
254     removeFromUsedLists(M, LocalVars);
255 
256     // Replace uses of ith variable with a constantexpr to the ith field of the
257     // instance that will be allocated by AMDGPUMachineFunction
258     Type *I32 = Type::getInt32Ty(Ctx);
259     for (size_t I = 0; I < LocalVars.size(); I++) {
260       GlobalVariable *GV = LocalVars[I];
261       Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)};
262       Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx);
263       if (F) {
264         GV->replaceUsesWithIf(GEP, [F](Use &U) {
265           return AMDGPU::isUsedOnlyFromFunction(U.getUser(), F);
266         });
267       } else {
268         GV->replaceAllUsesWith(GEP);
269       }
270       if (GV->use_empty()) {
271         UsedList.erase(GV);
272         GV->eraseFromParent();
273       }
274     }
275 
276     // Mark kernels with asm that reads the address of the allocated structure
277     // This is not necessary for lowering. This lets other passes, specifically
278     // PromoteAlloca, accurately calculate how much LDS will be used by the
279     // kernel after lowering.
280     if (!F) {
281       IRBuilder<> Builder(Ctx);
282       SmallPtrSet<Function *, 32> Kernels;
283       for (auto &I : M.functions()) {
284         Function *Func = &I;
285         if (AMDGPU::isKernelCC(Func) && !Kernels.contains(Func)) {
286           markUsedByKernel(Builder, Func, SGV);
287           Kernels.insert(Func);
288         }
289       }
290     }
291     return true;
292   }
293 };
294 
295 } // namespace
296 char AMDGPULowerModuleLDS::ID = 0;
297 
298 char &llvm::AMDGPULowerModuleLDSID = AMDGPULowerModuleLDS::ID;
299 
300 INITIALIZE_PASS(AMDGPULowerModuleLDS, DEBUG_TYPE,
301                 "Lower uses of LDS variables from non-kernel functions", false,
302                 false)
303 
304 ModulePass *llvm::createAMDGPULowerModuleLDSPass() {
305   return new AMDGPULowerModuleLDS();
306 }
307 
308 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M,
309                                                 ModuleAnalysisManager &) {
310   return AMDGPULowerModuleLDS().runOnModule(M) ? PreservedAnalyses::none()
311                                                : PreservedAnalyses::all();
312 }
313