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 simplication 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/AMDGPULDSUtils.h" 32 #include "llvm/ADT/STLExtras.h" 33 #include "llvm/IR/Constants.h" 34 #include "llvm/IR/DerivedTypes.h" 35 #include "llvm/IR/IRBuilder.h" 36 #include "llvm/IR/InlineAsm.h" 37 #include "llvm/IR/Instructions.h" 38 #include "llvm/InitializePasses.h" 39 #include "llvm/Pass.h" 40 #include "llvm/Support/Debug.h" 41 #include "llvm/Transforms/Utils/ModuleUtils.h" 42 #include <algorithm> 43 #include <vector> 44 45 #define DEBUG_TYPE "amdgpu-lower-module-lds" 46 47 using namespace llvm; 48 49 namespace { 50 51 class AMDGPULowerModuleLDS : public ModulePass { 52 53 static void removeFromUsedList(Module &M, StringRef Name, 54 SmallPtrSetImpl<Constant *> &ToRemove) { 55 GlobalVariable *GV = M.getNamedGlobal(Name); 56 if (!GV || ToRemove.empty()) { 57 return; 58 } 59 60 SmallVector<Constant *, 16> Init; 61 auto *CA = cast<ConstantArray>(GV->getInitializer()); 62 for (auto &Op : CA->operands()) { 63 // ModuleUtils::appendToUsed only inserts Constants 64 Constant *C = cast<Constant>(Op); 65 if (!ToRemove.contains(C->stripPointerCasts())) { 66 Init.push_back(C); 67 } 68 } 69 70 if (Init.size() == CA->getNumOperands()) { 71 return; // none to remove 72 } 73 74 GV->eraseFromParent(); 75 76 for (Constant *C : ToRemove) { 77 C->removeDeadConstantUsers(); 78 } 79 80 if (!Init.empty()) { 81 ArrayType *ATy = 82 ArrayType::get(Type::getInt8PtrTy(M.getContext()), Init.size()); 83 GV = 84 new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage, 85 ConstantArray::get(ATy, Init), Name); 86 GV->setSection("llvm.metadata"); 87 } 88 } 89 90 static void 91 removeFromUsedLists(Module &M, 92 const std::vector<GlobalVariable *> &LocalVars) { 93 SmallPtrSet<Constant *, 32> LocalVarsSet; 94 for (size_t I = 0; I < LocalVars.size(); I++) { 95 if (Constant *C = dyn_cast<Constant>(LocalVars[I]->stripPointerCasts())) { 96 LocalVarsSet.insert(C); 97 } 98 } 99 removeFromUsedList(M, "llvm.used", LocalVarsSet); 100 removeFromUsedList(M, "llvm.compiler.used", LocalVarsSet); 101 } 102 103 static void markUsedByKernel(IRBuilder<> &Builder, Function *Func, 104 GlobalVariable *SGV) { 105 // The llvm.amdgcn.module.lds instance is implicitly used by all kernels 106 // that might call a function which accesses a field within it. This is 107 // presently approximated to 'all kernels' if there are any such functions 108 // in the module. This implicit use is reified as an explicit use here so 109 // that later passes, specifically PromoteAlloca, account for the required 110 // memory without any knowledge of this transform. 111 112 // An operand bundle on llvm.donothing works because the call instruction 113 // survives until after the last pass that needs to account for LDS. It is 114 // better than inline asm as the latter survives until the end of codegen. A 115 // totally robust solution would be a function with the same semantics as 116 // llvm.donothing that takes a pointer to the instance and is lowered to a 117 // no-op after LDS is allocated, but that is not presently necessary. 118 119 LLVMContext &Ctx = Func->getContext(); 120 121 Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI()); 122 123 FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {}); 124 125 Function *Decl = 126 Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {}); 127 128 Value *UseInstance[1] = {Builder.CreateInBoundsGEP( 129 SGV->getValueType(), SGV, ConstantInt::get(Type::getInt32Ty(Ctx), 0))}; 130 131 Builder.CreateCall(FTy, Decl, {}, 132 {OperandBundleDefT<Value *>("ExplicitUse", UseInstance)}, 133 ""); 134 } 135 136 private: 137 SmallPtrSet<GlobalValue *, 32> UsedList; 138 139 public: 140 static char ID; 141 142 AMDGPULowerModuleLDS() : ModulePass(ID) { 143 initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry()); 144 } 145 146 bool runOnModule(Module &M) override { 147 UsedList = AMDGPU::getUsedList(M); 148 149 bool Changed = processUsedLDS(M); 150 151 for (Function &F : M.functions()) { 152 if (!AMDGPU::isKernelCC(&F)) 153 continue; 154 Changed |= processUsedLDS(M, &F); 155 } 156 157 UsedList.clear(); 158 return Changed; 159 } 160 161 private: 162 bool processUsedLDS(Module &M, Function *F = nullptr) { 163 LLVMContext &Ctx = M.getContext(); 164 const DataLayout &DL = M.getDataLayout(); 165 166 // Find variables to move into new struct instance 167 std::vector<GlobalVariable *> FoundLocalVars = 168 AMDGPU::findVariablesToLower(M, UsedList, F); 169 170 if (FoundLocalVars.empty()) { 171 // No variables to rewrite, no changes made. 172 return false; 173 } 174 175 // Increase the alignment of LDS globals if necessary to maximise the chance 176 // that we can use aligned LDS instructions to access them. 177 for (auto *GV : FoundLocalVars) { 178 Align Alignment(GV->getAlignment()); 179 TypeSize GVSize = DL.getTypeAllocSize(GV->getValueType()); 180 181 if (GVSize > 8) { 182 // We might want to use a b96 or b128 load/store 183 Alignment = std::max(Alignment, Align(16)); 184 } else if (GVSize > 4) { 185 // We might want to use a b64 load/store 186 Alignment = std::max(Alignment, Align(8)); 187 } else if (GVSize > 2) { 188 // We might want to use a b32 load/store 189 Alignment = std::max(Alignment, Align(4)); 190 } else if (GVSize > 1) { 191 // We might want to use a b16 load/store 192 Alignment = std::max(Alignment, Align(2)); 193 } 194 195 GV->setAlignment(Alignment); 196 } 197 198 // Sort by alignment, descending, to minimise padding. 199 // On ties, sort by size, descending, then by name, lexicographical. 200 llvm::stable_sort( 201 FoundLocalVars, 202 [&](const GlobalVariable *LHS, const GlobalVariable *RHS) -> bool { 203 Align ALHS = AMDGPU::getAlign(DL, LHS); 204 Align ARHS = AMDGPU::getAlign(DL, RHS); 205 if (ALHS != ARHS) { 206 return ALHS > ARHS; 207 } 208 209 TypeSize SLHS = DL.getTypeAllocSize(LHS->getValueType()); 210 TypeSize SRHS = DL.getTypeAllocSize(RHS->getValueType()); 211 if (SLHS != SRHS) { 212 return SLHS > SRHS; 213 } 214 215 // By variable name on tie for predictable order in test cases. 216 return LHS->getName() < RHS->getName(); 217 }); 218 219 std::vector<GlobalVariable *> LocalVars; 220 LocalVars.reserve(FoundLocalVars.size()); // will be at least this large 221 { 222 // This usually won't need to insert any padding, perhaps avoid the alloc 223 uint64_t CurrentOffset = 0; 224 for (size_t I = 0; I < FoundLocalVars.size(); I++) { 225 GlobalVariable *FGV = FoundLocalVars[I]; 226 Align DataAlign = AMDGPU::getAlign(DL, FGV); 227 228 uint64_t DataAlignV = DataAlign.value(); 229 if (uint64_t Rem = CurrentOffset % DataAlignV) { 230 uint64_t Padding = DataAlignV - Rem; 231 232 // Append an array of padding bytes to meet alignment requested 233 // Note (o + (a - (o % a)) ) % a == 0 234 // (offset + Padding ) % align == 0 235 236 Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding); 237 LocalVars.push_back(new GlobalVariable( 238 M, ATy, false, GlobalValue::InternalLinkage, UndefValue::get(ATy), 239 "", nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, 240 false)); 241 CurrentOffset += Padding; 242 } 243 244 LocalVars.push_back(FGV); 245 CurrentOffset += DL.getTypeAllocSize(FGV->getValueType()); 246 } 247 } 248 249 std::vector<Type *> LocalVarTypes; 250 LocalVarTypes.reserve(LocalVars.size()); 251 std::transform( 252 LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes), 253 [](const GlobalVariable *V) -> Type * { return V->getValueType(); }); 254 255 std::string VarName( 256 F ? (Twine("llvm.amdgcn.kernel.") + F->getName() + ".lds").str() 257 : "llvm.amdgcn.module.lds"); 258 StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t"); 259 260 Align MaxAlign = 261 AMDGPU::getAlign(DL, LocalVars[0]); // was sorted on alignment 262 263 GlobalVariable *SGV = new GlobalVariable( 264 M, LDSTy, false, GlobalValue::InternalLinkage, UndefValue::get(LDSTy), 265 VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, 266 false); 267 SGV->setAlignment(MaxAlign); 268 if (!F) { 269 appendToCompilerUsed( 270 M, {static_cast<GlobalValue *>( 271 ConstantExpr::getPointerBitCastOrAddrSpaceCast( 272 cast<Constant>(SGV), Type::getInt8PtrTy(Ctx)))}); 273 } 274 275 // The verifier rejects used lists containing an inttoptr of a constant 276 // so remove the variables from these lists before replaceAllUsesWith 277 removeFromUsedLists(M, LocalVars); 278 279 // Replace uses of ith variable with a constantexpr to the ith field of the 280 // instance that will be allocated by AMDGPUMachineFunction 281 Type *I32 = Type::getInt32Ty(Ctx); 282 for (size_t I = 0; I < LocalVars.size(); I++) { 283 GlobalVariable *GV = LocalVars[I]; 284 Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)}; 285 Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx); 286 if (F) { 287 GV->replaceUsesWithIf(GEP, [F](Use &U) { 288 return AMDGPU::isUsedOnlyFromFunction(U.getUser(), F); 289 }); 290 } else { 291 GV->replaceAllUsesWith(GEP); 292 } 293 if (GV->use_empty()) { 294 UsedList.erase(GV); 295 GV->eraseFromParent(); 296 } 297 } 298 299 // Mark kernels with asm that reads the address of the allocated structure 300 // This is not necessary for lowering. This lets other passes, specifically 301 // PromoteAlloca, accurately calculate how much LDS will be used by the 302 // kernel after lowering. 303 if (!F) { 304 IRBuilder<> Builder(Ctx); 305 SmallPtrSet<Function *, 32> Kernels; 306 for (auto &I : M.functions()) { 307 Function *Func = &I; 308 if (AMDGPU::isKernelCC(Func) && !Kernels.contains(Func)) { 309 markUsedByKernel(Builder, Func, SGV); 310 Kernels.insert(Func); 311 } 312 } 313 } 314 return true; 315 } 316 }; 317 318 } // namespace 319 char AMDGPULowerModuleLDS::ID = 0; 320 321 char &llvm::AMDGPULowerModuleLDSID = AMDGPULowerModuleLDS::ID; 322 323 INITIALIZE_PASS(AMDGPULowerModuleLDS, DEBUG_TYPE, 324 "Lower uses of LDS variables from non-kernel functions", false, 325 false) 326 327 ModulePass *llvm::createAMDGPULowerModuleLDSPass() { 328 return new AMDGPULowerModuleLDS(); 329 } 330 331 PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M, 332 ModuleAnalysisManager &) { 333 return AMDGPULowerModuleLDS().runOnModule(M) ? PreservedAnalyses::none() 334 : PreservedAnalyses::all(); 335 } 336