1 //===-- AMDGPULowerKernelArguments.cpp ------------------------------------------===// 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 /// \file This pass replaces accesses to kernel arguments with loads from 10 /// offsets from the kernarg base pointer. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "AMDGPU.h" 15 #include "GCNSubtarget.h" 16 #include "llvm/CodeGen/TargetPassConfig.h" 17 #include "llvm/IR/IRBuilder.h" 18 #include "llvm/IR/IntrinsicsAMDGPU.h" 19 #include "llvm/IR/MDBuilder.h" 20 #include "llvm/Target/TargetMachine.h" 21 22 #define DEBUG_TYPE "amdgpu-lower-kernel-arguments" 23 24 using namespace llvm; 25 26 namespace { 27 28 class PreloadKernelArgInfo { 29 private: 30 Function &F; 31 const GCNSubtarget &ST; 32 unsigned NumFreeUserSGPRs; 33 34 public: 35 SmallVector<llvm::Metadata *, 8> KernelArgMetadata; 36 37 PreloadKernelArgInfo(Function &F, const GCNSubtarget &ST) : F(F), ST(ST) { 38 setInitialFreeUserSGPRsCount(); 39 } 40 41 // Returns the maximum number of user SGPRs that we have available to preload 42 // arguments. 43 void setInitialFreeUserSGPRsCount() { 44 const unsigned MaxUserSGPRs = ST.getMaxNumUserSGPRs(); 45 GCNUserSGPRUsageInfo UserSGPRInfo(F, ST); 46 47 NumFreeUserSGPRs = MaxUserSGPRs - UserSGPRInfo.getNumUsedUserSGPRs(); 48 } 49 50 bool tryAllocPreloadSGPRs(unsigned AllocSize, uint64_t ArgOffset, 51 uint64_t LastExplicitArgOffset) { 52 // Check if this argument may be loaded into the same register as the 53 // previous argument. 54 if (!isAligned(Align(4), ArgOffset) && AllocSize < 4) 55 return true; 56 57 // Pad SGPRs for kernarg alignment. 58 unsigned Padding = ArgOffset - LastExplicitArgOffset; 59 unsigned PaddingSGPRs = alignTo(Padding, 4) / 4; 60 unsigned NumPreloadSGPRs = alignTo(AllocSize, 4) / 4; 61 if (NumPreloadSGPRs + PaddingSGPRs > NumFreeUserSGPRs) 62 return false; 63 64 NumFreeUserSGPRs -= (NumPreloadSGPRs + PaddingSGPRs); 65 return true; 66 } 67 }; 68 69 class AMDGPULowerKernelArguments : public FunctionPass { 70 public: 71 static char ID; 72 73 AMDGPULowerKernelArguments() : FunctionPass(ID) {} 74 75 bool runOnFunction(Function &F) override; 76 77 void getAnalysisUsage(AnalysisUsage &AU) const override { 78 AU.addRequired<TargetPassConfig>(); 79 AU.setPreservesAll(); 80 } 81 }; 82 83 } // end anonymous namespace 84 85 // skip allocas 86 static BasicBlock::iterator getInsertPt(BasicBlock &BB) { 87 BasicBlock::iterator InsPt = BB.getFirstInsertionPt(); 88 for (BasicBlock::iterator E = BB.end(); InsPt != E; ++InsPt) { 89 AllocaInst *AI = dyn_cast<AllocaInst>(&*InsPt); 90 91 // If this is a dynamic alloca, the value may depend on the loaded kernargs, 92 // so loads will need to be inserted before it. 93 if (!AI || !AI->isStaticAlloca()) 94 break; 95 } 96 97 return InsPt; 98 } 99 100 static bool lowerKernelArguments(Function &F, const TargetMachine &TM) { 101 CallingConv::ID CC = F.getCallingConv(); 102 if (CC != CallingConv::AMDGPU_KERNEL || F.arg_empty()) 103 return false; 104 105 const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F); 106 LLVMContext &Ctx = F.getParent()->getContext(); 107 const DataLayout &DL = F.getParent()->getDataLayout(); 108 BasicBlock &EntryBlock = *F.begin(); 109 IRBuilder<> Builder(&EntryBlock, getInsertPt(EntryBlock)); 110 111 const Align KernArgBaseAlign(16); // FIXME: Increase if necessary 112 const uint64_t BaseOffset = ST.getExplicitKernelArgOffset(); 113 114 Align MaxAlign; 115 // FIXME: Alignment is broken with explicit arg offset.; 116 const uint64_t TotalKernArgSize = ST.getKernArgSegmentSize(F, MaxAlign); 117 if (TotalKernArgSize == 0) 118 return false; 119 120 CallInst *KernArgSegment = 121 Builder.CreateIntrinsic(Intrinsic::amdgcn_kernarg_segment_ptr, {}, {}, 122 nullptr, F.getName() + ".kernarg.segment"); 123 124 KernArgSegment->addRetAttr(Attribute::NonNull); 125 KernArgSegment->addRetAttr( 126 Attribute::getWithDereferenceableBytes(Ctx, TotalKernArgSize)); 127 128 uint64_t ExplicitArgOffset = 0; 129 // Preloaded kernel arguments must be sequential. 130 bool InPreloadSequence = true; 131 PreloadKernelArgInfo PreloadInfo(F, ST); 132 133 for (Argument &Arg : F.args()) { 134 const bool IsByRef = Arg.hasByRefAttr(); 135 Type *ArgTy = IsByRef ? Arg.getParamByRefType() : Arg.getType(); 136 MaybeAlign ParamAlign = IsByRef ? Arg.getParamAlign() : std::nullopt; 137 Align ABITypeAlign = DL.getValueOrABITypeAlignment(ParamAlign, ArgTy); 138 139 uint64_t Size = DL.getTypeSizeInBits(ArgTy); 140 uint64_t AllocSize = DL.getTypeAllocSize(ArgTy); 141 142 uint64_t EltOffset = alignTo(ExplicitArgOffset, ABITypeAlign) + BaseOffset; 143 uint64_t LastExplicitArgOffset = ExplicitArgOffset; 144 ExplicitArgOffset = alignTo(ExplicitArgOffset, ABITypeAlign) + AllocSize; 145 146 // Try to preload this argument into user SGPRs. 147 if (Arg.hasInRegAttr() && InPreloadSequence && ST.hasKernargPreload() && 148 !ST.needsKernargPreloadBackwardsCompatibility() && 149 !Arg.getType()->isAggregateType()) 150 if (PreloadInfo.tryAllocPreloadSGPRs(AllocSize, EltOffset, 151 LastExplicitArgOffset)) 152 continue; 153 154 InPreloadSequence = false; 155 156 if (Arg.use_empty()) 157 continue; 158 159 // If this is byval, the loads are already explicit in the function. We just 160 // need to rewrite the pointer values. 161 if (IsByRef) { 162 Value *ArgOffsetPtr = Builder.CreateConstInBoundsGEP1_64( 163 Builder.getInt8Ty(), KernArgSegment, EltOffset, 164 Arg.getName() + ".byval.kernarg.offset"); 165 166 Value *CastOffsetPtr = 167 Builder.CreateAddrSpaceCast(ArgOffsetPtr, Arg.getType()); 168 Arg.replaceAllUsesWith(CastOffsetPtr); 169 continue; 170 } 171 172 if (PointerType *PT = dyn_cast<PointerType>(ArgTy)) { 173 // FIXME: Hack. We rely on AssertZext to be able to fold DS addressing 174 // modes on SI to know the high bits are 0 so pointer adds don't wrap. We 175 // can't represent this with range metadata because it's only allowed for 176 // integer types. 177 if ((PT->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS || 178 PT->getAddressSpace() == AMDGPUAS::REGION_ADDRESS) && 179 !ST.hasUsableDSOffset()) 180 continue; 181 182 // FIXME: We can replace this with equivalent alias.scope/noalias 183 // metadata, but this appears to be a lot of work. 184 if (Arg.hasNoAliasAttr()) 185 continue; 186 } 187 188 auto *VT = dyn_cast<FixedVectorType>(ArgTy); 189 bool IsV3 = VT && VT->getNumElements() == 3; 190 bool DoShiftOpt = Size < 32 && !ArgTy->isAggregateType(); 191 192 VectorType *V4Ty = nullptr; 193 194 int64_t AlignDownOffset = alignDown(EltOffset, 4); 195 int64_t OffsetDiff = EltOffset - AlignDownOffset; 196 Align AdjustedAlign = commonAlignment( 197 KernArgBaseAlign, DoShiftOpt ? AlignDownOffset : EltOffset); 198 199 Value *ArgPtr; 200 Type *AdjustedArgTy; 201 if (DoShiftOpt) { // FIXME: Handle aggregate types 202 // Since we don't have sub-dword scalar loads, avoid doing an extload by 203 // loading earlier than the argument address, and extracting the relevant 204 // bits. 205 // TODO: Update this for GFX12 which does have scalar sub-dword loads. 206 // 207 // Additionally widen any sub-dword load to i32 even if suitably aligned, 208 // so that CSE between different argument loads works easily. 209 ArgPtr = Builder.CreateConstInBoundsGEP1_64( 210 Builder.getInt8Ty(), KernArgSegment, AlignDownOffset, 211 Arg.getName() + ".kernarg.offset.align.down"); 212 AdjustedArgTy = Builder.getInt32Ty(); 213 } else { 214 ArgPtr = Builder.CreateConstInBoundsGEP1_64( 215 Builder.getInt8Ty(), KernArgSegment, EltOffset, 216 Arg.getName() + ".kernarg.offset"); 217 AdjustedArgTy = ArgTy; 218 } 219 220 if (IsV3 && Size >= 32) { 221 V4Ty = FixedVectorType::get(VT->getElementType(), 4); 222 // Use the hack that clang uses to avoid SelectionDAG ruining v3 loads 223 AdjustedArgTy = V4Ty; 224 } 225 226 LoadInst *Load = 227 Builder.CreateAlignedLoad(AdjustedArgTy, ArgPtr, AdjustedAlign); 228 Load->setMetadata(LLVMContext::MD_invariant_load, MDNode::get(Ctx, {})); 229 230 MDBuilder MDB(Ctx); 231 232 if (isa<PointerType>(ArgTy)) { 233 if (Arg.hasNonNullAttr()) 234 Load->setMetadata(LLVMContext::MD_nonnull, MDNode::get(Ctx, {})); 235 236 uint64_t DerefBytes = Arg.getDereferenceableBytes(); 237 if (DerefBytes != 0) { 238 Load->setMetadata( 239 LLVMContext::MD_dereferenceable, 240 MDNode::get(Ctx, 241 MDB.createConstant( 242 ConstantInt::get(Builder.getInt64Ty(), DerefBytes)))); 243 } 244 245 uint64_t DerefOrNullBytes = Arg.getDereferenceableOrNullBytes(); 246 if (DerefOrNullBytes != 0) { 247 Load->setMetadata( 248 LLVMContext::MD_dereferenceable_or_null, 249 MDNode::get(Ctx, 250 MDB.createConstant(ConstantInt::get(Builder.getInt64Ty(), 251 DerefOrNullBytes)))); 252 } 253 254 if (MaybeAlign ParamAlign = Arg.getParamAlign()) { 255 Load->setMetadata( 256 LLVMContext::MD_align, 257 MDNode::get(Ctx, MDB.createConstant(ConstantInt::get( 258 Builder.getInt64Ty(), ParamAlign->value())))); 259 } 260 } 261 262 // TODO: Convert noalias arg to !noalias 263 264 if (DoShiftOpt) { 265 Value *ExtractBits = OffsetDiff == 0 ? 266 Load : Builder.CreateLShr(Load, OffsetDiff * 8); 267 268 IntegerType *ArgIntTy = Builder.getIntNTy(Size); 269 Value *Trunc = Builder.CreateTrunc(ExtractBits, ArgIntTy); 270 Value *NewVal = Builder.CreateBitCast(Trunc, ArgTy, 271 Arg.getName() + ".load"); 272 Arg.replaceAllUsesWith(NewVal); 273 } else if (IsV3) { 274 Value *Shuf = Builder.CreateShuffleVector(Load, ArrayRef<int>{0, 1, 2}, 275 Arg.getName() + ".load"); 276 Arg.replaceAllUsesWith(Shuf); 277 } else { 278 Load->setName(Arg.getName() + ".load"); 279 Arg.replaceAllUsesWith(Load); 280 } 281 } 282 283 KernArgSegment->addRetAttr( 284 Attribute::getWithAlignment(Ctx, std::max(KernArgBaseAlign, MaxAlign))); 285 286 return true; 287 } 288 289 bool AMDGPULowerKernelArguments::runOnFunction(Function &F) { 290 auto &TPC = getAnalysis<TargetPassConfig>(); 291 const TargetMachine &TM = TPC.getTM<TargetMachine>(); 292 return lowerKernelArguments(F, TM); 293 } 294 295 INITIALIZE_PASS_BEGIN(AMDGPULowerKernelArguments, DEBUG_TYPE, 296 "AMDGPU Lower Kernel Arguments", false, false) 297 INITIALIZE_PASS_END(AMDGPULowerKernelArguments, DEBUG_TYPE, "AMDGPU Lower Kernel Arguments", 298 false, false) 299 300 char AMDGPULowerKernelArguments::ID = 0; 301 302 FunctionPass *llvm::createAMDGPULowerKernelArgumentsPass() { 303 return new AMDGPULowerKernelArguments(); 304 } 305 306 PreservedAnalyses 307 AMDGPULowerKernelArgumentsPass::run(Function &F, FunctionAnalysisManager &AM) { 308 bool Changed = lowerKernelArguments(F, TM); 309 if (Changed) { 310 // TODO: Preserves a lot more. 311 PreservedAnalyses PA; 312 PA.preserveSet<CFGAnalyses>(); 313 return PA; 314 } 315 316 return PreservedAnalyses::all(); 317 } 318