1 //===-- AMDGPULowerModuleLDSPass.cpp ------------------------------*- C++ -*-=// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This pass eliminates LDS uses from non-kernel functions. 10 // 11 // The strategy is to create a new struct with a field for each LDS variable 12 // and allocate that struct at the same address for every kernel. Uses of the 13 // original LDS variables are then replaced with compile time offsets from that 14 // known address. AMDGPUMachineFunction allocates the LDS global. 15 // 16 // Local variables with constant annotation or non-undef initializer are passed 17 // through unchanged for simplification or error diagnostics in later passes. 18 // 19 // To reduce the memory overhead variables that are only used by kernels are 20 // excluded from this transform. The analysis to determine whether a variable 21 // is only used by a kernel is cheap and conservative so this may allocate 22 // a variable in every kernel when it was not strictly necessary to do so. 23 // 24 // A possible future refinement is to specialise the structure per-kernel, so 25 // that fields can be elided based on more expensive analysis. 26 // 27 //===----------------------------------------------------------------------===// 28 29 #include "AMDGPU.h" 30 #include "Utils/AMDGPUBaseInfo.h" 31 #include "Utils/AMDGPUMemoryUtils.h" 32 #include "llvm/ADT/BitVector.h" 33 #include "llvm/ADT/DenseMap.h" 34 #include "llvm/ADT/STLExtras.h" 35 #include "llvm/Analysis/CallGraph.h" 36 #include "llvm/IR/Constants.h" 37 #include "llvm/IR/DerivedTypes.h" 38 #include "llvm/IR/IRBuilder.h" 39 #include "llvm/IR/InlineAsm.h" 40 #include "llvm/IR/Instructions.h" 41 #include "llvm/IR/MDBuilder.h" 42 #include "llvm/InitializePasses.h" 43 #include "llvm/Pass.h" 44 #include "llvm/Support/CommandLine.h" 45 #include "llvm/Support/Debug.h" 46 #include "llvm/Support/OptimizedStructLayout.h" 47 #include "llvm/Transforms/Utils/ModuleUtils.h" 48 #include <tuple> 49 #include <vector> 50 51 #define DEBUG_TYPE "amdgpu-lower-module-lds" 52 53 using namespace llvm; 54 55 static cl::opt<bool> SuperAlignLDSGlobals( 56 "amdgpu-super-align-lds-globals", 57 cl::desc("Increase alignment of LDS if it is not on align boundary"), 58 cl::init(true), cl::Hidden); 59 60 namespace { 61 class AMDGPULowerModuleLDS : public ModulePass { 62 63 static void removeFromUsedList(Module &M, StringRef Name, 64 SmallPtrSetImpl<Constant *> &ToRemove) { 65 GlobalVariable *GV = M.getNamedGlobal(Name); 66 if (!GV || ToRemove.empty()) { 67 return; 68 } 69 70 SmallVector<Constant *, 16> Init; 71 auto *CA = cast<ConstantArray>(GV->getInitializer()); 72 for (auto &Op : CA->operands()) { 73 // ModuleUtils::appendToUsed only inserts Constants 74 Constant *C = cast<Constant>(Op); 75 if (!ToRemove.contains(C->stripPointerCasts())) { 76 Init.push_back(C); 77 } 78 } 79 80 if (Init.size() == CA->getNumOperands()) { 81 return; // none to remove 82 } 83 84 GV->eraseFromParent(); 85 86 for (Constant *C : ToRemove) { 87 C->removeDeadConstantUsers(); 88 } 89 90 if (!Init.empty()) { 91 ArrayType *ATy = 92 ArrayType::get(Type::getInt8PtrTy(M.getContext()), Init.size()); 93 GV = 94 new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage, 95 ConstantArray::get(ATy, Init), Name); 96 GV->setSection("llvm.metadata"); 97 } 98 } 99 100 static void 101 removeFromUsedLists(Module &M, 102 const std::vector<GlobalVariable *> &LocalVars) { 103 // The verifier rejects used lists containing an inttoptr of a constant 104 // so remove the variables from these lists before replaceAllUsesWith 105 106 SmallPtrSet<Constant *, 32> LocalVarsSet; 107 for (GlobalVariable *LocalVar : LocalVars) 108 if (Constant *C = dyn_cast<Constant>(LocalVar->stripPointerCasts())) 109 LocalVarsSet.insert(C); 110 removeFromUsedList(M, "llvm.used", LocalVarsSet); 111 removeFromUsedList(M, "llvm.compiler.used", LocalVarsSet); 112 } 113 114 static void markUsedByKernel(IRBuilder<> &Builder, Function *Func, 115 GlobalVariable *SGV) { 116 // The llvm.amdgcn.module.lds instance is implicitly used by all kernels 117 // that might call a function which accesses a field within it. This is 118 // presently approximated to 'all kernels' if there are any such functions 119 // in the module. This implicit use is redefined as an explicit use here so 120 // that later passes, specifically PromoteAlloca, account for the required 121 // memory without any knowledge of this transform. 122 123 // An operand bundle on llvm.donothing works because the call instruction 124 // survives until after the last pass that needs to account for LDS. It is 125 // better than inline asm as the latter survives until the end of codegen. A 126 // totally robust solution would be a function with the same semantics as 127 // llvm.donothing that takes a pointer to the instance and is lowered to a 128 // no-op after LDS is allocated, but that is not presently necessary. 129 130 LLVMContext &Ctx = Func->getContext(); 131 132 Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI()); 133 134 FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {}); 135 136 Function *Decl = 137 Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {}); 138 139 Value *UseInstance[1] = {Builder.CreateInBoundsGEP( 140 SGV->getValueType(), SGV, ConstantInt::get(Type::getInt32Ty(Ctx), 0))}; 141 142 Builder.CreateCall(FTy, Decl, {}, 143 {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)}, 144 ""); 145 } 146 147 public: 148 static char ID; 149 150 AMDGPULowerModuleLDS() : ModulePass(ID) { 151 initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry()); 152 } 153 154 bool runOnModule(Module &M) override { 155 LLVMContext &Ctx = M.getContext(); 156 CallGraph CG = CallGraph(M); 157 bool Changed = superAlignLDSGlobals(M); 158 159 // Move variables used by functions into amdgcn.module.lds 160 std::vector<GlobalVariable *> ModuleScopeVariables = 161 AMDGPU::findLDSVariablesToLower(M, nullptr); 162 if (!ModuleScopeVariables.empty()) { 163 std::string VarName = "llvm.amdgcn.module.lds"; 164 165 GlobalVariable *SGV; 166 DenseMap<GlobalVariable *, Constant *> LDSVarsToConstantGEP; 167 std::tie(SGV, LDSVarsToConstantGEP) = 168 createLDSVariableReplacement(M, VarName, ModuleScopeVariables); 169 170 appendToCompilerUsed( 171 M, {static_cast<GlobalValue *>( 172 ConstantExpr::getPointerBitCastOrAddrSpaceCast( 173 cast<Constant>(SGV), Type::getInt8PtrTy(Ctx)))}); 174 175 removeFromUsedLists(M, ModuleScopeVariables); 176 replaceLDSVariablesWithStruct(M, ModuleScopeVariables, SGV, 177 LDSVarsToConstantGEP, 178 [](Use &) { return true; }); 179 180 // This ensures the variable is allocated when called functions access it. 181 // It also lets other passes, specifically PromoteAlloca, accurately 182 // calculate how much LDS will be used by the kernel after lowering. 183 184 IRBuilder<> Builder(Ctx); 185 for (Function &Func : M.functions()) { 186 if (!Func.isDeclaration() && AMDGPU::isKernelCC(&Func)) { 187 const CallGraphNode *N = CG[&Func]; 188 const bool CalleesRequireModuleLDS = N->size() > 0; 189 190 if (CalleesRequireModuleLDS) { 191 // If a function this kernel might call requires module LDS, 192 // annotate the kernel to let later passes know it will allocate 193 // this structure, even if not apparent from the IR. 194 markUsedByKernel(Builder, &Func, SGV); 195 } else { 196 // However if we are certain this kernel cannot call a function that 197 // requires module LDS, annotate the kernel so the backend can elide 198 // the allocation without repeating callgraph walks. 199 Func.addFnAttr("amdgpu-elide-module-lds"); 200 } 201 } 202 } 203 204 Changed = true; 205 } 206 207 // Move variables used by kernels into per-kernel instances 208 for (Function &F : M.functions()) { 209 if (F.isDeclaration()) 210 continue; 211 212 // Only lower compute kernels' LDS. 213 if (!AMDGPU::isKernel(F.getCallingConv())) 214 continue; 215 216 std::vector<GlobalVariable *> KernelUsedVariables = 217 AMDGPU::findLDSVariablesToLower(M, &F); 218 219 // Replace all constant uses with instructions if they belong to the 220 // current kernel. Unnecessary, removing will cause test churn. 221 for (GlobalVariable *GV : KernelUsedVariables) { 222 for (User *U : make_early_inc_range(GV->users())) { 223 if (ConstantExpr *C = dyn_cast<ConstantExpr>(U)) 224 AMDGPU::replaceConstantUsesInFunction(C, &F); 225 } 226 GV->removeDeadConstantUsers(); 227 } 228 229 if (!KernelUsedVariables.empty()) { 230 std::string VarName = 231 (Twine("llvm.amdgcn.kernel.") + F.getName() + ".lds").str(); 232 GlobalVariable *SGV; 233 DenseMap<GlobalVariable *, Constant *> LDSVarsToConstantGEP; 234 std::tie(SGV, LDSVarsToConstantGEP) = 235 createLDSVariableReplacement(M, VarName, KernelUsedVariables); 236 237 removeFromUsedLists(M, KernelUsedVariables); 238 replaceLDSVariablesWithStruct( 239 M, KernelUsedVariables, SGV, LDSVarsToConstantGEP, [&F](Use &U) { 240 Instruction *I = dyn_cast<Instruction>(U.getUser()); 241 return I && I->getFunction() == &F; 242 }); 243 Changed = true; 244 } 245 } 246 247 for (auto &GV : make_early_inc_range(M.globals())) 248 if (AMDGPU::isLDSVariableToLower(GV) && GV.use_empty()) 249 GV.eraseFromParent(); 250 251 return Changed; 252 } 253 254 private: 255 // Increase the alignment of LDS globals if necessary to maximise the chance 256 // that we can use aligned LDS instructions to access them. 257 static bool superAlignLDSGlobals(Module &M) { 258 const DataLayout &DL = M.getDataLayout(); 259 bool Changed = false; 260 if (!SuperAlignLDSGlobals) { 261 return Changed; 262 } 263 264 for (auto &GV : M.globals()) { 265 if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) { 266 // Only changing alignment of LDS variables 267 continue; 268 } 269 if (!GV.hasInitializer()) { 270 // cuda/hip extern __shared__ variable, leave alignment alone 271 continue; 272 } 273 274 Align Alignment = AMDGPU::getAlign(DL, &GV); 275 TypeSize GVSize = DL.getTypeAllocSize(GV.getValueType()); 276 277 if (GVSize > 8) { 278 // We might want to use a b96 or b128 load/store 279 Alignment = std::max(Alignment, Align(16)); 280 } else if (GVSize > 4) { 281 // We might want to use a b64 load/store 282 Alignment = std::max(Alignment, Align(8)); 283 } else if (GVSize > 2) { 284 // We might want to use a b32 load/store 285 Alignment = std::max(Alignment, Align(4)); 286 } else if (GVSize > 1) { 287 // We might want to use a b16 load/store 288 Alignment = std::max(Alignment, Align(2)); 289 } 290 291 if (Alignment != AMDGPU::getAlign(DL, &GV)) { 292 Changed = true; 293 GV.setAlignment(Alignment); 294 } 295 } 296 return Changed; 297 } 298 299 std::tuple<GlobalVariable *, DenseMap<GlobalVariable *, Constant *>> 300 createLDSVariableReplacement( 301 Module &M, std::string VarName, 302 std::vector<GlobalVariable *> const &LDSVarsToTransform) { 303 // Create a struct instance containing LDSVarsToTransform and map from those 304 // variables to ConstantExprGEP 305 // Variables may be introduced to meet alignment requirements. No aliasing 306 // metadata is useful for these as they have no uses. Erased before return. 307 308 LLVMContext &Ctx = M.getContext(); 309 const DataLayout &DL = M.getDataLayout(); 310 assert(!LDSVarsToTransform.empty()); 311 312 SmallVector<OptimizedStructLayoutField, 8> LayoutFields; 313 LayoutFields.reserve(LDSVarsToTransform.size()); 314 for (GlobalVariable *GV : LDSVarsToTransform) { 315 OptimizedStructLayoutField F(GV, DL.getTypeAllocSize(GV->getValueType()), 316 AMDGPU::getAlign(DL, GV)); 317 LayoutFields.emplace_back(F); 318 } 319 320 performOptimizedStructLayout(LayoutFields); 321 322 std::vector<GlobalVariable *> LocalVars; 323 BitVector IsPaddingField; 324 LocalVars.reserve(LDSVarsToTransform.size()); // will be at least this large 325 IsPaddingField.reserve(LDSVarsToTransform.size()); 326 { 327 uint64_t CurrentOffset = 0; 328 for (size_t I = 0; I < LayoutFields.size(); I++) { 329 GlobalVariable *FGV = static_cast<GlobalVariable *>( 330 const_cast<void *>(LayoutFields[I].Id)); 331 Align DataAlign = LayoutFields[I].Alignment; 332 333 uint64_t DataAlignV = DataAlign.value(); 334 if (uint64_t Rem = CurrentOffset % DataAlignV) { 335 uint64_t Padding = DataAlignV - Rem; 336 337 // Append an array of padding bytes to meet alignment requested 338 // Note (o + (a - (o % a)) ) % a == 0 339 // (offset + Padding ) % align == 0 340 341 Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding); 342 LocalVars.push_back(new GlobalVariable( 343 M, ATy, false, GlobalValue::InternalLinkage, UndefValue::get(ATy), 344 "", nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, 345 false)); 346 IsPaddingField.push_back(true); 347 CurrentOffset += Padding; 348 } 349 350 LocalVars.push_back(FGV); 351 IsPaddingField.push_back(false); 352 CurrentOffset += LayoutFields[I].Size; 353 } 354 } 355 356 std::vector<Type *> LocalVarTypes; 357 LocalVarTypes.reserve(LocalVars.size()); 358 std::transform( 359 LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes), 360 [](const GlobalVariable *V) -> Type * { return V->getValueType(); }); 361 362 StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t"); 363 364 Align StructAlign = 365 AMDGPU::getAlign(DL, LocalVars[0]); 366 367 GlobalVariable *SGV = new GlobalVariable( 368 M, LDSTy, false, GlobalValue::InternalLinkage, UndefValue::get(LDSTy), 369 VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, 370 false); 371 SGV->setAlignment(StructAlign); 372 373 DenseMap<GlobalVariable *, Constant *> Map; 374 Type *I32 = Type::getInt32Ty(Ctx); 375 for (size_t I = 0; I < LocalVars.size(); I++) { 376 GlobalVariable *GV = LocalVars[I]; 377 Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)}; 378 Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx, true); 379 if (IsPaddingField[I]) { 380 assert(GV->use_empty()); 381 GV->eraseFromParent(); 382 } else { 383 Map[GV] = GEP; 384 } 385 } 386 assert(Map.size() == LDSVarsToTransform.size()); 387 return std::make_tuple(SGV, std::move(Map)); 388 } 389 390 template <typename PredicateTy> 391 void replaceLDSVariablesWithStruct( 392 Module &M, std::vector<GlobalVariable *> const &LDSVarsToTransform, 393 GlobalVariable *SGV, 394 DenseMap<GlobalVariable *, Constant *> &LDSVarsToConstantGEP, 395 PredicateTy Predicate) { 396 LLVMContext &Ctx = M.getContext(); 397 const DataLayout &DL = M.getDataLayout(); 398 399 // Create alias.scope and their lists. Each field in the new structure 400 // does not alias with all other fields. 401 SmallVector<MDNode *> AliasScopes; 402 SmallVector<Metadata *> NoAliasList; 403 const size_t NumberVars = LDSVarsToTransform.size(); 404 if (NumberVars > 1) { 405 MDBuilder MDB(Ctx); 406 AliasScopes.reserve(NumberVars); 407 MDNode *Domain = MDB.createAnonymousAliasScopeDomain(); 408 for (size_t I = 0; I < NumberVars; I++) { 409 MDNode *Scope = MDB.createAnonymousAliasScope(Domain); 410 AliasScopes.push_back(Scope); 411 } 412 NoAliasList.append(&AliasScopes[1], AliasScopes.end()); 413 } 414 415 // Replace uses of ith variable with a constantexpr to the corresponding 416 // field of the instance that will be allocated by AMDGPUMachineFunction 417 for (size_t I = 0; I < NumberVars; I++) { 418 GlobalVariable *GV = LDSVarsToTransform[I]; 419 Constant *GEP = LDSVarsToConstantGEP[GV]; 420 421 GV->replaceUsesWithIf(GEP, Predicate); 422 if (GV->use_empty()) { 423 GV->eraseFromParent(); 424 } 425 426 APInt APOff(DL.getIndexTypeSizeInBits(GEP->getType()), 0); 427 GEP->stripAndAccumulateInBoundsConstantOffsets(DL, APOff); 428 uint64_t Offset = APOff.getZExtValue(); 429 430 Align A = commonAlignment(SGV->getAlign().valueOrOne(), Offset); 431 432 if (I) 433 NoAliasList[I - 1] = AliasScopes[I - 1]; 434 MDNode *NoAlias = 435 NoAliasList.empty() ? nullptr : MDNode::get(Ctx, NoAliasList); 436 MDNode *AliasScope = 437 AliasScopes.empty() ? nullptr : MDNode::get(Ctx, {AliasScopes[I]}); 438 439 refineUsesAlignmentAndAA(GEP, A, DL, AliasScope, NoAlias); 440 } 441 } 442 443 void refineUsesAlignmentAndAA(Value *Ptr, Align A, const DataLayout &DL, 444 MDNode *AliasScope, MDNode *NoAlias, 445 unsigned MaxDepth = 5) { 446 if (!MaxDepth || (A == 1 && !AliasScope)) 447 return; 448 449 for (User *U : Ptr->users()) { 450 if (auto *I = dyn_cast<Instruction>(U)) { 451 if (AliasScope && I->mayReadOrWriteMemory()) { 452 MDNode *AS = I->getMetadata(LLVMContext::MD_alias_scope); 453 AS = (AS ? MDNode::getMostGenericAliasScope(AS, AliasScope) 454 : AliasScope); 455 I->setMetadata(LLVMContext::MD_alias_scope, AS); 456 457 MDNode *NA = I->getMetadata(LLVMContext::MD_noalias); 458 NA = (NA ? MDNode::intersect(NA, NoAlias) : NoAlias); 459 I->setMetadata(LLVMContext::MD_noalias, NA); 460 } 461 } 462 463 if (auto *LI = dyn_cast<LoadInst>(U)) { 464 LI->setAlignment(std::max(A, LI->getAlign())); 465 continue; 466 } 467 if (auto *SI = dyn_cast<StoreInst>(U)) { 468 if (SI->getPointerOperand() == Ptr) 469 SI->setAlignment(std::max(A, SI->getAlign())); 470 continue; 471 } 472 if (auto *AI = dyn_cast<AtomicRMWInst>(U)) { 473 // None of atomicrmw operations can work on pointers, but let's 474 // check it anyway in case it will or we will process ConstantExpr. 475 if (AI->getPointerOperand() == Ptr) 476 AI->setAlignment(std::max(A, AI->getAlign())); 477 continue; 478 } 479 if (auto *AI = dyn_cast<AtomicCmpXchgInst>(U)) { 480 if (AI->getPointerOperand() == Ptr) 481 AI->setAlignment(std::max(A, AI->getAlign())); 482 continue; 483 } 484 if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) { 485 unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP->getType()); 486 APInt Off(BitWidth, 0); 487 if (GEP->getPointerOperand() == Ptr) { 488 Align GA; 489 if (GEP->accumulateConstantOffset(DL, Off)) 490 GA = commonAlignment(A, Off.getLimitedValue()); 491 refineUsesAlignmentAndAA(GEP, GA, DL, AliasScope, NoAlias, 492 MaxDepth - 1); 493 } 494 continue; 495 } 496 if (auto *I = dyn_cast<Instruction>(U)) { 497 if (I->getOpcode() == Instruction::BitCast || 498 I->getOpcode() == Instruction::AddrSpaceCast) 499 refineUsesAlignmentAndAA(I, A, DL, AliasScope, NoAlias, MaxDepth - 1); 500 } 501 } 502 } 503 }; 504 505 } // namespace 506 char AMDGPULowerModuleLDS::ID = 0; 507 508 char &llvm::AMDGPULowerModuleLDSID = AMDGPULowerModuleLDS::ID; 509 510 INITIALIZE_PASS(AMDGPULowerModuleLDS, DEBUG_TYPE, 511 "Lower uses of LDS variables from non-kernel functions", false, 512 false) 513 514 ModulePass *llvm::createAMDGPULowerModuleLDSPass() { 515 return new AMDGPULowerModuleLDS(); 516 } 517 518 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M, 519 ModuleAnalysisManager &) { 520 return AMDGPULowerModuleLDS().runOnModule(M) ? PreservedAnalyses::none() 521 : PreservedAnalyses::all(); 522 } 523