[AMDGPU] Scratch setup fix on AMDPAL gfx9+ merge shader

[tpr]

Summary

With OS type AMDPAL, the scratch descriptor is hardwired to be loaded
from offset 0 of the global information table, whose low pointer is
passed in s0. For a merge shader on gfx9+, it needs to be s8 instead, as
the hardware reserves s0-s7.

[arsenm]

lib/Target/AMDGPU/SIFrameLowering.cpp

397 ↗

(On Diff #130287)

Why aren't there corresponding changes in SIISelLowering? I would expect the argument lowering code would need to be aware of this

test/CodeGen/AMDGPU/amdpal_scratch_mergeshader.ll

1 ↗

(On Diff #130287)

Use GCN instead of CHECK with multiple prefixes

12 ↗

(On Diff #130287)

You don't need the linkage. Most of these arguments also look unused.

Why is most of this function body necessary? Shouldn't you only need a single volatile scratch access?

[tpr]

lib/Target/AMDGPU/SIFrameLowering.cpp

397 ↗

(On Diff #130287)

The frontend does actually set up args for s0-s7 as well as the user sgprs starting at s8. So arg lowering does not need to do anything different.

[tpr]

V2: Rebased. Addressed some review comments in the test.

[tpr]

test/CodeGen/AMDGPU/amdpal_scratch_mergeshader.ll

12 ↗

(On Diff #130287)

Removed the linkage and arg9 onwards. arg0-arg8 are needed. I got that body by using bugpoint on a real shader; any further manual reduction results in the scratch access disappearing. It needs to be an alloca to test what the fix is doing.

[tpr]

Ping. Can I land the revised fix?

[timcorringham]

Looks good to me.

[nhaehnle]

One tiny nitpick, apart from that LGTM.

lib/Target/AMDGPU/AMDGPUSubtarget.h

859 ↗

(On Diff #133031)

Merged, here and elsewhere?

llvm/trunk/lib/Target/AMDGPU/AMDGPUSubtarget.h

public:
   unsigned getOccupancyWithNumVGPRs(unsigned VGPRs) const;
 
   /// \returns true if the flat_scratch register should be initialized with the
   /// pointer to the wave's scratch memory rather than a size and offset.
   bool flatScratchIsPointer() const {
     return getGeneration() >= GFX9;
   }
 
+  /// \returns true if the machine has merged shaders in which s0-s7 are
+  /// reserved by the hardware and user SGPRs start at s8
+  bool hasMergedShaders() const {
+    return getGeneration() >= GFX9;
+  }
+
   /// \returns SGPR allocation granularity supported by the subtarget.
   unsigned getSGPRAllocGranule() const {
     return AMDGPU::IsaInfo::getSGPRAllocGranule(getFeatureBits());
   }
 
   /// \returns SGPR encoding granularity supported by the subtarget.
   unsigned getSGPREncodingGranule() const {
     return AMDGPU::IsaInfo::getSGPREncodingGranule(getFeatureBits());

llvm/trunk/lib/Target/AMDGPU/SIFrameLowering.cpp

if (ST.isAmdPalOS()) {
     if (MFI->getGITPtrHigh() != 0xffffffff) {
       BuildMI(MBB, I, DL, SMovB32, RsrcHi)
         .addImm(MFI->getGITPtrHigh())
         .addReg(ScratchRsrcReg, RegState::ImplicitDefine);
     } else {
       const MCInstrDesc &GetPC64 = TII->get(AMDGPU::S_GETPC_B64);
       BuildMI(MBB, I, DL, GetPC64, Rsrc01);
     }
+    auto GitPtrLo = AMDGPU::SGPR0; // Low GIT address passed in
+    if (ST.hasMergedShaders()) {
+      switch (MF.getFunction().getCallingConv()) {
+        case CallingConv::AMDGPU_HS:
+        case CallingConv::AMDGPU_GS:
+          // 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;
+      }
+    }
     BuildMI(MBB, I, DL, SMovB32, RsrcLo)
-      .addReg(AMDGPU::SGPR0) // Low address passed in
+      .addReg(GitPtrLo)
       .addReg(ScratchRsrcReg, RegState::ImplicitDefine);
 
     // We now have the GIT ptr - now get the scratch descriptor from the entry
     // at offset 0.
     PointerType *PtrTy =
       PointerType::get(Type::getInt64Ty(MF.getFunction().getContext()),
                        AMDGPUAS::CONSTANT_ADDRESS);
     MachinePointerInfo PtrInfo(UndefValue::get(PtrTy));

llvm/trunk/test/CodeGen/AMDGPU/amdpal_scratch_mergedshader.ll

+; RUN: llc -mtriple=amdgcn--amdpal -mcpu=gfx900 -verify-machineinstrs < %s | FileCheck -check-prefix=GCN -check-prefix=GFX9 -enable-var-scope %s
+
+; On gfx9 and later, a HS is a merged shader, in which s0-s7 are reserved by the
+; hardware, so the PAL puts the GIT (global information table) in s8 rather
+; than s0.
+
+; GCN-LABEL: {{^}}_amdgpu_hs_main:
+; GCN: s_getpc_b64 s{{\[}}[[GITPTR:[0-9]+]]:
+; PREGFX9: s_mov_b32 s[[GITPTR]], s0
+; GFX9: s_mov_b32 s[[GITPTR]], s8
+
+define amdgpu_hs void @_amdgpu_hs_main(i32 inreg %arg, i32 inreg %arg1, i32 inreg %arg2, i32 inreg %arg3, i32 inreg %arg4, i32 inreg %arg5, i32 inreg %arg6, i32 inreg %arg7, <6 x i32> inreg %arg8) {
+.entry:
+  %__llpc_global_proxy_7.i = alloca [3 x <4 x float>], align 16, addrspace(5)
+  %tmp = icmp ult i32 undef, undef
+  br i1 %tmp, label %.beginls, label %.endls
+
+.beginls:                                         ; preds = %.entry
+  %tmp15 = extractelement <6 x i32> %arg8, i32 3
+  %.0.vec.insert.i = insertelement <2 x i32> undef, i32 %tmp15, i32 0
+  %.4.vec.insert.i = shufflevector <2 x i32> %.0.vec.insert.i, <2 x i32> undef, <2 x i32> <i32 0, i32 3>
+  %tmp16 = bitcast <2 x i32> %.4.vec.insert.i to i64
+  br label %.endls
+
+.endls:                                           ; preds = %.beginls, %.entry
+  %.fca.2.gep120.i = getelementptr inbounds [3 x <4 x float>], [3 x <4 x float>] addrspace(5)* %__llpc_global_proxy_7.i, i64 0, i64 2
+  store <4 x float> <float 9.000000e+00, float 1.000000e+01, float 1.100000e+01, float 1.200000e+01>, <4 x float> addrspace(5)* %.fca.2.gep120.i, align 16
+  br label %bb
+
+bb:                                               ; preds = %bb, %.endls
+  %lsr.iv182 = phi [3 x <4 x float>] addrspace(5)* [ undef, %bb ], [ %__llpc_global_proxy_7.i, %.endls ]
+  %scevgep183 = getelementptr [3 x <4 x float>], [3 x <4 x float>] addrspace(5)* %lsr.iv182, i32 0, i32 1
+  br label %bb
+}