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