AMDGPU: AMDPAL scratch buffer support

[tpr]

Summary

(From TimC, with further modifications.)

Added support for scratch (including spilling) for OS type amdpal:
generates code to set up the scratch descriptor if it is needed.

With amdpal, the scratch resource descriptor is loaded from offset 0 of
the global information table. The low 32 bits of the address of the
global information table is passed in s0.

Added amdgpu-git-ptr-high function attribute to hard-wire the high 32
bits of the address of the global information table. If the function
attribute is not specified, or is 0xffffffff, then the backend generates
code to use the high 32 bits of pc.

The documentation for the AMDPAL ABI will be added in a later commit.

[arsenm]

lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

920

I don't think we actually care what order these are printed in. Can you just move all of these to one place?

lib/Target/AMDGPU/SIFrameLowering.cpp

358–416

This looks like it breaks mesa

395–396

Sort the add MMO to last

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

163–164

You can combine these and check if it's empty

test/CodeGen/AMDGPU/amdpal.ll

11

-enable-var-scope on new tests

16

mising -LABEL

[arsenm]

lib/Target/AMDGPU/SIFrameLowering.cpp

358–416

Oh, I see that just moved. Can you maybe extract this + the mesa setup to a separate function?

[tpr]

Addressed review comments. I'm a bit nervous though that the changes are now
more intrusive for the non-amdpal case, as I am not set up to test anything
other than amdpal.

[arsenm]

lib/Target/AMDGPU/SIFrameLowering.cpp

392

Why implicit def of s0? I don't think this should be necessary

[tpr]

Addressed review comment.

[arsenm]

LGTM

lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

899

This looks like it's already here, but while you're here can you remove this check in a separate patch?

lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

if (AMDGPU::isCompute(MF.getFunction()->getCallingConv())) {
     OutStreamer->EmitIntValue(S_00B860_WAVESIZE(CurrentProgramInfo.ScratchBlocks), 4);
 
     // TODO: Should probably note flat usage somewhere. SC emits a "FlatPtr32 =
     // 0" comment but I don't see a corresponding field in the register spec.
   } else {
     OutStreamer->EmitIntValue(RsrcReg, 4);
     OutStreamer->EmitIntValue(S_00B028_VGPRS(CurrentProgramInfo.VGPRBlocks) |
                               S_00B028_SGPRS(CurrentProgramInfo.SGPRBlocks), 4);
+    unsigned Rsrc2Val = 0;
     if (STM.isVGPRSpillingEnabled(*MF.getFunction())) {
       OutStreamer->EmitIntValue(R_0286E8_SPI_TMPRING_SIZE, 4);
       OutStreamer->EmitIntValue(S_0286E8_WAVESIZE(CurrentProgramInfo.ScratchBlocks), 4);
+      if (TM.getTargetTriple().getOS() == Triple::AMDPAL)
+        Rsrc2Val = S_00B84C_SCRATCH_EN(CurrentProgramInfo.ScratchBlocks > 0);
     }
-  }
-
     if (MF.getFunction()->getCallingConv() == CallingConv::AMDGPU_PS) {
-    OutStreamer->EmitIntValue(R_00B02C_SPI_SHADER_PGM_RSRC2_PS, 4);
-    OutStreamer->EmitIntValue(S_00B02C_EXTRA_LDS_SIZE(CurrentProgramInfo.LDSBlocks), 4);
       OutStreamer->EmitIntValue(R_0286CC_SPI_PS_INPUT_ENA, 4);
       OutStreamer->EmitIntValue(MFI->getPSInputEnable(), 4);
       OutStreamer->EmitIntValue(R_0286D0_SPI_PS_INPUT_ADDR, 4);
       OutStreamer->EmitIntValue(MFI->getPSInputAddr(), 4);
+      Rsrc2Val |= S_00B02C_EXTRA_LDS_SIZE(CurrentProgramInfo.LDSBlocks);
+    }
+    if (Rsrc2Val) {
+      OutStreamer->EmitIntValue(RsrcReg + 4 /*rsrc2*/, 4);
+      OutStreamer->EmitIntValue(Rsrc2Val, 4);
+    }
   }
 
   OutStreamer->EmitIntValue(R_SPILLED_SGPRS, 4);
   OutStreamer->EmitIntValue(MFI->getNumSpilledSGPRs(), 4);
   OutStreamer->EmitIntValue(R_SPILLED_VGPRS, 4);
   OutStreamer->EmitIntValue(MFI->getNumSpilledVGPRs(), 4);
 }
 
 // This is supposed to be log2(Size)
 static amd_element_byte_size_t getElementByteSizeValue(unsigned Size) {
   switch (Size) {
   case 4:
     return AMD_ELEMENT_4_BYTES;

lib/Target/AMDGPU/SIFrameLowering.h

std::pair<unsigned, unsigned> getReservedPrivateSegmentWaveByteOffsetReg(
     const SIInstrInfo *TII,
     const SIRegisterInfo *TRI,
     SIMachineFunctionInfo *MFI,
     MachineFunction &MF) const;
 
   /// \brief Emits debugger prologue.
   void emitDebuggerPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const;
 
+  // Emit scratch setup code for AMDPAL or Mesa, assuming ResourceRegUsed is set.
+  void emitEntryFunctionScratchSetup(const SISubtarget &ST, MachineFunction &MF,
+      MachineBasicBlock &MBB, SIMachineFunctionInfo *MFI,
+      MachineBasicBlock::iterator I, unsigned PreloadedPrivateBufferReg,
+      unsigned ScratchRsrcReg) const;
+
 public:
   bool hasFP(const MachineFunction &MF) const override;
   bool hasSP(const MachineFunction &MF) const;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_SIFRAMELOWERING_H

lib/Target/AMDGPU/SIFrameLowering.cpp

SIFrameLowering::getReservedPrivateSegmentWaveByteOffsetReg(
   return std::make_pair(ScratchWaveOffsetReg, SPReg);
 }
 
 void SIFrameLowering::emitEntryFunctionPrologue(MachineFunction &MF,
                                                 MachineBasicBlock &MBB) const {
   // Emit debugger prologue if "amdgpu-debugger-emit-prologue" attribute was
   // specified.
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
-  auto AMDGPUASI = ST.getAMDGPUAS();
   if (ST.debuggerEmitPrologue())
     emitDebuggerPrologue(MF, MBB);
 
   assert(&MF.front() == &MBB && "Shrink-wrapping not yet supported");
 
   SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
 
   // If we only have SGPR spills, we won't actually be using scratch memory
BuildMI(MBB, I, DL, TII->get(AMDGPU::COPY), ScratchWaveOffsetReg)
               MRI.isPhysRegUsed(ScratchWaveOffsetReg) ? 0 : RegState::Kill);
   }
 
   if (CopyBuffer && !CopyBufferFirst) {
     BuildMI(MBB, I, DL, TII->get(AMDGPU::COPY), ScratchRsrcReg)
       .addReg(PreloadedPrivateBufferReg, RegState::Kill);
   }
 
-  if (ResourceRegUsed && (ST.isMesaGfxShader(MF) || (PreloadedPrivateBufferReg == AMDGPU::NoRegister))) {
+  if (ResourceRegUsed)
+    emitEntryFunctionScratchSetup(ST, MF, MBB, MFI, I,
+        PreloadedPrivateBufferReg, ScratchRsrcReg);
+}
+
+// Emit scratch setup code for AMDPAL or Mesa, assuming ResourceRegUsed is set.
+void SIFrameLowering::emitEntryFunctionScratchSetup(const SISubtarget &ST,
+      MachineFunction &MF, MachineBasicBlock &MBB, SIMachineFunctionInfo *MFI,
+      MachineBasicBlock::iterator I, unsigned PreloadedPrivateBufferReg,
+      unsigned ScratchRsrcReg) const {
+
+  const SIInstrInfo *TII = ST.getInstrInfo();
+  const SIRegisterInfo *TRI = &TII->getRegisterInfo();
+  DebugLoc DL;
+  auto AMDGPUASI = ST.getAMDGPUAS();
+
+  if (ST.isAmdPalOS()) {
+    // The pointer to the GIT is formed from the offset passed in and either
+    // the amdgpu-git-ptr-high function attribute or the top part of the PC
+    unsigned RsrcLo = TRI->getSubReg(ScratchRsrcReg, AMDGPU::sub0);
+    unsigned RsrcHi = TRI->getSubReg(ScratchRsrcReg, AMDGPU::sub1);
+    unsigned Rsrc01 = TRI->getSubReg(ScratchRsrcReg, AMDGPU::sub0_sub1);
+
+    const MCInstrDesc &SMovB32 = TII->get(AMDGPU::S_MOV_B32);
+
+    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);
+    }
+    BuildMI(MBB, I, DL, SMovB32, RsrcLo)
+      .addReg(AMDGPU::SGPR0) // Low address passed in
+      .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));
+    const MCInstrDesc &LoadDwordX4 = TII->get(AMDGPU::S_LOAD_DWORDX4_IMM);
+    auto MMO = MF.getMachineMemOperand(PtrInfo,
+                                       MachineMemOperand::MOLoad |
+                                       MachineMemOperand::MOInvariant |
+                                       MachineMemOperand::MODereferenceable,
+                                       0, 0);
+    BuildMI(MBB, I, DL, LoadDwordX4, ScratchRsrcReg)
+      .addReg(Rsrc01)
+      .addImm(0) // offset
+      .addImm(0) // glc
+      .addReg(ScratchRsrcReg, RegState::ImplicitDefine)
+      .addMemOperand(MMO);
+    return;
+  }
+  if (ST.isMesaGfxShader(MF)
+      || (PreloadedPrivateBufferReg == AMDGPU::NoRegister)) {
     assert(!ST.isAmdCodeObjectV2(MF));
     const MCInstrDesc &SMovB32 = TII->get(AMDGPU::S_MOV_B32);
 
     unsigned Rsrc2 = TRI->getSubReg(ScratchRsrcReg, AMDGPU::sub2);
     unsigned Rsrc3 = TRI->getSubReg(ScratchRsrcReg, AMDGPU::sub3);
 
     // Use relocations to get the pointer, and setup the other bits manually.
     uint64_t Rsrc23 = TII->getScratchRsrcWords23();

lib/Target/AMDGPU/SIMachineFunctionInfo.h

private:
   // Compute directly in sgpr[0:1]
   // Other shaders indirect 64-bits at sgpr[0:1]
   bool ImplicitBufferPtr : 1;
 
   // Pointer to where the ABI inserts special kernel arguments separate from the
   // user arguments. This is an offset from the KernargSegmentPtr.
   bool ImplicitArgPtr : 1;
 
+  // The hard-wired high half of the address of the global information table
+  // for AMDPAL OS type. 0xffffffff represents no hard-wired high half, since
+  // current hardware only allows a 16 bit value.
+  unsigned GITPtrHigh;
+
   MCPhysReg getNextUserSGPR() const {
     assert(NumSystemSGPRs == 0 && "System SGPRs must be added after user SGPRs");
     return AMDGPU::SGPR0 + NumUserSGPRs;
   }
 
   MCPhysReg getNextSystemSGPR() const {
     return AMDGPU::SGPR0 + NumUserSGPRs + NumSystemSGPRs;
   }
public:
   getPreloadedValue(AMDGPUFunctionArgInfo::PreloadedValue Value) const {
     return ArgInfo.getPreloadedValue(Value);
   }
 
   unsigned getPreloadedReg(AMDGPUFunctionArgInfo::PreloadedValue Value) const {
     return ArgInfo.getPreloadedValue(Value).first->getRegister();
   }
 
+  unsigned getGITPtrHigh() const {
+    return GITPtrHigh;
+  }
+
   unsigned getNumUserSGPRs() const {
     return NumUserSGPRs;
   }
 
   unsigned getNumPreloadedSGPRs() const {
     return NumUserSGPRs + NumSystemSGPRs;
   }
 

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

: AMDGPUMachineFunction(MF),
     WorkGroupIDY(false),
     WorkGroupIDZ(false),
     WorkGroupInfo(false),
     PrivateSegmentWaveByteOffset(false),
     WorkItemIDX(false),
     WorkItemIDY(false),
     WorkItemIDZ(false),
     ImplicitBufferPtr(false),
-    ImplicitArgPtr(false) {
+    ImplicitArgPtr(false),
+    GITPtrHigh(0xffffffff) {
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
   const Function *F = MF.getFunction();
   FlatWorkGroupSizes = ST.getFlatWorkGroupSizes(*F);
   WavesPerEU = ST.getWavesPerEU(*F);
 
   if (!isEntryFunction()) {
     // Non-entry functions have no special inputs for now, other registers
     // required for scratch access.
if (F->hasFnAttribute("amdgpu-kernarg-segment-ptr"))
     KernargSegmentPtr = true;
 
   if (ST.hasFlatAddressSpace() && isEntryFunction() && IsCOV2) {
     // TODO: This could be refined a lot. The attribute is a poor way of
     // detecting calls that may require it before argument lowering.
     if (HasStackObjects || F->hasFnAttribute("amdgpu-flat-scratch"))
       FlatScratchInit = true;
   }
+
+  Attribute A = F->getFnAttribute("amdgpu-git-ptr-high");
+  StringRef S = A.getValueAsString();
+  if (!S.empty())
+    S.consumeInteger(0, GITPtrHigh);
 }
 
 unsigned SIMachineFunctionInfo::addPrivateSegmentBuffer(
   const SIRegisterInfo &TRI) {
   ArgInfo.PrivateSegmentBuffer =
     ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_128RegClass));
   NumUserSGPRs += 4;

test/CodeGen/AMDGPU/amdpal.ll

-; RUN: llc < %s -mtriple=amdgcn--amdpal -mcpu=tahiti | FileCheck --check-prefix=PAL %s
+; RUN: llc < %s -mtriple=amdgcn--amdpal -mcpu=tahiti | FileCheck --check-prefix=PAL --enable-var-scope %s
 
 ; PAL: .AMDGPU.config
 
 define amdgpu_kernel void @simple(i32 addrspace(1)* %out) {
 entry:
   store i32 0, i32 addrspace(1)* %out
   ret void
 }
 
+; Check code sequence for amdpal use of scratch for alloca. This is the case
+; where the high half of the address comes from s_getpc.
+
+; PAL-LABEL: {{^}}scratch:
+; PAL: s_getpc_b64 s{{\[}}[[GITPTR:[0-9]+]]:
+; PAL: s_mov_b32 s[[GITPTR]], s0
+; PAL: s_load_dwordx4 s{{\[}}[[SCRATCHDESC:[0-9]+]]:{{[0-9]+]}}, s{{\[}}[[GITPTR]]:
+; PAL: buffer_store{{.*}}, s{{\[}}[[SCRATCHDESC]]:
+
+define amdgpu_kernel void @scratch(<2 x i32> %in, i32 %idx, i32* %out) {
+entry:
+  %v = alloca [2 x i32]
+  %vv = bitcast [2 x i32]* %v to <2 x i32>*
+  store <2 x i32> %in, <2 x i32>* %vv
+  %e = getelementptr [2 x i32], [2 x i32]* %v, i32 0, i32 %idx
+  %x = load i32, i32* %e
+  store i32 %x, i32* %out
+  ret void
+}
+
+; Check code sequence for amdpal use of scratch for alloca. This is the case
+; where the amdgpu-git-ptr-high function attribute gives the high half of the
+; address to use.
+; Looks like you can't do arithmetic on a filecheck variable, so we can't test
+; that the s_movk_i32 is into a reg that is one more than the following
+; s_mov_b32.
+
+; PAL-LABEL: {{^}}scratch2:
+; 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]]:
+; PAL: buffer_store{{.*}}, s{{\[}}[[SCRATCHDESC]]:
+
+define amdgpu_kernel void @scratch2(<2 x i32> %in, i32 %idx, i32* %out) #0 {
+entry:
+  %v = alloca [2 x i32]
+  %vv = bitcast [2 x i32]* %v to <2 x i32>*
+  store <2 x i32> %in, <2 x i32>* %vv
+  %e = getelementptr [2 x i32], [2 x i32]* %v, i32 0, i32 %idx
+  %x = load i32, i32* %e
+  store i32 %x, i32* %out
+  ret void
+}
+
+attributes #0 = { nounwind "amdgpu-git-ptr-high"="0x1234" }