Skip to content

[AMDGPU] Architected SGPRs for GFX12 #76140

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.

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

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 4 additions & 0 deletions llvm/include/llvm/IR/IntrinsicsAMDGPU.td
Original file line number Diff line number Diff line change
Expand Up @@ -2491,6 +2491,10 @@ def int_amdgcn_flat_atomic_fmax_num : AMDGPUGlobalAtomicRtn<llvm_anyfloat_ty>;
def int_amdgcn_global_atomic_fmin_num : AMDGPUGlobalAtomicRtn<llvm_anyfloat_ty>;
def int_amdgcn_global_atomic_fmax_num : AMDGPUGlobalAtomicRtn<llvm_anyfloat_ty>;

// i32 @llvm.amdgcn.wave.id()
def int_amdgcn_wave_id :
DefaultAttrsIntrinsic<[llvm_i32_ty], [], [IntrNoMem, IntrSpeculatable]>;

//===----------------------------------------------------------------------===//
// Deep learning intrinsics.
//===----------------------------------------------------------------------===//
Expand Down
1 change: 1 addition & 0 deletions llvm/lib/Target/AMDGPU/AMDGPU.td
Original file line number Diff line number Diff line change
Expand Up @@ -1471,6 +1471,7 @@ def FeatureISAVersion12 : FeatureSet<
FeatureWavefrontSize32,
FeatureShaderCyclesRegister,
FeatureArchitectedFlatScratch,
FeatureArchitectedSGPRs,
FeatureAtomicFaddRtnInsts,
FeatureAtomicFaddNoRtnInsts,
FeatureFlatAtomicFaddF32Inst,
Expand Down
14 changes: 8 additions & 6 deletions llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -75,15 +75,14 @@ void AMDGPUArgumentUsageInfo::print(raw_ostream &OS, const Module *M) const {
<< " WorkGroupIDY: " << FI.second.WorkGroupIDY
<< " WorkGroupIDZ: " << FI.second.WorkGroupIDZ
<< " WorkGroupInfo: " << FI.second.WorkGroupInfo
<< " WaveID: " << FI.second.WaveID
<< " LDSKernelId: " << FI.second.LDSKernelId
<< " PrivateSegmentWaveByteOffset: "
<< FI.second.PrivateSegmentWaveByteOffset
<< FI.second.PrivateSegmentWaveByteOffset
<< " ImplicitBufferPtr: " << FI.second.ImplicitBufferPtr
<< " ImplicitArgPtr: " << FI.second.ImplicitArgPtr
<< " WorkItemIDX " << FI.second.WorkItemIDX
<< " WorkItemIDY " << FI.second.WorkItemIDY
<< " WorkItemIDZ " << FI.second.WorkItemIDZ
<< '\n';
<< " ImplicitArgPtr: " << FI.second.ImplicitArgPtr << " WorkItemIDX "
<< FI.second.WorkItemIDX << " WorkItemIDY " << FI.second.WorkItemIDY
<< " WorkItemIDZ " << FI.second.WorkItemIDZ << '\n';
}
}

Expand All @@ -108,6 +107,9 @@ AMDGPUFunctionArgInfo::getPreloadedValue(
case AMDGPUFunctionArgInfo::WORKGROUP_ID_Z:
return std::tuple(WorkGroupIDZ ? &WorkGroupIDZ : nullptr,
&AMDGPU::SGPR_32RegClass, LLT::scalar(32));
case AMDGPUFunctionArgInfo::WAVE_ID:
return std::tuple(WaveID ? &WaveID : nullptr, &AMDGPU::SGPR_32RegClass,
LLT::scalar(32));
case AMDGPUFunctionArgInfo::LDS_KERNEL_ID:
return std::tuple(LDSKernelId ? &LDSKernelId : nullptr,
&AMDGPU::SGPR_32RegClass, LLT::scalar(32));
Expand Down
2 changes: 2 additions & 0 deletions llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ struct AMDGPUFunctionArgInfo {
WORKGROUP_ID_X = 10,
WORKGROUP_ID_Y = 11,
WORKGROUP_ID_Z = 12,
WAVE_ID = 13,
PRIVATE_SEGMENT_WAVE_BYTE_OFFSET = 14,
IMPLICIT_BUFFER_PTR = 15,
IMPLICIT_ARG_PTR = 16,
Expand Down Expand Up @@ -141,6 +142,7 @@ struct AMDGPUFunctionArgInfo {
ArgDescriptor WorkGroupIDY;
ArgDescriptor WorkGroupIDZ;
ArgDescriptor WorkGroupInfo;
ArgDescriptor WaveID;
ArgDescriptor PrivateSegmentWaveByteOffset;

// Pointer with offset from kernargsegmentptr to where special ABI arguments
Expand Down
3 changes: 3 additions & 0 deletions llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -6922,6 +6922,9 @@ bool AMDGPULegalizerInfo::legalizeIntrinsic(LegalizerHelper &Helper,
case Intrinsic::amdgcn_workgroup_id_z:
return legalizePreloadedArgIntrin(MI, MRI, B,
AMDGPUFunctionArgInfo::WORKGROUP_ID_Z);
case Intrinsic::amdgcn_wave_id:
return legalizePreloadedArgIntrin(MI, MRI, B,
AMDGPUFunctionArgInfo::WAVE_ID);
case Intrinsic::amdgcn_lds_kernel_id:
return legalizePreloadedArgIntrin(MI, MRI, B,
AMDGPUFunctionArgInfo::LDS_KERNEL_ID);
Expand Down
7 changes: 6 additions & 1 deletion llvm/lib/Target/AMDGPU/SIISelLowering.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2431,6 +2431,9 @@ void SITargetLowering::allocateSystemSGPRs(CCState &CCInfo,
CCInfo.AllocateReg(Reg);
}

if (HasArchitectedSGPRs && Info.hasWaveID())
CCInfo.AllocateReg(Info.addWaveID());

if (Info.hasWorkGroupInfo()) {
Register Reg = Info.addWorkGroupInfo();
MF.addLiveIn(Reg, &AMDGPU::SGPR_32RegClass);
Expand Down Expand Up @@ -2635,7 +2638,7 @@ SDValue SITargetLowering::LowerFormalArguments(
assert(!UserSGPRInfo.hasFlatScratchInit());
if (CallConv != CallingConv::AMDGPU_CS || !Subtarget->hasArchitectedSGPRs())
assert(!Info->hasWorkGroupIDX() && !Info->hasWorkGroupIDY() &&
!Info->hasWorkGroupIDZ());
!Info->hasWorkGroupIDZ() && !Info->hasWaveID());
}

if (CallConv == CallingConv::AMDGPU_PS) {
Expand Down Expand Up @@ -7775,6 +7778,8 @@ SDValue SITargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
case Intrinsic::amdgcn_workgroup_id_z:
return getPreloadedValue(DAG, *MFI, VT,
AMDGPUFunctionArgInfo::WORKGROUP_ID_Z);
case Intrinsic::amdgcn_wave_id:
return getPreloadedValue(DAG, *MFI, VT, AMDGPUFunctionArgInfo::WAVE_ID);
case Intrinsic::amdgcn_lds_kernel_id: {
if (MFI->isEntryFunction())
return getLDSKernelId(DAG, DL);
Expand Down
17 changes: 11 additions & 6 deletions llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,10 @@ SIMachineFunctionInfo::SIMachineFunctionInfo(const Function &F,
const GCNSubtarget *STI)
: AMDGPUMachineFunction(F, *STI), Mode(F, *STI), GWSResourcePSV(getTM(STI)),
UserSGPRInfo(F, *STI), WorkGroupIDX(false), WorkGroupIDY(false),
WorkGroupIDZ(false), WorkGroupInfo(false), LDSKernelId(false),
PrivateSegmentWaveByteOffset(false), WorkItemIDX(false),
WorkItemIDY(false), WorkItemIDZ(false), ImplicitArgPtr(false),
GITPtrHigh(0xffffffff), HighBitsOf32BitAddress(0) {
WorkGroupIDZ(false), WorkGroupInfo(false), WaveID(false),
LDSKernelId(false), PrivateSegmentWaveByteOffset(false),
WorkItemIDX(false), WorkItemIDY(false), WorkItemIDZ(false),
ImplicitArgPtr(false), GITPtrHigh(0xffffffff), HighBitsOf32BitAddress(0) {
const GCNSubtarget &ST = *static_cast<const GCNSubtarget *>(STI);
FlatWorkGroupSizes = ST.getFlatWorkGroupSizes(F);
WavesPerEU = ST.getWavesPerEU(F);
Expand Down Expand Up @@ -107,8 +107,9 @@ SIMachineFunctionInfo::SIMachineFunctionInfo(const Function &F,
MayNeedAGPRs = false; // We will select all MAI with VGPR operands.
}

if (!AMDGPU::isGraphics(CC) ||
(CC == CallingConv::AMDGPU_CS && ST.hasArchitectedSGPRs())) {
bool HasArchitectedSGPRs =
CC == CallingConv::AMDGPU_CS && ST.hasArchitectedSGPRs();
if (!AMDGPU::isGraphics(CC) || HasArchitectedSGPRs) {
if (IsKernel || !F.hasFnAttribute("amdgpu-no-workgroup-id-x"))
WorkGroupIDX = true;

Expand All @@ -135,6 +136,10 @@ SIMachineFunctionInfo::SIMachineFunctionInfo(const Function &F,
LDSKernelId = true;
}

// For GFX12+.
if (HasArchitectedSGPRs)
WaveID = true;

if (isEntryFunction()) {
// X, XY, and XYZ are the only supported combinations, so make sure Y is
// enabled if Z is.
Expand Down
9 changes: 9 additions & 0 deletions llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
Original file line number Diff line number Diff line change
Expand Up @@ -447,6 +447,7 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
bool WorkGroupIDY : 1;
bool WorkGroupIDZ : 1;
bool WorkGroupInfo : 1;
bool WaveID : 1;
bool LDSKernelId : 1;
bool PrivateSegmentWaveByteOffset : 1;

Expand Down Expand Up @@ -782,6 +783,14 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
return ArgInfo.WorkGroupInfo.getRegister();
}

// Supported for GFX12+.
Register addWaveID() {
ArgInfo.WaveID = ArgDescriptor::createRegister(AMDGPU::TTMP8, 0x1f << 25);
return ArgInfo.WaveID.getRegister();
}

bool hasWaveID() const { return WaveID; }

bool hasLDSKernelId() const { return LDSKernelId; }

// Add special VGPR inputs
Expand Down
44 changes: 44 additions & 0 deletions llvm/test/CodeGen/AMDGPU/lower-work-group-id-intrinsics.ll
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx900 -mattr=+architected-sgprs --verify-machineinstrs < %s | FileCheck -check-prefix=GFX9-SDAG %s
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx900 -mattr=+architected-sgprs -global-isel --verify-machineinstrs < %s | FileCheck -check-prefix=GFX9-GISEL %s
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1200 --verify-machineinstrs < %s | FileCheck -check-prefix=GFX12-SDAG %s
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1200 -global-isel --verify-machineinstrs < %s | FileCheck -check-prefix=GFX12-GISEL %s

define amdgpu_cs void @_amdgpu_cs_main() {
; GFX9-SDAG-LABEL: _amdgpu_cs_main:
Expand All @@ -23,6 +25,30 @@ define amdgpu_cs void @_amdgpu_cs_main() {
; GFX9-GISEL-NEXT: v_mov_b32_e32 v2, s2
; GFX9-GISEL-NEXT: buffer_store_dwordx3 v[0:2], off, s[0:3], 0
; GFX9-GISEL-NEXT: s_endpgm
;
; GFX12-SDAG-LABEL: _amdgpu_cs_main:
; GFX12-SDAG: ; %bb.0: ; %.entry
; GFX12-SDAG-NEXT: s_lshr_b32 s2, ttmp7, 16
; GFX12-SDAG-NEXT: s_and_b32 s1, ttmp7, 0xffff
; GFX12-SDAG-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
; GFX12-SDAG-NEXT: v_dual_mov_b32 v0, ttmp9 :: v_dual_mov_b32 v1, s1
; GFX12-SDAG-NEXT: v_mov_b32_e32 v2, s2
; GFX12-SDAG-NEXT: buffer_store_b96 v[0:2], off, s[0:3], null
; GFX12-SDAG-NEXT: s_nop 0
; GFX12-SDAG-NEXT: s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
; GFX12-SDAG-NEXT: s_endpgm
;
; GFX12-GISEL-LABEL: _amdgpu_cs_main:
; GFX12-GISEL: ; %bb.0: ; %.entry
; GFX12-GISEL-NEXT: s_mov_b32 s0, ttmp9
; GFX12-GISEL-NEXT: s_and_b32 s1, ttmp7, 0xffff
; GFX12-GISEL-NEXT: s_lshr_b32 s2, ttmp7, 16
; GFX12-GISEL-NEXT: v_dual_mov_b32 v0, s0 :: v_dual_mov_b32 v1, s1
; GFX12-GISEL-NEXT: v_mov_b32_e32 v2, s2
; GFX12-GISEL-NEXT: buffer_store_b96 v[0:2], off, s[0:3], null
; GFX12-GISEL-NEXT: s_nop 0
; GFX12-GISEL-NEXT: s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
; GFX12-GISEL-NEXT: s_endpgm
.entry:
%idx = call i32 @llvm.amdgcn.workgroup.id.x()
%idy = call i32 @llvm.amdgcn.workgroup.id.y()
Expand Down Expand Up @@ -68,6 +94,24 @@ define amdgpu_cs void @caller() {
; GFX9-GISEL-NEXT: s_mov_b64 s[2:3], s[10:11]
; GFX9-GISEL-NEXT: s_swappc_b64 s[30:31], s[4:5]
; GFX9-GISEL-NEXT: s_endpgm
;
; GFX12-SDAG-LABEL: caller:
; GFX12-SDAG: ; %bb.0:
; GFX12-SDAG-NEXT: v_mov_b32_e32 v0, ttmp9
; GFX12-SDAG-NEXT: s_mov_b32 s1, callee@abs32@hi
; GFX12-SDAG-NEXT: s_mov_b32 s0, callee@abs32@lo
; GFX12-SDAG-NEXT: s_mov_b32 s32, 0
; GFX12-SDAG-NEXT: s_swappc_b64 s[30:31], s[0:1]
; GFX12-SDAG-NEXT: s_endpgm
;
; GFX12-GISEL-LABEL: caller:
; GFX12-GISEL: ; %bb.0:
; GFX12-GISEL-NEXT: v_mov_b32_e32 v0, ttmp9
; GFX12-GISEL-NEXT: s_mov_b32 s0, callee@abs32@lo
; GFX12-GISEL-NEXT: s_mov_b32 s1, callee@abs32@hi
; GFX12-GISEL-NEXT: s_mov_b32 s32, 0
; GFX12-GISEL-NEXT: s_swappc_b64 s[30:31], s[0:1]
; GFX12-GISEL-NEXT: s_endpgm
%idx = call i32 @llvm.amdgcn.workgroup.id.x()
call amdgpu_gfx void @callee(i32 %idx)
ret void
Expand Down
34 changes: 34 additions & 0 deletions llvm/test/CodeGen/AMDGPU/wave-id-intrinsic.ll
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx900 -mattr=+architected-sgprs --verify-machineinstrs < %s | FileCheck -check-prefixes=GFX9,GFX9-SDAG %s
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx900 -mattr=+architected-sgprs -global-isel --verify-machineinstrs < %s | FileCheck -check-prefixes=GFX9,GFX9-GISEL %s
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1200 --verify-machineinstrs < %s | FileCheck -check-prefixes=GFX12,GFX12-SDAG %s
; RUN: llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1200 -global-isel --verify-machineinstrs < %s | FileCheck -check-prefixes=GFX12,GFX12-GISEL %s

define amdgpu_cs void @test_wave_id(ptr addrspace(1) %out) {
; GFX9-LABEL: test_wave_id:
; GFX9: ; %bb.0:
; GFX9-NEXT: s_bfe_u32 s0, ttmp8, 0x50019
; GFX9-NEXT: v_mov_b32_e32 v2, s0
; GFX9-NEXT: global_store_dword v[0:1], v2, off
; GFX9-NEXT: s_endpgm
;
; GFX12-LABEL: test_wave_id:
; GFX12: ; %bb.0:
; GFX12-NEXT: s_bfe_u32 s0, ttmp8, 0x50019
; GFX12-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
; GFX12-NEXT: v_mov_b32_e32 v2, s0
; GFX12-NEXT: global_store_b32 v[0:1], v2, off
; GFX12-NEXT: s_nop 0
; GFX12-NEXT: s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
; GFX12-NEXT: s_endpgm
%waveid = call i32 @llvm.amdgcn.wave.id()
store i32 %waveid, ptr addrspace(1) %out
ret void
}

Copy link
Contributor

Choose a reason for hiding this comment

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

What about uses in a callable function?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@cdevadas @nhaehnle is that supported and how is it supposed to work?

Copy link
Collaborator

Choose a reason for hiding this comment

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

Yes, that ought to work in the exact same way since we never change ttmp registers from "user" (non-trap) code.

Copy link
Contributor

Choose a reason for hiding this comment

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

Should get corresponding tests

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Is it OK to commit this as-is and then work on supporting wave.id and workgroup.id in callable functions? @cdevadas is working on that part.

Copy link
Contributor

Choose a reason for hiding this comment

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

I don't think it requires any additional work. I would expect it to work as-is and just needs testing. The workgroup ID handling requires additional work to support, but that's only tangentially related

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It does not "just work". I think that's probably because we don't call allocateSystemSGPRs for non-entry functions, but that what CD is looking into.

Copy link
Contributor

Choose a reason for hiding this comment

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

OK, I see what you mean. But that means this is routing through a lot of complexity in order to support function ABI handling. This case is much simpler and can bypass all of it. The intrinsic lowering just needs to directly copy from the TTMP register. The existing inputs are a poor analogy if we have one dedicated, constant register

Copy link
Collaborator

Choose a reason for hiding this comment

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

I currently don't see any codegen for workgroupIDs and waveID for callable gfx functions (those with amdgpu_gfx CC). Should we select the code for all three workgroupIDs and waveID intrinsics when Architected SGPRs are enabled?
cc. @nhaehnle

Copy link
Collaborator

Choose a reason for hiding this comment

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

I think so, yes.

declare i32 @llvm.amdgcn.wave.id()
;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
; GFX12-GISEL: {{.*}}
; GFX12-SDAG: {{.*}}
; GFX9-GISEL: {{.*}}
; GFX9-SDAG: {{.*}}
Loading