AMDGPU/GlobalISel: Fix flat_scratch_init handling for shaders

[arsenm]

Summary

I don't think this is actually defined for mesa, but this is what we
were doing on the DAG path.

[sebastian-ne]

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

505–509

Isn’t this already handled here for HSA?
I didn’t hear of mesa using flat scratch (and we don’t have any flat-scratch+mesa test), so should we update the dag code instead? I wonder why we are calling allocateHSAUserSGPRs for non-HSA platforms in the SDAG.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

505–509

The GISel version is factored a bit differently with a cleaner break between the kernel and non-kernel handling. allocateHSAUserSGPRs isn't called on the shader path (it's a bit ugly that it is for the DAG and it happens to work out).

Mesa should have some kind of ABI definition for flat scratch, but in the absence of one, we still need to produce something. This matches the DAG behavior. Alternatively I could switch mesa to follow what PAL does, but that would be a separate potentially breaking change

[sebastian-ne]

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

505–509

I’d prefer to change the SDAG to do the same for mesa as is done for PAL.
Duplicating the HSA init code for mesa seems wrong to me. I don’t think this would work correctly nor does the code look particularly appealing.
From what I can see, mesa still does not use flat scratch.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

505–509

I'm not attempting to address mesa's flat usage. if that change goes ahead, both would update in the future. For now gisel should just match behavior

[arsenm]

f482e86980a8580bae4ecef1b1dcd447ecb4484c

llvm/lib/Target/AMDGPU/AMDGPUCallLowering.cpp

static void allocateHSAUserSGPRs(CCState &CCInfo,
   }
 
   if (Info.hasDispatchID()) {
     Register DispatchIDReg = Info.addDispatchID(TRI);
     MF.addLiveIn(DispatchIDReg, &AMDGPU::SGPR_64RegClass);
     CCInfo.AllocateReg(DispatchIDReg);
   }
 
   if (Info.hasFlatScratchInit()) {
     Register FlatScratchInitReg = Info.addFlatScratchInit(TRI);
     MF.addLiveIn(FlatScratchInitReg, &AMDGPU::SGPR_64RegClass);
     CCInfo.AllocateReg(FlatScratchInitReg);
   }
 
   // TODO: Add GridWorkGroupCount user SGPRs when used. For now with HSA we read
   // these from the dispatch pointer.
 }
 
 bool AMDGPUCallLowering::lowerFormalArgumentsKernel(
     MachineIRBuilder &B, const Function &F,
     ArrayRef<ArrayRef<Register>> VRegs) const {
bool AMDGPUCallLowering::lowerFormalArguments(
   }
 
   if (Info->hasImplicitBufferPtr()) {
     Register ImplicitBufferPtrReg = Info->addImplicitBufferPtr(*TRI);
     MF.addLiveIn(ImplicitBufferPtrReg, &AMDGPU::SGPR_64RegClass);
     CCInfo.AllocateReg(ImplicitBufferPtrReg);
   }
 
+  // FIXME: This probably isn't defined for mesa
+  if (Info->hasFlatScratchInit() && !Subtarget.isAmdPalOS()) {
+    Register FlatScratchInitReg = Info->addFlatScratchInit(*TRI);
+    MF.addLiveIn(FlatScratchInitReg, &AMDGPU::SGPR_64RegClass);
+    CCInfo.AllocateReg(FlatScratchInitReg);
+  }
+
   SmallVector<ArgInfo, 32> SplitArgs;
   unsigned Idx = 0;
   unsigned PSInputNum = 0;
 
   // Insert the hidden sret parameter if the return value won't fit in the
   // return registers.
   if (!FLI.CanLowerReturn)
     insertSRetIncomingArgument(F, SplitArgs, FLI.DemoteRegister, MRI, DL);

llvm/test/CodeGen/AMDGPU/GlobalISel/flat-scratch-init.gfx.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -global-isel -amdgpu-enable-flat-scratch -mtriple=amdgcn-mesa-mesa3d -mcpu=gfx900 < %s | FileCheck -check-prefix=MESA %s
+; RUN: llc -global-isel -amdgpu-enable-flat-scratch -mtriple=amdgcn-amd-amdpal -mcpu=gfx900 < %s | FileCheck -check-prefix=PAL %s
+
+; Test that the initialization for flat_scratch doesn't crash.  PAL
+; doesn't add a user SGPR for initializing flat_scratch, mesa does
+; (although this probably isn't actually defined).
+
+define amdgpu_ps void @amdgpu_ps() {
+; MESA-LABEL: amdgpu_ps:
+; MESA:       ; %bb.0:
+; MESA-NEXT:    s_add_u32 flat_scratch_lo, s2, s4
+; MESA-NEXT:    s_getreg_b32 s0, hwreg(HW_REG_SH_MEM_BASES, 0, 16)
+; MESA-NEXT:    s_addc_u32 flat_scratch_hi, s3, 0
+; MESA-NEXT:    s_lshl_b32 s0, s0, 16
+; MESA-NEXT:    v_mov_b32_e32 v0, 4
+; MESA-NEXT:    v_mov_b32_e32 v1, s0
+; MESA-NEXT:    v_mov_b32_e32 v2, 0
+; MESA-NEXT:    flat_store_dword v[0:1], v2
+; MESA-NEXT:    s_waitcnt vmcnt(0)
+; MESA-NEXT:    s_endpgm
+;
+; PAL-LABEL: amdgpu_ps:
+; PAL:       ; %bb.0:
+; PAL-NEXT:    s_getpc_b64 s[2:3]
+; PAL-NEXT:    s_mov_b32 s2, s0
+; PAL-NEXT:    s_load_dwordx2 s[2:3], s[2:3], 0x0
+; PAL-NEXT:    v_mov_b32_e32 v0, 4
+; PAL-NEXT:    v_mov_b32_e32 v2, 0
+; PAL-NEXT:    s_waitcnt lgkmcnt(0)
+; PAL-NEXT:    s_and_b32 s3, s3, 0xffff
+; PAL-NEXT:    s_add_u32 flat_scratch_lo, s2, s0
+; PAL-NEXT:    s_getreg_b32 s0, hwreg(HW_REG_SH_MEM_BASES, 0, 16)
+; PAL-NEXT:    s_addc_u32 flat_scratch_hi, s3, 0
+; PAL-NEXT:    s_lshl_b32 s0, s0, 16
+; PAL-NEXT:    v_mov_b32_e32 v1, s0
+; PAL-NEXT:    flat_store_dword v[0:1], v2
+; PAL-NEXT:    s_waitcnt vmcnt(0)
+; PAL-NEXT:    s_endpgm
+  %alloca = alloca i32, addrspace(5)
+  %cast = addrspacecast i32 addrspace(5)* %alloca to i32*
+  store volatile i32 0, i32* %cast
+  ret void
+}
+