[AMDGPU] For OS type AMDPAL, fixed scratch on compute shader

[tpr]

Summary

For OS type AMDPAL, the scratch descriptor is loaded from offset 0 of
the GIT, whose 32 bit pointer is in s0 (s8 for gfx9 merged shaders).

This commit fixes that to use offset 0x10 instead of offset 0 for a
compute shader, per the PAL ABI spec.

Change-Id: I93dffa647758e37f613bb5e0dfca840d82e6d26f

[arsenm]

LGTM

[dstuttard]

LGTM

[tpr]

I had to revert this because it caused a test failure only with LLVM_ENABLE_EXPENSIVE_CHECKS.

[tpr]

V2: Ensure livein set for s0 (s8 in gfx9 merged shader) when using it to load
scratch decriptor.

[tpr]

What's the protocol for a fix to a change that I had to revert due to test failure? Does it need to come back to review here? Or should I just have landed it?

[nhaehnle]

LGTM.

In D44468#1050348, @tpr wrote:

What's the protocol for a fix to a change that I had to revert due to test failure? Does it need to come back to review here? Or should I just have landed it?

I don't believe there are completely hard rules. Personally, I'd say it depends on (a) how invasive the fix is and (b) how quickly you find it. In this particular case, I'd say the change is small enough that you could have just landed it. I'd go further and say that if the problem is so obvious that you realize it within a very short time of receiving the failure notifications, you might just commit the fix without reverting. Generally you should run the lit tests yourself locally though, to reduce the chances of this happening in the first place.

lib/Target/AMDGPU/SIFrameLowering.cpp

if (ST.hasMergedShaders()) {
           // Low GIT address is passed in s8 rather than s0 for an LS+HS or
           // ES+GS merged shader on gfx9+.
           GitPtrLo = AMDGPU::SGPR8;
           break;
         default:
           break;
       }
     }
+    MF.getRegInfo().addLiveIn(GitPtrLo);
+    MF.front().addLiveIn(GitPtrLo);
     BuildMI(MBB, I, DL, SMovB32, RsrcLo)
       .addReg(GitPtrLo)
       .addReg(ScratchRsrcReg, RegState::ImplicitDefine);
 
     // We now have the GIT ptr - now get the scratch descriptor from the entry
-    // at offset 0.
+    // at offset 0 (or offset 16 for a compute shader).
     PointerType *PtrTy =
       PointerType::get(Type::getInt64Ty(MF.getFunction().getContext()),
                        AMDGPUAS::CONSTANT_ADDRESS);
     MachinePointerInfo PtrInfo(UndefValue::get(PtrTy));
     const MCInstrDesc &LoadDwordX4 = TII->get(AMDGPU::S_LOAD_DWORDX4_IMM);
     auto MMO = MF.getMachineMemOperand(PtrInfo,
                                        MachineMemOperand::MOLoad |
                                        MachineMemOperand::MOInvariant |
                                        MachineMemOperand::MODereferenceable,
                                        0, 0);
+    unsigned Offset
+        = MF.getFunction().getCallingConv() == CallingConv::AMDGPU_CS ? 16 : 0;
     BuildMI(MBB, I, DL, LoadDwordX4, ScratchRsrcReg)
       .addReg(Rsrc01)
-      .addImm(0) // offset
+      .addImm(Offset) // offset
       .addImm(0) // glc
       .addReg(ScratchRsrcReg, RegState::ImplicitDefine)
       .addMemOperand(MMO);
     return;
   }
   if (ST.isMesaGfxShader(MF)
       || (PreloadedPrivateBufferReg == AMDGPU::NoRegister)) {
     assert(!ST.isAmdCodeObjectV2(MF));

test/CodeGen/AMDGPU/amdpal.ll

entry:
   %vv = bitcast [2 x i32] addrspace(5)* %v to <2 x i32> addrspace(5)*
   store <2 x i32> %in, <2 x i32> addrspace(5)* %vv
   %e = getelementptr [2 x i32], [2 x i32] addrspace(5)* %v, i32 0, i32 %idx
   %x = load i32, i32 addrspace(5)* %e
   store i32 %x, i32 addrspace(5)* %out
   ret void
 }
 
+; Check code sequence for amdpal use of scratch for alloca in a compute shader.
+; The scratch descriptor is loaded from offset 0x10 of the GIT, rather than offset
+; 0 in a graphics shader.
+
+; PAL-LABEL: {{^}}scratch2_cs:
+; PAL: s_movk_i32 s{{[0-9]+}}, 0x1234
+; PAL: s_mov_b32 s[[GITPTR:[0-9]+]], s0
+; PAL: s_load_dwordx4 s{{\[}}[[SCRATCHDESC:[0-9]+]]:{{[0-9]+]}}, s{{\[}}[[GITPTR]]:{{[0-9]+\]}}, 0x10
+; PAL: buffer_store{{.*}}, s{{\[}}[[SCRATCHDESC]]:
+
+define amdgpu_cs void @scratch2_cs(i32 inreg, i32 inreg, i32 inreg, <3 x i32> inreg, i32 inreg, <3 x i32> %coord, <2 x i32> %in, i32 %extra, i32 %idx) #0 {
+entry:
+  %v = alloca [3 x i32], addrspace(5)
+  %v0 = getelementptr [3 x i32], [3 x i32] addrspace(5)* %v, i32 0, i32 0
+  %v1 = getelementptr [3 x i32], [3 x i32] addrspace(5)* %v, i32 0, i32 1
+  store i32 %extra, i32 addrspace(5)* %v0
+  %v1a = bitcast i32 addrspace(5)* %v1 to [2 x i32] addrspace(5)*
+  %vv = bitcast [2 x i32] addrspace(5)* %v1a to <2 x i32> addrspace(5)*
+  store <2 x i32> %in, <2 x i32> addrspace(5)* %vv
+  %e = getelementptr [2 x i32], [2 x i32] addrspace(5)* %v1a, i32 0, i32 %idx
+  %x = load i32, i32 addrspace(5)* %e
+  %xf = bitcast i32 %x to float
+  call void @llvm.amdgcn.buffer.store.f32(float %xf, <4 x i32> undef, i32 0, i32 0, i1 0, i1 0)
+  ret void
+}
+
 attributes #0 = { nounwind "amdgpu-git-ptr-high"="0x1234" }
 
+declare void @llvm.amdgcn.buffer.store.f32(float, <4 x i32>, i32, i32, i1, i1)
+
+
 ; Check we have CS_NUM_USED_VGPRS in PAL metadata.
 ; PAL: .amd_amdgpu_pal_metadata {{.*}},0x10000027,