xref: /llvm-project/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp (revision ccf5d624f9a30911923b2cb3963cacb8076835d8)
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 local data store, LDS, uses from non-kernel functions.
10 // LDS is contiguous memory allocated per kernel execution.
11 //
12 // Background.
13 //
14 // The programming model is global variables, or equivalently function local
15 // static variables, accessible from kernels or other functions. For uses from
16 // kernels this is straightforward - assign an integer to the kernel for the
17 // memory required by all the variables combined, allocate them within that.
18 // For uses from functions there are performance tradeoffs to choose between.
19 //
20 // This model means the GPU runtime can specify the amount of memory allocated.
21 // If this is more than the kernel assumed, the excess can be made available
22 // using a language specific feature, which IR represents as a variable with
23 // no initializer. This feature is referred to here as "Dynamic LDS" and is
24 // lowered slightly differently to the normal case.
25 //
26 // Consequences of this GPU feature:
27 // - memory is limited and exceeding it halts compilation
28 // - a global accessed by one kernel exists independent of other kernels
29 // - a global exists independent of simultaneous execution of the same kernel
30 // - the address of the global may be different from different kernels as they
31 //   do not alias, which permits only allocating variables they use
32 // - if the address is allowed to differ, functions need help to find it
33 //
34 // Uses from kernels are implemented here by grouping them in a per-kernel
35 // struct instance. This duplicates the variables, accurately modelling their
36 // aliasing properties relative to a single global representation. It also
37 // permits control over alignment via padding.
38 //
39 // Uses from functions are more complicated and the primary purpose of this
40 // IR pass. Several different lowering are chosen between to meet requirements
41 // to avoid allocating any LDS where it is not necessary, as that impacts
42 // occupancy and may fail the compilation, while not imposing overhead on a
43 // feature whose primary advantage over global memory is performance. The basic
44 // design goal is to avoid one kernel imposing overhead on another.
45 //
46 // Implementation.
47 //
48 // LDS variables with constant annotation or non-undef initializer are passed
49 // through unchanged for simplification or error diagnostics in later passes.
50 // Non-undef initializers are not yet implemented for LDS.
51 //
52 // LDS variables that are always allocated at the same address can be found
53 // by lookup at that address. Otherwise runtime information/cost is required.
54 //
55 // The simplest strategy possible is to group all LDS variables in a single
56 // struct and allocate that struct in every kernel such that the original
57 // variables are always at the same address. LDS is however a limited resource
58 // so this strategy is unusable in practice. It is not implemented here.
59 //
60 // Strategy | Precise allocation | Zero runtime cost | General purpose |
61 //  --------+--------------------+-------------------+-----------------+
62 //   Module |                 No |               Yes |             Yes |
63 //    Table |                Yes |                No |             Yes |
64 //   Kernel |                Yes |               Yes |              No |
65 //   Hybrid |                Yes |           Partial |             Yes |
66 //
67 // "Module" spends LDS memory to save cycles. "Table" spends cycles and global
68 // memory to save LDS. "Kernel" is as fast as kernel allocation but only works
69 // for variables that are known reachable from a single kernel. "Hybrid" picks
70 // between all three. When forced to choose between LDS and cycles we minimise
71 // LDS use.
72 
73 // The "module" lowering implemented here finds LDS variables which are used by
74 // non-kernel functions and creates a new struct with a field for each of those
75 // LDS variables. Variables that are only used from kernels are excluded.
76 //
77 // The "table" lowering implemented here has three components.
78 // First kernels are assigned a unique integer identifier which is available in
79 // functions it calls through the intrinsic amdgcn_lds_kernel_id. The integer
80 // is passed through a specific SGPR, thus works with indirect calls.
81 // Second, each kernel allocates LDS variables independent of other kernels and
82 // writes the addresses it chose for each variable into an array in consistent
83 // order. If the kernel does not allocate a given variable, it writes undef to
84 // the corresponding array location. These arrays are written to a constant
85 // table in the order matching the kernel unique integer identifier.
86 // Third, uses from non-kernel functions are replaced with a table lookup using
87 // the intrinsic function to find the address of the variable.
88 //
89 // "Kernel" lowering is only applicable for variables that are unambiguously
90 // reachable from exactly one kernel. For those cases, accesses to the variable
91 // can be lowered to ConstantExpr address of a struct instance specific to that
92 // one kernel. This is zero cost in space and in compute. It will raise a fatal
93 // error on any variable that might be reachable from multiple kernels and is
94 // thus most easily used as part of the hybrid lowering strategy.
95 //
96 // Hybrid lowering is a mixture of the above. It uses the zero cost kernel
97 // lowering where it can. It lowers the variable accessed by the greatest
98 // number of kernels using the module strategy as that is free for the first
99 // variable. Any futher variables that can be lowered with the module strategy
100 // without incurring LDS memory overhead are. The remaining ones are lowered
101 // via table.
102 //
103 // Consequences
104 // - No heuristics or user controlled magic numbers, hybrid is the right choice
105 // - Kernels that don't use functions (or have had them all inlined) are not
106 //   affected by any lowering for kernels that do.
107 // - Kernels that don't make indirect function calls are not affected by those
108 //   that do.
109 // - Variables which are used by lots of kernels, e.g. those injected by a
110 //   language runtime in most kernels, are expected to have no overhead
111 // - Implementations that instantiate templates per-kernel where those templates
112 //   use LDS are expected to hit the "Kernel" lowering strategy
113 // - The runtime properties impose a cost in compiler implementation complexity
114 //
115 // Dynamic LDS implementation
116 // Dynamic LDS is lowered similarly to the "table" strategy above and uses the
117 // same intrinsic to identify which kernel is at the root of the dynamic call
118 // graph. This relies on the specified behaviour that all dynamic LDS variables
119 // alias one another, i.e. are at the same address, with respect to a given
120 // kernel. Therefore this pass creates new dynamic LDS variables for each kernel
121 // that allocates any dynamic LDS and builds a table of addresses out of those.
122 // The AMDGPUPromoteAlloca pass skips kernels that use dynamic LDS.
123 // The corresponding optimisation for "kernel" lowering where the table lookup
124 // is elided is not implemented.
125 //
126 //
127 // Implementation notes / limitations
128 // A single LDS global variable represents an instance per kernel that can reach
129 // said variables. This pass essentially specialises said variables per kernel.
130 // Handling ConstantExpr during the pass complicated this significantly so now
131 // all ConstantExpr uses of LDS variables are expanded to instructions. This
132 // may need amending when implementing non-undef initialisers.
133 //
134 // Lowering is split between this IR pass and the back end. This pass chooses
135 // where given variables should be allocated and marks them with metadata,
136 // MD_absolute_symbol. The backend places the variables in coincidentally the
137 // same location and raises a fatal error if something has gone awry. This works
138 // in practice because the only pass between this one and the backend that
139 // changes LDS is PromoteAlloca and the changes it makes do not conflict.
140 //
141 // Addresses are written to constant global arrays based on the same metadata.
142 //
143 // The backend lowers LDS variables in the order of traversal of the function.
144 // This is at odds with the deterministic layout required. The workaround is to
145 // allocate the fixed-address variables immediately upon starting the function
146 // where they can be placed as intended. This requires a means of mapping from
147 // the function to the variables that it allocates. For the module scope lds,
148 // this is via metadata indicating whether the variable is not required. If a
149 // pass deletes that metadata, a fatal error on disagreement with the absolute
150 // symbol metadata will occur. For kernel scope and dynamic, this is by _name_
151 // correspondence between the function and the variable. It requires the
152 // kernel to have a name (which is only a limitation for tests in practice) and
153 // for nothing to rename the corresponding symbols. This is a hazard if the pass
154 // is run multiple times during debugging. Alternative schemes considered all
155 // involve bespoke metadata.
156 //
157 // If the name correspondence can be replaced, multiple distinct kernels that
158 // have the same memory layout can map to the same kernel id (as the address
159 // itself is handled by the absolute symbol metadata) and that will allow more
160 // uses of the "kernel" style faster lowering and reduce the size of the lookup
161 // tables.
162 //
163 // There is a test that checks this does not fire for a graphics shader. This
164 // lowering is expected to work for graphics if the isKernel test is changed.
165 //
166 // The current markUsedByKernel is sufficient for PromoteAlloca but is elided
167 // before codegen. Replacing this with an equivalent intrinsic which lasts until
168 // shortly after the machine function lowering of LDS would help break the name
169 // mapping. The other part needed is probably to amend PromoteAlloca to embed
170 // the LDS variables it creates in the same struct created here. That avoids the
171 // current hazard where a PromoteAlloca LDS variable might be allocated before
172 // the kernel scope (and thus error on the address check). Given a new invariant
173 // that no LDS variables exist outside of the structs managed here, and an
174 // intrinsic that lasts until after the LDS frame lowering, it should be
175 // possible to drop the name mapping and fold equivalent memory layouts.
176 //
177 //===----------------------------------------------------------------------===//
178 
179 #include "AMDGPU.h"
180 #include "AMDGPUMemoryUtils.h"
181 #include "AMDGPUTargetMachine.h"
182 #include "Utils/AMDGPUBaseInfo.h"
183 #include "llvm/ADT/BitVector.h"
184 #include "llvm/ADT/DenseMap.h"
185 #include "llvm/ADT/DenseSet.h"
186 #include "llvm/ADT/STLExtras.h"
187 #include "llvm/ADT/SetOperations.h"
188 #include "llvm/Analysis/CallGraph.h"
189 #include "llvm/CodeGen/TargetPassConfig.h"
190 #include "llvm/IR/Constants.h"
191 #include "llvm/IR/DerivedTypes.h"
192 #include "llvm/IR/IRBuilder.h"
193 #include "llvm/IR/InlineAsm.h"
194 #include "llvm/IR/Instructions.h"
195 #include "llvm/IR/IntrinsicsAMDGPU.h"
196 #include "llvm/IR/MDBuilder.h"
197 #include "llvm/IR/ReplaceConstant.h"
198 #include "llvm/InitializePasses.h"
199 #include "llvm/Pass.h"
200 #include "llvm/Support/CommandLine.h"
201 #include "llvm/Support/Debug.h"
202 #include "llvm/Support/Format.h"
203 #include "llvm/Support/OptimizedStructLayout.h"
204 #include "llvm/Support/raw_ostream.h"
205 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
206 #include "llvm/Transforms/Utils/ModuleUtils.h"
207 
208 #include <vector>
209 
210 #include <cstdio>
211 
212 #define DEBUG_TYPE "amdgpu-lower-module-lds"
213 
214 using namespace llvm;
215 using namespace AMDGPU;
216 
217 namespace {
218 
219 cl::opt<bool> SuperAlignLDSGlobals(
220     "amdgpu-super-align-lds-globals",
221     cl::desc("Increase alignment of LDS if it is not on align boundary"),
222     cl::init(true), cl::Hidden);
223 
224 enum class LoweringKind { module, table, kernel, hybrid };
225 cl::opt<LoweringKind> LoweringKindLoc(
226     "amdgpu-lower-module-lds-strategy",
227     cl::desc("Specify lowering strategy for function LDS access:"), cl::Hidden,
228     cl::init(LoweringKind::hybrid),
229     cl::values(
230         clEnumValN(LoweringKind::table, "table", "Lower via table lookup"),
231         clEnumValN(LoweringKind::module, "module", "Lower via module struct"),
232         clEnumValN(
233             LoweringKind::kernel, "kernel",
234             "Lower variables reachable from one kernel, otherwise abort"),
235         clEnumValN(LoweringKind::hybrid, "hybrid",
236                    "Lower via mixture of above strategies")));
237 
238 template <typename T> std::vector<T> sortByName(std::vector<T> &&V) {
239   llvm::sort(V.begin(), V.end(), [](const auto *L, const auto *R) {
240     return L->getName() < R->getName();
241   });
242   return {std::move(V)};
243 }
244 
245 class AMDGPULowerModuleLDS {
246   const AMDGPUTargetMachine &TM;
247 
248   static void
249   removeLocalVarsFromUsedLists(Module &M,
250                                const DenseSet<GlobalVariable *> &LocalVars) {
251     // The verifier rejects used lists containing an inttoptr of a constant
252     // so remove the variables from these lists before replaceAllUsesWith
253     SmallPtrSet<Constant *, 8> LocalVarsSet;
254     for (GlobalVariable *LocalVar : LocalVars)
255       LocalVarsSet.insert(cast<Constant>(LocalVar->stripPointerCasts()));
256 
257     removeFromUsedLists(
258         M, [&LocalVarsSet](Constant *C) { return LocalVarsSet.count(C); });
259 
260     for (GlobalVariable *LocalVar : LocalVars)
261       LocalVar->removeDeadConstantUsers();
262   }
263 
264   static void markUsedByKernel(Function *Func, GlobalVariable *SGV) {
265     // The llvm.amdgcn.module.lds instance is implicitly used by all kernels
266     // that might call a function which accesses a field within it. This is
267     // presently approximated to 'all kernels' if there are any such functions
268     // in the module. This implicit use is redefined as an explicit use here so
269     // that later passes, specifically PromoteAlloca, account for the required
270     // memory without any knowledge of this transform.
271 
272     // An operand bundle on llvm.donothing works because the call instruction
273     // survives until after the last pass that needs to account for LDS. It is
274     // better than inline asm as the latter survives until the end of codegen. A
275     // totally robust solution would be a function with the same semantics as
276     // llvm.donothing that takes a pointer to the instance and is lowered to a
277     // no-op after LDS is allocated, but that is not presently necessary.
278 
279     // This intrinsic is eliminated shortly before instruction selection. It
280     // does not suffice to indicate to ISel that a given global which is not
281     // immediately used by the kernel must still be allocated by it. An
282     // equivalent target specific intrinsic which lasts until immediately after
283     // codegen would suffice for that, but one would still need to ensure that
284     // the variables are allocated in the anticipated order.
285     BasicBlock *Entry = &Func->getEntryBlock();
286     IRBuilder<> Builder(Entry, Entry->getFirstNonPHIIt());
287 
288     Function *Decl = Intrinsic::getOrInsertDeclaration(
289         Func->getParent(), Intrinsic::donothing, {});
290 
291     Value *UseInstance[1] = {
292         Builder.CreateConstInBoundsGEP1_32(SGV->getValueType(), SGV, 0)};
293 
294     Builder.CreateCall(
295         Decl, {}, {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)});
296   }
297 
298 public:
299   AMDGPULowerModuleLDS(const AMDGPUTargetMachine &TM_) : TM(TM_) {}
300 
301   struct LDSVariableReplacement {
302     GlobalVariable *SGV = nullptr;
303     DenseMap<GlobalVariable *, Constant *> LDSVarsToConstantGEP;
304   };
305 
306   // remap from lds global to a constantexpr gep to where it has been moved to
307   // for each kernel
308   // an array with an element for each kernel containing where the corresponding
309   // variable was remapped to
310 
311   static Constant *getAddressesOfVariablesInKernel(
312       LLVMContext &Ctx, ArrayRef<GlobalVariable *> Variables,
313       const DenseMap<GlobalVariable *, Constant *> &LDSVarsToConstantGEP) {
314     // Create a ConstantArray containing the address of each Variable within the
315     // kernel corresponding to LDSVarsToConstantGEP, or poison if that kernel
316     // does not allocate it
317     // TODO: Drop the ptrtoint conversion
318 
319     Type *I32 = Type::getInt32Ty(Ctx);
320 
321     ArrayType *KernelOffsetsType = ArrayType::get(I32, Variables.size());
322 
323     SmallVector<Constant *> Elements;
324     for (GlobalVariable *GV : Variables) {
325       auto ConstantGepIt = LDSVarsToConstantGEP.find(GV);
326       if (ConstantGepIt != LDSVarsToConstantGEP.end()) {
327         auto *elt = ConstantExpr::getPtrToInt(ConstantGepIt->second, I32);
328         Elements.push_back(elt);
329       } else {
330         Elements.push_back(PoisonValue::get(I32));
331       }
332     }
333     return ConstantArray::get(KernelOffsetsType, Elements);
334   }
335 
336   static GlobalVariable *buildLookupTable(
337       Module &M, ArrayRef<GlobalVariable *> Variables,
338       ArrayRef<Function *> kernels,
339       DenseMap<Function *, LDSVariableReplacement> &KernelToReplacement) {
340     if (Variables.empty()) {
341       return nullptr;
342     }
343     LLVMContext &Ctx = M.getContext();
344 
345     const size_t NumberVariables = Variables.size();
346     const size_t NumberKernels = kernels.size();
347 
348     ArrayType *KernelOffsetsType =
349         ArrayType::get(Type::getInt32Ty(Ctx), NumberVariables);
350 
351     ArrayType *AllKernelsOffsetsType =
352         ArrayType::get(KernelOffsetsType, NumberKernels);
353 
354     Constant *Missing = PoisonValue::get(KernelOffsetsType);
355     std::vector<Constant *> overallConstantExprElts(NumberKernels);
356     for (size_t i = 0; i < NumberKernels; i++) {
357       auto Replacement = KernelToReplacement.find(kernels[i]);
358       overallConstantExprElts[i] =
359           (Replacement == KernelToReplacement.end())
360               ? Missing
361               : getAddressesOfVariablesInKernel(
362                     Ctx, Variables, Replacement->second.LDSVarsToConstantGEP);
363     }
364 
365     Constant *init =
366         ConstantArray::get(AllKernelsOffsetsType, overallConstantExprElts);
367 
368     return new GlobalVariable(
369         M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
370         "llvm.amdgcn.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
371         AMDGPUAS::CONSTANT_ADDRESS);
372   }
373 
374   void replaceUseWithTableLookup(Module &M, IRBuilder<> &Builder,
375                                  GlobalVariable *LookupTable,
376                                  GlobalVariable *GV, Use &U,
377                                  Value *OptionalIndex) {
378     // Table is a constant array of the same length as OrderedKernels
379     LLVMContext &Ctx = M.getContext();
380     Type *I32 = Type::getInt32Ty(Ctx);
381     auto *I = cast<Instruction>(U.getUser());
382 
383     Value *tableKernelIndex = getTableLookupKernelIndex(M, I->getFunction());
384 
385     if (auto *Phi = dyn_cast<PHINode>(I)) {
386       BasicBlock *BB = Phi->getIncomingBlock(U);
387       Builder.SetInsertPoint(&(*(BB->getFirstInsertionPt())));
388     } else {
389       Builder.SetInsertPoint(I);
390     }
391 
392     SmallVector<Value *, 3> GEPIdx = {
393         ConstantInt::get(I32, 0),
394         tableKernelIndex,
395     };
396     if (OptionalIndex)
397       GEPIdx.push_back(OptionalIndex);
398 
399     Value *Address = Builder.CreateInBoundsGEP(
400         LookupTable->getValueType(), LookupTable, GEPIdx, GV->getName());
401 
402     Value *loaded = Builder.CreateLoad(I32, Address);
403 
404     Value *replacement =
405         Builder.CreateIntToPtr(loaded, GV->getType(), GV->getName());
406 
407     U.set(replacement);
408   }
409 
410   void replaceUsesInInstructionsWithTableLookup(
411       Module &M, ArrayRef<GlobalVariable *> ModuleScopeVariables,
412       GlobalVariable *LookupTable) {
413 
414     LLVMContext &Ctx = M.getContext();
415     IRBuilder<> Builder(Ctx);
416     Type *I32 = Type::getInt32Ty(Ctx);
417 
418     for (size_t Index = 0; Index < ModuleScopeVariables.size(); Index++) {
419       auto *GV = ModuleScopeVariables[Index];
420 
421       for (Use &U : make_early_inc_range(GV->uses())) {
422         auto *I = dyn_cast<Instruction>(U.getUser());
423         if (!I)
424           continue;
425 
426         replaceUseWithTableLookup(M, Builder, LookupTable, GV, U,
427                                   ConstantInt::get(I32, Index));
428       }
429     }
430   }
431 
432   static DenseSet<Function *> kernelsThatIndirectlyAccessAnyOfPassedVariables(
433       Module &M, LDSUsesInfoTy &LDSUsesInfo,
434       DenseSet<GlobalVariable *> const &VariableSet) {
435 
436     DenseSet<Function *> KernelSet;
437 
438     if (VariableSet.empty())
439       return KernelSet;
440 
441     for (Function &Func : M.functions()) {
442       if (Func.isDeclaration() || !isKernelLDS(&Func))
443         continue;
444       for (GlobalVariable *GV : LDSUsesInfo.indirect_access[&Func]) {
445         if (VariableSet.contains(GV)) {
446           KernelSet.insert(&Func);
447           break;
448         }
449       }
450     }
451 
452     return KernelSet;
453   }
454 
455   static GlobalVariable *
456   chooseBestVariableForModuleStrategy(const DataLayout &DL,
457                                       VariableFunctionMap &LDSVars) {
458     // Find the global variable with the most indirect uses from kernels
459 
460     struct CandidateTy {
461       GlobalVariable *GV = nullptr;
462       size_t UserCount = 0;
463       size_t Size = 0;
464 
465       CandidateTy() = default;
466 
467       CandidateTy(GlobalVariable *GV, uint64_t UserCount, uint64_t AllocSize)
468           : GV(GV), UserCount(UserCount), Size(AllocSize) {}
469 
470       bool operator<(const CandidateTy &Other) const {
471         // Fewer users makes module scope variable less attractive
472         if (UserCount < Other.UserCount) {
473           return true;
474         }
475         if (UserCount > Other.UserCount) {
476           return false;
477         }
478 
479         // Bigger makes module scope variable less attractive
480         if (Size < Other.Size) {
481           return false;
482         }
483 
484         if (Size > Other.Size) {
485           return true;
486         }
487 
488         // Arbitrary but consistent
489         return GV->getName() < Other.GV->getName();
490       }
491     };
492 
493     CandidateTy MostUsed;
494 
495     for (auto &K : LDSVars) {
496       GlobalVariable *GV = K.first;
497       if (K.second.size() <= 1) {
498         // A variable reachable by only one kernel is best lowered with kernel
499         // strategy
500         continue;
501       }
502       CandidateTy Candidate(
503           GV, K.second.size(),
504           DL.getTypeAllocSize(GV->getValueType()).getFixedValue());
505       if (MostUsed < Candidate)
506         MostUsed = Candidate;
507     }
508 
509     return MostUsed.GV;
510   }
511 
512   static void recordLDSAbsoluteAddress(Module *M, GlobalVariable *GV,
513                                        uint32_t Address) {
514     // Write the specified address into metadata where it can be retrieved by
515     // the assembler. Format is a half open range, [Address Address+1)
516     LLVMContext &Ctx = M->getContext();
517     auto *IntTy =
518         M->getDataLayout().getIntPtrType(Ctx, AMDGPUAS::LOCAL_ADDRESS);
519     auto *MinC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address));
520     auto *MaxC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address + 1));
521     GV->setMetadata(LLVMContext::MD_absolute_symbol,
522                     MDNode::get(Ctx, {MinC, MaxC}));
523   }
524 
525   DenseMap<Function *, Value *> tableKernelIndexCache;
526   Value *getTableLookupKernelIndex(Module &M, Function *F) {
527     // Accesses from a function use the amdgcn_lds_kernel_id intrinsic which
528     // lowers to a read from a live in register. Emit it once in the entry
529     // block to spare deduplicating it later.
530     auto [It, Inserted] = tableKernelIndexCache.try_emplace(F);
531     if (Inserted) {
532       auto InsertAt = F->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
533       IRBuilder<> Builder(&*InsertAt);
534 
535       It->second =
536           Builder.CreateIntrinsic(Intrinsic::amdgcn_lds_kernel_id, {}, {});
537     }
538 
539     return It->second;
540   }
541 
542   static std::vector<Function *> assignLDSKernelIDToEachKernel(
543       Module *M, DenseSet<Function *> const &KernelsThatAllocateTableLDS,
544       DenseSet<Function *> const &KernelsThatIndirectlyAllocateDynamicLDS) {
545     // Associate kernels in the set with an arbitrary but reproducible order and
546     // annotate them with that order in metadata. This metadata is recognised by
547     // the backend and lowered to a SGPR which can be read from using
548     // amdgcn_lds_kernel_id.
549 
550     std::vector<Function *> OrderedKernels;
551     if (!KernelsThatAllocateTableLDS.empty() ||
552         !KernelsThatIndirectlyAllocateDynamicLDS.empty()) {
553 
554       for (Function &Func : M->functions()) {
555         if (Func.isDeclaration())
556           continue;
557         if (!isKernelLDS(&Func))
558           continue;
559 
560         if (KernelsThatAllocateTableLDS.contains(&Func) ||
561             KernelsThatIndirectlyAllocateDynamicLDS.contains(&Func)) {
562           assert(Func.hasName()); // else fatal error earlier
563           OrderedKernels.push_back(&Func);
564         }
565       }
566 
567       // Put them in an arbitrary but reproducible order
568       OrderedKernels = sortByName(std::move(OrderedKernels));
569 
570       // Annotate the kernels with their order in this vector
571       LLVMContext &Ctx = M->getContext();
572       IRBuilder<> Builder(Ctx);
573 
574       if (OrderedKernels.size() > UINT32_MAX) {
575         // 32 bit keeps it in one SGPR. > 2**32 kernels won't fit on the GPU
576         report_fatal_error("Unimplemented LDS lowering for > 2**32 kernels");
577       }
578 
579       for (size_t i = 0; i < OrderedKernels.size(); i++) {
580         Metadata *AttrMDArgs[1] = {
581             ConstantAsMetadata::get(Builder.getInt32(i)),
582         };
583         OrderedKernels[i]->setMetadata("llvm.amdgcn.lds.kernel.id",
584                                        MDNode::get(Ctx, AttrMDArgs));
585       }
586     }
587     return OrderedKernels;
588   }
589 
590   static void partitionVariablesIntoIndirectStrategies(
591       Module &M, LDSUsesInfoTy const &LDSUsesInfo,
592       VariableFunctionMap &LDSToKernelsThatNeedToAccessItIndirectly,
593       DenseSet<GlobalVariable *> &ModuleScopeVariables,
594       DenseSet<GlobalVariable *> &TableLookupVariables,
595       DenseSet<GlobalVariable *> &KernelAccessVariables,
596       DenseSet<GlobalVariable *> &DynamicVariables) {
597 
598     GlobalVariable *HybridModuleRoot =
599         LoweringKindLoc != LoweringKind::hybrid
600             ? nullptr
601             : chooseBestVariableForModuleStrategy(
602                   M.getDataLayout(), LDSToKernelsThatNeedToAccessItIndirectly);
603 
604     DenseSet<Function *> const EmptySet;
605     DenseSet<Function *> const &HybridModuleRootKernels =
606         HybridModuleRoot
607             ? LDSToKernelsThatNeedToAccessItIndirectly[HybridModuleRoot]
608             : EmptySet;
609 
610     for (auto &K : LDSToKernelsThatNeedToAccessItIndirectly) {
611       // Each iteration of this loop assigns exactly one global variable to
612       // exactly one of the implementation strategies.
613 
614       GlobalVariable *GV = K.first;
615       assert(AMDGPU::isLDSVariableToLower(*GV));
616       assert(K.second.size() != 0);
617 
618       if (AMDGPU::isDynamicLDS(*GV)) {
619         DynamicVariables.insert(GV);
620         continue;
621       }
622 
623       switch (LoweringKindLoc) {
624       case LoweringKind::module:
625         ModuleScopeVariables.insert(GV);
626         break;
627 
628       case LoweringKind::table:
629         TableLookupVariables.insert(GV);
630         break;
631 
632       case LoweringKind::kernel:
633         if (K.second.size() == 1) {
634           KernelAccessVariables.insert(GV);
635         } else {
636           report_fatal_error(
637               "cannot lower LDS '" + GV->getName() +
638               "' to kernel access as it is reachable from multiple kernels");
639         }
640         break;
641 
642       case LoweringKind::hybrid: {
643         if (GV == HybridModuleRoot) {
644           assert(K.second.size() != 1);
645           ModuleScopeVariables.insert(GV);
646         } else if (K.second.size() == 1) {
647           KernelAccessVariables.insert(GV);
648         } else if (set_is_subset(K.second, HybridModuleRootKernels)) {
649           ModuleScopeVariables.insert(GV);
650         } else {
651           TableLookupVariables.insert(GV);
652         }
653         break;
654       }
655       }
656     }
657 
658     // All LDS variables accessed indirectly have now been partitioned into
659     // the distinct lowering strategies.
660     assert(ModuleScopeVariables.size() + TableLookupVariables.size() +
661                KernelAccessVariables.size() + DynamicVariables.size() ==
662            LDSToKernelsThatNeedToAccessItIndirectly.size());
663   }
664 
665   static GlobalVariable *lowerModuleScopeStructVariables(
666       Module &M, DenseSet<GlobalVariable *> const &ModuleScopeVariables,
667       DenseSet<Function *> const &KernelsThatAllocateModuleLDS) {
668     // Create a struct to hold the ModuleScopeVariables
669     // Replace all uses of those variables from non-kernel functions with the
670     // new struct instance Replace only the uses from kernel functions that will
671     // allocate this instance. That is a space optimisation - kernels that use a
672     // subset of the module scope struct and do not need to allocate it for
673     // indirect calls will only allocate the subset they use (they do so as part
674     // of the per-kernel lowering).
675     if (ModuleScopeVariables.empty()) {
676       return nullptr;
677     }
678 
679     LLVMContext &Ctx = M.getContext();
680 
681     LDSVariableReplacement ModuleScopeReplacement =
682         createLDSVariableReplacement(M, "llvm.amdgcn.module.lds",
683                                      ModuleScopeVariables);
684 
685     appendToCompilerUsed(M, {static_cast<GlobalValue *>(
686                                 ConstantExpr::getPointerBitCastOrAddrSpaceCast(
687                                     cast<Constant>(ModuleScopeReplacement.SGV),
688                                     PointerType::getUnqual(Ctx)))});
689 
690     // module.lds will be allocated at zero in any kernel that allocates it
691     recordLDSAbsoluteAddress(&M, ModuleScopeReplacement.SGV, 0);
692 
693     // historic
694     removeLocalVarsFromUsedLists(M, ModuleScopeVariables);
695 
696     // Replace all uses of module scope variable from non-kernel functions
697     replaceLDSVariablesWithStruct(
698         M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) {
699           Instruction *I = dyn_cast<Instruction>(U.getUser());
700           if (!I) {
701             return false;
702           }
703           Function *F = I->getFunction();
704           return !isKernelLDS(F);
705         });
706 
707     // Replace uses of module scope variable from kernel functions that
708     // allocate the module scope variable, otherwise leave them unchanged
709     // Record on each kernel whether the module scope global is used by it
710 
711     for (Function &Func : M.functions()) {
712       if (Func.isDeclaration() || !isKernelLDS(&Func))
713         continue;
714 
715       if (KernelsThatAllocateModuleLDS.contains(&Func)) {
716         replaceLDSVariablesWithStruct(
717             M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) {
718               Instruction *I = dyn_cast<Instruction>(U.getUser());
719               if (!I) {
720                 return false;
721               }
722               Function *F = I->getFunction();
723               return F == &Func;
724             });
725 
726         markUsedByKernel(&Func, ModuleScopeReplacement.SGV);
727       }
728     }
729 
730     return ModuleScopeReplacement.SGV;
731   }
732 
733   static DenseMap<Function *, LDSVariableReplacement>
734   lowerKernelScopeStructVariables(
735       Module &M, LDSUsesInfoTy &LDSUsesInfo,
736       DenseSet<GlobalVariable *> const &ModuleScopeVariables,
737       DenseSet<Function *> const &KernelsThatAllocateModuleLDS,
738       GlobalVariable *MaybeModuleScopeStruct) {
739 
740     // Create a struct for each kernel for the non-module-scope variables.
741 
742     DenseMap<Function *, LDSVariableReplacement> KernelToReplacement;
743     for (Function &Func : M.functions()) {
744       if (Func.isDeclaration() || !isKernelLDS(&Func))
745         continue;
746 
747       DenseSet<GlobalVariable *> KernelUsedVariables;
748       // Allocating variables that are used directly in this struct to get
749       // alignment aware allocation and predictable frame size.
750       for (auto &v : LDSUsesInfo.direct_access[&Func]) {
751         if (!AMDGPU::isDynamicLDS(*v)) {
752           KernelUsedVariables.insert(v);
753         }
754       }
755 
756       // Allocating variables that are accessed indirectly so that a lookup of
757       // this struct instance can find them from nested functions.
758       for (auto &v : LDSUsesInfo.indirect_access[&Func]) {
759         if (!AMDGPU::isDynamicLDS(*v)) {
760           KernelUsedVariables.insert(v);
761         }
762       }
763 
764       // Variables allocated in module lds must all resolve to that struct,
765       // not to the per-kernel instance.
766       if (KernelsThatAllocateModuleLDS.contains(&Func)) {
767         for (GlobalVariable *v : ModuleScopeVariables) {
768           KernelUsedVariables.erase(v);
769         }
770       }
771 
772       if (KernelUsedVariables.empty()) {
773         // Either used no LDS, or the LDS it used was all in the module struct
774         // or dynamically sized
775         continue;
776       }
777 
778       // The association between kernel function and LDS struct is done by
779       // symbol name, which only works if the function in question has a
780       // name This is not expected to be a problem in practice as kernels
781       // are called by name making anonymous ones (which are named by the
782       // backend) difficult to use. This does mean that llvm test cases need
783       // to name the kernels.
784       if (!Func.hasName()) {
785         report_fatal_error("Anonymous kernels cannot use LDS variables");
786       }
787 
788       std::string VarName =
789           (Twine("llvm.amdgcn.kernel.") + Func.getName() + ".lds").str();
790 
791       auto Replacement =
792           createLDSVariableReplacement(M, VarName, KernelUsedVariables);
793 
794       // If any indirect uses, create a direct use to ensure allocation
795       // TODO: Simpler to unconditionally mark used but that regresses
796       // codegen in test/CodeGen/AMDGPU/noclobber-barrier.ll
797       auto Accesses = LDSUsesInfo.indirect_access.find(&Func);
798       if ((Accesses != LDSUsesInfo.indirect_access.end()) &&
799           !Accesses->second.empty())
800         markUsedByKernel(&Func, Replacement.SGV);
801 
802       // remove preserves existing codegen
803       removeLocalVarsFromUsedLists(M, KernelUsedVariables);
804       KernelToReplacement[&Func] = Replacement;
805 
806       // Rewrite uses within kernel to the new struct
807       replaceLDSVariablesWithStruct(
808           M, KernelUsedVariables, Replacement, [&Func](Use &U) {
809             Instruction *I = dyn_cast<Instruction>(U.getUser());
810             return I && I->getFunction() == &Func;
811           });
812     }
813     return KernelToReplacement;
814   }
815 
816   static GlobalVariable *
817   buildRepresentativeDynamicLDSInstance(Module &M, LDSUsesInfoTy &LDSUsesInfo,
818                                         Function *func) {
819     // Create a dynamic lds variable with a name associated with the passed
820     // function that has the maximum alignment of any dynamic lds variable
821     // reachable from this kernel. Dynamic LDS is allocated after the static LDS
822     // allocation, possibly after alignment padding. The representative variable
823     // created here has the maximum alignment of any other dynamic variable
824     // reachable by that kernel. All dynamic LDS variables are allocated at the
825     // same address in each kernel in order to provide the documented aliasing
826     // semantics. Setting the alignment here allows this IR pass to accurately
827     // predict the exact constant at which it will be allocated.
828 
829     assert(isKernelLDS(func));
830 
831     LLVMContext &Ctx = M.getContext();
832     const DataLayout &DL = M.getDataLayout();
833     Align MaxDynamicAlignment(1);
834 
835     auto UpdateMaxAlignment = [&MaxDynamicAlignment, &DL](GlobalVariable *GV) {
836       if (AMDGPU::isDynamicLDS(*GV)) {
837         MaxDynamicAlignment =
838             std::max(MaxDynamicAlignment, AMDGPU::getAlign(DL, GV));
839       }
840     };
841 
842     for (GlobalVariable *GV : LDSUsesInfo.indirect_access[func]) {
843       UpdateMaxAlignment(GV);
844     }
845 
846     for (GlobalVariable *GV : LDSUsesInfo.direct_access[func]) {
847       UpdateMaxAlignment(GV);
848     }
849 
850     assert(func->hasName()); // Checked by caller
851     auto *emptyCharArray = ArrayType::get(Type::getInt8Ty(Ctx), 0);
852     GlobalVariable *N = new GlobalVariable(
853         M, emptyCharArray, false, GlobalValue::ExternalLinkage, nullptr,
854         Twine("llvm.amdgcn." + func->getName() + ".dynlds"), nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
855         false);
856     N->setAlignment(MaxDynamicAlignment);
857 
858     assert(AMDGPU::isDynamicLDS(*N));
859     return N;
860   }
861 
862   DenseMap<Function *, GlobalVariable *> lowerDynamicLDSVariables(
863       Module &M, LDSUsesInfoTy &LDSUsesInfo,
864       DenseSet<Function *> const &KernelsThatIndirectlyAllocateDynamicLDS,
865       DenseSet<GlobalVariable *> const &DynamicVariables,
866       std::vector<Function *> const &OrderedKernels) {
867     DenseMap<Function *, GlobalVariable *> KernelToCreatedDynamicLDS;
868     if (!KernelsThatIndirectlyAllocateDynamicLDS.empty()) {
869       LLVMContext &Ctx = M.getContext();
870       IRBuilder<> Builder(Ctx);
871       Type *I32 = Type::getInt32Ty(Ctx);
872 
873       std::vector<Constant *> newDynamicLDS;
874 
875       // Table is built in the same order as OrderedKernels
876       for (auto &func : OrderedKernels) {
877 
878         if (KernelsThatIndirectlyAllocateDynamicLDS.contains(func)) {
879           assert(isKernelLDS(func));
880           if (!func->hasName()) {
881             report_fatal_error("Anonymous kernels cannot use LDS variables");
882           }
883 
884           GlobalVariable *N =
885               buildRepresentativeDynamicLDSInstance(M, LDSUsesInfo, func);
886 
887           KernelToCreatedDynamicLDS[func] = N;
888 
889           markUsedByKernel(func, N);
890 
891           auto *emptyCharArray = ArrayType::get(Type::getInt8Ty(Ctx), 0);
892           auto *GEP = ConstantExpr::getGetElementPtr(
893               emptyCharArray, N, ConstantInt::get(I32, 0), true);
894           newDynamicLDS.push_back(ConstantExpr::getPtrToInt(GEP, I32));
895         } else {
896           newDynamicLDS.push_back(PoisonValue::get(I32));
897         }
898       }
899       assert(OrderedKernels.size() == newDynamicLDS.size());
900 
901       ArrayType *t = ArrayType::get(I32, newDynamicLDS.size());
902       Constant *init = ConstantArray::get(t, newDynamicLDS);
903       GlobalVariable *table = new GlobalVariable(
904           M, t, true, GlobalValue::InternalLinkage, init,
905           "llvm.amdgcn.dynlds.offset.table", nullptr,
906           GlobalValue::NotThreadLocal, AMDGPUAS::CONSTANT_ADDRESS);
907 
908       for (GlobalVariable *GV : DynamicVariables) {
909         for (Use &U : make_early_inc_range(GV->uses())) {
910           auto *I = dyn_cast<Instruction>(U.getUser());
911           if (!I)
912             continue;
913           if (isKernelLDS(I->getFunction()))
914             continue;
915 
916           replaceUseWithTableLookup(M, Builder, table, GV, U, nullptr);
917         }
918       }
919     }
920     return KernelToCreatedDynamicLDS;
921   }
922 
923   static GlobalVariable *uniquifyGVPerKernel(Module &M, GlobalVariable *GV,
924                                              Function *KF) {
925     bool NeedsReplacement = false;
926     for (Use &U : GV->uses()) {
927       if (auto *I = dyn_cast<Instruction>(U.getUser())) {
928         Function *F = I->getFunction();
929         if (isKernelLDS(F) && F != KF) {
930           NeedsReplacement = true;
931           break;
932         }
933       }
934     }
935     if (!NeedsReplacement)
936       return GV;
937     // Create a new GV used only by this kernel and its function
938     GlobalVariable *NewGV = new GlobalVariable(
939         M, GV->getValueType(), GV->isConstant(), GV->getLinkage(),
940         GV->getInitializer(), GV->getName() + "." + KF->getName(), nullptr,
941         GV->getThreadLocalMode(), GV->getType()->getAddressSpace());
942     NewGV->copyAttributesFrom(GV);
943     for (Use &U : make_early_inc_range(GV->uses())) {
944       if (auto *I = dyn_cast<Instruction>(U.getUser())) {
945         Function *F = I->getFunction();
946         if (!isKernelLDS(F) || F == KF) {
947           U.getUser()->replaceUsesOfWith(GV, NewGV);
948         }
949       }
950     }
951     return NewGV;
952   }
953 
954   bool lowerSpecialLDSVariables(
955       Module &M, LDSUsesInfoTy &LDSUsesInfo,
956       VariableFunctionMap &LDSToKernelsThatNeedToAccessItIndirectly) {
957     bool Changed = false;
958     // The 1st round: give module-absolute assignments
959     int NumAbsolutes = 0;
960     std::vector<GlobalVariable *> OrderedGVs;
961     for (auto &K : LDSToKernelsThatNeedToAccessItIndirectly) {
962       GlobalVariable *GV = K.first;
963       if (!isNamedBarrier(*GV))
964         continue;
965       // give a module-absolute assignment if it is indirectly accessed by
966       // multiple kernels. This is not precise, but we don't want to duplicate
967       // a function when it is called by multiple kernels.
968       if (LDSToKernelsThatNeedToAccessItIndirectly[GV].size() > 1) {
969         OrderedGVs.push_back(GV);
970       } else {
971         // leave it to the 2nd round, which will give a kernel-relative
972         // assignment if it is only indirectly accessed by one kernel
973         LDSUsesInfo.direct_access[*K.second.begin()].insert(GV);
974       }
975       LDSToKernelsThatNeedToAccessItIndirectly.erase(GV);
976     }
977     OrderedGVs = sortByName(std::move(OrderedGVs));
978     for (GlobalVariable *GV : OrderedGVs) {
979       int BarId = ++NumAbsolutes;
980       unsigned BarrierScope = llvm::AMDGPU::Barrier::BARRIER_SCOPE_WORKGROUP;
981       // 4 bits for alignment, 5 bits for the barrier num,
982       // 3 bits for the barrier scope
983       unsigned Offset = 0x802000u | BarrierScope << 9 | BarId << 4;
984       recordLDSAbsoluteAddress(&M, GV, Offset);
985     }
986     OrderedGVs.clear();
987 
988     // The 2nd round: give a kernel-relative assignment for GV that
989     // either only indirectly accessed by single kernel or only directly
990     // accessed by multiple kernels.
991     std::vector<Function *> OrderedKernels;
992     for (auto &K : LDSUsesInfo.direct_access) {
993       Function *F = K.first;
994       assert(isKernelLDS(F));
995       OrderedKernels.push_back(F);
996     }
997     OrderedKernels = sortByName(std::move(OrderedKernels));
998 
999     llvm::DenseMap<Function *, uint32_t> Kernel2BarId;
1000     for (Function *F : OrderedKernels) {
1001       for (GlobalVariable *GV : LDSUsesInfo.direct_access[F]) {
1002         if (!isNamedBarrier(*GV))
1003           continue;
1004 
1005         LDSUsesInfo.direct_access[F].erase(GV);
1006         if (GV->isAbsoluteSymbolRef()) {
1007           // already assigned
1008           continue;
1009         }
1010         OrderedGVs.push_back(GV);
1011       }
1012       OrderedGVs = sortByName(std::move(OrderedGVs));
1013       for (GlobalVariable *GV : OrderedGVs) {
1014         // GV could also be used directly by other kernels. If so, we need to
1015         // create a new GV used only by this kernel and its function.
1016         auto NewGV = uniquifyGVPerKernel(M, GV, F);
1017         Changed |= (NewGV != GV);
1018         int BarId = (NumAbsolutes + 1);
1019         if (Kernel2BarId.find(F) != Kernel2BarId.end()) {
1020           BarId = (Kernel2BarId[F] + 1);
1021         }
1022         Kernel2BarId[F] = BarId;
1023         unsigned BarrierScope = llvm::AMDGPU::Barrier::BARRIER_SCOPE_WORKGROUP;
1024         unsigned Offset = 0x802000u | BarrierScope << 9 | BarId << 4;
1025         recordLDSAbsoluteAddress(&M, NewGV, Offset);
1026       }
1027       OrderedGVs.clear();
1028     }
1029     // Also erase those special LDS variables from indirect_access.
1030     for (auto &K : LDSUsesInfo.indirect_access) {
1031       assert(isKernelLDS(K.first));
1032       for (GlobalVariable *GV : K.second) {
1033         if (isNamedBarrier(*GV))
1034           K.second.erase(GV);
1035       }
1036     }
1037     return Changed;
1038   }
1039 
1040   bool runOnModule(Module &M) {
1041     CallGraph CG = CallGraph(M);
1042     bool Changed = superAlignLDSGlobals(M);
1043 
1044     Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M);
1045 
1046     Changed = true; // todo: narrow this down
1047 
1048     // For each kernel, what variables does it access directly or through
1049     // callees
1050     LDSUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDS(CG, M);
1051 
1052     // For each variable accessed through callees, which kernels access it
1053     VariableFunctionMap LDSToKernelsThatNeedToAccessItIndirectly;
1054     for (auto &K : LDSUsesInfo.indirect_access) {
1055       Function *F = K.first;
1056       assert(isKernelLDS(F));
1057       for (GlobalVariable *GV : K.second) {
1058         LDSToKernelsThatNeedToAccessItIndirectly[GV].insert(F);
1059       }
1060     }
1061 
1062     if (LDSUsesInfo.HasSpecialGVs) {
1063       // Special LDS variables need special address assignment
1064       Changed |= lowerSpecialLDSVariables(
1065           M, LDSUsesInfo, LDSToKernelsThatNeedToAccessItIndirectly);
1066     }
1067 
1068     // Partition variables accessed indirectly into the different strategies
1069     DenseSet<GlobalVariable *> ModuleScopeVariables;
1070     DenseSet<GlobalVariable *> TableLookupVariables;
1071     DenseSet<GlobalVariable *> KernelAccessVariables;
1072     DenseSet<GlobalVariable *> DynamicVariables;
1073     partitionVariablesIntoIndirectStrategies(
1074         M, LDSUsesInfo, LDSToKernelsThatNeedToAccessItIndirectly,
1075         ModuleScopeVariables, TableLookupVariables, KernelAccessVariables,
1076         DynamicVariables);
1077 
1078     // If the kernel accesses a variable that is going to be stored in the
1079     // module instance through a call then that kernel needs to allocate the
1080     // module instance
1081     const DenseSet<Function *> KernelsThatAllocateModuleLDS =
1082         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
1083                                                         ModuleScopeVariables);
1084     const DenseSet<Function *> KernelsThatAllocateTableLDS =
1085         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
1086                                                         TableLookupVariables);
1087 
1088     const DenseSet<Function *> KernelsThatIndirectlyAllocateDynamicLDS =
1089         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
1090                                                         DynamicVariables);
1091 
1092     GlobalVariable *MaybeModuleScopeStruct = lowerModuleScopeStructVariables(
1093         M, ModuleScopeVariables, KernelsThatAllocateModuleLDS);
1094 
1095     DenseMap<Function *, LDSVariableReplacement> KernelToReplacement =
1096         lowerKernelScopeStructVariables(M, LDSUsesInfo, ModuleScopeVariables,
1097                                         KernelsThatAllocateModuleLDS,
1098                                         MaybeModuleScopeStruct);
1099 
1100     // Lower zero cost accesses to the kernel instances just created
1101     for (auto &GV : KernelAccessVariables) {
1102       auto &funcs = LDSToKernelsThatNeedToAccessItIndirectly[GV];
1103       assert(funcs.size() == 1); // Only one kernel can access it
1104       LDSVariableReplacement Replacement =
1105           KernelToReplacement[*(funcs.begin())];
1106 
1107       DenseSet<GlobalVariable *> Vec;
1108       Vec.insert(GV);
1109 
1110       replaceLDSVariablesWithStruct(M, Vec, Replacement, [](Use &U) {
1111         return isa<Instruction>(U.getUser());
1112       });
1113     }
1114 
1115     // The ith element of this vector is kernel id i
1116     std::vector<Function *> OrderedKernels =
1117         assignLDSKernelIDToEachKernel(&M, KernelsThatAllocateTableLDS,
1118                                       KernelsThatIndirectlyAllocateDynamicLDS);
1119 
1120     if (!KernelsThatAllocateTableLDS.empty()) {
1121       LLVMContext &Ctx = M.getContext();
1122       IRBuilder<> Builder(Ctx);
1123 
1124       // The order must be consistent between lookup table and accesses to
1125       // lookup table
1126       auto TableLookupVariablesOrdered =
1127           sortByName(std::vector<GlobalVariable *>(TableLookupVariables.begin(),
1128                                                    TableLookupVariables.end()));
1129 
1130       GlobalVariable *LookupTable = buildLookupTable(
1131           M, TableLookupVariablesOrdered, OrderedKernels, KernelToReplacement);
1132       replaceUsesInInstructionsWithTableLookup(M, TableLookupVariablesOrdered,
1133                                                LookupTable);
1134     }
1135 
1136     DenseMap<Function *, GlobalVariable *> KernelToCreatedDynamicLDS =
1137         lowerDynamicLDSVariables(M, LDSUsesInfo,
1138                                  KernelsThatIndirectlyAllocateDynamicLDS,
1139                                  DynamicVariables, OrderedKernels);
1140 
1141     // Strip amdgpu-no-lds-kernel-id from all functions reachable from the
1142     // kernel. We may have inferred this wasn't used prior to the pass.
1143     // TODO: We could filter out subgraphs that do not access LDS globals.
1144     for (auto *KernelSet : {&KernelsThatIndirectlyAllocateDynamicLDS,
1145                             &KernelsThatAllocateTableLDS})
1146       for (Function *F : *KernelSet)
1147         removeFnAttrFromReachable(CG, F, {"amdgpu-no-lds-kernel-id"});
1148 
1149     // All kernel frames have been allocated. Calculate and record the
1150     // addresses.
1151     {
1152       const DataLayout &DL = M.getDataLayout();
1153 
1154       for (Function &Func : M.functions()) {
1155         if (Func.isDeclaration() || !isKernelLDS(&Func))
1156           continue;
1157 
1158         // All three of these are optional. The first variable is allocated at
1159         // zero. They are allocated by AMDGPUMachineFunction as one block.
1160         // Layout:
1161         //{
1162         //  module.lds
1163         //  alignment padding
1164         //  kernel instance
1165         //  alignment padding
1166         //  dynamic lds variables
1167         //}
1168 
1169         const bool AllocateModuleScopeStruct =
1170             MaybeModuleScopeStruct &&
1171             KernelsThatAllocateModuleLDS.contains(&Func);
1172 
1173         auto Replacement = KernelToReplacement.find(&Func);
1174         const bool AllocateKernelScopeStruct =
1175             Replacement != KernelToReplacement.end();
1176 
1177         const bool AllocateDynamicVariable =
1178             KernelToCreatedDynamicLDS.contains(&Func);
1179 
1180         uint32_t Offset = 0;
1181 
1182         if (AllocateModuleScopeStruct) {
1183           // Allocated at zero, recorded once on construction, not once per
1184           // kernel
1185           Offset += DL.getTypeAllocSize(MaybeModuleScopeStruct->getValueType());
1186         }
1187 
1188         if (AllocateKernelScopeStruct) {
1189           GlobalVariable *KernelStruct = Replacement->second.SGV;
1190           Offset = alignTo(Offset, AMDGPU::getAlign(DL, KernelStruct));
1191           recordLDSAbsoluteAddress(&M, KernelStruct, Offset);
1192           Offset += DL.getTypeAllocSize(KernelStruct->getValueType());
1193         }
1194 
1195         // If there is dynamic allocation, the alignment needed is included in
1196         // the static frame size. There may be no reference to the dynamic
1197         // variable in the kernel itself, so without including it here, that
1198         // alignment padding could be missed.
1199         if (AllocateDynamicVariable) {
1200           GlobalVariable *DynamicVariable = KernelToCreatedDynamicLDS[&Func];
1201           Offset = alignTo(Offset, AMDGPU::getAlign(DL, DynamicVariable));
1202           recordLDSAbsoluteAddress(&M, DynamicVariable, Offset);
1203         }
1204 
1205         if (Offset != 0) {
1206           (void)TM; // TODO: Account for target maximum LDS
1207           std::string Buffer;
1208           raw_string_ostream SS{Buffer};
1209           SS << format("%u", Offset);
1210 
1211           // Instead of explicitly marking kernels that access dynamic variables
1212           // using special case metadata, annotate with min-lds == max-lds, i.e.
1213           // that there is no more space available for allocating more static
1214           // LDS variables. That is the right condition to prevent allocating
1215           // more variables which would collide with the addresses assigned to
1216           // dynamic variables.
1217           if (AllocateDynamicVariable)
1218             SS << format(",%u", Offset);
1219 
1220           Func.addFnAttr("amdgpu-lds-size", Buffer);
1221         }
1222       }
1223     }
1224 
1225     for (auto &GV : make_early_inc_range(M.globals()))
1226       if (AMDGPU::isLDSVariableToLower(GV)) {
1227         // probably want to remove from used lists
1228         GV.removeDeadConstantUsers();
1229         if (GV.use_empty())
1230           GV.eraseFromParent();
1231       }
1232 
1233     return Changed;
1234   }
1235 
1236 private:
1237   // Increase the alignment of LDS globals if necessary to maximise the chance
1238   // that we can use aligned LDS instructions to access them.
1239   static bool superAlignLDSGlobals(Module &M) {
1240     const DataLayout &DL = M.getDataLayout();
1241     bool Changed = false;
1242     if (!SuperAlignLDSGlobals) {
1243       return Changed;
1244     }
1245 
1246     for (auto &GV : M.globals()) {
1247       if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) {
1248         // Only changing alignment of LDS variables
1249         continue;
1250       }
1251       if (!GV.hasInitializer()) {
1252         // cuda/hip extern __shared__ variable, leave alignment alone
1253         continue;
1254       }
1255 
1256       if (GV.isAbsoluteSymbolRef()) {
1257         // If the variable is already allocated, don't change the alignment
1258         continue;
1259       }
1260 
1261       Align Alignment = AMDGPU::getAlign(DL, &GV);
1262       TypeSize GVSize = DL.getTypeAllocSize(GV.getValueType());
1263 
1264       if (GVSize > 8) {
1265         // We might want to use a b96 or b128 load/store
1266         Alignment = std::max(Alignment, Align(16));
1267       } else if (GVSize > 4) {
1268         // We might want to use a b64 load/store
1269         Alignment = std::max(Alignment, Align(8));
1270       } else if (GVSize > 2) {
1271         // We might want to use a b32 load/store
1272         Alignment = std::max(Alignment, Align(4));
1273       } else if (GVSize > 1) {
1274         // We might want to use a b16 load/store
1275         Alignment = std::max(Alignment, Align(2));
1276       }
1277 
1278       if (Alignment != AMDGPU::getAlign(DL, &GV)) {
1279         Changed = true;
1280         GV.setAlignment(Alignment);
1281       }
1282     }
1283     return Changed;
1284   }
1285 
1286   static LDSVariableReplacement createLDSVariableReplacement(
1287       Module &M, std::string VarName,
1288       DenseSet<GlobalVariable *> const &LDSVarsToTransform) {
1289     // Create a struct instance containing LDSVarsToTransform and map from those
1290     // variables to ConstantExprGEP
1291     // Variables may be introduced to meet alignment requirements. No aliasing
1292     // metadata is useful for these as they have no uses. Erased before return.
1293 
1294     LLVMContext &Ctx = M.getContext();
1295     const DataLayout &DL = M.getDataLayout();
1296     assert(!LDSVarsToTransform.empty());
1297 
1298     SmallVector<OptimizedStructLayoutField, 8> LayoutFields;
1299     LayoutFields.reserve(LDSVarsToTransform.size());
1300     {
1301       // The order of fields in this struct depends on the order of
1302       // variables in the argument which varies when changing how they
1303       // are identified, leading to spurious test breakage.
1304       auto Sorted = sortByName(std::vector<GlobalVariable *>(
1305           LDSVarsToTransform.begin(), LDSVarsToTransform.end()));
1306 
1307       for (GlobalVariable *GV : Sorted) {
1308         OptimizedStructLayoutField F(GV,
1309                                      DL.getTypeAllocSize(GV->getValueType()),
1310                                      AMDGPU::getAlign(DL, GV));
1311         LayoutFields.emplace_back(F);
1312       }
1313     }
1314 
1315     performOptimizedStructLayout(LayoutFields);
1316 
1317     std::vector<GlobalVariable *> LocalVars;
1318     BitVector IsPaddingField;
1319     LocalVars.reserve(LDSVarsToTransform.size()); // will be at least this large
1320     IsPaddingField.reserve(LDSVarsToTransform.size());
1321     {
1322       uint64_t CurrentOffset = 0;
1323       for (auto &F : LayoutFields) {
1324         GlobalVariable *FGV =
1325             static_cast<GlobalVariable *>(const_cast<void *>(F.Id));
1326         Align DataAlign = F.Alignment;
1327 
1328         uint64_t DataAlignV = DataAlign.value();
1329         if (uint64_t Rem = CurrentOffset % DataAlignV) {
1330           uint64_t Padding = DataAlignV - Rem;
1331 
1332           // Append an array of padding bytes to meet alignment requested
1333           // Note (o +      (a - (o % a)) ) % a == 0
1334           //      (offset + Padding       ) % align == 0
1335 
1336           Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding);
1337           LocalVars.push_back(new GlobalVariable(
1338               M, ATy, false, GlobalValue::InternalLinkage,
1339               PoisonValue::get(ATy), "", nullptr, GlobalValue::NotThreadLocal,
1340               AMDGPUAS::LOCAL_ADDRESS, false));
1341           IsPaddingField.push_back(true);
1342           CurrentOffset += Padding;
1343         }
1344 
1345         LocalVars.push_back(FGV);
1346         IsPaddingField.push_back(false);
1347         CurrentOffset += F.Size;
1348       }
1349     }
1350 
1351     std::vector<Type *> LocalVarTypes;
1352     LocalVarTypes.reserve(LocalVars.size());
1353     std::transform(
1354         LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes),
1355         [](const GlobalVariable *V) -> Type * { return V->getValueType(); });
1356 
1357     StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t");
1358 
1359     Align StructAlign = AMDGPU::getAlign(DL, LocalVars[0]);
1360 
1361     GlobalVariable *SGV = new GlobalVariable(
1362         M, LDSTy, false, GlobalValue::InternalLinkage, PoisonValue::get(LDSTy),
1363         VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
1364         false);
1365     SGV->setAlignment(StructAlign);
1366 
1367     DenseMap<GlobalVariable *, Constant *> Map;
1368     Type *I32 = Type::getInt32Ty(Ctx);
1369     for (size_t I = 0; I < LocalVars.size(); I++) {
1370       GlobalVariable *GV = LocalVars[I];
1371       Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)};
1372       Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx, true);
1373       if (IsPaddingField[I]) {
1374         assert(GV->use_empty());
1375         GV->eraseFromParent();
1376       } else {
1377         Map[GV] = GEP;
1378       }
1379     }
1380     assert(Map.size() == LDSVarsToTransform.size());
1381     return {SGV, std::move(Map)};
1382   }
1383 
1384   template <typename PredicateTy>
1385   static void replaceLDSVariablesWithStruct(
1386       Module &M, DenseSet<GlobalVariable *> const &LDSVarsToTransformArg,
1387       const LDSVariableReplacement &Replacement, PredicateTy Predicate) {
1388     LLVMContext &Ctx = M.getContext();
1389     const DataLayout &DL = M.getDataLayout();
1390 
1391     // A hack... we need to insert the aliasing info in a predictable order for
1392     // lit tests. Would like to have them in a stable order already, ideally the
1393     // same order they get allocated, which might mean an ordered set container
1394     auto LDSVarsToTransform = sortByName(std::vector<GlobalVariable *>(
1395         LDSVarsToTransformArg.begin(), LDSVarsToTransformArg.end()));
1396 
1397     // Create alias.scope and their lists. Each field in the new structure
1398     // does not alias with all other fields.
1399     SmallVector<MDNode *> AliasScopes;
1400     SmallVector<Metadata *> NoAliasList;
1401     const size_t NumberVars = LDSVarsToTransform.size();
1402     if (NumberVars > 1) {
1403       MDBuilder MDB(Ctx);
1404       AliasScopes.reserve(NumberVars);
1405       MDNode *Domain = MDB.createAnonymousAliasScopeDomain();
1406       for (size_t I = 0; I < NumberVars; I++) {
1407         MDNode *Scope = MDB.createAnonymousAliasScope(Domain);
1408         AliasScopes.push_back(Scope);
1409       }
1410       NoAliasList.append(&AliasScopes[1], AliasScopes.end());
1411     }
1412 
1413     // Replace uses of ith variable with a constantexpr to the corresponding
1414     // field of the instance that will be allocated by AMDGPUMachineFunction
1415     for (size_t I = 0; I < NumberVars; I++) {
1416       GlobalVariable *GV = LDSVarsToTransform[I];
1417       Constant *GEP = Replacement.LDSVarsToConstantGEP.at(GV);
1418 
1419       GV->replaceUsesWithIf(GEP, Predicate);
1420 
1421       APInt APOff(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
1422       GEP->stripAndAccumulateInBoundsConstantOffsets(DL, APOff);
1423       uint64_t Offset = APOff.getZExtValue();
1424 
1425       Align A =
1426           commonAlignment(Replacement.SGV->getAlign().valueOrOne(), Offset);
1427 
1428       if (I)
1429         NoAliasList[I - 1] = AliasScopes[I - 1];
1430       MDNode *NoAlias =
1431           NoAliasList.empty() ? nullptr : MDNode::get(Ctx, NoAliasList);
1432       MDNode *AliasScope =
1433           AliasScopes.empty() ? nullptr : MDNode::get(Ctx, {AliasScopes[I]});
1434 
1435       refineUsesAlignmentAndAA(GEP, A, DL, AliasScope, NoAlias);
1436     }
1437   }
1438 
1439   static void refineUsesAlignmentAndAA(Value *Ptr, Align A,
1440                                        const DataLayout &DL, MDNode *AliasScope,
1441                                        MDNode *NoAlias, unsigned MaxDepth = 5) {
1442     if (!MaxDepth || (A == 1 && !AliasScope))
1443       return;
1444 
1445     for (User *U : Ptr->users()) {
1446       if (auto *I = dyn_cast<Instruction>(U)) {
1447         if (AliasScope && I->mayReadOrWriteMemory()) {
1448           MDNode *AS = I->getMetadata(LLVMContext::MD_alias_scope);
1449           AS = (AS ? MDNode::getMostGenericAliasScope(AS, AliasScope)
1450                    : AliasScope);
1451           I->setMetadata(LLVMContext::MD_alias_scope, AS);
1452 
1453           MDNode *NA = I->getMetadata(LLVMContext::MD_noalias);
1454           NA = (NA ? MDNode::intersect(NA, NoAlias) : NoAlias);
1455           I->setMetadata(LLVMContext::MD_noalias, NA);
1456         }
1457       }
1458 
1459       if (auto *LI = dyn_cast<LoadInst>(U)) {
1460         LI->setAlignment(std::max(A, LI->getAlign()));
1461         continue;
1462       }
1463       if (auto *SI = dyn_cast<StoreInst>(U)) {
1464         if (SI->getPointerOperand() == Ptr)
1465           SI->setAlignment(std::max(A, SI->getAlign()));
1466         continue;
1467       }
1468       if (auto *AI = dyn_cast<AtomicRMWInst>(U)) {
1469         // None of atomicrmw operations can work on pointers, but let's
1470         // check it anyway in case it will or we will process ConstantExpr.
1471         if (AI->getPointerOperand() == Ptr)
1472           AI->setAlignment(std::max(A, AI->getAlign()));
1473         continue;
1474       }
1475       if (auto *AI = dyn_cast<AtomicCmpXchgInst>(U)) {
1476         if (AI->getPointerOperand() == Ptr)
1477           AI->setAlignment(std::max(A, AI->getAlign()));
1478         continue;
1479       }
1480       if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
1481         unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP->getType());
1482         APInt Off(BitWidth, 0);
1483         if (GEP->getPointerOperand() == Ptr) {
1484           Align GA;
1485           if (GEP->accumulateConstantOffset(DL, Off))
1486             GA = commonAlignment(A, Off.getLimitedValue());
1487           refineUsesAlignmentAndAA(GEP, GA, DL, AliasScope, NoAlias,
1488                                    MaxDepth - 1);
1489         }
1490         continue;
1491       }
1492       if (auto *I = dyn_cast<Instruction>(U)) {
1493         if (I->getOpcode() == Instruction::BitCast ||
1494             I->getOpcode() == Instruction::AddrSpaceCast)
1495           refineUsesAlignmentAndAA(I, A, DL, AliasScope, NoAlias, MaxDepth - 1);
1496       }
1497     }
1498   }
1499 };
1500 
1501 class AMDGPULowerModuleLDSLegacy : public ModulePass {
1502 public:
1503   const AMDGPUTargetMachine *TM;
1504   static char ID;
1505 
1506   AMDGPULowerModuleLDSLegacy(const AMDGPUTargetMachine *TM_ = nullptr)
1507       : ModulePass(ID), TM(TM_) {
1508     initializeAMDGPULowerModuleLDSLegacyPass(*PassRegistry::getPassRegistry());
1509   }
1510 
1511   void getAnalysisUsage(AnalysisUsage &AU) const override {
1512     if (!TM)
1513       AU.addRequired<TargetPassConfig>();
1514   }
1515 
1516   bool runOnModule(Module &M) override {
1517     if (!TM) {
1518       auto &TPC = getAnalysis<TargetPassConfig>();
1519       TM = &TPC.getTM<AMDGPUTargetMachine>();
1520     }
1521 
1522     return AMDGPULowerModuleLDS(*TM).runOnModule(M);
1523   }
1524 };
1525 
1526 } // namespace
1527 char AMDGPULowerModuleLDSLegacy::ID = 0;
1528 
1529 char &llvm::AMDGPULowerModuleLDSLegacyPassID = AMDGPULowerModuleLDSLegacy::ID;
1530 
1531 INITIALIZE_PASS_BEGIN(AMDGPULowerModuleLDSLegacy, DEBUG_TYPE,
1532                       "Lower uses of LDS variables from non-kernel functions",
1533                       false, false)
1534 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
1535 INITIALIZE_PASS_END(AMDGPULowerModuleLDSLegacy, DEBUG_TYPE,
1536                     "Lower uses of LDS variables from non-kernel functions",
1537                     false, false)
1538 
1539 ModulePass *
1540 llvm::createAMDGPULowerModuleLDSLegacyPass(const AMDGPUTargetMachine *TM) {
1541   return new AMDGPULowerModuleLDSLegacy(TM);
1542 }
1543 
1544 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M,
1545                                                 ModuleAnalysisManager &) {
1546   return AMDGPULowerModuleLDS(TM).runOnModule(M) ? PreservedAnalyses::none()
1547                                                  : PreservedAnalyses::all();
1548 }
1549