LLVM 19.0.0git
AMDGPULowerKernelArguments.cpp
Go to the documentation of this file.
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"
17#include "llvm/IR/IRBuilder.h"
18#include "llvm/IR/IntrinsicsAMDGPU.h"
19#include "llvm/IR/MDBuilder.h"
21
22#define DEBUG_TYPE "amdgpu-lower-kernel-arguments"
23
24using namespace llvm;
25
26namespace {
27
28class PreloadKernelArgInfo {
29private:
30 Function &F;
31 const GCNSubtarget &ST;
32 unsigned NumFreeUserSGPRs;
33
34public:
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
69class AMDGPULowerKernelArguments : public FunctionPass {
70public:
71 static char ID;
72
73 AMDGPULowerKernelArguments() : FunctionPass(ID) {}
74
75 bool runOnFunction(Function &F) override;
76
77 void getAnalysisUsage(AnalysisUsage &AU) const override {
79 AU.setPreservesAll();
80 }
81};
82
83} // end anonymous namespace
84
85// skip allocas
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
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 !Arg.getType()->isAggregateType())
149 if (PreloadInfo.tryAllocPreloadSGPRs(AllocSize, EltOffset,
150 LastExplicitArgOffset))
151 continue;
152
153 InPreloadSequence = false;
154
155 if (Arg.use_empty())
156 continue;
157
158 // If this is byval, the loads are already explicit in the function. We just
159 // need to rewrite the pointer values.
160 if (IsByRef) {
161 Value *ArgOffsetPtr = Builder.CreateConstInBoundsGEP1_64(
162 Builder.getInt8Ty(), KernArgSegment, EltOffset,
163 Arg.getName() + ".byval.kernarg.offset");
164
165 Value *CastOffsetPtr =
166 Builder.CreateAddrSpaceCast(ArgOffsetPtr, Arg.getType());
167 Arg.replaceAllUsesWith(CastOffsetPtr);
168 continue;
169 }
170
171 if (PointerType *PT = dyn_cast<PointerType>(ArgTy)) {
172 // FIXME: Hack. We rely on AssertZext to be able to fold DS addressing
173 // modes on SI to know the high bits are 0 so pointer adds don't wrap. We
174 // can't represent this with range metadata because it's only allowed for
175 // integer types.
176 if ((PT->getAddressSpace() == AMDGPUAS::LOCAL_ADDRESS ||
177 PT->getAddressSpace() == AMDGPUAS::REGION_ADDRESS) &&
178 !ST.hasUsableDSOffset())
179 continue;
180
181 // FIXME: We can replace this with equivalent alias.scope/noalias
182 // metadata, but this appears to be a lot of work.
183 if (Arg.hasNoAliasAttr())
184 continue;
185 }
186
187 auto *VT = dyn_cast<FixedVectorType>(ArgTy);
188 bool IsV3 = VT && VT->getNumElements() == 3;
189 bool DoShiftOpt = Size < 32 && !ArgTy->isAggregateType();
190
191 VectorType *V4Ty = nullptr;
192
193 int64_t AlignDownOffset = alignDown(EltOffset, 4);
194 int64_t OffsetDiff = EltOffset - AlignDownOffset;
195 Align AdjustedAlign = commonAlignment(
196 KernArgBaseAlign, DoShiftOpt ? AlignDownOffset : EltOffset);
197
198 Value *ArgPtr;
199 Type *AdjustedArgTy;
200 if (DoShiftOpt) { // FIXME: Handle aggregate types
201 // Since we don't have sub-dword scalar loads, avoid doing an extload by
202 // loading earlier than the argument address, and extracting the relevant
203 // bits.
204 // TODO: Update this for GFX12 which does have scalar sub-dword loads.
205 //
206 // Additionally widen any sub-dword load to i32 even if suitably aligned,
207 // so that CSE between different argument loads works easily.
208 ArgPtr = Builder.CreateConstInBoundsGEP1_64(
209 Builder.getInt8Ty(), KernArgSegment, AlignDownOffset,
210 Arg.getName() + ".kernarg.offset.align.down");
211 AdjustedArgTy = Builder.getInt32Ty();
212 } else {
213 ArgPtr = Builder.CreateConstInBoundsGEP1_64(
214 Builder.getInt8Ty(), KernArgSegment, EltOffset,
215 Arg.getName() + ".kernarg.offset");
216 AdjustedArgTy = ArgTy;
217 }
218
219 if (IsV3 && Size >= 32) {
220 V4Ty = FixedVectorType::get(VT->getElementType(), 4);
221 // Use the hack that clang uses to avoid SelectionDAG ruining v3 loads
222 AdjustedArgTy = V4Ty;
223 }
224
225 LoadInst *Load =
226 Builder.CreateAlignedLoad(AdjustedArgTy, ArgPtr, AdjustedAlign);
227 Load->setMetadata(LLVMContext::MD_invariant_load, MDNode::get(Ctx, {}));
228
229 MDBuilder MDB(Ctx);
230
231 if (isa<PointerType>(ArgTy)) {
232 if (Arg.hasNonNullAttr())
233 Load->setMetadata(LLVMContext::MD_nonnull, MDNode::get(Ctx, {}));
234
235 uint64_t DerefBytes = Arg.getDereferenceableBytes();
236 if (DerefBytes != 0) {
237 Load->setMetadata(
238 LLVMContext::MD_dereferenceable,
239 MDNode::get(Ctx,
240 MDB.createConstant(
241 ConstantInt::get(Builder.getInt64Ty(), DerefBytes))));
242 }
243
244 uint64_t DerefOrNullBytes = Arg.getDereferenceableOrNullBytes();
245 if (DerefOrNullBytes != 0) {
246 Load->setMetadata(
247 LLVMContext::MD_dereferenceable_or_null,
248 MDNode::get(Ctx,
249 MDB.createConstant(ConstantInt::get(Builder.getInt64Ty(),
250 DerefOrNullBytes))));
251 }
252
253 if (MaybeAlign ParamAlign = Arg.getParamAlign()) {
254 Load->setMetadata(
255 LLVMContext::MD_align,
256 MDNode::get(Ctx, MDB.createConstant(ConstantInt::get(
257 Builder.getInt64Ty(), ParamAlign->value()))));
258 }
259 }
260
261 // TODO: Convert noalias arg to !noalias
262
263 if (DoShiftOpt) {
264 Value *ExtractBits = OffsetDiff == 0 ?
265 Load : Builder.CreateLShr(Load, OffsetDiff * 8);
266
267 IntegerType *ArgIntTy = Builder.getIntNTy(Size);
268 Value *Trunc = Builder.CreateTrunc(ExtractBits, ArgIntTy);
269 Value *NewVal = Builder.CreateBitCast(Trunc, ArgTy,
270 Arg.getName() + ".load");
271 Arg.replaceAllUsesWith(NewVal);
272 } else if (IsV3) {
273 Value *Shuf = Builder.CreateShuffleVector(Load, ArrayRef<int>{0, 1, 2},
274 Arg.getName() + ".load");
275 Arg.replaceAllUsesWith(Shuf);
276 } else {
277 Load->setName(Arg.getName() + ".load");
278 Arg.replaceAllUsesWith(Load);
279 }
280 }
281
282 KernArgSegment->addRetAttr(
283 Attribute::getWithAlignment(Ctx, std::max(KernArgBaseAlign, MaxAlign)));
284
285 return true;
286}
287
288bool AMDGPULowerKernelArguments::runOnFunction(Function &F) {
289 auto &TPC = getAnalysis<TargetPassConfig>();
290 const TargetMachine &TM = TPC.getTM<TargetMachine>();
291 return lowerKernelArguments(F, TM);
292}
293
294INITIALIZE_PASS_BEGIN(AMDGPULowerKernelArguments, DEBUG_TYPE,
295 "AMDGPU Lower Kernel Arguments", false, false)
296INITIALIZE_PASS_END(AMDGPULowerKernelArguments, DEBUG_TYPE, "AMDGPU Lower Kernel Arguments",
298
299char AMDGPULowerKernelArguments::ID = 0;
300
302 return new AMDGPULowerKernelArguments();
303}
304
307 bool Changed = lowerKernelArguments(F, TM);
308 if (Changed) {
309 // TODO: Preserves a lot more.
312 return PA;
313 }
314
315 return PreservedAnalyses::all();
316}
MachineBasicBlock MachineBasicBlock::iterator DebugLoc DL
AMDGPU Lower Kernel Arguments
static BasicBlock::iterator getInsertPt(BasicBlock &BB)
static bool lowerKernelArguments(Function &F, const TargetMachine &TM)
#define DEBUG_TYPE
uint64_t Size
AMD GCN specific subclass of TargetSubtarget.
#define F(x, y, z)
Definition: MD5.cpp:55
const char LLVMTargetMachineRef TM
#define INITIALIZE_PASS_END(passName, arg, name, cfg, analysis)
Definition: PassSupport.h:59
#define INITIALIZE_PASS_BEGIN(passName, arg, name, cfg, analysis)
Definition: PassSupport.h:52
Target-Independent Code Generator Pass Configuration Options pass.
PreservedAnalyses run(Function &, FunctionAnalysisManager &)
an instruction to allocate memory on the stack
Definition: Instructions.h:59
bool isStaticAlloca() const
Return true if this alloca is in the entry block of the function and is a constant size.
A container for analyses that lazily runs them and caches their results.
Definition: PassManager.h:348
Represent the analysis usage information of a pass.
AnalysisUsage & addRequired()
void setPreservesAll()
Set by analyses that do not transform their input at all.
This class represents an incoming formal argument to a Function.
Definition: Argument.h:31
ArrayRef - Represent a constant reference to an array (0 or more elements consecutively in memory),...
Definition: ArrayRef.h:41
static Attribute getWithDereferenceableBytes(LLVMContext &Context, uint64_t Bytes)
Definition: Attributes.cpp:204
static Attribute getWithAlignment(LLVMContext &Context, Align Alignment)
Return a uniquified Attribute object that has the specific alignment set.
Definition: Attributes.cpp:194
LLVM Basic Block Representation.
Definition: BasicBlock.h:60
iterator end()
Definition: BasicBlock.h:442
const_iterator getFirstInsertionPt() const
Returns an iterator to the first instruction in this block that is suitable for inserting a non-PHI i...
Definition: BasicBlock.cpp:398
InstListType::iterator iterator
Instruction iterators...
Definition: BasicBlock.h:164
Represents analyses that only rely on functions' control flow.
Definition: Analysis.h:70
void addRetAttr(Attribute::AttrKind Kind)
Adds the attribute to the return value.
Definition: InstrTypes.h:1828
This class represents a function call, abstracting a target machine's calling convention.
A parsed version of the target data layout string in and methods for querying it.
Definition: DataLayout.h:110
static FixedVectorType * get(Type *ElementType, unsigned NumElts)
Definition: Type.cpp:692
FunctionPass class - This class is used to implement most global optimizations.
Definition: Pass.h:311
virtual bool runOnFunction(Function &F)=0
runOnFunction - Virtual method overriden by subclasses to do the per-function processing of the pass.
unsigned getNumUsedUserSGPRs() const
Value * CreateTrunc(Value *V, Type *DestTy, const Twine &Name="")
Definition: IRBuilder.h:2001
IntegerType * getIntNTy(unsigned N)
Fetch the type representing an N-bit integer.
Definition: IRBuilder.h:533
LoadInst * CreateAlignedLoad(Type *Ty, Value *Ptr, MaybeAlign Align, const char *Name)
Definition: IRBuilder.h:1801
CallInst * CreateIntrinsic(Intrinsic::ID ID, ArrayRef< Type * > Types, ArrayRef< Value * > Args, Instruction *FMFSource=nullptr, const Twine &Name="")
Create a call to intrinsic ID with Args, mangled using Types.
Definition: IRBuilder.cpp:932
Value * CreateLShr(Value *LHS, Value *RHS, const Twine &Name="", bool isExact=false)
Definition: IRBuilder.h:1431
IntegerType * getInt32Ty()
Fetch the type representing a 32-bit integer.
Definition: IRBuilder.h:520
IntegerType * getInt64Ty()
Fetch the type representing a 64-bit integer.
Definition: IRBuilder.h:525
Value * CreateBitCast(Value *V, Type *DestTy, const Twine &Name="")
Definition: IRBuilder.h:2105
Value * CreateShuffleVector(Value *V1, Value *V2, Value *Mask, const Twine &Name="")
Definition: IRBuilder.h:2472
Value * CreateConstInBoundsGEP1_64(Type *Ty, Value *Ptr, uint64_t Idx0, const Twine &Name="")
Definition: IRBuilder.h:1931
IntegerType * getInt8Ty()
Fetch the type representing an 8-bit integer.
Definition: IRBuilder.h:510
Value * CreateAddrSpaceCast(Value *V, Type *DestTy, const Twine &Name="")
Definition: IRBuilder.h:2110
This provides a uniform API for creating instructions and inserting them into a basic block: either a...
Definition: IRBuilder.h:2644
Class to represent integer types.
Definition: DerivedTypes.h:40
This is an important class for using LLVM in a threaded context.
Definition: LLVMContext.h:67
An instruction for reading from memory.
Definition: Instructions.h:184
ConstantAsMetadata * createConstant(Constant *C)
Return the given constant as metadata.
Definition: MDBuilder.cpp:24
static MDTuple * get(LLVMContext &Context, ArrayRef< Metadata * > MDs)
Definition: Metadata.h:1541
virtual void getAnalysisUsage(AnalysisUsage &) const
getAnalysisUsage - This function should be overriden by passes that need analysis information to do t...
Definition: Pass.cpp:98
A set of analyses that are preserved following a run of a transformation pass.
Definition: Analysis.h:109
static PreservedAnalyses all()
Construct a special preserved set that preserves all passes.
Definition: Analysis.h:115
void preserveSet()
Mark an analysis set as preserved.
Definition: Analysis.h:144
This is a 'vector' (really, a variable-sized array), optimized for the case when the array is small.
Definition: SmallVector.h:1209
Primary interface to the complete machine description for the target machine.
Definition: TargetMachine.h:76
Target-Independent Code Generator Pass Configuration Options.
The instances of the Type class are immutable: once they are created, they are never changed.
Definition: Type.h:45
bool isAggregateType() const
Return true if the type is an aggregate type.
Definition: Type.h:295
LLVM Value Representation.
Definition: Value.h:74
Type * getType() const
All values are typed, get the type of this value.
Definition: Value.h:255
void replaceAllUsesWith(Value *V)
Change all uses of this to point to a new Value.
Definition: Value.cpp:534
StringRef getName() const
Return a constant reference to the value's name.
Definition: Value.cpp:309
@ REGION_ADDRESS
Address space for region memory. (GDS)
@ LOCAL_ADDRESS
Address space for local memory.
unsigned ID
LLVM IR allows to use arbitrary numbers as calling convention identifiers.
Definition: CallingConv.h:24
@ AMDGPU_KERNEL
Used for AMDGPU code object kernels.
Definition: CallingConv.h:200
This is an optimization pass for GlobalISel generic memory operations.
Definition: AddressRanges.h:18
bool isAligned(Align Lhs, uint64_t SizeInBytes)
Checks that SizeInBytes is a multiple of the alignment.
Definition: Alignment.h:145
FunctionPass * createAMDGPULowerKernelArgumentsPass()
uint64_t alignTo(uint64_t Size, Align A)
Returns a multiple of A needed to store Size bytes.
Definition: Alignment.h:155
Align commonAlignment(Align A, uint64_t Offset)
Returns the alignment that satisfies both alignments.
Definition: Alignment.h:212
uint64_t alignDown(uint64_t Value, uint64_t Align, uint64_t Skew=0)
Returns the largest uint64_t less than or equal to Value and is Skew mod Align.
Definition: MathExtras.h:428
This struct is a compact representation of a valid (non-zero power of two) alignment.
Definition: Alignment.h:39
This struct is a compact representation of a valid (power of two) or undefined (0) alignment.
Definition: Alignment.h:117