Skip to content

[AMDGPU] Enable serializing of allocated preload kernarg SGPRs info #168374

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Nov 22, 2025
+382 −2
Merged

[AMDGPU] Enable serializing of allocated preload kernarg SGPRs info #168374

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
63f624e
[AMDGPU] Enable serializing of allocated preload kernarg SGPRs info
tyb0807 Nov 17, 2025
13a6ac7
Address comments
tyb0807 Nov 19, 2025
d6b6f6a
More tests + address comments
tyb0807 Nov 20, 2025
2272203
Format + review comments
tyb0807 Nov 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2027,6 +2027,35 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
MFI->ArgInfo.WorkItemIDZ, 0, 0)))
return true;

// Parse FirstKernArgPreloadReg separately, since it's a Register,
// not ArgDescriptor.
if (YamlMFI.ArgInfo && YamlMFI.ArgInfo->FirstKernArgPreloadReg) {
const auto &A = *YamlMFI.ArgInfo->FirstKernArgPreloadReg;

if (!A.IsRegister) {
const MemoryBuffer &Buffer =
*PFS.SM->getMemoryBuffer(PFS.SM->getMainFileID());
Error =
SMDiagnostic(*PFS.SM, SMLoc(), Buffer.getBufferIdentifier(), 1, 0,
SourceMgr::DK_Error,
"firstKernArgPreloadReg must be a register", "", {}, {});
return true;
}

Register Reg;
if (parseNamedRegisterReference(PFS, Reg, A.RegisterName.Value, Error)) {
SourceRange = A.RegisterName.SourceRange;
return true;
}

if (!AMDGPU::SGPR_32RegClass.contains(Reg))
return diagnoseRegisterClass(A.RegisterName);

MFI->ArgInfo.FirstKernArgPreloadReg = Reg;

MFI->NumUserSGPRs += YamlMFI.NumKernargPreloadSGPRs;
}

if (ST.hasIEEEMode())
MFI->Mode.IEEE = YamlMFI.Mode.IEEE;
if (ST.hasDX10ClampMode())
Expand Down
20 changes: 18 additions & 2 deletions llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,7 +696,6 @@ convertArgumentInfo(const AMDGPUFunctionArgInfo &ArgInfo,
return true;
};

// TODO: Need to serialize kernarg preloads.
bool Any = false;
Any |= convertArg(AI.PrivateSegmentBuffer, ArgInfo.PrivateSegmentBuffer);
Any |= convertArg(AI.DispatchPtr, ArgInfo.DispatchPtr);
Expand All @@ -718,6 +717,20 @@ convertArgumentInfo(const AMDGPUFunctionArgInfo &ArgInfo,
Any |= convertArg(AI.WorkItemIDY, ArgInfo.WorkItemIDY);
Any |= convertArg(AI.WorkItemIDZ, ArgInfo.WorkItemIDZ);

// Write FirstKernArgPreloadReg separately, since it's a Register,
// not ArgDescriptor.
if (ArgInfo.FirstKernArgPreloadReg) {
Register Reg = ArgInfo.FirstKernArgPreloadReg;
if (Reg.isPhysical()) {
yaml::SIArgument SA = yaml::SIArgument::createArgument(true);
raw_string_ostream OS(SA.RegisterName.Value);
OS << printReg(Reg, &TRI);
Copy link
Collaborator

@qcolombet qcolombet Nov 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reuse convertArg here instead of inlining the logic.

Copy link
Contributor Author

@tyb0807 tyb0807 Nov 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we can because ArgInfo.FirstKernArgPreloadReg is a Register and not ArgDescriptor


AI.FirstKernArgPreloadReg = SA;
Any = true;
}
}

if (Any)
return AI;

Expand Down Expand Up @@ -750,7 +763,8 @@ yaml::SIMachineFunctionInfo::SIMachineFunctionInfo(
Mode(MFI.getMode()), HasInitWholeWave(MFI.hasInitWholeWave()),
IsWholeWaveFunction(MFI.isWholeWaveFunction()),
DynamicVGPRBlockSize(MFI.getDynamicVGPRBlockSize()),
ScratchReservedForDynamicVGPRs(MFI.getScratchReservedForDynamicVGPRs()) {
ScratchReservedForDynamicVGPRs(MFI.getScratchReservedForDynamicVGPRs()),
NumKernargPreloadSGPRs(MFI.getNumKernargPreloadedSGPRs()) {
for (Register Reg : MFI.getSGPRSpillPhysVGPRs())
SpillPhysVGPRS.push_back(regToString(Reg, TRI));

Expand Down Expand Up @@ -799,6 +813,8 @@ bool SIMachineFunctionInfo::initializeBaseYamlFields(
ReturnsVoid = YamlMFI.ReturnsVoid;
IsWholeWaveFunction = YamlMFI.IsWholeWaveFunction;

UserSGPRInfo.allocKernargPreloadSGPRs(YamlMFI.NumKernargPreloadSGPRs);

if (YamlMFI.ScavengeFI) {
auto FIOrErr = YamlMFI.ScavengeFI->getFI(MF.getFrameInfo());
if (!FIOrErr) {
Expand Down
5 changes: 5 additions & 0 deletions llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,7 @@ struct SIArgumentInfo {
std::optional<SIArgument> DispatchID;
std::optional<SIArgument> FlatScratchInit;
std::optional<SIArgument> PrivateSegmentSize;
std::optional<SIArgument> FirstKernArgPreloadReg;

std::optional<SIArgument> WorkGroupIDX;
std::optional<SIArgument> WorkGroupIDY;
Expand All @@ -195,6 +196,7 @@ template <> struct MappingTraits<SIArgumentInfo> {
YamlIO.mapOptional("dispatchID", AI.DispatchID);
YamlIO.mapOptional("flatScratchInit", AI.FlatScratchInit);
YamlIO.mapOptional("privateSegmentSize", AI.PrivateSegmentSize);
YamlIO.mapOptional("firstKernArgPreloadReg", AI.FirstKernArgPreloadReg);

YamlIO.mapOptional("workGroupIDX", AI.WorkGroupIDX);
YamlIO.mapOptional("workGroupIDY", AI.WorkGroupIDY);
Expand Down Expand Up @@ -305,6 +307,8 @@ struct SIMachineFunctionInfo final : public yaml::MachineFunctionInfo {
unsigned DynamicVGPRBlockSize = 0;
unsigned ScratchReservedForDynamicVGPRs = 0;

unsigned NumKernargPreloadSGPRs = 0;

SIMachineFunctionInfo() = default;
SIMachineFunctionInfo(const llvm::SIMachineFunctionInfo &,
const TargetRegisterInfo &TRI,
Expand Down Expand Up @@ -361,6 +365,7 @@ template <> struct MappingTraits<SIMachineFunctionInfo> {
YamlIO.mapOptional("dynamicVGPRBlockSize", MFI.DynamicVGPRBlockSize, false);
YamlIO.mapOptional("scratchReservedForDynamicVGPRs",
MFI.ScratchReservedForDynamicVGPRs, 0);
YamlIO.mapOptional("numKernargPreloadSGPRs", MFI.NumKernargPreloadSGPRs, 0);
YamlIO.mapOptional("isWholeWaveFunction", MFI.IsWholeWaveFunction, false);
}
};
Expand Down
2 changes: 2 additions & 0 deletions llvm/test/CodeGen/MIR/AMDGPU/long-branch-reg-all-sgpr-used.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define amdgpu_kernel void @long_branch_used_all_sgprs(ptr addrspace(1) %arg, i32 %cnd) #0 {
Expand Down Expand Up @@ -320,6 +321,7 @@
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define amdgpu_kernel void @long_branch_high_num_sgprs_used(ptr addrspace(1) %arg, i32 %cnd) #0 {
Expand Down
1 change: 1 addition & 0 deletions llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-after-pei.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
; AFTER-PEI-NEXT: hasInitWholeWave: false
; AFTER-PEI-NEXT: dynamicVGPRBlockSize: 0
; AFTER-PEI-NEXT: scratchReservedForDynamicVGPRs: 0
; AFTER-PEI-NEXT: numKernargPreloadSGPRs: 0
; AFTER-PEI-NEXT: isWholeWaveFunction: false
; AFTER-PEI-NEXT: body:
define amdgpu_kernel void @scavenge_fi(ptr addrspace(1) %out, i32 %in) #0 {
Expand Down
1 change: 1 addition & 0 deletions llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-long-branch-reg-debug.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define amdgpu_kernel void @uniform_long_forward_branch_debug(ptr addrspace(1) %arg, i32 %arg1) #0 !dbg !5 {
Expand Down
1 change: 1 addition & 0 deletions llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-long-branch-reg.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define amdgpu_kernel void @uniform_long_forward_branch(ptr addrspace(1) %arg, i32 %arg1) #0 {
Expand Down
4 changes: 4 additions & 0 deletions llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
# FULL-NEXT: hasInitWholeWave: false
# FULL-NEXT: dynamicVGPRBlockSize: 0
# FULL-NEXT: scratchReservedForDynamicVGPRs: 0
# FULL-NEXT: numKernargPreloadSGPRs: 0
# FULL-NEXT: isWholeWaveFunction: false
# FULL-NEXT: body:

Expand Down Expand Up @@ -167,6 +168,7 @@ body: |
# FULL-NEXT: hasInitWholeWave: false
# FULL-NEXT: dynamicVGPRBlockSize: 0
# FULL-NEXT: scratchReservedForDynamicVGPRs: 0
# FULL-NEXT: numKernargPreloadSGPRs: 0
# FULL-NEXT: isWholeWaveFunction: false
# FULL-NEXT: body:

Expand Down Expand Up @@ -248,6 +250,7 @@ body: |
# FULL-NEXT: hasInitWholeWave: false
# FULL-NEXT: dynamicVGPRBlockSize: 0
# FULL-NEXT: scratchReservedForDynamicVGPRs: 0
# FULL-NEXT: numKernargPreloadSGPRs: 0
# FULL-NEXT: isWholeWaveFunction: false
# FULL-NEXT: body:

Expand Down Expand Up @@ -330,6 +333,7 @@ body: |
# FULL-NEXT: hasInitWholeWave: false
# FULL-NEXT: dynamicVGPRBlockSize: 0
# FULL-NEXT: scratchReservedForDynamicVGPRs: 0
# FULL-NEXT: numKernargPreloadSGPRs: 0
# FULL-NEXT: isWholeWaveFunction: false
# FULL-NEXT: body:

Expand Down
4 changes: 4 additions & 0 deletions llvm/test/CodeGen/MIR/AMDGPU/machine-function-info.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define amdgpu_kernel void @kernel(i32 %arg0, i64 %arg1, <16 x i32> %arg2) {
Expand Down Expand Up @@ -110,6 +111,7 @@ define amdgpu_kernel void @kernel(i32 %arg0, i64 %arg1, <16 x i32> %arg2) {
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define amdgpu_ps void @ps_shader(i32 %arg0, i32 inreg %arg1) {
Expand Down Expand Up @@ -186,6 +188,7 @@ define amdgpu_ps void @gds_size_shader(i32 %arg0, i32 inreg %arg1) #5 {
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define void @function() {
Expand Down Expand Up @@ -244,6 +247,7 @@ define void @function() {
; CHECK-NEXT: hasInitWholeWave: false
; CHECK-NEXT: dynamicVGPRBlockSize: 0
; CHECK-NEXT: scratchReservedForDynamicVGPRs: 0
; CHECK-NEXT: numKernargPreloadSGPRs: 0
; CHECK-NEXT: isWholeWaveFunction: false
; CHECK-NEXT: body:
define void @function_nsz() #0 {
Expand Down
Loading