[AMDGPU] Allow architected SGPRs for workgroup IDs

[cdevadas]

Summary

Some subtargets use architected SGPRs for workgroup
IDs instead of the regular SGPRs. This patch enables
the support for the same and is guarded under the
subtarget feature FeatureArchitectedSGPRs.

[foad]

llvm/test/CodeGen/AMDGPU/workgroup-id-in-arch-sgprs.ll

27

Why are all these stores volatile?

54

Doesn't this value need to be ANDed with 0xffff?

[cdevadas]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

2260

Also, should I guard this check with HasArchitectedSGPRS?
The numPreloaded regular SGPRs will vary if TTMP registers are used for workgroup IDs.

llvm/test/CodeGen/AMDGPU/workgroup-id-in-arch-sgprs.ll

27

Will change.

54

If there is no WorkGroupIDZ, I thought I can skip that part.
Just like the packedTID here.
https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SIISelLowering.cpp#L1934

[foad]

llvm/test/CodeGen/AMDGPU/workgroup-id-in-arch-sgprs.ll

54

I don't understand. How do you know "there is no WorkGroupIDZ"? Just because this function does not call @llvm.amdgcn.workgroup.id.z?

[cdevadas]

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

707

The hasWorkGroupIDZ() will be set to 1 if there is a workgroup IDZ used. Otherwise, this function would return false.
The WorkGroupIDX, WorkGroupIDY, and WorkGroupIDZ fields in the MachineFunctionInfo would have set/reset it by now.

[foad]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

2260

@arsenm?

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

707

OK, I understand now, if the function does not call @llvm.amdgcn.workgroup.id.z then the compiler will not set COMPUTE_PGM_RSRC2:TGID_Z_EN so the hardware will initialize the top half of TTMP7 to 0.

[cdevadas]

Adjusted NumRequiredSystemSGPRs before force initializing the 16 input SGPRs.

[foad]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

2200

It is not at all clear how hasArchitectedSGPRs and hasUserSGPRInit16Bug should interact. Instead of guessing, how about asserting !HasArchitectedSGPRs here unless/until there is an ASIC that has both features?

[cdevadas]

Added an assertion as per review comment.

[foad]

LGTM, thanks.

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

if (PVID == AMDGPUFunctionArgInfo::PreloadedValue::KERNARG_SEGMENT_PTR) {
       return DAG.getConstant(0, SDLoc(), VT);
     }
 
     // It's undefined behavior if a function marked with the amdgpu-no-*
     // attributes uses the corresponding intrinsic.
     return DAG.getUNDEF(VT);
   }
 
-  return CreateLiveInRegister(DAG, RC, Reg->getRegister(), VT);
+  return loadInputValue(DAG, RC, VT, SDLoc(DAG.getEntryNode()), *Reg);
 }
 
 static void processPSInputArgs(SmallVectorImpl<ISD::InputArg> &Splits,
                                CallingConv::ID CallConv,
                                ArrayRef<ISD::InputArg> Ins, BitVector &Skipped,
                                FunctionType *FType,
                                SIMachineFunctionInfo *Info) {
   for (unsigned I = 0, E = Ins.size(), PSInputNum = 0; I != E; ++I) {

 }
 
 // Allocate special input registers that are initialized per-wave.
 void SITargetLowering::allocateSystemSGPRs(CCState &CCInfo,
                                            MachineFunction &MF,
                                            SIMachineFunctionInfo &Info,
                                            CallingConv::ID CallConv,
                                            bool IsShader) const {
+  bool HasArchitectedSGPRs = Subtarget->hasArchitectedSGPRs();
   if (Subtarget->hasUserSGPRInit16Bug() && !IsShader) {
     // Note: user SGPRs are handled by the front-end for graphics shaders
     // Pad up the used user SGPRs with dead inputs.
-    unsigned CurrentUserSGPRs = Info.getNumUserSGPRs();
 
+    // TODO: NumRequiredSystemSGPRs computation should be adjusted appropriately
+    // before enabling architected SGPRs for workgroup IDs.
+    assert(!HasArchitectedSGPRs && "Unhandled feature for the subtarget");
+
+    unsigned CurrentUserSGPRs = Info.getNumUserSGPRs();
     // Note we do not count the PrivateSegmentWaveByteOffset. We do not want to
     // rely on it to reach 16 since if we end up having no stack usage, it will
     // not really be added.
     unsigned NumRequiredSystemSGPRs = Info.hasWorkGroupIDX() +
                                       Info.hasWorkGroupIDY() +
                                       Info.hasWorkGroupIDZ() +
                                       Info.hasWorkGroupInfo();
     for (unsigned i = NumRequiredSystemSGPRs + CurrentUserSGPRs; i < 16; ++i) {
       Register Reg = Info.addReservedUserSGPR();
       MF.addLiveIn(Reg, &AMDGPU::SGPR_32RegClass);
       CCInfo.AllocateReg(Reg);
     }
   }
 
   if (Info.hasWorkGroupIDX()) {
-    Register Reg = Info.addWorkGroupIDX();
+    Register Reg = Info.addWorkGroupIDX(HasArchitectedSGPRs);
+    if (!HasArchitectedSGPRs)
       MF.addLiveIn(Reg, &AMDGPU::SGPR_32RegClass);
+
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info.hasWorkGroupIDY()) {
-    Register Reg = Info.addWorkGroupIDY();
+    Register Reg = Info.addWorkGroupIDY(HasArchitectedSGPRs);
+    if (!HasArchitectedSGPRs)
       MF.addLiveIn(Reg, &AMDGPU::SGPR_32RegClass);
+
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info.hasWorkGroupIDZ()) {
-    Register Reg = Info.addWorkGroupIDZ();
+    Register Reg = Info.addWorkGroupIDZ(HasArchitectedSGPRs);
+    if (!HasArchitectedSGPRs)
       MF.addLiveIn(Reg, &AMDGPU::SGPR_32RegClass);
+
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info.hasWorkGroupInfo()) {
     Register Reg = Info.addWorkGroupInfo();
     MF.addLiveIn(Reg, &AMDGPU::SGPR_32RegClass);
     CCInfo.AllocateReg(Reg);
   }
if (Info.hasPrivateSegmentWaveByteOffset()) {
     } else
       PrivateSegmentWaveByteOffsetReg = Info.addPrivateSegmentWaveByteOffset();
 
     MF.addLiveIn(PrivateSegmentWaveByteOffsetReg, &AMDGPU::SGPR_32RegClass);
     CCInfo.AllocateReg(PrivateSegmentWaveByteOffsetReg);
   }
 
   assert(!Subtarget->hasUserSGPRInit16Bug() || IsShader ||
          Info.getNumPreloadedSGPRs() >= 16);
 }
 
 static void reservePrivateMemoryRegs(const TargetMachine &TM,
                                      MachineFunction &MF,
                                      const SIRegisterInfo &TRI,
                                      SIMachineFunctionInfo &Info) {
   // Now that we've figured out where the scratch register inputs are, see if
   // should reserve the arguments and use them directly.

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

public:
   /// Increment user SGPRs used for padding the argument list only.
   Register addReservedUserSGPR() {
     Register Next = getNextUserSGPR();
     ++NumUserSGPRs;
     return Next;
   }
 
   // Add system SGPRs.
-  Register addWorkGroupIDX() {
-    ArgInfo.WorkGroupIDX = ArgDescriptor::createRegister(getNextSystemSGPR());
+  Register addWorkGroupIDX(bool HasArchitectedSGPRs) {
+    Register Reg = HasArchitectedSGPRs ? AMDGPU::TTMP9 : getNextSystemSGPR();
+    ArgInfo.WorkGroupIDX = ArgDescriptor::createRegister(Reg);
+    if (!HasArchitectedSGPRs)
       NumSystemSGPRs += 1;
+
     return ArgInfo.WorkGroupIDX.getRegister();
   }
 
-  Register addWorkGroupIDY() {
-    ArgInfo.WorkGroupIDY = ArgDescriptor::createRegister(getNextSystemSGPR());
+  Register addWorkGroupIDY(bool HasArchitectedSGPRs) {
+    Register Reg = HasArchitectedSGPRs ? AMDGPU::TTMP7 : getNextSystemSGPR();
+    unsigned Mask = HasArchitectedSGPRs && hasWorkGroupIDZ() ? 0xffff : ~0u;
+    ArgInfo.WorkGroupIDY = ArgDescriptor::createRegister(Reg, Mask);
+    if (!HasArchitectedSGPRs)
       NumSystemSGPRs += 1;
+
     return ArgInfo.WorkGroupIDY.getRegister();
   }
 
-  Register addWorkGroupIDZ() {
-    ArgInfo.WorkGroupIDZ = ArgDescriptor::createRegister(getNextSystemSGPR());
+  Register addWorkGroupIDZ(bool HasArchitectedSGPRs) {
+    Register Reg = HasArchitectedSGPRs ? AMDGPU::TTMP7 : getNextSystemSGPR();
+    unsigned Mask = HasArchitectedSGPRs ? 0xffff << 16 : ~0u;
+    ArgInfo.WorkGroupIDZ = ArgDescriptor::createRegister(Reg, Mask);
+    if (!HasArchitectedSGPRs)
       NumSystemSGPRs += 1;
+
     return ArgInfo.WorkGroupIDZ.getRegister();
   }
 
   Register addWorkGroupInfo() {
     ArgInfo.WorkGroupInfo = ArgDescriptor::createRegister(getNextSystemSGPR());
     NumSystemSGPRs += 1;
     return ArgInfo.WorkGroupInfo.getRegister();
   }

llvm/test/CodeGen/AMDGPU/workgroup-id-in-arch-sgprs.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -mattr=+architected-sgprs --verify-machineinstrs < %s | FileCheck -check-prefix=GCN-SDAG %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -mattr=+architected-sgprs -global-isel --verify-machineinstrs < %s | FileCheck -check-prefix=GCN-GISEL %s
+
+define amdgpu_kernel void @workgroup_id_x(ptr addrspace(1) %ptrx) {
+; GCN-SDAG-LABEL: workgroup_id_x:
+; GCN-SDAG:       ; %bb.0:
+; GCN-SDAG-NEXT:    s_load_dwordx2 s[0:1], s[4:5], 0x0
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v0, 0
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v1, ttmp9
+; GCN-SDAG-NEXT:    s_waitcnt lgkmcnt(0)
+; GCN-SDAG-NEXT:    global_store_dword v0, v1, s[0:1]
+; GCN-SDAG-NEXT:    s_endpgm
+;
+; GCN-GISEL-LABEL: workgroup_id_x:
+; GCN-GISEL:       ; %bb.0:
+; GCN-GISEL-NEXT:    s_load_dwordx2 s[0:1], s[4:5], 0x0
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v0, ttmp9
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v1, 0
+; GCN-GISEL-NEXT:    s_waitcnt lgkmcnt(0)
+; GCN-GISEL-NEXT:    global_store_dword v1, v0, s[0:1]
+; GCN-GISEL-NEXT:    s_endpgm
+  %idx = call i32 @llvm.amdgcn.workgroup.id.x()
+  store i32 %idx, ptr addrspace(1) %ptrx
+
+  ret void
+}
+
+define amdgpu_kernel void @workgroup_id_xy(ptr addrspace(1) %ptrx, ptr addrspace(1) %ptry) {
+; GCN-SDAG-LABEL: workgroup_id_xy:
+; GCN-SDAG:       ; %bb.0:
+; GCN-SDAG-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v0, 0
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v1, ttmp9
+; GCN-SDAG-NEXT:    s_waitcnt lgkmcnt(0)
+; GCN-SDAG-NEXT:    global_store_dword v0, v1, s[0:1]
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v1, ttmp7
+; GCN-SDAG-NEXT:    global_store_dword v0, v1, s[2:3]
+; GCN-SDAG-NEXT:    s_endpgm
+;
+; GCN-GISEL-LABEL: workgroup_id_xy:
+; GCN-GISEL:       ; %bb.0:
+; GCN-GISEL-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v0, 0
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v1, ttmp9
+; GCN-GISEL-NEXT:    s_waitcnt lgkmcnt(0)
+; GCN-GISEL-NEXT:    global_store_dword v0, v1, s[0:1]
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v1, ttmp7
+; GCN-GISEL-NEXT:    global_store_dword v0, v1, s[2:3]
+; GCN-GISEL-NEXT:    s_endpgm
+  %idx = call i32 @llvm.amdgcn.workgroup.id.x()
+  store i32 %idx, ptr addrspace(1) %ptrx
+  %idy = call i32 @llvm.amdgcn.workgroup.id.y()
+  store i32 %idy, ptr addrspace(1) %ptry
+
+  ret void
+}
+
+define amdgpu_kernel void @workgroup_id_xyz(ptr addrspace(1) %ptrx, ptr addrspace(1) %ptry, ptr addrspace(1) %ptrz) {
+; GCN-SDAG-LABEL: workgroup_id_xyz:
+; GCN-SDAG:       ; %bb.0:
+; GCN-SDAG-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
+; GCN-SDAG-NEXT:    s_load_dwordx2 s[6:7], s[4:5], 0x10
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v0, 0
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v1, ttmp9
+; GCN-SDAG-NEXT:    s_waitcnt lgkmcnt(0)
+; GCN-SDAG-NEXT:    global_store_dword v0, v1, s[0:1]
+; GCN-SDAG-NEXT:    s_and_b32 s0, ttmp7, 0xffff
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v1, s0
+; GCN-SDAG-NEXT:    s_lshr_b32 s0, ttmp7, 16
+; GCN-SDAG-NEXT:    global_store_dword v0, v1, s[2:3]
+; GCN-SDAG-NEXT:    v_mov_b32_e32 v1, s0
+; GCN-SDAG-NEXT:    global_store_dword v0, v1, s[6:7]
+; GCN-SDAG-NEXT:    s_endpgm
+;
+; GCN-GISEL-LABEL: workgroup_id_xyz:
+; GCN-GISEL:       ; %bb.0:
+; GCN-GISEL-NEXT:    s_load_dwordx4 s[0:3], s[4:5], 0x0
+; GCN-GISEL-NEXT:    s_load_dwordx2 s[6:7], s[4:5], 0x10
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v0, ttmp9
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v1, 0
+; GCN-GISEL-NEXT:    s_waitcnt lgkmcnt(0)
+; GCN-GISEL-NEXT:    global_store_dword v1, v0, s[0:1]
+; GCN-GISEL-NEXT:    s_and_b32 s0, ttmp7, 0xffff
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v0, s0
+; GCN-GISEL-NEXT:    s_lshr_b32 s0, ttmp7, 16
+; GCN-GISEL-NEXT:    global_store_dword v1, v0, s[2:3]
+; GCN-GISEL-NEXT:    v_mov_b32_e32 v0, s0
+; GCN-GISEL-NEXT:    global_store_dword v1, v0, s[6:7]
+; GCN-GISEL-NEXT:    s_endpgm
+  %idx = call i32 @llvm.amdgcn.workgroup.id.x()
+  store i32 %idx, ptr addrspace(1) %ptrx
+  %idy = call i32 @llvm.amdgcn.workgroup.id.y()
+  store i32 %idy, ptr addrspace(1) %ptry
+  %idz = call i32 @llvm.amdgcn.workgroup.id.z()
+  store i32 %idz, ptr addrspace(1) %ptrz
+
+  ret void
+}
+
+declare i32 @llvm.amdgcn.workgroup.id.x()
+declare i32 @llvm.amdgcn.workgroup.id.y()
+declare i32 @llvm.amdgcn.workgroup.id.z()