xref: /llvm-project/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp (revision 62951784f04ecd640240d7b33fcddc2ee6862bb7)
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 not yet implemented for non-kernel functions.
24 // This lowering could be extended to handle that use case, but would probably
25 // require closer integration with promoteAllocaToLDS.
26 //
27 // Consequences of this GPU feature:
28 // - memory is limited and exceeding it halts compilation
29 // - a global accessed by one kernel exists independent of other kernels
30 // - a global exists independent of simultaneous execution of the same kernel
31 // - the address of the global may be different from different kernels as they
32 //   do not alias, which permits only allocating variables they use
33 // - if the address is allowed to differ, functions need help to find it
34 //
35 // Uses from kernels are implemented here by grouping them in a per-kernel
36 // struct instance. This duplicates the variables, accurately modelling their
37 // aliasing properties relative to a single global representation. It also
38 // permits control over alignment via padding.
39 //
40 // Uses from functions are more complicated and the primary purpose of this
41 // IR pass. Several different lowering are chosen between to meet requirements
42 // to avoid allocating any LDS where it is not necessary, as that impacts
43 // occupancy and may fail the compilation, while not imposing overhead on a
44 // feature whose primary advantage over global memory is performance. The basic
45 // design goal is to avoid one kernel imposing overhead on another.
46 //
47 // Implementation.
48 //
49 // LDS variables with constant annotation or non-undef initializer are passed
50 // through unchanged for simplification or error diagnostics in later passes.
51 // Non-undef initializers are not yet implemented for LDS.
52 //
53 // LDS variables that are always allocated at the same address can be found
54 // by lookup at that address. Otherwise runtime information/cost is required.
55 //
56 // The simplest strategy possible is to group all LDS variables in a single
57 // struct and allocate that struct in every kernel such that the original
58 // variables are always at the same address. LDS is however a limited resource
59 // so this strategy is unusable in practice. It is not implemented here.
60 //
61 // Strategy | Precise allocation | Zero runtime cost | General purpose |
62 //  --------+--------------------+-------------------+-----------------+
63 //   Module |                 No |               Yes |             Yes |
64 //    Table |                Yes |                No |             Yes |
65 //   Kernel |                Yes |               Yes |              No |
66 //   Hybrid |                Yes |           Partial |             Yes |
67 //
68 // Module spends LDS memory to save cycles. Table spends cycles and global
69 // memory to save LDS. Kernel is as fast as kernel allocation but only works
70 // for variables that are known reachable from a single kernel. Hybrid picks
71 // between all three. When forced to choose between LDS and cycles it minimises
72 // LDS use.
73 
74 // The "module" lowering implemented here finds LDS variables which are used by
75 // non-kernel functions and creates a new struct with a field for each of those
76 // LDS variables. Variables that are only used from kernels are excluded.
77 // Kernels that do not use this struct are annoteated with the attribute
78 // amdgpu-elide-module-lds which allows the back end to elide the allocation.
79 //
80 // The "table" lowering implemented here has three components.
81 // First kernels are assigned a unique integer identifier which is available in
82 // functions it calls through the intrinsic amdgcn_lds_kernel_id. The integer
83 // is passed through a specific SGPR, thus works with indirect calls.
84 // Second, each kernel allocates LDS variables independent of other kernels and
85 // writes the addresses it chose for each variable into an array in consistent
86 // order. If the kernel does not allocate a given variable, it writes undef to
87 // the corresponding array location. These arrays are written to a constant
88 // table in the order matching the kernel unique integer identifier.
89 // Third, uses from non-kernel functions are replaced with a table lookup using
90 // the intrinsic function to find the address of the variable.
91 //
92 // "Kernel" lowering is only applicable for variables that are unambiguously
93 // reachable from exactly one kernel. For those cases, accesses to the variable
94 // can be lowered to ConstantExpr address of a struct instance specific to that
95 // one kernel. This is zero cost in space and in compute. It will raise a fatal
96 // error on any variable that might be reachable from multiple kernels and is
97 // thus most easily used as part of the hybrid lowering strategy.
98 //
99 // Hybrid lowering is a mixture of the above. It uses the zero cost kernel
100 // lowering where it can. It lowers the variable accessed by the greatest
101 // number of kernels using the module strategy as that is free for the first
102 // variable. Any futher variables that can be lowered with the module strategy
103 // without incurring LDS memory overhead are. The remaining ones are lowered
104 // via table.
105 //
106 // Consequences
107 // - No heuristics or user controlled magic numbers, hybrid is the right choice
108 // - Kernels that don't use functions (or have had them all inlined) are not
109 //   affected by any lowering for kernels that do.
110 // - Kernels that don't make indirect function calls are not affected by those
111 //   that do.
112 // - Variables which are used by lots of kernels, e.g. those injected by a
113 //   language runtime in most kernels, are expected to have no overhead
114 // - Implementations that instantiate templates per-kernel where those templates
115 //   use LDS are expected to hit the "Kernel" lowering strategy
116 // - The runtime properties impose a cost in compiler implementation complexity
117 //
118 //===----------------------------------------------------------------------===//
119 
120 #include "AMDGPU.h"
121 #include "Utils/AMDGPUBaseInfo.h"
122 #include "Utils/AMDGPUMemoryUtils.h"
123 #include "llvm/ADT/BitVector.h"
124 #include "llvm/ADT/DenseMap.h"
125 #include "llvm/ADT/DenseSet.h"
126 #include "llvm/ADT/STLExtras.h"
127 #include "llvm/ADT/SetOperations.h"
128 #include "llvm/ADT/SetVector.h"
129 #include "llvm/Analysis/CallGraph.h"
130 #include "llvm/IR/Constants.h"
131 #include "llvm/IR/DerivedTypes.h"
132 #include "llvm/IR/IRBuilder.h"
133 #include "llvm/IR/InlineAsm.h"
134 #include "llvm/IR/Instructions.h"
135 #include "llvm/IR/IntrinsicsAMDGPU.h"
136 #include "llvm/IR/MDBuilder.h"
137 #include "llvm/IR/ReplaceConstant.h"
138 #include "llvm/InitializePasses.h"
139 #include "llvm/Pass.h"
140 #include "llvm/Support/CommandLine.h"
141 #include "llvm/Support/Debug.h"
142 #include "llvm/Support/OptimizedStructLayout.h"
143 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
144 #include "llvm/Transforms/Utils/ModuleUtils.h"
145 
146 #include <tuple>
147 #include <vector>
148 
149 #include <cstdio>
150 
151 #define DEBUG_TYPE "amdgpu-lower-module-lds"
152 
153 using namespace llvm;
154 
155 namespace {
156 
157 cl::opt<bool> SuperAlignLDSGlobals(
158     "amdgpu-super-align-lds-globals",
159     cl::desc("Increase alignment of LDS if it is not on align boundary"),
160     cl::init(true), cl::Hidden);
161 
162 enum class LoweringKind { module, table, kernel, hybrid };
163 cl::opt<LoweringKind> LoweringKindLoc(
164     "amdgpu-lower-module-lds-strategy",
165     cl::desc("Specify lowering strategy for function LDS access:"), cl::Hidden,
166     cl::init(LoweringKind::hybrid),
167     cl::values(
168         clEnumValN(LoweringKind::table, "table", "Lower via table lookup"),
169         clEnumValN(LoweringKind::module, "module", "Lower via module struct"),
170         clEnumValN(
171             LoweringKind::kernel, "kernel",
172             "Lower variables reachable from one kernel, otherwise abort"),
173         clEnumValN(LoweringKind::hybrid, "hybrid",
174                    "Lower via mixture of above strategies")));
175 
176 bool isKernelLDS(const Function *F) {
177   // Some weirdness here. AMDGPU::isKernelCC does not call into
178   // AMDGPU::isKernel with the calling conv, it instead calls into
179   // isModuleEntryFunction which returns true for more calling conventions
180   // than AMDGPU::isKernel does. There's a FIXME on AMDGPU::isKernel.
181   // There's also a test that checks that the LDS lowering does not hit on
182   // a graphics shader, denoted amdgpu_ps, so stay with the limited case.
183   // Putting LDS in the name of the function to draw attention to this.
184   return AMDGPU::isKernel(F->getCallingConv());
185 }
186 
187 class AMDGPULowerModuleLDS : public ModulePass {
188 
189   static void
190   removeLocalVarsFromUsedLists(Module &M,
191                                const DenseSet<GlobalVariable *> &LocalVars) {
192     // The verifier rejects used lists containing an inttoptr of a constant
193     // so remove the variables from these lists before replaceAllUsesWith
194     SmallPtrSet<Constant *, 8> LocalVarsSet;
195     for (GlobalVariable *LocalVar : LocalVars)
196       LocalVarsSet.insert(cast<Constant>(LocalVar->stripPointerCasts()));
197 
198     removeFromUsedLists(
199         M, [&LocalVarsSet](Constant *C) { return LocalVarsSet.count(C); });
200 
201     for (GlobalVariable *LocalVar : LocalVars)
202       LocalVar->removeDeadConstantUsers();
203   }
204 
205   static void markUsedByKernel(IRBuilder<> &Builder, Function *Func,
206                                GlobalVariable *SGV) {
207     // The llvm.amdgcn.module.lds instance is implicitly used by all kernels
208     // that might call a function which accesses a field within it. This is
209     // presently approximated to 'all kernels' if there are any such functions
210     // in the module. This implicit use is redefined as an explicit use here so
211     // that later passes, specifically PromoteAlloca, account for the required
212     // memory without any knowledge of this transform.
213 
214     // An operand bundle on llvm.donothing works because the call instruction
215     // survives until after the last pass that needs to account for LDS. It is
216     // better than inline asm as the latter survives until the end of codegen. A
217     // totally robust solution would be a function with the same semantics as
218     // llvm.donothing that takes a pointer to the instance and is lowered to a
219     // no-op after LDS is allocated, but that is not presently necessary.
220 
221     // This intrinsic is eliminated shortly before instruction selection. It
222     // does not suffice to indicate to ISel that a given global which is not
223     // immediately used by the kernel must still be allocated by it. An
224     // equivalent target specific intrinsic which lasts until immediately after
225     // codegen would suffice for that, but one would still need to ensure that
226     // the variables are allocated in the anticpated order.
227 
228     LLVMContext &Ctx = Func->getContext();
229 
230     Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI());
231 
232     FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {});
233 
234     Function *Decl =
235         Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {});
236 
237     Value *UseInstance[1] = {Builder.CreateInBoundsGEP(
238         SGV->getValueType(), SGV, ConstantInt::get(Type::getInt32Ty(Ctx), 0))};
239 
240     Builder.CreateCall(FTy, Decl, {},
241                        {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)},
242                        "");
243   }
244 
245   static bool eliminateConstantExprUsesOfLDSFromAllInstructions(Module &M) {
246     // Constants are uniqued within LLVM. A ConstantExpr referring to a LDS
247     // global may have uses from multiple different functions as a result.
248     // This pass specialises LDS variables with respect to the kernel that
249     // allocates them.
250 
251     // This is semantically equivalent to (the unimplemented as slow):
252     // for (auto &F : M.functions())
253     //   for (auto &BB : F)
254     //     for (auto &I : BB)
255     //       for (Use &Op : I.operands())
256     //         if (constantExprUsesLDS(Op))
257     //           replaceConstantExprInFunction(I, Op);
258 
259     SmallVector<Constant *> LDSGlobals;
260     for (auto &GV : M.globals())
261       if (AMDGPU::isLDSVariableToLower(GV))
262         LDSGlobals.push_back(&GV);
263 
264     return convertUsersOfConstantsToInstructions(LDSGlobals);
265   }
266 
267 public:
268   static char ID;
269 
270   AMDGPULowerModuleLDS() : ModulePass(ID) {
271     initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry());
272   }
273 
274   using FunctionVariableMap = DenseMap<Function *, DenseSet<GlobalVariable *>>;
275 
276   using VariableFunctionMap = DenseMap<GlobalVariable *, DenseSet<Function *>>;
277 
278   static void getUsesOfLDSByFunction(CallGraph const &CG, Module &M,
279                                      FunctionVariableMap &kernels,
280                                      FunctionVariableMap &functions) {
281 
282     // Get uses from the current function, excluding uses by called functions
283     // Two output variables to avoid walking the globals list twice
284     for (auto &GV : M.globals()) {
285       if (!AMDGPU::isLDSVariableToLower(GV)) {
286         continue;
287       }
288 
289       SmallVector<User *, 16> Stack(GV.users());
290       for (User *V : GV.users()) {
291         if (auto *I = dyn_cast<Instruction>(V)) {
292           Function *F = I->getFunction();
293           if (isKernelLDS(F)) {
294             kernels[F].insert(&GV);
295           } else {
296             functions[F].insert(&GV);
297           }
298         }
299       }
300     }
301   }
302 
303   struct LDSUsesInfoTy {
304     FunctionVariableMap direct_access;
305     FunctionVariableMap indirect_access;
306   };
307 
308   static LDSUsesInfoTy getTransitiveUsesOfLDS(CallGraph const &CG, Module &M) {
309 
310     FunctionVariableMap direct_map_kernel;
311     FunctionVariableMap direct_map_function;
312     getUsesOfLDSByFunction(CG, M, direct_map_kernel, direct_map_function);
313 
314     // Collect variables that are used by functions whose address has escaped
315     DenseSet<GlobalVariable *> VariablesReachableThroughFunctionPointer;
316     for (Function &F : M.functions()) {
317       if (!isKernelLDS(&F))
318         if (F.hasAddressTaken(nullptr,
319                               /* IgnoreCallbackUses */ false,
320                               /* IgnoreAssumeLikeCalls */ false,
321                               /* IgnoreLLVMUsed */ true,
322                               /* IgnoreArcAttachedCall */ false)) {
323           set_union(VariablesReachableThroughFunctionPointer,
324                     direct_map_function[&F]);
325         }
326     }
327 
328     auto functionMakesUnknownCall = [&](const Function *F) -> bool {
329       assert(!F->isDeclaration());
330       for (CallGraphNode::CallRecord R : *CG[F]) {
331         if (!R.second->getFunction()) {
332           return true;
333         }
334       }
335       return false;
336     };
337 
338     // Work out which variables are reachable through function calls
339     FunctionVariableMap transitive_map_function = direct_map_function;
340 
341     // If the function makes any unknown call, assume the worst case that it can
342     // access all variables accessed by functions whose address escaped
343     for (Function &F : M.functions()) {
344       if (!F.isDeclaration() && functionMakesUnknownCall(&F)) {
345         if (!isKernelLDS(&F)) {
346           set_union(transitive_map_function[&F],
347                     VariablesReachableThroughFunctionPointer);
348         }
349       }
350     }
351 
352     // Direct implementation of collecting all variables reachable from each
353     // function
354     for (Function &Func : M.functions()) {
355       if (Func.isDeclaration() || isKernelLDS(&Func))
356         continue;
357 
358       DenseSet<Function *> seen; // catches cycles
359       SmallVector<Function *, 4> wip{&Func};
360 
361       while (!wip.empty()) {
362         Function *F = wip.pop_back_val();
363 
364         // Can accelerate this by referring to transitive map for functions that
365         // have already been computed, with more care than this
366         set_union(transitive_map_function[&Func], direct_map_function[F]);
367 
368         for (CallGraphNode::CallRecord R : *CG[F]) {
369           Function *ith = R.second->getFunction();
370           if (ith) {
371             if (!seen.contains(ith)) {
372               seen.insert(ith);
373               wip.push_back(ith);
374             }
375           }
376         }
377       }
378     }
379 
380     // direct_map_kernel lists which variables are used by the kernel
381     // find the variables which are used through a function call
382     FunctionVariableMap indirect_map_kernel;
383 
384     for (Function &Func : M.functions()) {
385       if (Func.isDeclaration() || !isKernelLDS(&Func))
386         continue;
387 
388       for (CallGraphNode::CallRecord R : *CG[&Func]) {
389         Function *ith = R.second->getFunction();
390         if (ith) {
391           set_union(indirect_map_kernel[&Func], transitive_map_function[ith]);
392         } else {
393           set_union(indirect_map_kernel[&Func],
394                     VariablesReachableThroughFunctionPointer);
395         }
396       }
397     }
398 
399     return {std::move(direct_map_kernel), std::move(indirect_map_kernel)};
400   }
401 
402   struct LDSVariableReplacement {
403     GlobalVariable *SGV = nullptr;
404     DenseMap<GlobalVariable *, Constant *> LDSVarsToConstantGEP;
405   };
406 
407   // remap from lds global to a constantexpr gep to where it has been moved to
408   // for each kernel
409   // an array with an element for each kernel containing where the corresponding
410   // variable was remapped to
411 
412   static Constant *getAddressesOfVariablesInKernel(
413       LLVMContext &Ctx, ArrayRef<GlobalVariable *> Variables,
414       DenseMap<GlobalVariable *, Constant *> &LDSVarsToConstantGEP) {
415     // Create a ConstantArray containing the address of each Variable within the
416     // kernel corresponding to LDSVarsToConstantGEP, or poison if that kernel
417     // does not allocate it
418     // TODO: Drop the ptrtoint conversion
419 
420     Type *I32 = Type::getInt32Ty(Ctx);
421 
422     ArrayType *KernelOffsetsType = ArrayType::get(I32, Variables.size());
423 
424     SmallVector<Constant *> Elements;
425     for (size_t i = 0; i < Variables.size(); i++) {
426       GlobalVariable *GV = Variables[i];
427       if (LDSVarsToConstantGEP.count(GV) != 0) {
428         auto elt = ConstantExpr::getPtrToInt(LDSVarsToConstantGEP[GV], I32);
429         Elements.push_back(elt);
430       } else {
431         Elements.push_back(PoisonValue::get(I32));
432       }
433     }
434     return ConstantArray::get(KernelOffsetsType, Elements);
435   }
436 
437   static GlobalVariable *buildLookupTable(
438       Module &M, ArrayRef<GlobalVariable *> Variables,
439       ArrayRef<Function *> kernels,
440       DenseMap<Function *, LDSVariableReplacement> &KernelToReplacement) {
441     if (Variables.empty()) {
442       return nullptr;
443     }
444     LLVMContext &Ctx = M.getContext();
445 
446     const size_t NumberVariables = Variables.size();
447     const size_t NumberKernels = kernels.size();
448 
449     ArrayType *KernelOffsetsType =
450         ArrayType::get(Type::getInt32Ty(Ctx), NumberVariables);
451 
452     ArrayType *AllKernelsOffsetsType =
453         ArrayType::get(KernelOffsetsType, NumberKernels);
454 
455     std::vector<Constant *> overallConstantExprElts(NumberKernels);
456     for (size_t i = 0; i < NumberKernels; i++) {
457       LDSVariableReplacement Replacement = KernelToReplacement[kernels[i]];
458       overallConstantExprElts[i] = getAddressesOfVariablesInKernel(
459           Ctx, Variables, Replacement.LDSVarsToConstantGEP);
460     }
461 
462     Constant *init =
463         ConstantArray::get(AllKernelsOffsetsType, overallConstantExprElts);
464 
465     return new GlobalVariable(
466         M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
467         "llvm.amdgcn.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
468         AMDGPUAS::CONSTANT_ADDRESS);
469   }
470 
471   void replaceUseWithTableLookup(Module &M, IRBuilder<> &Builder,
472                                  GlobalVariable *LookupTable,
473                                  GlobalVariable *GV, Use &U, Value *Index) {
474     // Table is a constant array of the same length as OrderedKernels
475     LLVMContext &Ctx = M.getContext();
476     Type *I32 = Type::getInt32Ty(Ctx);
477     auto *I = cast<Instruction>(U.getUser());
478 
479     Value *tableKernelIndex = getTableLookupKernelIndex(M, I->getFunction());
480 
481     if (auto *Phi = dyn_cast<PHINode>(I)) {
482       BasicBlock *BB = Phi->getIncomingBlock(U);
483       Builder.SetInsertPoint(&(*(BB->getFirstInsertionPt())));
484     } else {
485       Builder.SetInsertPoint(I);
486     }
487 
488     Value *GEPIdx[3] = {
489         ConstantInt::get(I32, 0),
490         tableKernelIndex,
491         Index,
492     };
493 
494     Value *Address = Builder.CreateInBoundsGEP(
495         LookupTable->getValueType(), LookupTable, GEPIdx, GV->getName());
496 
497     Value *loaded = Builder.CreateLoad(I32, Address);
498 
499     Value *replacement =
500         Builder.CreateIntToPtr(loaded, GV->getType(), GV->getName());
501 
502     U.set(replacement);
503   }
504 
505   void replaceUsesInInstructionsWithTableLookup(
506       Module &M, ArrayRef<GlobalVariable *> ModuleScopeVariables,
507       GlobalVariable *LookupTable) {
508 
509     LLVMContext &Ctx = M.getContext();
510     IRBuilder<> Builder(Ctx);
511     Type *I32 = Type::getInt32Ty(Ctx);
512 
513     for (size_t Index = 0; Index < ModuleScopeVariables.size(); Index++) {
514       auto *GV = ModuleScopeVariables[Index];
515 
516       for (Use &U : make_early_inc_range(GV->uses())) {
517         auto *I = dyn_cast<Instruction>(U.getUser());
518         if (!I)
519           continue;
520 
521         replaceUseWithTableLookup(M, Builder, LookupTable, GV, U,
522                                   ConstantInt::get(I32, Index));
523       }
524     }
525   }
526 
527   static DenseSet<Function *> kernelsThatIndirectlyAccessAnyOfPassedVariables(
528       Module &M, LDSUsesInfoTy &LDSUsesInfo,
529       DenseSet<GlobalVariable *> const &VariableSet) {
530 
531     DenseSet<Function *> KernelSet;
532 
533     if (VariableSet.empty())
534       return KernelSet;
535 
536     for (Function &Func : M.functions()) {
537       if (Func.isDeclaration() || !isKernelLDS(&Func))
538         continue;
539       for (GlobalVariable *GV : LDSUsesInfo.indirect_access[&Func]) {
540         if (VariableSet.contains(GV)) {
541           KernelSet.insert(&Func);
542           break;
543         }
544       }
545     }
546 
547     return KernelSet;
548   }
549 
550   static GlobalVariable *
551   chooseBestVariableForModuleStrategy(const DataLayout &DL,
552                                       VariableFunctionMap &LDSVars) {
553     // Find the global variable with the most indirect uses from kernels
554 
555     struct CandidateTy {
556       GlobalVariable *GV = nullptr;
557       size_t UserCount = 0;
558       size_t Size = 0;
559 
560       CandidateTy() = default;
561 
562       CandidateTy(GlobalVariable *GV, uint64_t UserCount, uint64_t AllocSize)
563           : GV(GV), UserCount(UserCount), Size(AllocSize) {}
564 
565       bool operator<(const CandidateTy &Other) const {
566         // Fewer users makes module scope variable less attractive
567         if (UserCount < Other.UserCount) {
568           return true;
569         }
570         if (UserCount > Other.UserCount) {
571           return false;
572         }
573 
574         // Bigger makes module scope variable less attractive
575         if (Size < Other.Size) {
576           return false;
577         }
578 
579         if (Size > Other.Size) {
580           return true;
581         }
582 
583         // Arbitrary but consistent
584         return GV->getName() < Other.GV->getName();
585       }
586     };
587 
588     CandidateTy MostUsed;
589 
590     for (auto &K : LDSVars) {
591       GlobalVariable *GV = K.first;
592       if (K.second.size() <= 1) {
593         // A variable reachable by only one kernel is best lowered with kernel
594         // strategy
595         continue;
596       }
597       CandidateTy Candidate(
598           GV, K.second.size(),
599           DL.getTypeAllocSize(GV->getValueType()).getFixedValue());
600       if (MostUsed < Candidate)
601         MostUsed = Candidate;
602     }
603 
604     return MostUsed.GV;
605   }
606 
607   static void recordLDSAbsoluteAddress(Module *M, GlobalVariable *GV,
608                                        uint32_t Address) {
609     // Write the specified address into metadata where it can be retrieved by
610     // the assembler. Format is a half open range, [Address Address+1)
611     LLVMContext &Ctx = M->getContext();
612     auto *IntTy =
613         M->getDataLayout().getIntPtrType(Ctx, AMDGPUAS::LOCAL_ADDRESS);
614     auto *MinC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address));
615     auto *MaxC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address + 1));
616     GV->setMetadata(LLVMContext::MD_absolute_symbol,
617                     MDNode::get(Ctx, {MinC, MaxC}));
618   }
619 
620   DenseMap<Function *, Value *> tableKernelIndexCache;
621   Value *getTableLookupKernelIndex(Module &M, Function *F) {
622     // Accesses from a function use the amdgcn_lds_kernel_id intrinsic which
623     // lowers to a read from a live in register. Emit it once in the entry
624     // block to spare deduplicating it later.
625     if (tableKernelIndexCache.count(F) == 0) {
626       LLVMContext &Ctx = M.getContext();
627       IRBuilder<> Builder(Ctx);
628       FunctionType *FTy = FunctionType::get(Type::getInt32Ty(Ctx), {});
629       Function *Decl =
630           Intrinsic::getDeclaration(&M, Intrinsic::amdgcn_lds_kernel_id, {});
631 
632       BasicBlock::iterator it =
633           F->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
634       Instruction &i = *it;
635       Builder.SetInsertPoint(&i);
636 
637       tableKernelIndexCache[F] = Builder.CreateCall(FTy, Decl, {});
638     }
639 
640     return tableKernelIndexCache[F];
641   }
642 
643   static std::vector<Function *> assignLDSKernelIDToEachKernel(
644       Module *M, DenseSet<Function *> const &KernelsThatAllocateTableLDS) {
645     // Associate kernels in the set with an arbirary but reproducible order and
646     // annotate them with that order in metadata. This metadata is recognised by
647     // the backend and lowered to a SGPR which can be read from using
648     // amdgcn_lds_kernel_id.
649 
650     std::vector<Function *> OrderedKernels;
651     {
652       for (Function &Func : M->functions()) {
653         if (Func.isDeclaration())
654           continue;
655         if (!isKernelLDS(&Func))
656           continue;
657 
658         if (KernelsThatAllocateTableLDS.contains(&Func)) {
659           assert(Func.hasName()); // else fatal error earlier
660           OrderedKernels.push_back(&Func);
661         }
662       }
663 
664       // Put them in an arbitrary but reproducible order
665       llvm::sort(OrderedKernels.begin(), OrderedKernels.end(),
666                  [](const Function *lhs, const Function *rhs) -> bool {
667                    return lhs->getName() < rhs->getName();
668                  });
669 
670       // Annotate the kernels with their order in this vector
671       LLVMContext &Ctx = M->getContext();
672       IRBuilder<> Builder(Ctx);
673 
674       if (OrderedKernels.size() > UINT32_MAX) {
675         // 32 bit keeps it in one SGPR. > 2**32 kernels won't fit on the GPU
676         report_fatal_error("Unimplemented LDS lowering for > 2**32 kernels");
677       }
678 
679       for (size_t i = 0; i < OrderedKernels.size(); i++) {
680         Metadata *AttrMDArgs[1] = {
681             ConstantAsMetadata::get(Builder.getInt32(i)),
682         };
683         OrderedKernels[i]->setMetadata("llvm.amdgcn.lds.kernel.id",
684                                        MDNode::get(Ctx, AttrMDArgs));
685       }
686     }
687     return OrderedKernels;
688   }
689 
690   static void partitionVariablesIntoIndirectStrategies(
691       Module &M, LDSUsesInfoTy const &LDSUsesInfo,
692       VariableFunctionMap &LDSToKernelsThatNeedToAccessItIndirectly,
693       DenseSet<GlobalVariable *> &ModuleScopeVariables,
694       DenseSet<GlobalVariable *> &TableLookupVariables,
695       DenseSet<GlobalVariable *> &KernelAccessVariables) {
696 
697     GlobalVariable *HybridModuleRoot =
698         LoweringKindLoc != LoweringKind::hybrid
699             ? nullptr
700             : chooseBestVariableForModuleStrategy(
701                   M.getDataLayout(), LDSToKernelsThatNeedToAccessItIndirectly);
702 
703     DenseSet<Function *> const EmptySet;
704     DenseSet<Function *> const &HybridModuleRootKernels =
705         HybridModuleRoot
706             ? LDSToKernelsThatNeedToAccessItIndirectly[HybridModuleRoot]
707             : EmptySet;
708 
709     for (auto &K : LDSToKernelsThatNeedToAccessItIndirectly) {
710       // Each iteration of this loop assigns exactly one global variable to
711       // exactly one of the implementation strategies.
712 
713       GlobalVariable *GV = K.first;
714       assert(AMDGPU::isLDSVariableToLower(*GV));
715       assert(K.second.size() != 0);
716 
717       switch (LoweringKindLoc) {
718       case LoweringKind::module:
719         ModuleScopeVariables.insert(GV);
720         break;
721 
722       case LoweringKind::table:
723         TableLookupVariables.insert(GV);
724         break;
725 
726       case LoweringKind::kernel:
727         if (K.second.size() == 1) {
728           KernelAccessVariables.insert(GV);
729         } else {
730           report_fatal_error(
731               "cannot lower LDS '" + GV->getName() +
732               "' to kernel access as it is reachable from multiple kernels");
733         }
734         break;
735 
736       case LoweringKind::hybrid: {
737         if (GV == HybridModuleRoot) {
738           assert(K.second.size() != 1);
739           ModuleScopeVariables.insert(GV);
740         } else if (K.second.size() == 1) {
741           KernelAccessVariables.insert(GV);
742         } else if (set_is_subset(K.second, HybridModuleRootKernels)) {
743           ModuleScopeVariables.insert(GV);
744         } else {
745           TableLookupVariables.insert(GV);
746         }
747         break;
748       }
749       }
750     }
751 
752     // All LDS variables accessed indirectly have now been partitioned into
753     // the distinct lowering strategies.
754     assert(ModuleScopeVariables.size() + TableLookupVariables.size() +
755                KernelAccessVariables.size() ==
756            LDSToKernelsThatNeedToAccessItIndirectly.size());
757   }
758 
759   static GlobalVariable *lowerModuleScopeStructVariables(
760       Module &M, DenseSet<GlobalVariable *> const &ModuleScopeVariables,
761       DenseSet<Function *> const &KernelsThatAllocateModuleLDS) {
762     // Create a struct to hold the ModuleScopeVariables
763     // Replace all uses of those variables from non-kernel functions with the
764     // new struct instance Replace only the uses from kernel functions that will
765     // allocate this instance. That is a space optimisation - kernels that use a
766     // subset of the module scope struct and do not need to allocate it for
767     // indirect calls will only allocate the subset they use (they do so as part
768     // of the per-kernel lowering).
769     if (ModuleScopeVariables.empty()) {
770       return nullptr;
771     }
772 
773     LLVMContext &Ctx = M.getContext();
774 
775     LDSVariableReplacement ModuleScopeReplacement =
776         createLDSVariableReplacement(M, "llvm.amdgcn.module.lds",
777                                      ModuleScopeVariables);
778 
779     appendToCompilerUsed(M, {static_cast<GlobalValue *>(
780                                 ConstantExpr::getPointerBitCastOrAddrSpaceCast(
781                                     cast<Constant>(ModuleScopeReplacement.SGV),
782                                     Type::getInt8PtrTy(Ctx)))});
783 
784     // module.lds will be allocated at zero in any kernel that allocates it
785     recordLDSAbsoluteAddress(&M, ModuleScopeReplacement.SGV, 0);
786 
787     // historic
788     removeLocalVarsFromUsedLists(M, ModuleScopeVariables);
789 
790     // Replace all uses of module scope variable from non-kernel functions
791     replaceLDSVariablesWithStruct(
792         M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) {
793           Instruction *I = dyn_cast<Instruction>(U.getUser());
794           if (!I) {
795             return false;
796           }
797           Function *F = I->getFunction();
798           return !isKernelLDS(F);
799         });
800 
801     // Replace uses of module scope variable from kernel functions that
802     // allocate the module scope variable, otherwise leave them unchanged
803     // Record on each kernel whether the module scope global is used by it
804 
805     IRBuilder<> Builder(Ctx);
806 
807     for (Function &Func : M.functions()) {
808       if (Func.isDeclaration() || !isKernelLDS(&Func))
809         continue;
810 
811       if (KernelsThatAllocateModuleLDS.contains(&Func)) {
812         replaceLDSVariablesWithStruct(
813             M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) {
814               Instruction *I = dyn_cast<Instruction>(U.getUser());
815               if (!I) {
816                 return false;
817               }
818               Function *F = I->getFunction();
819               return F == &Func;
820             });
821 
822         markUsedByKernel(Builder, &Func, ModuleScopeReplacement.SGV);
823 
824       } else {
825         Func.addFnAttr("amdgpu-elide-module-lds");
826       }
827     }
828 
829     return ModuleScopeReplacement.SGV;
830   }
831 
832   static DenseMap<Function *, LDSVariableReplacement>
833   lowerKernelScopeStructVariables(
834       Module &M, LDSUsesInfoTy &LDSUsesInfo,
835       DenseSet<GlobalVariable *> const &ModuleScopeVariables,
836       DenseSet<Function *> const &KernelsThatAllocateModuleLDS,
837       GlobalVariable *MaybeModuleScopeStruct) {
838 
839     // Create a struct for each kernel for the non-module-scope variables.
840 
841     DenseMap<Function *, LDSVariableReplacement> KernelToReplacement;
842     for (Function &Func : M.functions()) {
843       if (Func.isDeclaration() || !isKernelLDS(&Func))
844         continue;
845 
846       DenseSet<GlobalVariable *> KernelUsedVariables;
847       // Allocating variables that are used directly in this struct to get
848       // alignment aware allocation and predictable frame size.
849       for (auto &v : LDSUsesInfo.direct_access[&Func]) {
850         KernelUsedVariables.insert(v);
851       }
852 
853       // Allocating variables that are accessed indirectly so that a lookup of
854       // this struct instance can find them from nested functions.
855       for (auto &v : LDSUsesInfo.indirect_access[&Func]) {
856         KernelUsedVariables.insert(v);
857       }
858 
859       // Variables allocated in module lds must all resolve to that struct,
860       // not to the per-kernel instance.
861       if (KernelsThatAllocateModuleLDS.contains(&Func)) {
862         for (GlobalVariable *v : ModuleScopeVariables) {
863           KernelUsedVariables.erase(v);
864         }
865       }
866 
867       if (KernelUsedVariables.empty()) {
868         // Either used no LDS, or the LDS it used was all in the module struct
869         continue;
870       }
871 
872       // The association between kernel function and LDS struct is done by
873       // symbol name, which only works if the function in question has a
874       // name This is not expected to be a problem in practice as kernels
875       // are called by name making anonymous ones (which are named by the
876       // backend) difficult to use. This does mean that llvm test cases need
877       // to name the kernels.
878       if (!Func.hasName()) {
879         report_fatal_error("Anonymous kernels cannot use LDS variables");
880       }
881 
882       std::string VarName =
883           (Twine("llvm.amdgcn.kernel.") + Func.getName() + ".lds").str();
884 
885       auto Replacement =
886           createLDSVariableReplacement(M, VarName, KernelUsedVariables);
887 
888       // This struct is allocated at a predictable address that can be
889       // calculated now, recorded in metadata then used to lower references to
890       // it during codegen.
891       {
892         // frame layout, starting from 0
893         //{
894         //  module.lds
895         //  alignment padding
896         //  kernel instance
897         //}
898 
899         if (!MaybeModuleScopeStruct ||
900             Func.hasFnAttribute("amdgpu-elide-module-lds")) {
901           // There's no module.lds for this kernel so this replacement struct
902           // goes first
903           recordLDSAbsoluteAddress(&M, Replacement.SGV, 0);
904         } else {
905           const DataLayout &DL = M.getDataLayout();
906           TypeSize ModuleSize =
907               DL.getTypeAllocSize(MaybeModuleScopeStruct->getValueType());
908           GlobalVariable *KernelStruct = Replacement.SGV;
909           Align KernelAlign = AMDGPU::getAlign(DL, KernelStruct);
910           recordLDSAbsoluteAddress(&M, Replacement.SGV,
911                                    alignTo(ModuleSize, KernelAlign));
912         }
913       }
914 
915       // remove preserves existing codegen
916       removeLocalVarsFromUsedLists(M, KernelUsedVariables);
917       KernelToReplacement[&Func] = Replacement;
918 
919       // Rewrite uses within kernel to the new struct
920       replaceLDSVariablesWithStruct(
921           M, KernelUsedVariables, Replacement, [&Func](Use &U) {
922             Instruction *I = dyn_cast<Instruction>(U.getUser());
923             return I && I->getFunction() == &Func;
924           });
925     }
926     return KernelToReplacement;
927   }
928 
929   bool runOnModule(Module &M) override {
930     CallGraph CG = CallGraph(M);
931     bool Changed = superAlignLDSGlobals(M);
932 
933     Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M);
934 
935     Changed = true; // todo: narrow this down
936 
937     // For each kernel, what variables does it access directly or through
938     // callees
939     LDSUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDS(CG, M);
940 
941     // For each variable accessed through callees, which kernels access it
942     VariableFunctionMap LDSToKernelsThatNeedToAccessItIndirectly;
943     for (auto &K : LDSUsesInfo.indirect_access) {
944       Function *F = K.first;
945       assert(isKernelLDS(F));
946       for (GlobalVariable *GV : K.second) {
947         LDSToKernelsThatNeedToAccessItIndirectly[GV].insert(F);
948       }
949     }
950 
951     DenseSet<GlobalVariable *> ModuleScopeVariables;
952     DenseSet<GlobalVariable *> TableLookupVariables;
953     DenseSet<GlobalVariable *> KernelAccessVariables;
954     partitionVariablesIntoIndirectStrategies(
955         M, LDSUsesInfo, LDSToKernelsThatNeedToAccessItIndirectly,
956         ModuleScopeVariables, TableLookupVariables, KernelAccessVariables);
957 
958     // If the kernel accesses a variable that is going to be stored in the
959     // module instance through a call then that kernel needs to allocate the
960     // module instance
961     DenseSet<Function *> KernelsThatAllocateModuleLDS =
962         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
963                                                         ModuleScopeVariables);
964     DenseSet<Function *> KernelsThatAllocateTableLDS =
965         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
966                                                         TableLookupVariables);
967 
968     GlobalVariable *MaybeModuleScopeStruct = lowerModuleScopeStructVariables(
969         M, ModuleScopeVariables, KernelsThatAllocateModuleLDS);
970 
971     DenseMap<Function *, LDSVariableReplacement> KernelToReplacement =
972         lowerKernelScopeStructVariables(M, LDSUsesInfo, ModuleScopeVariables,
973                                         KernelsThatAllocateModuleLDS,
974                                         MaybeModuleScopeStruct);
975 
976     // Lower zero cost accesses to the kernel instances just created
977     for (auto &GV : KernelAccessVariables) {
978       auto &funcs = LDSToKernelsThatNeedToAccessItIndirectly[GV];
979       assert(funcs.size() == 1); // Only one kernel can access it
980       LDSVariableReplacement Replacement =
981           KernelToReplacement[*(funcs.begin())];
982 
983       DenseSet<GlobalVariable *> Vec;
984       Vec.insert(GV);
985 
986       // TODO: Looks like a latent bug, Replacement may not be marked
987       // UsedByKernel here
988       replaceLDSVariablesWithStruct(M, Vec, Replacement, [](Use &U) {
989         return isa<Instruction>(U.getUser());
990       });
991     }
992 
993     if (!KernelsThatAllocateTableLDS.empty()) {
994       LLVMContext &Ctx = M.getContext();
995       IRBuilder<> Builder(Ctx);
996 
997       // The ith element of this vector is kernel id i
998       std::vector<Function *> OrderedKernels =
999           assignLDSKernelIDToEachKernel(&M, KernelsThatAllocateTableLDS);
1000 
1001       for (size_t i = 0; i < OrderedKernels.size(); i++) {
1002         markUsedByKernel(Builder, OrderedKernels[i],
1003                          KernelToReplacement[OrderedKernels[i]].SGV);
1004       }
1005 
1006       // The order must be consistent between lookup table and accesses to
1007       // lookup table
1008       std::vector<GlobalVariable *> TableLookupVariablesOrdered(
1009           TableLookupVariables.begin(), TableLookupVariables.end());
1010       llvm::sort(TableLookupVariablesOrdered.begin(),
1011                  TableLookupVariablesOrdered.end(),
1012                  [](const GlobalVariable *lhs, const GlobalVariable *rhs) {
1013                    return lhs->getName() < rhs->getName();
1014                  });
1015 
1016       GlobalVariable *LookupTable = buildLookupTable(
1017           M, TableLookupVariablesOrdered, OrderedKernels, KernelToReplacement);
1018       replaceUsesInInstructionsWithTableLookup(M, TableLookupVariablesOrdered,
1019                                                LookupTable);
1020     }
1021 
1022     for (auto &GV : make_early_inc_range(M.globals()))
1023       if (AMDGPU::isLDSVariableToLower(GV)) {
1024         // probably want to remove from used lists
1025         GV.removeDeadConstantUsers();
1026         if (GV.use_empty())
1027           GV.eraseFromParent();
1028       }
1029 
1030     return Changed;
1031   }
1032 
1033 private:
1034   // Increase the alignment of LDS globals if necessary to maximise the chance
1035   // that we can use aligned LDS instructions to access them.
1036   static bool superAlignLDSGlobals(Module &M) {
1037     const DataLayout &DL = M.getDataLayout();
1038     bool Changed = false;
1039     if (!SuperAlignLDSGlobals) {
1040       return Changed;
1041     }
1042 
1043     for (auto &GV : M.globals()) {
1044       if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) {
1045         // Only changing alignment of LDS variables
1046         continue;
1047       }
1048       if (!GV.hasInitializer()) {
1049         // cuda/hip extern __shared__ variable, leave alignment alone
1050         continue;
1051       }
1052 
1053       Align Alignment = AMDGPU::getAlign(DL, &GV);
1054       TypeSize GVSize = DL.getTypeAllocSize(GV.getValueType());
1055 
1056       if (GVSize > 8) {
1057         // We might want to use a b96 or b128 load/store
1058         Alignment = std::max(Alignment, Align(16));
1059       } else if (GVSize > 4) {
1060         // We might want to use a b64 load/store
1061         Alignment = std::max(Alignment, Align(8));
1062       } else if (GVSize > 2) {
1063         // We might want to use a b32 load/store
1064         Alignment = std::max(Alignment, Align(4));
1065       } else if (GVSize > 1) {
1066         // We might want to use a b16 load/store
1067         Alignment = std::max(Alignment, Align(2));
1068       }
1069 
1070       if (Alignment != AMDGPU::getAlign(DL, &GV)) {
1071         Changed = true;
1072         GV.setAlignment(Alignment);
1073       }
1074     }
1075     return Changed;
1076   }
1077 
1078   static LDSVariableReplacement createLDSVariableReplacement(
1079       Module &M, std::string VarName,
1080       DenseSet<GlobalVariable *> const &LDSVarsToTransform) {
1081     // Create a struct instance containing LDSVarsToTransform and map from those
1082     // variables to ConstantExprGEP
1083     // Variables may be introduced to meet alignment requirements. No aliasing
1084     // metadata is useful for these as they have no uses. Erased before return.
1085 
1086     LLVMContext &Ctx = M.getContext();
1087     const DataLayout &DL = M.getDataLayout();
1088     assert(!LDSVarsToTransform.empty());
1089 
1090     SmallVector<OptimizedStructLayoutField, 8> LayoutFields;
1091     LayoutFields.reserve(LDSVarsToTransform.size());
1092     {
1093       // The order of fields in this struct depends on the order of
1094       // varables in the argument which varies when changing how they
1095       // are identified, leading to spurious test breakage.
1096       std::vector<GlobalVariable *> Sorted(LDSVarsToTransform.begin(),
1097                                            LDSVarsToTransform.end());
1098       llvm::sort(Sorted.begin(), Sorted.end(),
1099                  [](const GlobalVariable *lhs, const GlobalVariable *rhs) {
1100                    return lhs->getName() < rhs->getName();
1101                  });
1102       for (GlobalVariable *GV : Sorted) {
1103         OptimizedStructLayoutField F(GV,
1104                                      DL.getTypeAllocSize(GV->getValueType()),
1105                                      AMDGPU::getAlign(DL, GV));
1106         LayoutFields.emplace_back(F);
1107       }
1108     }
1109 
1110     performOptimizedStructLayout(LayoutFields);
1111 
1112     std::vector<GlobalVariable *> LocalVars;
1113     BitVector IsPaddingField;
1114     LocalVars.reserve(LDSVarsToTransform.size()); // will be at least this large
1115     IsPaddingField.reserve(LDSVarsToTransform.size());
1116     {
1117       uint64_t CurrentOffset = 0;
1118       for (size_t I = 0; I < LayoutFields.size(); I++) {
1119         GlobalVariable *FGV = static_cast<GlobalVariable *>(
1120             const_cast<void *>(LayoutFields[I].Id));
1121         Align DataAlign = LayoutFields[I].Alignment;
1122 
1123         uint64_t DataAlignV = DataAlign.value();
1124         if (uint64_t Rem = CurrentOffset % DataAlignV) {
1125           uint64_t Padding = DataAlignV - Rem;
1126 
1127           // Append an array of padding bytes to meet alignment requested
1128           // Note (o +      (a - (o % a)) ) % a == 0
1129           //      (offset + Padding       ) % align == 0
1130 
1131           Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding);
1132           LocalVars.push_back(new GlobalVariable(
1133               M, ATy, false, GlobalValue::InternalLinkage, UndefValue::get(ATy),
1134               "", nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
1135               false));
1136           IsPaddingField.push_back(true);
1137           CurrentOffset += Padding;
1138         }
1139 
1140         LocalVars.push_back(FGV);
1141         IsPaddingField.push_back(false);
1142         CurrentOffset += LayoutFields[I].Size;
1143       }
1144     }
1145 
1146     std::vector<Type *> LocalVarTypes;
1147     LocalVarTypes.reserve(LocalVars.size());
1148     std::transform(
1149         LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes),
1150         [](const GlobalVariable *V) -> Type * { return V->getValueType(); });
1151 
1152     StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t");
1153 
1154     Align StructAlign = AMDGPU::getAlign(DL, LocalVars[0]);
1155 
1156     GlobalVariable *SGV = new GlobalVariable(
1157         M, LDSTy, false, GlobalValue::InternalLinkage, UndefValue::get(LDSTy),
1158         VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
1159         false);
1160     SGV->setAlignment(StructAlign);
1161 
1162     DenseMap<GlobalVariable *, Constant *> Map;
1163     Type *I32 = Type::getInt32Ty(Ctx);
1164     for (size_t I = 0; I < LocalVars.size(); I++) {
1165       GlobalVariable *GV = LocalVars[I];
1166       Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)};
1167       Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx, true);
1168       if (IsPaddingField[I]) {
1169         assert(GV->use_empty());
1170         GV->eraseFromParent();
1171       } else {
1172         Map[GV] = GEP;
1173       }
1174     }
1175     assert(Map.size() == LDSVarsToTransform.size());
1176     return {SGV, std::move(Map)};
1177   }
1178 
1179   template <typename PredicateTy>
1180   static void replaceLDSVariablesWithStruct(
1181       Module &M, DenseSet<GlobalVariable *> const &LDSVarsToTransformArg,
1182       LDSVariableReplacement Replacement, PredicateTy Predicate) {
1183     LLVMContext &Ctx = M.getContext();
1184     const DataLayout &DL = M.getDataLayout();
1185 
1186     // A hack... we need to insert the aliasing info in a predictable order for
1187     // lit tests. Would like to have them in a stable order already, ideally the
1188     // same order they get allocated, which might mean an ordered set container
1189     std::vector<GlobalVariable *> LDSVarsToTransform(
1190         LDSVarsToTransformArg.begin(), LDSVarsToTransformArg.end());
1191     llvm::sort(LDSVarsToTransform.begin(), LDSVarsToTransform.end(),
1192                [](const GlobalVariable *lhs, const GlobalVariable *rhs) {
1193                  return lhs->getName() < rhs->getName();
1194                });
1195 
1196     // Create alias.scope and their lists. Each field in the new structure
1197     // does not alias with all other fields.
1198     SmallVector<MDNode *> AliasScopes;
1199     SmallVector<Metadata *> NoAliasList;
1200     const size_t NumberVars = LDSVarsToTransform.size();
1201     if (NumberVars > 1) {
1202       MDBuilder MDB(Ctx);
1203       AliasScopes.reserve(NumberVars);
1204       MDNode *Domain = MDB.createAnonymousAliasScopeDomain();
1205       for (size_t I = 0; I < NumberVars; I++) {
1206         MDNode *Scope = MDB.createAnonymousAliasScope(Domain);
1207         AliasScopes.push_back(Scope);
1208       }
1209       NoAliasList.append(&AliasScopes[1], AliasScopes.end());
1210     }
1211 
1212     // Replace uses of ith variable with a constantexpr to the corresponding
1213     // field of the instance that will be allocated by AMDGPUMachineFunction
1214     for (size_t I = 0; I < NumberVars; I++) {
1215       GlobalVariable *GV = LDSVarsToTransform[I];
1216       Constant *GEP = Replacement.LDSVarsToConstantGEP[GV];
1217 
1218       GV->replaceUsesWithIf(GEP, Predicate);
1219 
1220       APInt APOff(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
1221       GEP->stripAndAccumulateInBoundsConstantOffsets(DL, APOff);
1222       uint64_t Offset = APOff.getZExtValue();
1223 
1224       Align A =
1225           commonAlignment(Replacement.SGV->getAlign().valueOrOne(), Offset);
1226 
1227       if (I)
1228         NoAliasList[I - 1] = AliasScopes[I - 1];
1229       MDNode *NoAlias =
1230           NoAliasList.empty() ? nullptr : MDNode::get(Ctx, NoAliasList);
1231       MDNode *AliasScope =
1232           AliasScopes.empty() ? nullptr : MDNode::get(Ctx, {AliasScopes[I]});
1233 
1234       refineUsesAlignmentAndAA(GEP, A, DL, AliasScope, NoAlias);
1235     }
1236   }
1237 
1238   static void refineUsesAlignmentAndAA(Value *Ptr, Align A,
1239                                        const DataLayout &DL, MDNode *AliasScope,
1240                                        MDNode *NoAlias, unsigned MaxDepth = 5) {
1241     if (!MaxDepth || (A == 1 && !AliasScope))
1242       return;
1243 
1244     for (User *U : Ptr->users()) {
1245       if (auto *I = dyn_cast<Instruction>(U)) {
1246         if (AliasScope && I->mayReadOrWriteMemory()) {
1247           MDNode *AS = I->getMetadata(LLVMContext::MD_alias_scope);
1248           AS = (AS ? MDNode::getMostGenericAliasScope(AS, AliasScope)
1249                    : AliasScope);
1250           I->setMetadata(LLVMContext::MD_alias_scope, AS);
1251 
1252           MDNode *NA = I->getMetadata(LLVMContext::MD_noalias);
1253           NA = (NA ? MDNode::intersect(NA, NoAlias) : NoAlias);
1254           I->setMetadata(LLVMContext::MD_noalias, NA);
1255         }
1256       }
1257 
1258       if (auto *LI = dyn_cast<LoadInst>(U)) {
1259         LI->setAlignment(std::max(A, LI->getAlign()));
1260         continue;
1261       }
1262       if (auto *SI = dyn_cast<StoreInst>(U)) {
1263         if (SI->getPointerOperand() == Ptr)
1264           SI->setAlignment(std::max(A, SI->getAlign()));
1265         continue;
1266       }
1267       if (auto *AI = dyn_cast<AtomicRMWInst>(U)) {
1268         // None of atomicrmw operations can work on pointers, but let's
1269         // check it anyway in case it will or we will process ConstantExpr.
1270         if (AI->getPointerOperand() == Ptr)
1271           AI->setAlignment(std::max(A, AI->getAlign()));
1272         continue;
1273       }
1274       if (auto *AI = dyn_cast<AtomicCmpXchgInst>(U)) {
1275         if (AI->getPointerOperand() == Ptr)
1276           AI->setAlignment(std::max(A, AI->getAlign()));
1277         continue;
1278       }
1279       if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
1280         unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP->getType());
1281         APInt Off(BitWidth, 0);
1282         if (GEP->getPointerOperand() == Ptr) {
1283           Align GA;
1284           if (GEP->accumulateConstantOffset(DL, Off))
1285             GA = commonAlignment(A, Off.getLimitedValue());
1286           refineUsesAlignmentAndAA(GEP, GA, DL, AliasScope, NoAlias,
1287                                    MaxDepth - 1);
1288         }
1289         continue;
1290       }
1291       if (auto *I = dyn_cast<Instruction>(U)) {
1292         if (I->getOpcode() == Instruction::BitCast ||
1293             I->getOpcode() == Instruction::AddrSpaceCast)
1294           refineUsesAlignmentAndAA(I, A, DL, AliasScope, NoAlias, MaxDepth - 1);
1295       }
1296     }
1297   }
1298 };
1299 
1300 } // namespace
1301 char AMDGPULowerModuleLDS::ID = 0;
1302 
1303 char &llvm::AMDGPULowerModuleLDSID = AMDGPULowerModuleLDS::ID;
1304 
1305 INITIALIZE_PASS(AMDGPULowerModuleLDS, DEBUG_TYPE,
1306                 "Lower uses of LDS variables from non-kernel functions", false,
1307                 false)
1308 
1309 ModulePass *llvm::createAMDGPULowerModuleLDSPass() {
1310   return new AMDGPULowerModuleLDS();
1311 }
1312 
1313 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M,
1314                                                 ModuleAnalysisManager &) {
1315   return AMDGPULowerModuleLDS().runOnModule(M) ? PreservedAnalyses::none()
1316                                                : PreservedAnalyses::all();
1317 }
1318