| //===-- AMDGPULowerModuleLDSPass.cpp ------------------------------*- C++ -*-=// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This pass eliminates local data store, LDS, uses from non-kernel functions. |
| // LDS is contiguous memory allocated per kernel execution. |
| // |
| // Background. |
| // |
| // The programming model is global variables, or equivalently function local |
| // static variables, accessible from kernels or other functions. For uses from |
| // kernels this is straightforward - assign an integer to the kernel for the |
| // memory required by all the variables combined, allocate them within that. |
| // For uses from functions there are performance tradeoffs to choose between. |
| // |
| // This model means the GPU runtime can specify the amount of memory allocated. |
| // If this is more than the kernel assumed, the excess can be made available |
| // using a language specific feature, which IR represents as a variable with |
| // no initializer. This feature is referred to here as "Dynamic LDS" and is |
| // lowered slightly differently to the normal case. |
| // |
| // Consequences of this GPU feature: |
| // - memory is limited and exceeding it halts compilation |
| // - a global accessed by one kernel exists independent of other kernels |
| // - a global exists independent of simultaneous execution of the same kernel |
| // - the address of the global may be different from different kernels as they |
| // do not alias, which permits only allocating variables they use |
| // - if the address is allowed to differ, functions need help to find it |
| // |
| // Uses from kernels are implemented here by grouping them in a per-kernel |
| // struct instance. This duplicates the variables, accurately modelling their |
| // aliasing properties relative to a single global representation. It also |
| // permits control over alignment via padding. |
| // |
| // Uses from functions are more complicated and the primary purpose of this |
| // IR pass. Several different lowering are chosen between to meet requirements |
| // to avoid allocating any LDS where it is not necessary, as that impacts |
| // occupancy and may fail the compilation, while not imposing overhead on a |
| // feature whose primary advantage over global memory is performance. The basic |
| // design goal is to avoid one kernel imposing overhead on another. |
| // |
| // Implementation. |
| // |
| // LDS variables with constant annotation or non-undef initializer are passed |
| // through unchanged for simplification or error diagnostics in later passes. |
| // Non-undef initializers are not yet implemented for LDS. |
| // |
| // LDS variables that are always allocated at the same address can be found |
| // by lookup at that address. Otherwise runtime information/cost is required. |
| // |
| // The simplest strategy possible is to group all LDS variables in a single |
| // struct and allocate that struct in every kernel such that the original |
| // variables are always at the same address. LDS is however a limited resource |
| // so this strategy is unusable in practice. It is not implemented here. |
| // |
| // Strategy | Precise allocation | Zero runtime cost | General purpose | |
| // --------+--------------------+-------------------+-----------------+ |
| // Module | No | Yes | Yes | |
| // Table | Yes | No | Yes | |
| // Kernel | Yes | Yes | No | |
| // Hybrid | Yes | Partial | Yes | |
| // |
| // "Module" spends LDS memory to save cycles. "Table" spends cycles and global |
| // memory to save LDS. "Kernel" is as fast as kernel allocation but only works |
| // for variables that are known reachable from a single kernel. "Hybrid" picks |
| // between all three. When forced to choose between LDS and cycles we minimise |
| // LDS use. |
| |
| // The "module" lowering implemented here finds LDS variables which are used by |
| // non-kernel functions and creates a new struct with a field for each of those |
| // LDS variables. Variables that are only used from kernels are excluded. |
| // |
| // The "table" lowering implemented here has three components. |
| // First kernels are assigned a unique integer identifier which is available in |
| // functions it calls through the intrinsic amdgcn_lds_kernel_id. The integer |
| // is passed through a specific SGPR, thus works with indirect calls. |
| // Second, each kernel allocates LDS variables independent of other kernels and |
| // writes the addresses it chose for each variable into an array in consistent |
| // order. If the kernel does not allocate a given variable, it writes undef to |
| // the corresponding array location. These arrays are written to a constant |
| // table in the order matching the kernel unique integer identifier. |
| // Third, uses from non-kernel functions are replaced with a table lookup using |
| // the intrinsic function to find the address of the variable. |
| // |
| // "Kernel" lowering is only applicable for variables that are unambiguously |
| // reachable from exactly one kernel. For those cases, accesses to the variable |
| // can be lowered to ConstantExpr address of a struct instance specific to that |
| // one kernel. This is zero cost in space and in compute. It will raise a fatal |
| // error on any variable that might be reachable from multiple kernels and is |
| // thus most easily used as part of the hybrid lowering strategy. |
| // |
| // Hybrid lowering is a mixture of the above. It uses the zero cost kernel |
| // lowering where it can. It lowers the variable accessed by the greatest |
| // number of kernels using the module strategy as that is free for the first |
| // variable. Any futher variables that can be lowered with the module strategy |
| // without incurring LDS memory overhead are. The remaining ones are lowered |
| // via table. |
| // |
| // Consequences |
| // - No heuristics or user controlled magic numbers, hybrid is the right choice |
| // - Kernels that don't use functions (or have had them all inlined) are not |
| // affected by any lowering for kernels that do. |
| // - Kernels that don't make indirect function calls are not affected by those |
| // that do. |
| // - Variables which are used by lots of kernels, e.g. those injected by a |
| // language runtime in most kernels, are expected to have no overhead |
| // - Implementations that instantiate templates per-kernel where those templates |
| // use LDS are expected to hit the "Kernel" lowering strategy |
| // - The runtime properties impose a cost in compiler implementation complexity |
| // |
| // Dynamic LDS implementation |
| // Dynamic LDS is lowered similarly to the "table" strategy above and uses the |
| // same intrinsic to identify which kernel is at the root of the dynamic call |
| // graph. This relies on the specified behaviour that all dynamic LDS variables |
| // alias one another, i.e. are at the same address, with respect to a given |
| // kernel. Therefore this pass creates new dynamic LDS variables for each kernel |
| // that allocates any dynamic LDS and builds a table of addresses out of those. |
| // The AMDGPUPromoteAlloca pass skips kernels that use dynamic LDS. |
| // The corresponding optimisation for "kernel" lowering where the table lookup |
| // is elided is not implemented. |
| // |
| // |
| // Implementation notes / limitations |
| // A single LDS global variable represents an instance per kernel that can reach |
| // said variables. This pass essentially specialises said variables per kernel. |
| // Handling ConstantExpr during the pass complicated this significantly so now |
| // all ConstantExpr uses of LDS variables are expanded to instructions. This |
| // may need amending when implementing non-undef initialisers. |
| // |
| // Lowering is split between this IR pass and the back end. This pass chooses |
| // where given variables should be allocated and marks them with metadata, |
| // MD_absolute_symbol. The backend places the variables in coincidentally the |
| // same location and raises a fatal error if something has gone awry. This works |
| // in practice because the only pass between this one and the backend that |
| // changes LDS is PromoteAlloca and the changes it makes do not conflict. |
| // |
| // Addresses are written to constant global arrays based on the same metadata. |
| // |
| // The backend lowers LDS variables in the order of traversal of the function. |
| // This is at odds with the deterministic layout required. The workaround is to |
| // allocate the fixed-address variables immediately upon starting the function |
| // where they can be placed as intended. This requires a means of mapping from |
| // the function to the variables that it allocates. For the module scope lds, |
| // this is via metadata indicating whether the variable is not required. If a |
| // pass deletes that metadata, a fatal error on disagreement with the absolute |
| // symbol metadata will occur. For kernel scope and dynamic, this is by _name_ |
| // correspondence between the function and the variable. It requires the |
| // kernel to have a name (which is only a limitation for tests in practice) and |
| // for nothing to rename the corresponding symbols. This is a hazard if the pass |
| // is run multiple times during debugging. Alternative schemes considered all |
| // involve bespoke metadata. |
| // |
| // If the name correspondence can be replaced, multiple distinct kernels that |
| // have the same memory layout can map to the same kernel id (as the address |
| // itself is handled by the absolute symbol metadata) and that will allow more |
| // uses of the "kernel" style faster lowering and reduce the size of the lookup |
| // tables. |
| // |
| // There is a test that checks this does not fire for a graphics shader. This |
| // lowering is expected to work for graphics if the isKernel test is changed. |
| // |
| // The current markUsedByKernel is sufficient for PromoteAlloca but is elided |
| // before codegen. Replacing this with an equivalent intrinsic which lasts until |
| // shortly after the machine function lowering of LDS would help break the name |
| // mapping. The other part needed is probably to amend PromoteAlloca to embed |
| // the LDS variables it creates in the same struct created here. That avoids the |
| // current hazard where a PromoteAlloca LDS variable might be allocated before |
| // the kernel scope (and thus error on the address check). Given a new invariant |
| // that no LDS variables exist outside of the structs managed here, and an |
| // intrinsic that lasts until after the LDS frame lowering, it should be |
| // possible to drop the name mapping and fold equivalent memory layouts. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "AMDGPU.h" |
| #include "AMDGPUTargetMachine.h" |
| #include "Utils/AMDGPUBaseInfo.h" |
| #include "Utils/AMDGPUMemoryUtils.h" |
| #include "llvm/ADT/BitVector.h" |
| #include "llvm/ADT/DenseMap.h" |
| #include "llvm/ADT/DenseSet.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "llvm/ADT/SetOperations.h" |
| #include "llvm/Analysis/CallGraph.h" |
| #include "llvm/CodeGen/TargetPassConfig.h" |
| #include "llvm/IR/Constants.h" |
| #include "llvm/IR/DerivedTypes.h" |
| #include "llvm/IR/IRBuilder.h" |
| #include "llvm/IR/InlineAsm.h" |
| #include "llvm/IR/Instructions.h" |
| #include "llvm/IR/IntrinsicsAMDGPU.h" |
| #include "llvm/IR/MDBuilder.h" |
| #include "llvm/IR/ReplaceConstant.h" |
| #include "llvm/InitializePasses.h" |
| #include "llvm/Pass.h" |
| #include "llvm/Support/CommandLine.h" |
| #include "llvm/Support/Debug.h" |
| #include "llvm/Support/Format.h" |
| #include "llvm/Support/OptimizedStructLayout.h" |
| #include "llvm/Support/raw_ostream.h" |
| #include "llvm/Transforms/Utils/BasicBlockUtils.h" |
| #include "llvm/Transforms/Utils/ModuleUtils.h" |
| |
| #include <vector> |
| |
| #include <cstdio> |
| |
| #define DEBUG_TYPE "amdgpu-lower-module-lds" |
| |
| using namespace llvm; |
| |
| namespace { |
| |
| cl::opt<bool> SuperAlignLDSGlobals( |
| "amdgpu-super-align-lds-globals", |
| cl::desc("Increase alignment of LDS if it is not on align boundary"), |
| cl::init(true), cl::Hidden); |
| |
| enum class LoweringKind { module, table, kernel, hybrid }; |
| cl::opt<LoweringKind> LoweringKindLoc( |
| "amdgpu-lower-module-lds-strategy", |
| cl::desc("Specify lowering strategy for function LDS access:"), cl::Hidden, |
| cl::init(LoweringKind::hybrid), |
| cl::values( |
| clEnumValN(LoweringKind::table, "table", "Lower via table lookup"), |
| clEnumValN(LoweringKind::module, "module", "Lower via module struct"), |
| clEnumValN( |
| LoweringKind::kernel, "kernel", |
| "Lower variables reachable from one kernel, otherwise abort"), |
| clEnumValN(LoweringKind::hybrid, "hybrid", |
| "Lower via mixture of above strategies"))); |
| |
| bool isKernelLDS(const Function *F) { |
| // Some weirdness here. AMDGPU::isKernelCC does not call into |
| // AMDGPU::isKernel with the calling conv, it instead calls into |
| // isModuleEntryFunction which returns true for more calling conventions |
| // than AMDGPU::isKernel does. There's a FIXME on AMDGPU::isKernel. |
| // There's also a test that checks that the LDS lowering does not hit on |
| // a graphics shader, denoted amdgpu_ps, so stay with the limited case. |
| // Putting LDS in the name of the function to draw attention to this. |
| return AMDGPU::isKernel(F->getCallingConv()); |
| } |
| |
| template <typename T> std::vector<T> sortByName(std::vector<T> &&V) { |
| llvm::sort(V.begin(), V.end(), [](const auto *L, const auto *R) { |
| return L->getName() < R->getName(); |
| }); |
| return {std::move(V)}; |
| } |
| |
| class AMDGPULowerModuleLDS { |
| const AMDGPUTargetMachine &TM; |
| |
| static void |
| removeLocalVarsFromUsedLists(Module &M, |
| const DenseSet<GlobalVariable *> &LocalVars) { |
| // The verifier rejects used lists containing an inttoptr of a constant |
| // so remove the variables from these lists before replaceAllUsesWith |
| SmallPtrSet<Constant *, 8> LocalVarsSet; |
| for (GlobalVariable *LocalVar : LocalVars) |
| LocalVarsSet.insert(cast<Constant>(LocalVar->stripPointerCasts())); |
| |
| removeFromUsedLists( |
| M, [&LocalVarsSet](Constant *C) { return LocalVarsSet.count(C); }); |
| |
| for (GlobalVariable *LocalVar : LocalVars) |
| LocalVar->removeDeadConstantUsers(); |
| } |
| |
| static void markUsedByKernel(Function *Func, GlobalVariable *SGV) { |
| // The llvm.amdgcn.module.lds instance is implicitly used by all kernels |
| // that might call a function which accesses a field within it. This is |
| // presently approximated to 'all kernels' if there are any such functions |
| // in the module. This implicit use is redefined as an explicit use here so |
| // that later passes, specifically PromoteAlloca, account for the required |
| // memory without any knowledge of this transform. |
| |
| // An operand bundle on llvm.donothing works because the call instruction |
| // survives until after the last pass that needs to account for LDS. It is |
| // better than inline asm as the latter survives until the end of codegen. A |
| // totally robust solution would be a function with the same semantics as |
| // llvm.donothing that takes a pointer to the instance and is lowered to a |
| // no-op after LDS is allocated, but that is not presently necessary. |
| |
| // This intrinsic is eliminated shortly before instruction selection. It |
| // does not suffice to indicate to ISel that a given global which is not |
| // immediately used by the kernel must still be allocated by it. An |
| // equivalent target specific intrinsic which lasts until immediately after |
| // codegen would suffice for that, but one would still need to ensure that |
| // the variables are allocated in the anticpated order. |
| BasicBlock *Entry = &Func->getEntryBlock(); |
| IRBuilder<> Builder(Entry, Entry->getFirstNonPHIIt()); |
| |
| Function *Decl = |
| Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {}); |
| |
| Value *UseInstance[1] = { |
| Builder.CreateConstInBoundsGEP1_32(SGV->getValueType(), SGV, 0)}; |
| |
| Builder.CreateCall( |
| Decl, {}, {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)}); |
| } |
| |
| static bool eliminateConstantExprUsesOfLDSFromAllInstructions(Module &M) { |
| // Constants are uniqued within LLVM. A ConstantExpr referring to a LDS |
| // global may have uses from multiple different functions as a result. |
| // This pass specialises LDS variables with respect to the kernel that |
| // allocates them. |
| |
| // This is semantically equivalent to (the unimplemented as slow): |
| // for (auto &F : M.functions()) |
| // for (auto &BB : F) |
| // for (auto &I : BB) |
| // for (Use &Op : I.operands()) |
| // if (constantExprUsesLDS(Op)) |
| // replaceConstantExprInFunction(I, Op); |
| |
| SmallVector<Constant *> LDSGlobals; |
| for (auto &GV : M.globals()) |
| if (AMDGPU::isLDSVariableToLower(GV)) |
| LDSGlobals.push_back(&GV); |
| |
| return convertUsersOfConstantsToInstructions(LDSGlobals); |
| } |
| |
| public: |
| AMDGPULowerModuleLDS(const AMDGPUTargetMachine &TM_) : TM(TM_) {} |
| |
| using FunctionVariableMap = DenseMap<Function *, DenseSet<GlobalVariable *>>; |
| |
| using VariableFunctionMap = DenseMap<GlobalVariable *, DenseSet<Function *>>; |
| |
| static void getUsesOfLDSByFunction(CallGraph const &CG, Module &M, |
| FunctionVariableMap &kernels, |
| FunctionVariableMap &functions) { |
| |
| // Get uses from the current function, excluding uses by called functions |
| // Two output variables to avoid walking the globals list twice |
| for (auto &GV : M.globals()) { |
| if (!AMDGPU::isLDSVariableToLower(GV)) { |
| continue; |
| } |
| |
| for (User *V : GV.users()) { |
| if (auto *I = dyn_cast<Instruction>(V)) { |
| Function *F = I->getFunction(); |
| if (isKernelLDS(F)) { |
| kernels[F].insert(&GV); |
| } else { |
| functions[F].insert(&GV); |
| } |
| } |
| } |
| } |
| } |
| |
| struct LDSUsesInfoTy { |
| FunctionVariableMap direct_access; |
| FunctionVariableMap indirect_access; |
| }; |
| |
| static LDSUsesInfoTy getTransitiveUsesOfLDS(CallGraph const &CG, Module &M) { |
| |
| FunctionVariableMap direct_map_kernel; |
| FunctionVariableMap direct_map_function; |
| getUsesOfLDSByFunction(CG, M, direct_map_kernel, direct_map_function); |
| |
| // Collect variables that are used by functions whose address has escaped |
| DenseSet<GlobalVariable *> VariablesReachableThroughFunctionPointer; |
| for (Function &F : M.functions()) { |
| if (!isKernelLDS(&F)) |
| if (F.hasAddressTaken(nullptr, |
| /* IgnoreCallbackUses */ false, |
| /* IgnoreAssumeLikeCalls */ false, |
| /* IgnoreLLVMUsed */ true, |
| /* IgnoreArcAttachedCall */ false)) { |
| set_union(VariablesReachableThroughFunctionPointer, |
| direct_map_function[&F]); |
| } |
| } |
| |
| auto functionMakesUnknownCall = [&](const Function *F) -> bool { |
| assert(!F->isDeclaration()); |
| for (const CallGraphNode::CallRecord &R : *CG[F]) { |
| if (!R.second->getFunction()) { |
| return true; |
| } |
| } |
| return false; |
| }; |
| |
| // Work out which variables are reachable through function calls |
| FunctionVariableMap transitive_map_function = direct_map_function; |
| |
| // If the function makes any unknown call, assume the worst case that it can |
| // access all variables accessed by functions whose address escaped |
| for (Function &F : M.functions()) { |
| if (!F.isDeclaration() && functionMakesUnknownCall(&F)) { |
| if (!isKernelLDS(&F)) { |
| set_union(transitive_map_function[&F], |
| VariablesReachableThroughFunctionPointer); |
| } |
| } |
| } |
| |
| // Direct implementation of collecting all variables reachable from each |
| // function |
| for (Function &Func : M.functions()) { |
| if (Func.isDeclaration() || isKernelLDS(&Func)) |
| continue; |
| |
| DenseSet<Function *> seen; // catches cycles |
| SmallVector<Function *, 4> wip{&Func}; |
| |
| while (!wip.empty()) { |
| Function *F = wip.pop_back_val(); |
| |
| // Can accelerate this by referring to transitive map for functions that |
| // have already been computed, with more care than this |
| set_union(transitive_map_function[&Func], direct_map_function[F]); |
| |
| for (const CallGraphNode::CallRecord &R : *CG[F]) { |
| Function *ith = R.second->getFunction(); |
| if (ith) { |
| if (!seen.contains(ith)) { |
| seen.insert(ith); |
| wip.push_back(ith); |
| } |
| } |
| } |
| } |
| } |
| |
| // direct_map_kernel lists which variables are used by the kernel |
| // find the variables which are used through a function call |
| FunctionVariableMap indirect_map_kernel; |
| |
| for (Function &Func : M.functions()) { |
| if (Func.isDeclaration() || !isKernelLDS(&Func)) |
| continue; |
| |
| for (const CallGraphNode::CallRecord &R : *CG[&Func]) { |
| Function *ith = R.second->getFunction(); |
| if (ith) { |
| set_union(indirect_map_kernel[&Func], transitive_map_function[ith]); |
| } else { |
| set_union(indirect_map_kernel[&Func], |
| VariablesReachableThroughFunctionPointer); |
| } |
| } |
| } |
| |
| // Verify that we fall into one of 2 cases: |
| // - All variables are absolute: this is a re-run of the pass |
| // so we don't have anything to do. |
| // - No variables are absolute. |
| std::optional<bool> HasAbsoluteGVs; |
| for (auto &Map : {direct_map_kernel, indirect_map_kernel}) { |
| for (auto &[Fn, GVs] : Map) { |
| for (auto *GV : GVs) { |
| bool IsAbsolute = GV->isAbsoluteSymbolRef(); |
| if (HasAbsoluteGVs.has_value()) { |
| if (*HasAbsoluteGVs != IsAbsolute) { |
| report_fatal_error( |
| "Module cannot mix absolute and non-absolute LDS GVs"); |
| } |
| } else |
| HasAbsoluteGVs = IsAbsolute; |
| } |
| } |
| } |
| |
| // If we only had absolute GVs, we have nothing to do, return an empty |
| // result. |
| if (HasAbsoluteGVs && *HasAbsoluteGVs) |
| return {FunctionVariableMap(), FunctionVariableMap()}; |
| |
| return {std::move(direct_map_kernel), std::move(indirect_map_kernel)}; |
| } |
| |
| struct LDSVariableReplacement { |
| GlobalVariable *SGV = nullptr; |
| DenseMap<GlobalVariable *, Constant *> LDSVarsToConstantGEP; |
| }; |
| |
| // remap from lds global to a constantexpr gep to where it has been moved to |
| // for each kernel |
| // an array with an element for each kernel containing where the corresponding |
| // variable was remapped to |
| |
| static Constant *getAddressesOfVariablesInKernel( |
| LLVMContext &Ctx, ArrayRef<GlobalVariable *> Variables, |
| const DenseMap<GlobalVariable *, Constant *> &LDSVarsToConstantGEP) { |
| // Create a ConstantArray containing the address of each Variable within the |
| // kernel corresponding to LDSVarsToConstantGEP, or poison if that kernel |
| // does not allocate it |
| // TODO: Drop the ptrtoint conversion |
| |
| Type *I32 = Type::getInt32Ty(Ctx); |
| |
| ArrayType *KernelOffsetsType = ArrayType::get(I32, Variables.size()); |
| |
| SmallVector<Constant *> Elements; |
| for (size_t i = 0; i < Variables.size(); i++) { |
| GlobalVariable *GV = Variables[i]; |
| auto ConstantGepIt = LDSVarsToConstantGEP.find(GV); |
| if (ConstantGepIt != LDSVarsToConstantGEP.end()) { |
| auto elt = ConstantExpr::getPtrToInt(ConstantGepIt->second, I32); |
| Elements.push_back(elt); |
| } else { |
| Elements.push_back(PoisonValue::get(I32)); |
| } |
| } |
| return ConstantArray::get(KernelOffsetsType, Elements); |
| } |
| |
| static GlobalVariable *buildLookupTable( |
| Module &M, ArrayRef<GlobalVariable *> Variables, |
| ArrayRef<Function *> kernels, |
| DenseMap<Function *, LDSVariableReplacement> &KernelToReplacement) { |
| if (Variables.empty()) { |
| return nullptr; |
| } |
| LLVMContext &Ctx = M.getContext(); |
| |
| const size_t NumberVariables = Variables.size(); |
| const size_t NumberKernels = kernels.size(); |
| |
| ArrayType *KernelOffsetsType = |
| ArrayType::get(Type::getInt32Ty(Ctx), NumberVariables); |
| |
| ArrayType *AllKernelsOffsetsType = |
| ArrayType::get(KernelOffsetsType, NumberKernels); |
| |
| Constant *Missing = PoisonValue::get(KernelOffsetsType); |
| std::vector<Constant *> overallConstantExprElts(NumberKernels); |
| for (size_t i = 0; i < NumberKernels; i++) { |
| auto Replacement = KernelToReplacement.find(kernels[i]); |
| overallConstantExprElts[i] = |
| (Replacement == KernelToReplacement.end()) |
| ? Missing |
| : getAddressesOfVariablesInKernel( |
| Ctx, Variables, Replacement->second.LDSVarsToConstantGEP); |
| } |
| |
| Constant *init = |
| ConstantArray::get(AllKernelsOffsetsType, overallConstantExprElts); |
| |
| return new GlobalVariable( |
| M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init, |
| "llvm.amdgcn.lds.offset.table", nullptr, GlobalValue::NotThreadLocal, |
| AMDGPUAS::CONSTANT_ADDRESS); |
| } |
| |
| void replaceUseWithTableLookup(Module &M, IRBuilder<> &Builder, |
| GlobalVariable *LookupTable, |
| GlobalVariable *GV, Use &U, |
| Value *OptionalIndex) { |
| // Table is a constant array of the same length as OrderedKernels |
| LLVMContext &Ctx = M.getContext(); |
| Type *I32 = Type::getInt32Ty(Ctx); |
| auto *I = cast<Instruction>(U.getUser()); |
| |
| Value *tableKernelIndex = getTableLookupKernelIndex(M, I->getFunction()); |
| |
| if (auto *Phi = dyn_cast<PHINode>(I)) { |
| BasicBlock *BB = Phi->getIncomingBlock(U); |
| Builder.SetInsertPoint(&(*(BB->getFirstInsertionPt()))); |
| } else { |
| Builder.SetInsertPoint(I); |
| } |
| |
| SmallVector<Value *, 3> GEPIdx = { |
| ConstantInt::get(I32, 0), |
| tableKernelIndex, |
| }; |
| if (OptionalIndex) |
| GEPIdx.push_back(OptionalIndex); |
| |
| Value *Address = Builder.CreateInBoundsGEP( |
| LookupTable->getValueType(), LookupTable, GEPIdx, GV->getName()); |
| |
| Value *loaded = Builder.CreateLoad(I32, Address); |
| |
| Value *replacement = |
| Builder.CreateIntToPtr(loaded, GV->getType(), GV->getName()); |
| |
| U.set(replacement); |
| } |
| |
| void replaceUsesInInstructionsWithTableLookup( |
| Module &M, ArrayRef<GlobalVariable *> ModuleScopeVariables, |
| GlobalVariable *LookupTable) { |
| |
| LLVMContext &Ctx = M.getContext(); |
| IRBuilder<> Builder(Ctx); |
| Type *I32 = Type::getInt32Ty(Ctx); |
| |
| for (size_t Index = 0; Index < ModuleScopeVariables.size(); Index++) { |
| auto *GV = ModuleScopeVariables[Index]; |
| |
| for (Use &U : make_early_inc_range(GV->uses())) { |
| auto *I = dyn_cast<Instruction>(U.getUser()); |
| if (!I) |
| continue; |
| |
| replaceUseWithTableLookup(M, Builder, LookupTable, GV, U, |
| ConstantInt::get(I32, Index)); |
| } |
| } |
| } |
| |
| static DenseSet<Function *> kernelsThatIndirectlyAccessAnyOfPassedVariables( |
| Module &M, LDSUsesInfoTy &LDSUsesInfo, |
| DenseSet<GlobalVariable *> const &VariableSet) { |
| |
| DenseSet<Function *> KernelSet; |
| |
| if (VariableSet.empty()) |
| return KernelSet; |
| |
| for (Function &Func : M.functions()) { |
| if (Func.isDeclaration() || !isKernelLDS(&Func)) |
| continue; |
| for (GlobalVariable *GV : LDSUsesInfo.indirect_access[&Func]) { |
| if (VariableSet.contains(GV)) { |
| KernelSet.insert(&Func); |
| break; |
| } |
| } |
| } |
| |
| return KernelSet; |
| } |
| |
| static GlobalVariable * |
| chooseBestVariableForModuleStrategy(const DataLayout &DL, |
| VariableFunctionMap &LDSVars) { |
| // Find the global variable with the most indirect uses from kernels |
| |
| struct CandidateTy { |
| GlobalVariable *GV = nullptr; |
| size_t UserCount = 0; |
| size_t Size = 0; |
| |
| CandidateTy() = default; |
| |
| CandidateTy(GlobalVariable *GV, uint64_t UserCount, uint64_t AllocSize) |
| : GV(GV), UserCount(UserCount), Size(AllocSize) {} |
| |
| bool operator<(const CandidateTy &Other) const { |
| // Fewer users makes module scope variable less attractive |
| if (UserCount < Other.UserCount) { |
| return true; |
| } |
| if (UserCount > Other.UserCount) { |
| return false; |
| } |
| |
| // Bigger makes module scope variable less attractive |
| if (Size < Other.Size) { |
| return false; |
| } |
| |
| if (Size > Other.Size) { |
| return true; |
| } |
| |
| // Arbitrary but consistent |
| return GV->getName() < Other.GV->getName(); |
| } |
| }; |
| |
| CandidateTy MostUsed; |
| |
| for (auto &K : LDSVars) { |
| GlobalVariable *GV = K.first; |
| if (K.second.size() <= 1) { |
| // A variable reachable by only one kernel is best lowered with kernel |
| // strategy |
| continue; |
| } |
| CandidateTy Candidate( |
| GV, K.second.size(), |
| DL.getTypeAllocSize(GV->getValueType()).getFixedValue()); |
| if (MostUsed < Candidate) |
| MostUsed = Candidate; |
| } |
| |
| return MostUsed.GV; |
| } |
| |
| static void recordLDSAbsoluteAddress(Module *M, GlobalVariable *GV, |
| uint32_t Address) { |
| // Write the specified address into metadata where it can be retrieved by |
| // the assembler. Format is a half open range, [Address Address+1) |
| LLVMContext &Ctx = M->getContext(); |
| auto *IntTy = |
| M->getDataLayout().getIntPtrType(Ctx, AMDGPUAS::LOCAL_ADDRESS); |
| auto *MinC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address)); |
| auto *MaxC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address + 1)); |
| GV->setMetadata(LLVMContext::MD_absolute_symbol, |
| MDNode::get(Ctx, {MinC, MaxC})); |
| } |
| |
| DenseMap<Function *, Value *> tableKernelIndexCache; |
| Value *getTableLookupKernelIndex(Module &M, Function *F) { |
| // Accesses from a function use the amdgcn_lds_kernel_id intrinsic which |
| // lowers to a read from a live in register. Emit it once in the entry |
| // block to spare deduplicating it later. |
| auto [It, Inserted] = tableKernelIndexCache.try_emplace(F); |
| if (Inserted) { |
| Function *Decl = |
| Intrinsic::getDeclaration(&M, Intrinsic::amdgcn_lds_kernel_id, {}); |
| |
| auto InsertAt = F->getEntryBlock().getFirstNonPHIOrDbgOrAlloca(); |
| IRBuilder<> Builder(&*InsertAt); |
| |
| It->second = Builder.CreateCall(Decl, {}); |
| } |
| |
| return It->second; |
| } |
| |
| static std::vector<Function *> assignLDSKernelIDToEachKernel( |
| Module *M, DenseSet<Function *> const &KernelsThatAllocateTableLDS, |
| DenseSet<Function *> const &KernelsThatIndirectlyAllocateDynamicLDS) { |
| // Associate kernels in the set with an arbirary but reproducible order and |
| // annotate them with that order in metadata. This metadata is recognised by |
| // the backend and lowered to a SGPR which can be read from using |
| // amdgcn_lds_kernel_id. |
| |
| std::vector<Function *> OrderedKernels; |
| if (!KernelsThatAllocateTableLDS.empty() || |
| !KernelsThatIndirectlyAllocateDynamicLDS.empty()) { |
| |
| for (Function &Func : M->functions()) { |
| if (Func.isDeclaration()) |
| continue; |
| if (!isKernelLDS(&Func)) |
| continue; |
| |
| if (KernelsThatAllocateTableLDS.contains(&Func) || |
| KernelsThatIndirectlyAllocateDynamicLDS.contains(&Func)) { |
| assert(Func.hasName()); // else fatal error earlier |
| OrderedKernels.push_back(&Func); |
| } |
| } |
| |
| // Put them in an arbitrary but reproducible order |
| OrderedKernels = sortByName(std::move(OrderedKernels)); |
| |
| // Annotate the kernels with their order in this vector |
| LLVMContext &Ctx = M->getContext(); |
| IRBuilder<> Builder(Ctx); |
| |
| if (OrderedKernels.size() > UINT32_MAX) { |
| // 32 bit keeps it in one SGPR. > 2**32 kernels won't fit on the GPU |
| report_fatal_error("Unimplemented LDS lowering for > 2**32 kernels"); |
| } |
| |
| for (size_t i = 0; i < OrderedKernels.size(); i++) { |
| Metadata *AttrMDArgs[1] = { |
| ConstantAsMetadata::get(Builder.getInt32(i)), |
| }; |
| OrderedKernels[i]->setMetadata("llvm.amdgcn.lds.kernel.id", |
| MDNode::get(Ctx, AttrMDArgs)); |
| } |
| } |
| return OrderedKernels; |
| } |
| |
| static void partitionVariablesIntoIndirectStrategies( |
| Module &M, LDSUsesInfoTy const &LDSUsesInfo, |
| VariableFunctionMap &LDSToKernelsThatNeedToAccessItIndirectly, |
| DenseSet<GlobalVariable *> &ModuleScopeVariables, |
| DenseSet<GlobalVariable *> &TableLookupVariables, |
| DenseSet<GlobalVariable *> &KernelAccessVariables, |
| DenseSet<GlobalVariable *> &DynamicVariables) { |
| |
| GlobalVariable *HybridModuleRoot = |
| LoweringKindLoc != LoweringKind::hybrid |
| ? nullptr |
| : chooseBestVariableForModuleStrategy( |
| M.getDataLayout(), LDSToKernelsThatNeedToAccessItIndirectly); |
| |
| DenseSet<Function *> const EmptySet; |
| DenseSet<Function *> const &HybridModuleRootKernels = |
| HybridModuleRoot |
| ? LDSToKernelsThatNeedToAccessItIndirectly[HybridModuleRoot] |
| : EmptySet; |
| |
| for (auto &K : LDSToKernelsThatNeedToAccessItIndirectly) { |
| // Each iteration of this loop assigns exactly one global variable to |
| // exactly one of the implementation strategies. |
| |
| GlobalVariable *GV = K.first; |
| assert(AMDGPU::isLDSVariableToLower(*GV)); |
| assert(K.second.size() != 0); |
| |
| if (AMDGPU::isDynamicLDS(*GV)) { |
| DynamicVariables.insert(GV); |
| continue; |
| } |
| |
| switch (LoweringKindLoc) { |
| case LoweringKind::module: |
| ModuleScopeVariables.insert(GV); |
| break; |
| |
| case LoweringKind::table: |
| TableLookupVariables.insert(GV); |
| break; |
| |
| case LoweringKind::kernel: |
| if (K.second.size() == 1) { |
| KernelAccessVariables.insert(GV); |
| } else { |
| report_fatal_error( |
| "cannot lower LDS '" + GV->getName() + |
| "' to kernel access as it is reachable from multiple kernels"); |
| } |
| break; |
| |
| case LoweringKind::hybrid: { |
| if (GV == HybridModuleRoot) { |
| assert(K.second.size() != 1); |
| ModuleScopeVariables.insert(GV); |
| } else if (K.second.size() == 1) { |
| KernelAccessVariables.insert(GV); |
| } else if (set_is_subset(K.second, HybridModuleRootKernels)) { |
| ModuleScopeVariables.insert(GV); |
| } else { |
| TableLookupVariables.insert(GV); |
| } |
| break; |
| } |
| } |
| } |
| |
| // All LDS variables accessed indirectly have now been partitioned into |
| // the distinct lowering strategies. |
| assert(ModuleScopeVariables.size() + TableLookupVariables.size() + |
| KernelAccessVariables.size() + DynamicVariables.size() == |
| LDSToKernelsThatNeedToAccessItIndirectly.size()); |
| } |
| |
| static GlobalVariable *lowerModuleScopeStructVariables( |
| Module &M, DenseSet<GlobalVariable *> const &ModuleScopeVariables, |
| DenseSet<Function *> const &KernelsThatAllocateModuleLDS) { |
| // Create a struct to hold the ModuleScopeVariables |
| // Replace all uses of those variables from non-kernel functions with the |
| // new struct instance Replace only the uses from kernel functions that will |
| // allocate this instance. That is a space optimisation - kernels that use a |
| // subset of the module scope struct and do not need to allocate it for |
| // indirect calls will only allocate the subset they use (they do so as part |
| // of the per-kernel lowering). |
| if (ModuleScopeVariables.empty()) { |
| return nullptr; |
| } |
| |
| LLVMContext &Ctx = M.getContext(); |
| |
| LDSVariableReplacement ModuleScopeReplacement = |
| createLDSVariableReplacement(M, "llvm.amdgcn.module.lds", |
| ModuleScopeVariables); |
| |
| appendToCompilerUsed(M, {static_cast<GlobalValue *>( |
| ConstantExpr::getPointerBitCastOrAddrSpaceCast( |
| cast<Constant>(ModuleScopeReplacement.SGV), |
| PointerType::getUnqual(Ctx)))}); |
| |
| // module.lds will be allocated at zero in any kernel that allocates it |
| recordLDSAbsoluteAddress(&M, ModuleScopeReplacement.SGV, 0); |
| |
| // historic |
| removeLocalVarsFromUsedLists(M, ModuleScopeVariables); |
| |
| // Replace all uses of module scope variable from non-kernel functions |
| replaceLDSVariablesWithStruct( |
| M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) { |
| Instruction *I = dyn_cast<Instruction>(U.getUser()); |
| if (!I) { |
| return false; |
| } |
| Function *F = I->getFunction(); |
| return !isKernelLDS(F); |
| }); |
| |
| // Replace uses of module scope variable from kernel functions that |
| // allocate the module scope variable, otherwise leave them unchanged |
| // Record on each kernel whether the module scope global is used by it |
| |
| for (Function &Func : M.functions()) { |
| if (Func.isDeclaration() || !isKernelLDS(&Func)) |
| continue; |
| |
| if (KernelsThatAllocateModuleLDS.contains(&Func)) { |
| replaceLDSVariablesWithStruct( |
| M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) { |
| Instruction *I = dyn_cast<Instruction>(U.getUser()); |
| if (!I) { |
| return false; |
| } |
| Function *F = I->getFunction(); |
| return F == &Func; |
| }); |
| |
| markUsedByKernel(&Func, ModuleScopeReplacement.SGV); |
| } |
| } |
| |
| return ModuleScopeReplacement.SGV; |
| } |
| |
| static DenseMap<Function *, LDSVariableReplacement> |
| lowerKernelScopeStructVariables( |
| Module &M, LDSUsesInfoTy &LDSUsesInfo, |
| DenseSet<GlobalVariable *> const &ModuleScopeVariables, |
| DenseSet<Function *> const &KernelsThatAllocateModuleLDS, |
| GlobalVariable *MaybeModuleScopeStruct) { |
| |
| // Create a struct for each kernel for the non-module-scope variables. |
| |
| DenseMap<Function *, LDSVariableReplacement> KernelToReplacement; |
| for (Function &Func : M.functions()) { |
| if (Func.isDeclaration() || !isKernelLDS(&Func)) |
| continue; |
| |
| DenseSet<GlobalVariable *> KernelUsedVariables; |
| // Allocating variables that are used directly in this struct to get |
| // alignment aware allocation and predictable frame size. |
| for (auto &v : LDSUsesInfo.direct_access[&Func]) { |
| if (!AMDGPU::isDynamicLDS(*v)) { |
| KernelUsedVariables.insert(v); |
| } |
| } |
| |
| // Allocating variables that are accessed indirectly so that a lookup of |
| // this struct instance can find them from nested functions. |
| for (auto &v : LDSUsesInfo.indirect_access[&Func]) { |
| if (!AMDGPU::isDynamicLDS(*v)) { |
| KernelUsedVariables.insert(v); |
| } |
| } |
| |
| // Variables allocated in module lds must all resolve to that struct, |
| // not to the per-kernel instance. |
| if (KernelsThatAllocateModuleLDS.contains(&Func)) { |
| for (GlobalVariable *v : ModuleScopeVariables) { |
| KernelUsedVariables.erase(v); |
| } |
| } |
| |
| if (KernelUsedVariables.empty()) { |
| // Either used no LDS, or the LDS it used was all in the module struct |
| // or dynamically sized |
| continue; |
| } |
| |
| // The association between kernel function and LDS struct is done by |
| // symbol name, which only works if the function in question has a |
| // name This is not expected to be a problem in practice as kernels |
| // are called by name making anonymous ones (which are named by the |
| // backend) difficult to use. This does mean that llvm test cases need |
| // to name the kernels. |
| if (!Func.hasName()) { |
| report_fatal_error("Anonymous kernels cannot use LDS variables"); |
| } |
| |
| std::string VarName = |
| (Twine("llvm.amdgcn.kernel.") + Func.getName() + ".lds").str(); |
| |
| auto Replacement = |
| createLDSVariableReplacement(M, VarName, KernelUsedVariables); |
| |
| // If any indirect uses, create a direct use to ensure allocation |
| // TODO: Simpler to unconditionally mark used but that regresses |
| // codegen in test/CodeGen/AMDGPU/noclobber-barrier.ll |
| auto Accesses = LDSUsesInfo.indirect_access.find(&Func); |
| if ((Accesses != LDSUsesInfo.indirect_access.end()) && |
| !Accesses->second.empty()) |
| markUsedByKernel(&Func, Replacement.SGV); |
| |
| // remove preserves existing codegen |
| removeLocalVarsFromUsedLists(M, KernelUsedVariables); |
| KernelToReplacement[&Func] = Replacement; |
| |
| // Rewrite uses within kernel to the new struct |
| replaceLDSVariablesWithStruct( |
| M, KernelUsedVariables, Replacement, [&Func](Use &U) { |
| Instruction *I = dyn_cast<Instruction>(U.getUser()); |
| return I && I->getFunction() == &Func; |
| }); |
| } |
| return KernelToReplacement; |
| } |
| |
| static GlobalVariable * |
| buildRepresentativeDynamicLDSInstance(Module &M, LDSUsesInfoTy &LDSUsesInfo, |
| Function *func) { |
| // Create a dynamic lds variable with a name associated with the passed |
| // function that has the maximum alignment of any dynamic lds variable |
| // reachable from this kernel. Dynamic LDS is allocated after the static LDS |
| // allocation, possibly after alignment padding. The representative variable |
| // created here has the maximum alignment of any other dynamic variable |
| // reachable by that kernel. All dynamic LDS variables are allocated at the |
| // same address in each kernel in order to provide the documented aliasing |
| // semantics. Setting the alignment here allows this IR pass to accurately |
| // predict the exact constant at which it will be allocated. |
| |
| assert(isKernelLDS(func)); |
| |
| LLVMContext &Ctx = M.getContext(); |
| const DataLayout &DL = M.getDataLayout(); |
| Align MaxDynamicAlignment(1); |
| |
| auto UpdateMaxAlignment = [&MaxDynamicAlignment, &DL](GlobalVariable *GV) { |
| if (AMDGPU::isDynamicLDS(*GV)) { |
| MaxDynamicAlignment = |
| std::max(MaxDynamicAlignment, AMDGPU::getAlign(DL, GV)); |
| } |
| }; |
| |
| for (GlobalVariable *GV : LDSUsesInfo.indirect_access[func]) { |
| UpdateMaxAlignment(GV); |
| } |
| |
| for (GlobalVariable *GV : LDSUsesInfo.direct_access[func]) { |
| UpdateMaxAlignment(GV); |
| } |
| |
| assert(func->hasName()); // Checked by caller |
| auto emptyCharArray = ArrayType::get(Type::getInt8Ty(Ctx), 0); |
| GlobalVariable *N = new GlobalVariable( |
| M, emptyCharArray, false, GlobalValue::ExternalLinkage, nullptr, |
| Twine("llvm.amdgcn." + func->getName() + ".dynlds"), nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, |
| false); |
| N->setAlignment(MaxDynamicAlignment); |
| |
| assert(AMDGPU::isDynamicLDS(*N)); |
| return N; |
| } |
| |
| /// Strip "amdgpu-no-lds-kernel-id" from any functions where we may have |
| /// introduced its use. If AMDGPUAttributor ran prior to the pass, we inferred |
| /// the lack of llvm.amdgcn.lds.kernel.id calls. |
| void removeNoLdsKernelIdFromReachable(CallGraph &CG, Function *KernelRoot) { |
| KernelRoot->removeFnAttr("amdgpu-no-lds-kernel-id"); |
| |
| SmallVector<Function *> WorkList({CG[KernelRoot]->getFunction()}); |
| SmallPtrSet<Function *, 8> Visited; |
| bool SeenUnknownCall = false; |
| |
| while (!WorkList.empty()) { |
| Function *F = WorkList.pop_back_val(); |
| |
| for (auto &CallRecord : *CG[F]) { |
| if (!CallRecord.second) |
| continue; |
| |
| Function *Callee = CallRecord.second->getFunction(); |
| if (!Callee) { |
| if (!SeenUnknownCall) { |
| SeenUnknownCall = true; |
| |
| // If we see any indirect calls, assume nothing about potential |
| // targets. |
| // TODO: This could be refined to possible LDS global users. |
| for (auto &ExternalCallRecord : *CG.getExternalCallingNode()) { |
| Function *PotentialCallee = |
| ExternalCallRecord.second->getFunction(); |
| assert(PotentialCallee); |
| if (!isKernelLDS(PotentialCallee)) |
| PotentialCallee->removeFnAttr("amdgpu-no-lds-kernel-id"); |
| } |
| } |
| } else { |
| Callee->removeFnAttr("amdgpu-no-lds-kernel-id"); |
| if (Visited.insert(Callee).second) |
| WorkList.push_back(Callee); |
| } |
| } |
| } |
| } |
| |
| DenseMap<Function *, GlobalVariable *> lowerDynamicLDSVariables( |
| Module &M, LDSUsesInfoTy &LDSUsesInfo, |
| DenseSet<Function *> const &KernelsThatIndirectlyAllocateDynamicLDS, |
| DenseSet<GlobalVariable *> const &DynamicVariables, |
| std::vector<Function *> const &OrderedKernels) { |
| DenseMap<Function *, GlobalVariable *> KernelToCreatedDynamicLDS; |
| if (!KernelsThatIndirectlyAllocateDynamicLDS.empty()) { |
| LLVMContext &Ctx = M.getContext(); |
| IRBuilder<> Builder(Ctx); |
| Type *I32 = Type::getInt32Ty(Ctx); |
| |
| std::vector<Constant *> newDynamicLDS; |
| |
| // Table is built in the same order as OrderedKernels |
| for (auto &func : OrderedKernels) { |
| |
| if (KernelsThatIndirectlyAllocateDynamicLDS.contains(func)) { |
| assert(isKernelLDS(func)); |
| if (!func->hasName()) { |
| report_fatal_error("Anonymous kernels cannot use LDS variables"); |
| } |
| |
| GlobalVariable *N = |
| buildRepresentativeDynamicLDSInstance(M, LDSUsesInfo, func); |
| |
| KernelToCreatedDynamicLDS[func] = N; |
| |
| markUsedByKernel(func, N); |
| |
| auto emptyCharArray = ArrayType::get(Type::getInt8Ty(Ctx), 0); |
| auto GEP = ConstantExpr::getGetElementPtr( |
| emptyCharArray, N, ConstantInt::get(I32, 0), true); |
| newDynamicLDS.push_back(ConstantExpr::getPtrToInt(GEP, I32)); |
| } else { |
| newDynamicLDS.push_back(PoisonValue::get(I32)); |
| } |
| } |
| assert(OrderedKernels.size() == newDynamicLDS.size()); |
| |
| ArrayType *t = ArrayType::get(I32, newDynamicLDS.size()); |
| Constant *init = ConstantArray::get(t, newDynamicLDS); |
| GlobalVariable *table = new GlobalVariable( |
| M, t, true, GlobalValue::InternalLinkage, init, |
| "llvm.amdgcn.dynlds.offset.table", nullptr, |
| GlobalValue::NotThreadLocal, AMDGPUAS::CONSTANT_ADDRESS); |
| |
| for (GlobalVariable *GV : DynamicVariables) { |
| for (Use &U : make_early_inc_range(GV->uses())) { |
| auto *I = dyn_cast<Instruction>(U.getUser()); |
| if (!I) |
| continue; |
| if (isKernelLDS(I->getFunction())) |
| continue; |
| |
| replaceUseWithTableLookup(M, Builder, table, GV, U, nullptr); |
| } |
| } |
| } |
| return KernelToCreatedDynamicLDS; |
| } |
| |
| bool runOnModule(Module &M) { |
| CallGraph CG = CallGraph(M); |
| bool Changed = superAlignLDSGlobals(M); |
| |
| Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M); |
| |
| Changed = true; // todo: narrow this down |
| |
| // For each kernel, what variables does it access directly or through |
| // callees |
| LDSUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDS(CG, M); |
| |
| // For each variable accessed through callees, which kernels access it |
| VariableFunctionMap LDSToKernelsThatNeedToAccessItIndirectly; |
| for (auto &K : LDSUsesInfo.indirect_access) { |
| Function *F = K.first; |
| assert(isKernelLDS(F)); |
| for (GlobalVariable *GV : K.second) { |
| LDSToKernelsThatNeedToAccessItIndirectly[GV].insert(F); |
| } |
| } |
| |
| // Partition variables accessed indirectly into the different strategies |
| DenseSet<GlobalVariable *> ModuleScopeVariables; |
| DenseSet<GlobalVariable *> TableLookupVariables; |
| DenseSet<GlobalVariable *> KernelAccessVariables; |
| DenseSet<GlobalVariable *> DynamicVariables; |
| partitionVariablesIntoIndirectStrategies( |
| M, LDSUsesInfo, LDSToKernelsThatNeedToAccessItIndirectly, |
| ModuleScopeVariables, TableLookupVariables, KernelAccessVariables, |
| DynamicVariables); |
| |
| // If the kernel accesses a variable that is going to be stored in the |
| // module instance through a call then that kernel needs to allocate the |
| // module instance |
| const DenseSet<Function *> KernelsThatAllocateModuleLDS = |
| kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo, |
| ModuleScopeVariables); |
| const DenseSet<Function *> KernelsThatAllocateTableLDS = |
| kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo, |
| TableLookupVariables); |
| |
| const DenseSet<Function *> KernelsThatIndirectlyAllocateDynamicLDS = |
| kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo, |
| DynamicVariables); |
| |
| GlobalVariable *MaybeModuleScopeStruct = lowerModuleScopeStructVariables( |
| M, ModuleScopeVariables, KernelsThatAllocateModuleLDS); |
| |
| DenseMap<Function *, LDSVariableReplacement> KernelToReplacement = |
| lowerKernelScopeStructVariables(M, LDSUsesInfo, ModuleScopeVariables, |
| KernelsThatAllocateModuleLDS, |
| MaybeModuleScopeStruct); |
| |
| // Lower zero cost accesses to the kernel instances just created |
| for (auto &GV : KernelAccessVariables) { |
| auto &funcs = LDSToKernelsThatNeedToAccessItIndirectly[GV]; |
| assert(funcs.size() == 1); // Only one kernel can access it |
| LDSVariableReplacement Replacement = |
| KernelToReplacement[*(funcs.begin())]; |
| |
| DenseSet<GlobalVariable *> Vec; |
| Vec.insert(GV); |
| |
| replaceLDSVariablesWithStruct(M, Vec, Replacement, [](Use &U) { |
| return isa<Instruction>(U.getUser()); |
| }); |
| } |
| |
| // The ith element of this vector is kernel id i |
| std::vector<Function *> OrderedKernels = |
| assignLDSKernelIDToEachKernel(&M, KernelsThatAllocateTableLDS, |
| KernelsThatIndirectlyAllocateDynamicLDS); |
| |
| if (!KernelsThatAllocateTableLDS.empty()) { |
| LLVMContext &Ctx = M.getContext(); |
| IRBuilder<> Builder(Ctx); |
| |
| // The order must be consistent between lookup table and accesses to |
| // lookup table |
| auto TableLookupVariablesOrdered = |
| sortByName(std::vector<GlobalVariable *>(TableLookupVariables.begin(), |
| TableLookupVariables.end())); |
| |
| GlobalVariable *LookupTable = buildLookupTable( |
| M, TableLookupVariablesOrdered, OrderedKernels, KernelToReplacement); |
| replaceUsesInInstructionsWithTableLookup(M, TableLookupVariablesOrdered, |
| LookupTable); |
| |
| // Strip amdgpu-no-lds-kernel-id from all functions reachable from the |
| // kernel. We may have inferred this wasn't used prior to the pass. |
| // |
| // TODO: We could filter out subgraphs that do not access LDS globals. |
| for (Function *F : KernelsThatAllocateTableLDS) |
| removeNoLdsKernelIdFromReachable(CG, F); |
| } |
| |
| DenseMap<Function *, GlobalVariable *> KernelToCreatedDynamicLDS = |
| lowerDynamicLDSVariables(M, LDSUsesInfo, |
| KernelsThatIndirectlyAllocateDynamicLDS, |
| DynamicVariables, OrderedKernels); |
| |
| // All kernel frames have been allocated. Calculate and record the |
| // addresses. |
| { |
| const DataLayout &DL = M.getDataLayout(); |
| |
| for (Function &Func : M.functions()) { |
| if (Func.isDeclaration() || !isKernelLDS(&Func)) |
| continue; |
| |
| // All three of these are optional. The first variable is allocated at |
| // zero. They are allocated by AMDGPUMachineFunction as one block. |
| // Layout: |
| //{ |
| // module.lds |
| // alignment padding |
| // kernel instance |
| // alignment padding |
| // dynamic lds variables |
| //} |
| |
| const bool AllocateModuleScopeStruct = |
| MaybeModuleScopeStruct && |
| KernelsThatAllocateModuleLDS.contains(&Func); |
| |
| auto Replacement = KernelToReplacement.find(&Func); |
| const bool AllocateKernelScopeStruct = |
| Replacement != KernelToReplacement.end(); |
| |
| const bool AllocateDynamicVariable = |
| KernelToCreatedDynamicLDS.contains(&Func); |
| |
| uint32_t Offset = 0; |
| |
| if (AllocateModuleScopeStruct) { |
| // Allocated at zero, recorded once on construction, not once per |
| // kernel |
| Offset += DL.getTypeAllocSize(MaybeModuleScopeStruct->getValueType()); |
| } |
| |
| if (AllocateKernelScopeStruct) { |
| GlobalVariable *KernelStruct = Replacement->second.SGV; |
| Offset = alignTo(Offset, AMDGPU::getAlign(DL, KernelStruct)); |
| recordLDSAbsoluteAddress(&M, KernelStruct, Offset); |
| Offset += DL.getTypeAllocSize(KernelStruct->getValueType()); |
| } |
| |
| // If there is dynamic allocation, the alignment needed is included in |
| // the static frame size. There may be no reference to the dynamic |
| // variable in the kernel itself, so without including it here, that |
| // alignment padding could be missed. |
| if (AllocateDynamicVariable) { |
| GlobalVariable *DynamicVariable = KernelToCreatedDynamicLDS[&Func]; |
| Offset = alignTo(Offset, AMDGPU::getAlign(DL, DynamicVariable)); |
| recordLDSAbsoluteAddress(&M, DynamicVariable, Offset); |
| } |
| |
| if (Offset != 0) { |
| (void)TM; // TODO: Account for target maximum LDS |
| std::string Buffer; |
| raw_string_ostream SS{Buffer}; |
| SS << format("%u", Offset); |
| |
| // Instead of explictly marking kernels that access dynamic variables |
| // using special case metadata, annotate with min-lds == max-lds, i.e. |
| // that there is no more space available for allocating more static |
| // LDS variables. That is the right condition to prevent allocating |
| // more variables which would collide with the addresses assigned to |
| // dynamic variables. |
| if (AllocateDynamicVariable) |
| SS << format(",%u", Offset); |
| |
| Func.addFnAttr("amdgpu-lds-size", Buffer); |
| } |
| } |
| } |
| |
| for (auto &GV : make_early_inc_range(M.globals())) |
| if (AMDGPU::isLDSVariableToLower(GV)) { |
| // probably want to remove from used lists |
| GV.removeDeadConstantUsers(); |
| if (GV.use_empty()) |
| GV.eraseFromParent(); |
| } |
| |
| return Changed; |
| } |
| |
| private: |
| // Increase the alignment of LDS globals if necessary to maximise the chance |
| // that we can use aligned LDS instructions to access them. |
| static bool superAlignLDSGlobals(Module &M) { |
| const DataLayout &DL = M.getDataLayout(); |
| bool Changed = false; |
| if (!SuperAlignLDSGlobals) { |
| return Changed; |
| } |
| |
| for (auto &GV : M.globals()) { |
| if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) { |
| // Only changing alignment of LDS variables |
| continue; |
| } |
| if (!GV.hasInitializer()) { |
| // cuda/hip extern __shared__ variable, leave alignment alone |
| continue; |
| } |
| |
| Align Alignment = AMDGPU::getAlign(DL, &GV); |
| TypeSize GVSize = DL.getTypeAllocSize(GV.getValueType()); |
| |
| if (GVSize > 8) { |
| // We might want to use a b96 or b128 load/store |
| Alignment = std::max(Alignment, Align(16)); |
| } else if (GVSize > 4) { |
| // We might want to use a b64 load/store |
| Alignment = std::max(Alignment, Align(8)); |
| } else if (GVSize > 2) { |
| // We might want to use a b32 load/store |
| Alignment = std::max(Alignment, Align(4)); |
| } else if (GVSize > 1) { |
| // We might want to use a b16 load/store |
| Alignment = std::max(Alignment, Align(2)); |
| } |
| |
| if (Alignment != AMDGPU::getAlign(DL, &GV)) { |
| Changed = true; |
| GV.setAlignment(Alignment); |
| } |
| } |
| return Changed; |
| } |
| |
| static LDSVariableReplacement createLDSVariableReplacement( |
| Module &M, std::string VarName, |
| DenseSet<GlobalVariable *> const &LDSVarsToTransform) { |
| // Create a struct instance containing LDSVarsToTransform and map from those |
| // variables to ConstantExprGEP |
| // Variables may be introduced to meet alignment requirements. No aliasing |
| // metadata is useful for these as they have no uses. Erased before return. |
| |
| LLVMContext &Ctx = M.getContext(); |
| const DataLayout &DL = M.getDataLayout(); |
| assert(!LDSVarsToTransform.empty()); |
| |
| SmallVector<OptimizedStructLayoutField, 8> LayoutFields; |
| LayoutFields.reserve(LDSVarsToTransform.size()); |
| { |
| // The order of fields in this struct depends on the order of |
| // varables in the argument which varies when changing how they |
| // are identified, leading to spurious test breakage. |
| auto Sorted = sortByName(std::vector<GlobalVariable *>( |
| LDSVarsToTransform.begin(), LDSVarsToTransform.end())); |
| |
| for (GlobalVariable *GV : Sorted) { |
| OptimizedStructLayoutField F(GV, |
| DL.getTypeAllocSize(GV->getValueType()), |
| AMDGPU::getAlign(DL, GV)); |
| LayoutFields.emplace_back(F); |
| } |
| } |
| |
| performOptimizedStructLayout(LayoutFields); |
| |
| std::vector<GlobalVariable *> LocalVars; |
| BitVector IsPaddingField; |
| LocalVars.reserve(LDSVarsToTransform.size()); // will be at least this large |
| IsPaddingField.reserve(LDSVarsToTransform.size()); |
| { |
| uint64_t CurrentOffset = 0; |
| for (size_t I = 0; I < LayoutFields.size(); I++) { |
| GlobalVariable *FGV = static_cast<GlobalVariable *>( |
| const_cast<void *>(LayoutFields[I].Id)); |
| Align DataAlign = LayoutFields[I].Alignment; |
| |
| uint64_t DataAlignV = DataAlign.value(); |
| if (uint64_t Rem = CurrentOffset % DataAlignV) { |
| uint64_t Padding = DataAlignV - Rem; |
| |
| // Append an array of padding bytes to meet alignment requested |
| // Note (o + (a - (o % a)) ) % a == 0 |
| // (offset + Padding ) % align == 0 |
| |
| Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding); |
| LocalVars.push_back(new GlobalVariable( |
| M, ATy, false, GlobalValue::InternalLinkage, |
| PoisonValue::get(ATy), "", nullptr, GlobalValue::NotThreadLocal, |
| AMDGPUAS::LOCAL_ADDRESS, false)); |
| IsPaddingField.push_back(true); |
| CurrentOffset += Padding; |
| } |
| |
| LocalVars.push_back(FGV); |
| IsPaddingField.push_back(false); |
| CurrentOffset += LayoutFields[I].Size; |
| } |
| } |
| |
| std::vector<Type *> LocalVarTypes; |
| LocalVarTypes.reserve(LocalVars.size()); |
| std::transform( |
| LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes), |
| [](const GlobalVariable *V) -> Type * { return V->getValueType(); }); |
| |
| StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t"); |
| |
| Align StructAlign = AMDGPU::getAlign(DL, LocalVars[0]); |
| |
| GlobalVariable *SGV = new GlobalVariable( |
| M, LDSTy, false, GlobalValue::InternalLinkage, PoisonValue::get(LDSTy), |
| VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, |
| false); |
| SGV->setAlignment(StructAlign); |
| |
| DenseMap<GlobalVariable *, Constant *> Map; |
| Type *I32 = Type::getInt32Ty(Ctx); |
| for (size_t I = 0; I < LocalVars.size(); I++) { |
| GlobalVariable *GV = LocalVars[I]; |
| Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)}; |
| Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx, true); |
| if (IsPaddingField[I]) { |
| assert(GV->use_empty()); |
| GV->eraseFromParent(); |
| } else { |
| Map[GV] = GEP; |
| } |
| } |
| assert(Map.size() == LDSVarsToTransform.size()); |
| return {SGV, std::move(Map)}; |
| } |
| |
| template <typename PredicateTy> |
| static void replaceLDSVariablesWithStruct( |
| Module &M, DenseSet<GlobalVariable *> const &LDSVarsToTransformArg, |
| const LDSVariableReplacement &Replacement, PredicateTy Predicate) { |
| LLVMContext &Ctx = M.getContext(); |
| const DataLayout &DL = M.getDataLayout(); |
| |
| // A hack... we need to insert the aliasing info in a predictable order for |
| // lit tests. Would like to have them in a stable order already, ideally the |
| // same order they get allocated, which might mean an ordered set container |
| auto LDSVarsToTransform = sortByName(std::vector<GlobalVariable *>( |
| LDSVarsToTransformArg.begin(), LDSVarsToTransformArg.end())); |
| |
| // Create alias.scope and their lists. Each field in the new structure |
| // does not alias with all other fields. |
| SmallVector<MDNode *> AliasScopes; |
| SmallVector<Metadata *> NoAliasList; |
| const size_t NumberVars = LDSVarsToTransform.size(); |
| if (NumberVars > 1) { |
| MDBuilder MDB(Ctx); |
| AliasScopes.reserve(NumberVars); |
| MDNode *Domain = MDB.createAnonymousAliasScopeDomain(); |
| for (size_t I = 0; I < NumberVars; I++) { |
| MDNode *Scope = MDB.createAnonymousAliasScope(Domain); |
| AliasScopes.push_back(Scope); |
| } |
| NoAliasList.append(&AliasScopes[1], AliasScopes.end()); |
| } |
| |
| // Replace uses of ith variable with a constantexpr to the corresponding |
| // field of the instance that will be allocated by AMDGPUMachineFunction |
| for (size_t I = 0; I < NumberVars; I++) { |
| GlobalVariable *GV = LDSVarsToTransform[I]; |
| Constant *GEP = Replacement.LDSVarsToConstantGEP.at(GV); |
| |
| GV->replaceUsesWithIf(GEP, Predicate); |
| |
| APInt APOff(DL.getIndexTypeSizeInBits(GEP->getType()), 0); |
| GEP->stripAndAccumulateInBoundsConstantOffsets(DL, APOff); |
| uint64_t Offset = APOff.getZExtValue(); |
| |
| Align A = |
| commonAlignment(Replacement.SGV->getAlign().valueOrOne(), Offset); |
| |
| if (I) |
| NoAliasList[I - 1] = AliasScopes[I - 1]; |
| MDNode *NoAlias = |
| NoAliasList.empty() ? nullptr : MDNode::get(Ctx, NoAliasList); |
| MDNode *AliasScope = |
| AliasScopes.empty() ? nullptr : MDNode::get(Ctx, {AliasScopes[I]}); |
| |
| refineUsesAlignmentAndAA(GEP, A, DL, AliasScope, NoAlias); |
| } |
| } |
| |
| static void refineUsesAlignmentAndAA(Value *Ptr, Align A, |
| const DataLayout &DL, MDNode *AliasScope, |
| MDNode *NoAlias, unsigned MaxDepth = 5) { |
| if (!MaxDepth || (A == 1 && !AliasScope)) |
| return; |
| |
| for (User *U : Ptr->users()) { |
| if (auto *I = dyn_cast<Instruction>(U)) { |
| if (AliasScope && I->mayReadOrWriteMemory()) { |
| MDNode *AS = I->getMetadata(LLVMContext::MD_alias_scope); |
| AS = (AS ? MDNode::getMostGenericAliasScope(AS, AliasScope) |
| : AliasScope); |
| I->setMetadata(LLVMContext::MD_alias_scope, AS); |
| |
| MDNode *NA = I->getMetadata(LLVMContext::MD_noalias); |
| NA = (NA ? MDNode::intersect(NA, NoAlias) : NoAlias); |
| I->setMetadata(LLVMContext::MD_noalias, NA); |
| } |
| } |
| |
| if (auto *LI = dyn_cast<LoadInst>(U)) { |
| LI->setAlignment(std::max(A, LI->getAlign())); |
| continue; |
| } |
| if (auto *SI = dyn_cast<StoreInst>(U)) { |
| if (SI->getPointerOperand() == Ptr) |
| SI->setAlignment(std::max(A, SI->getAlign())); |
| continue; |
| } |
| if (auto *AI = dyn_cast<AtomicRMWInst>(U)) { |
| // None of atomicrmw operations can work on pointers, but let's |
| // check it anyway in case it will or we will process ConstantExpr. |
| if (AI->getPointerOperand() == Ptr) |
| AI->setAlignment(std::max(A, AI->getAlign())); |
| continue; |
| } |
| if (auto *AI = dyn_cast<AtomicCmpXchgInst>(U)) { |
| if (AI->getPointerOperand() == Ptr) |
| AI->setAlignment(std::max(A, AI->getAlign())); |
| continue; |
| } |
| if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) { |
| unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP->getType()); |
| APInt Off(BitWidth, 0); |
| if (GEP->getPointerOperand() == Ptr) { |
| Align GA; |
| if (GEP->accumulateConstantOffset(DL, Off)) |
| GA = commonAlignment(A, Off.getLimitedValue()); |
| refineUsesAlignmentAndAA(GEP, GA, DL, AliasScope, NoAlias, |
| MaxDepth - 1); |
| } |
| continue; |
| } |
| if (auto *I = dyn_cast<Instruction>(U)) { |
| if (I->getOpcode() == Instruction::BitCast || |
| I->getOpcode() == Instruction::AddrSpaceCast) |
| refineUsesAlignmentAndAA(I, A, DL, AliasScope, NoAlias, MaxDepth - 1); |
| } |
| } |
| } |
| }; |
| |
| class AMDGPULowerModuleLDSLegacy : public ModulePass { |
| public: |
| const AMDGPUTargetMachine *TM; |
| static char ID; |
| |
| AMDGPULowerModuleLDSLegacy(const AMDGPUTargetMachine *TM_ = nullptr) |
| : ModulePass(ID), TM(TM_) { |
| initializeAMDGPULowerModuleLDSLegacyPass(*PassRegistry::getPassRegistry()); |
| } |
| |
| void getAnalysisUsage(AnalysisUsage &AU) const override { |
| if (!TM) |
| AU.addRequired<TargetPassConfig>(); |
| } |
| |
| bool runOnModule(Module &M) override { |
| if (!TM) { |
| auto &TPC = getAnalysis<TargetPassConfig>(); |
| TM = &TPC.getTM<AMDGPUTargetMachine>(); |
| } |
| |
| return AMDGPULowerModuleLDS(*TM).runOnModule(M); |
| } |
| }; |
| |
| } // namespace |
| char AMDGPULowerModuleLDSLegacy::ID = 0; |
| |
| char &llvm::AMDGPULowerModuleLDSLegacyPassID = AMDGPULowerModuleLDSLegacy::ID; |
| |
| INITIALIZE_PASS_BEGIN(AMDGPULowerModuleLDSLegacy, DEBUG_TYPE, |
| "Lower uses of LDS variables from non-kernel functions", |
| false, false) |
| INITIALIZE_PASS_DEPENDENCY(TargetPassConfig) |
| INITIALIZE_PASS_END(AMDGPULowerModuleLDSLegacy, DEBUG_TYPE, |
| "Lower uses of LDS variables from non-kernel functions", |
| false, false) |
| |
| ModulePass * |
| llvm::createAMDGPULowerModuleLDSLegacyPass(const AMDGPUTargetMachine *TM) { |
| return new AMDGPULowerModuleLDSLegacy(TM); |
| } |
| |
| PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M, |
| ModuleAnalysisManager &) { |
| return AMDGPULowerModuleLDS(TM).runOnModule(M) ? PreservedAnalyses::none() |
| : PreservedAnalyses::all(); |
| } |
