AMDGPU/SI: Use the correct scratch wave offset register for shaders.

[bnieuwenhuizen]

Summary

The code previously always used s1 as it was using the user + system SGPR
information for compute kernels. This is incorrect for Mesa shaders though,

The register should be the next SGPR after all user and system SGPR's.
We use that Mesa adds arguments for all input and system SGPR's and
take the next available SGPR for the scratch wave offset register.

Signed-off-by: Bas Nieuwenhuizen <[e-mail]>

[arsenm]

lib/Target/AMDGPU/SIISelLowering.cpp

735–748

You should avoid having the isShader repeated all of these times, and the others should all be in the same assert. Maybe move under the isShader above?

[nhaehnle]

This looks like an improvement, but I think it is still incorrect for shaders with non-void returns, where the scratch information may have to be passed through.

[arsenm]

In D18941#399137, @nhaehnle wrote:

This looks like an improvement, but I think it is still incorrect for shaders with non-void returns, where the scratch information may have to be passed through.

The assert can still be cleaned up to only one

[nhaehnle]

In D18941#399137, @nhaehnle wrote:

This looks like an improvement, but I think it is still incorrect for shaders with non-void returns, where the scratch information may have to be passed through.

As discussed on IRC, it's fine to go ahead with this patch without fixing this, since it isn't a regression. We really need to make sure that non-mono shaders work correctly with spills though.

[nhaehnle]

And I again forgot to mention: I agree with Matt about the assert.

[bnieuwenhuizen]

Changed the asserts. Still used two though as we also want to assert that some get enabled for kernels.

[• tstellarAMD]

You should add the AMDGPU/SI: Prefix to the commit message.

lib/Target/AMDGPU/SIISelLowering.cpp

 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/ADT/SmallString.h"
 
 using namespace llvm;
 
+static unsigned findFirstFreeSGPR(CCState &CCInfo) {
+  unsigned NumSGPRs = AMDGPU::SGPR_32RegClass.getNumRegs();
+  for (unsigned Reg = 0; Reg < NumSGPRs; ++Reg) {
+    if (!CCInfo.isAllocated(AMDGPU::SGPR0 + Reg)) {
+      return AMDGPU::SGPR0 + Reg;
+    }
+  }
+  llvm_unreachable("Cannot allocate sgpr");
+}
+
 SITargetLowering::SITargetLowering(TargetMachine &TM,
                                    const AMDGPUSubtarget &STI)
     : AMDGPUTargetLowering(TM, STI) {
   addRegisterClass(MVT::i1, &AMDGPU::VReg_1RegClass);
   addRegisterClass(MVT::i64, &AMDGPU::SReg_64RegClass);
 
   addRegisterClass(MVT::i32, &AMDGPU::SReg_32RegClass);
   addRegisterClass(MVT::f32, &AMDGPU::VGPR_32RegClass);
Info->isPSInputAllocated(11)))) {
     CCInfo.AllocateReg(AMDGPU::VGPR1);
     Info->markPSInputAllocated(0);
     Info->PSInputEna |= 1;
   }
 
   if (!AMDGPU::isShader(CallConv)) {
     getOriginalFunctionArgs(DAG, DAG.getMachineFunction().getFunction(), Ins,
                             Splits);
+
+    assert(Info->hasWorkGroupIDX() && Info->hasWorkItemIDX());
+  } else {
+    assert(!Info->hasPrivateSegmentBuffer() && !Info->hasDispatchPtr() &&
+           !Info->hasKernargSegmentPtr() && !Info->hasFlatScratchInit() &&
+           !Info->hasWorkGroupIDX() && !Info->hasWorkGroupIDY() &&
+           !Info->hasWorkGroupIDZ() && !Info->hasWorkGroupInfo() &&
+           !Info->hasWorkItemIDX() && !Info->hasWorkItemIDY() &&
+           !Info->hasWorkItemIDZ());
   }
 
   // FIXME: How should these inputs interact with inreg / custom SGPR inputs?
   if (Info->hasPrivateSegmentBuffer()) {
     unsigned PrivateSegmentBufferReg = Info->addPrivateSegmentBuffer(*TRI);
     MF.addLiveIn(PrivateSegmentBufferReg, &AMDGPU::SReg_128RegClass);
     CCInfo.AllocateReg(PrivateSegmentBufferReg);
   }
 
   if (Info->hasDispatchPtr()) {
     unsigned DispatchPtrReg = Info->addDispatchPtr(*TRI);
     MF.addLiveIn(DispatchPtrReg, &AMDGPU::SReg_64RegClass);
     CCInfo.AllocateReg(DispatchPtrReg);
   }
 
   if (Info->hasKernargSegmentPtr()) {
     unsigned InputPtrReg = Info->addKernargSegmentPtr(*TRI);
     MF.addLiveIn(InputPtrReg, &AMDGPU::SReg_64RegClass);
     CCInfo.AllocateReg(InputPtrReg);
   }
 
   if (Info->hasFlatScratchInit()) {
     unsigned FlatScratchInitReg = Info->addFlatScratchInit(*TRI);
     MF.addLiveIn(FlatScratchInitReg, &AMDGPU::SReg_64RegClass);
SDValue SITargetLowering::LowerFormalArguments(
   // TODO: Add GridWorkGroupCount user SGPRs when used. For now with HSA we read
   // these from the dispatch pointer.
 
   // Start adding system SGPRs.
   if (Info->hasWorkGroupIDX()) {
     unsigned Reg = Info->addWorkGroupIDX();
     MF.addLiveIn(Reg, &AMDGPU::SReg_32RegClass);
     CCInfo.AllocateReg(Reg);
-  } else
-    llvm_unreachable("work group id x is always enabled");
+  }
 
   if (Info->hasWorkGroupIDY()) {
     unsigned Reg = Info->addWorkGroupIDY();
     MF.addLiveIn(Reg, &AMDGPU::SReg_32RegClass);
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info->hasWorkGroupIDZ()) {
     unsigned Reg = Info->addWorkGroupIDZ();
     MF.addLiveIn(Reg, &AMDGPU::SReg_32RegClass);
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info->hasWorkGroupInfo()) {
     unsigned Reg = Info->addWorkGroupInfo();
     MF.addLiveIn(Reg, &AMDGPU::SReg_32RegClass);
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info->hasPrivateSegmentWaveByteOffset()) {
     // Scratch wave offset passed in system SGPR.
-    unsigned PrivateSegmentWaveByteOffsetReg
-      = Info->addPrivateSegmentWaveByteOffset();
+    unsigned PrivateSegmentWaveByteOffsetReg;
+
+    if (AMDGPU::isShader(CallConv)) {
+      PrivateSegmentWaveByteOffsetReg = findFirstFreeSGPR(CCInfo);
+      Info->setPrivateSegmentWaveByteOffset(PrivateSegmentWaveByteOffsetReg);
+    } else
+      PrivateSegmentWaveByteOffsetReg = Info->addPrivateSegmentWaveByteOffset();
 
     MF.addLiveIn(PrivateSegmentWaveByteOffsetReg, &AMDGPU::SGPR_32RegClass);
     CCInfo.AllocateReg(PrivateSegmentWaveByteOffsetReg);
   }
 
   // Now that we've figured out where the scratch register inputs are, see if
   // should reserve the arguments and use them directly.
   bool HasStackObjects = MF.getFrameInfo()->hasStackObjects();
if (HasStackObjects) {
       Info->setScratchWaveOffsetReg(ReservedOffsetReg);
     }
   }
 
   if (Info->hasWorkItemIDX()) {
     unsigned Reg = TRI->getPreloadedValue(MF, SIRegisterInfo::WORKITEM_ID_X);
     MF.addLiveIn(Reg, &AMDGPU::VGPR_32RegClass);
     CCInfo.AllocateReg(Reg);
-  } else
-    llvm_unreachable("workitem id x should always be enabled");
+  }
 
   if (Info->hasWorkItemIDY()) {
     unsigned Reg = TRI->getPreloadedValue(MF, SIRegisterInfo::WORKITEM_ID_Y);
     MF.addLiveIn(Reg, &AMDGPU::VGPR_32RegClass);
     CCInfo.AllocateReg(Reg);
   }
 
   if (Info->hasWorkItemIDZ()) {

lib/Target/AMDGPU/SIMachineFunctionInfo.h

public:
   }
 
   unsigned addPrivateSegmentWaveByteOffset() {
     PrivateSegmentWaveByteOffsetSystemSGPR = getNextSystemSGPR();
     NumSystemSGPRs += 1;
     return PrivateSegmentWaveByteOffsetSystemSGPR;
   }
 
+  void setPrivateSegmentWaveByteOffset(unsigned Reg) {
+    PrivateSegmentWaveByteOffsetSystemSGPR = Reg;
+  }
+
   bool hasPrivateSegmentBuffer() const {
     return PrivateSegmentBuffer;
   }
 
   bool hasDispatchPtr() const {
     return DispatchPtr;
   }
 

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

: AMDGPUMachineFunction(MF),
     DispatchPtr(false),
     QueuePtr(false),
     DispatchID(false),
     KernargSegmentPtr(false),
     FlatScratchInit(false),
     GridWorkgroupCountX(false),
     GridWorkgroupCountY(false),
     GridWorkgroupCountZ(false),
-    WorkGroupIDX(true),
+    WorkGroupIDX(false),
     WorkGroupIDY(false),
     WorkGroupIDZ(false),
     WorkGroupInfo(false),
     PrivateSegmentWaveByteOffset(false),
-    WorkItemIDX(true),
+    WorkItemIDX(false),
     WorkItemIDY(false),
     WorkItemIDZ(false) {
   const AMDGPUSubtarget &ST = MF.getSubtarget<AMDGPUSubtarget>();
   const Function *F = MF.getFunction();
 
   PSInputAddr = AMDGPU::getInitialPSInputAddr(*F);
 
   const MachineFrameInfo *FrameInfo = MF.getFrameInfo();
 
-  if (!AMDGPU::isShader(F->getCallingConv()))
+  if (!AMDGPU::isShader(F->getCallingConv())) {
     KernargSegmentPtr = true;
+    WorkGroupIDX = true;
+    WorkItemIDX = true;
+  }
 
   if (F->hasFnAttribute("amdgpu-work-group-id-y"))
     WorkGroupIDY = true;
 
   if (F->hasFnAttribute("amdgpu-work-group-id-z"))
     WorkGroupIDZ = true;
 
   if (F->hasFnAttribute("amdgpu-work-item-id-y"))

test/CodeGen/AMDGPU/large-alloca-graphics.ll

 ; RUN: llc -march=amdgcn -mcpu=bonaire < %s | FileCheck -check-prefix=GCN -check-prefix=CI -check-prefix=ALL %s
 ; RUN: llc -march=amdgcn -mcpu=carrizo < %s | FileCheck -check-prefix=GCN -check-prefix=VI -check-prefix=ALL %s
 
 ; ALL-LABEL: {{^}}large_alloca_pixel_shader:
 ; GCN: s_mov_b32 s8, SCRATCH_RSRC_DWORD0
 ; GCN: s_mov_b32 s9, SCRATCH_RSRC_DWORD1
 ; GCN: s_mov_b32 s10, -1
 ; CI: s_mov_b32 s11, 0x98f000
 ; VI: s_mov_b32 s11, 0x980000
 
-; GCN: buffer_store_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s1 offen
-; GCN: buffer_load_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s1 offen
+; GCN: buffer_store_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s0 offen
+; GCN: buffer_load_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s0 offen
 
 ; ALL: ; ScratchSize: 32772
 define amdgpu_ps void @large_alloca_pixel_shader(i32 %x, i32 %y) #0 {
   %large = alloca [8192 x i32], align 4
   %gep = getelementptr [8192 x i32], [8192 x i32]* %large, i32 0, i32 8191
   store volatile i32 %x, i32* %gep
   %gep1 = getelementptr [8192 x i32], [8192 x i32]* %large, i32 0, i32 %y
   %val = load volatile i32, i32* %gep1
   store volatile i32 %val, i32 addrspace(1)* undef
   ret void
 }
 
 ; ALL-LABEL: {{^}}large_alloca_pixel_shader_inreg:
 ; GCN: s_mov_b32 s8, SCRATCH_RSRC_DWORD0
 ; GCN: s_mov_b32 s9, SCRATCH_RSRC_DWORD1
 ; GCN: s_mov_b32 s10, -1
 ; CI: s_mov_b32 s11, 0x98f000
 ; VI: s_mov_b32 s11, 0x980000
 
-; GCN: buffer_store_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s1 offen
-; GCN: buffer_load_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s1 offen
+; GCN: buffer_store_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s2 offen
+; GCN: buffer_load_dword {{v[0-9]+}}, {{v[0-9]+}}, s[8:11], s2 offen
 
 ; ALL: ; ScratchSize: 32772
 define amdgpu_ps void @large_alloca_pixel_shader_inreg(i32 inreg %x, i32 inreg %y) #0 {
   %large = alloca [8192 x i32], align 4
   %gep = getelementptr [8192 x i32], [8192 x i32]* %large, i32 0, i32 8191
   store volatile i32 %x, i32* %gep
   %gep1 = getelementptr [8192 x i32], [8192 x i32]* %large, i32 0, i32 %y
   %val = load volatile i32, i32* %gep1
   store volatile i32 %val, i32 addrspace(1)* undef
   ret void
 }
 
 attributes #0 = { nounwind  }

test/CodeGen/AMDGPU/vgpr-spill-emergency-stack-slot.ll

 ; RUN: llc -march=amdgcn -mcpu=tahiti -mattr=+vgpr-spilling -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=SI %s
 ; RUN: llc -march=amdgcn -mcpu=fiji -mattr=+vgpr-spilling -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=VI %s
 
 ; This ends up using all 255 registers and requires register
 ; scavenging which will fail to find an unsued register.
 
 ; Check the ScratchSize to avoid regressions from spilling
 ; intermediate register class copies.
 
 ; FIXME: The same register is initialized to 0 for every spill.
 
 ; GCN-LABEL: {{^}}main:
 
+; GCN: s_mov_b32 s11, s12
 ; GCN: s_mov_b32 s12, SCRATCH_RSRC_DWORD0
 ; GCN-NEXT: s_mov_b32 s13, SCRATCH_RSRC_DWORD1
 ; GCN-NEXT: s_mov_b32 s14, -1
 ; SI-NEXT: s_mov_b32 s15, 0x98f000
 ; VI-NEXT: s_mov_b32 s15, 0x980000
 
 ; s12 is offset user SGPR
 ; GCN: buffer_store_dword {{v[0-9]+}}, s[12:15], s11 offset:{{[0-9]+}} ; 16-byte Folded Spill