[AMDGPU] Fix VGPR spills where offset doesn't fit in 12 bits

[scott.linder]

Summary

Scale the offset of VGPR spills by the wave size when it cannot fit in the 12-bit offset immediate field and so is added to the soffset SGPR. This accounts for hardware swizzling of scratch memory.

[kzhuravl]

Needs a test

[t-tye]

lib/Target/AMDGPU/SIRegisterInfo.cpp

543 ↗

(On Diff #155954)

Setting this to Offset does not seem correct as the real value needs to be the scaled offset which is computed below. Here we just need the variable declared and it only has a valid value when RanOutOfSGPRs is true. So perhaps initialize to 0?

Probably a better approach is to eliminate RanOutOfSGPRs and instead just have ScratchOffsetRegDelta which is initialized to 0. Then set it when it is actually used (as already done). Then change the "if (RanOutOfSGPRs)" below to "if (ScratchOffsetRegDelta != 0)".

548 ↗

(On Diff #155954)

If this is introduced then seems it should be used as well in the NumSubRegs and Size calculations above which are also in terms of the EltSize.

550 ↗

(On Diff #155954)

Should this be Offset + Size - EltSize? Note that the loop is 0 based and increments by EltSize.

[arsenm]

lib/Target/AMDGPU/SIRegisterInfo.cpp

550 ↗

(On Diff #155954)

Yes, and this needs a very specific test to stress it

[scott.linder]

Addressed feedback, and added at least one test to exercise the fix and the condition for putting the offset in an SGPR.

I would like to add more tests for Offset + Size - EltSize where Size != EltSize but I have had some trouble getting a ValueReg with subregisters to survive until the spill occurs. E.g. if I load and store a <2 x i32> it is spilled as two distinct i32.

[arsenm]

In D49448#1166974, @scott.linder wrote:

Addressed feedback, and added at least one test to exercise the fix and the condition for putting the offset in an SGPR.

I would like to add more tests for Offset + Size - EltSize where Size != EltSize but I have had some trouble getting a ValueReg with subregisters to survive until the spill occurs. E.g. if I load and store a <2 x i32> it is spilled as two distinct i32.

You may have better success making all of these register clobbers in a single asm statement. Also -stress-regalloc will probably help make this simpler.

This is also a candidate for just making it a MIR test, but in this case I would probably prefer if the IR test worked

test/CodeGen/AMDGPU/spill-offset-calculation.ll

89 ↗

(On Diff #156125)

You can check the literal SGPR used for this

152 ↗

(On Diff #156125)

Should also have a test where the SGPR needs to be increment and restored (though I thought we already had one)

Also could use versions for non-kernels

[scott.linder]

test/CodeGen/AMDGPU/spill-offset-calculation.ll

152 ↗

(On Diff #156125)

This case is actually not scavenging another SGPR, it is just incrementing/decrementing the scratch offset SGPR. I don't know how to coerce CodeGen to do the spilling when the register scavenger is available (i.e. not in scavengeFrameVirtualRegs).

[arsenm]

test/CodeGen/AMDGPU/spill-offset-calculation.ll

152 ↗

(On Diff #156125)

OK, I assumed this was scavenging because you used a variable for it, as well as there being no check line for the restore. You should check the literal chosen register here because it will be fixed for the function, and need to check the restore.

[scott.linder]

Address some more feedback; use -stress-regalloc to cut down on the clobbers needed, add non-kernel tests, explicitly test the increment/decrement case, including the scratch offset SGPR.

I will try adding a MIR test to exercise the subregister condition and loop; with an IR test I don't know how to get a VGPR with subregs to survive "AMDGPU DAG->DAG Pattern Instruction Selection"

[arsenm]

In D49448#1168395, @scott.linder wrote:

Address some more feedback; use -stress-regalloc to cut down on the clobbers needed, add non-kernel tests, explicitly test the increment/decrement case, including the scratch offset SGPR.

I will try adding a MIR test to exercise the subregister condition and loop; with an IR test I don't know how to get a VGPR with subregs to survive "AMDGPU DAG->DAG Pattern Instruction Selection"

A physical register tuple inline asm use should work

[scott.linder]

Add tests for the the boundary condition for subregs.

Thank you for the help @arsenm; are there more cases I should be testing?

[arsenm]

test/CodeGen/AMDGPU/spill-offset-calculation.ll

179 ↗

(On Diff #156347)

This inline assembly is invalid, since it doesn't list the memory or physical register uses.

You don't need something meaningful here anyway, you should just need a comment with the register name to model the constraint

210 ↗

(On Diff #156347)

Ditto

[scott.linder]

Address feedback

[arsenm]

LGTM with the 2 remaining asms fixed

test/CodeGen/AMDGPU/spill-offset-calculation.ll

75 ↗

(On Diff #157513)

Missed one

106 ↗

(On Diff #157513)

Missed one

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

void SIRegisterInfo::buildSpillLoadStore(MachineBasicBlock::iterator MI,
   const GCNSubtarget &ST =  MF->getSubtarget<GCNSubtarget>();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const MachineFrameInfo &MFI = MF->getFrameInfo();
 
   const MCInstrDesc &Desc = TII->get(LoadStoreOp);
   const DebugLoc &DL = MI->getDebugLoc();
   bool IsStore = Desc.mayStore();
 
-  bool RanOutOfSGPRs = false;
   bool Scavenged = false;
   unsigned SOffset = ScratchOffsetReg;
 
+  const unsigned EltSize = 4;
   const TargetRegisterClass *RC = getRegClassForReg(MF->getRegInfo(), ValueReg);
-  unsigned NumSubRegs = AMDGPU::getRegBitWidth(RC->getID()) / 32;
-  unsigned Size = NumSubRegs * 4;
+  unsigned NumSubRegs = AMDGPU::getRegBitWidth(RC->getID()) / (EltSize * CHAR_BIT);
+  unsigned Size = NumSubRegs * EltSize;
   int64_t Offset = InstOffset + MFI.getObjectOffset(Index);
-  const int64_t OriginalImmOffset = Offset;
+  int64_t ScratchOffsetRegDelta = 0;
 
   unsigned Align = MFI.getObjectAlignment(Index);
   const MachinePointerInfo &BasePtrInfo = MMO->getPointerInfo();
 
-  if (!isUInt<12>(Offset + Size)) {
+  assert((Offset % EltSize) == 0 && "unexpected VGPR spill offset");
+
+  if (!isUInt<12>(Offset + Size - EltSize)) {
     SOffset = AMDGPU::NoRegister;
 
+    // We currently only support spilling VGPRs to EltSize boundaries, meaning
+    // we can simplify the adjustment of Offset here to just scale with
+    // WavefrontSize.
+    Offset *= ST.getWavefrontSize();
+
     // We don't have access to the register scavenger if this function is called
     // during  PEI::scavengeFrameVirtualRegs().
     if (RS)
       SOffset = RS->FindUnusedReg(&AMDGPU::SGPR_32RegClass);
 
     if (SOffset == AMDGPU::NoRegister) {
       // There are no free SGPRs, and since we are in the process of spilling
       // VGPRs too.  Since we need a VGPR in order to spill SGPRs (this is true
       // on SI/CI and on VI it is true until we implement spilling using scalar
       // stores), we have no way to free up an SGPR.  Our solution here is to
       // add the offset directly to the ScratchOffset register, and then
       // subtract the offset after the spill to return ScratchOffset to it's
       // original value.
-      RanOutOfSGPRs = true;
       SOffset = ScratchOffsetReg;
+      ScratchOffsetRegDelta = Offset;
     } else {
       Scavenged = true;
     }
 
     BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_ADD_U32), SOffset)
       .addReg(ScratchOffsetReg)
       .addImm(Offset);
 
     Offset = 0;
   }
 
-  const unsigned EltSize = 4;
-
   for (unsigned i = 0, e = NumSubRegs; i != e; ++i, Offset += EltSize) {
     unsigned SubReg = NumSubRegs == 1 ?
       ValueReg : getSubReg(ValueReg, getSubRegFromChannel(i));
 
     unsigned SOffsetRegState = 0;
     unsigned SrcDstRegState = getDefRegState(!IsStore);
     if (i + 1 == e) {
       SOffsetRegState |= getKillRegState(Scavenged);
auto MIB = BuildMI(*MBB, MI, DL, Desc)
       .addImm(0) // slc
       .addImm(0) // tfe
       .addMemOperand(NewMMO);
 
     if (NumSubRegs > 1)
       MIB.addReg(ValueReg, RegState::Implicit | SrcDstRegState);
   }
 
-  if (RanOutOfSGPRs) {
+  if (ScratchOffsetRegDelta != 0) {
     // Subtract the offset we added to the ScratchOffset register.
     BuildMI(*MBB, MI, DL, TII->get(AMDGPU::S_SUB_U32), ScratchOffsetReg)
         .addReg(ScratchOffsetReg)
-      .addImm(OriginalImmOffset);
+        .addImm(ScratchOffsetRegDelta);
   }
 }
 
 static std::pair<unsigned, unsigned> getSpillEltSize(unsigned SuperRegSize,
                                                      bool Store) {
   if (SuperRegSize % 16 == 0) {
     return { 16, Store ? AMDGPU::S_BUFFER_STORE_DWORDX4_SGPR :
                          AMDGPU::S_BUFFER_LOAD_DWORDX4_SGPR };

llvm/trunk/test/CodeGen/AMDGPU/spill-offset-calculation.ll

+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -enable-misched=0 -post-RA-scheduler=0 -stress-regalloc=8 < %s | FileCheck %s
+
+; Test that the VGPR spiller correctly switches to SGPR offsets when the
+; instruction offset field would overflow, and that it accounts for memory
+; swizzling.
+
+; CHECK-LABEL: test_inst_offset_kernel
+define amdgpu_kernel void @test_inst_offset_kernel() {
+entry:
+  ; Occupy 4092 bytes of scratch, so the offset of the spill of %a just fits in
+  ; the instruction offset field.
+  %alloca = alloca i8, i32 4088, align 4, addrspace(5)
+  %buf = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+
+  %aptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s{{[0-9]+}} offset:4092 ; 4-byte Folded Spill
+  %a = load volatile i32, i32 addrspace(5)* %aptr
+
+  ; Force %a to spill.
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  %outptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  store volatile i32 %a, i32 addrspace(5)* %outptr
+
+  ret void
+}
+
+; CHECK-LABEL: test_sgpr_offset_kernel
+define amdgpu_kernel void @test_sgpr_offset_kernel() {
+entry:
+  ; Occupy 4096 bytes of scratch, so the offset of the spill of %a does not
+  ; fit in the instruction, and has to live in the SGPR offset.
+  %alloca = alloca i8, i32 4092, align 4, addrspace(5)
+  %buf = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+
+  %aptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  ; 0x40000 / 64 = 4096 (for wave64)
+  ; CHECK: s_add_u32 s7, s7, 0x40000
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s7 ; 4-byte Folded Spill
+  ; CHECK: s_sub_u32 s7, s7, 0x40000
+  %a = load volatile i32, i32 addrspace(5)* %aptr
+
+  ; Force %a to spill
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  %outptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  store volatile i32 %a, i32 addrspace(5)* %outptr
+
+  ret void
+}
+
+; CHECK-LABEL: test_sgpr_offset_subregs_kernel
+define amdgpu_kernel void @test_sgpr_offset_subregs_kernel() {
+entry:
+  ; Occupy 4088 bytes of scratch, so that the spill of the last subreg of %a
+  ; still fits below offset 4096 (4088 + 8 - 4 = 4092), and can be placed in
+  ; the instruction offset field.
+  %alloca = alloca i8, i32 4084, align 4, addrspace(5)
+  %bufv1 = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+  %bufv2 = bitcast i8 addrspace(5)* %alloca to <2 x i32> addrspace(5)*
+
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s{{[0-9]+}} offset:4088 ; 4-byte Folded Spill
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s{{[0-9]+}} offset:4092 ; 4-byte Folded Spill
+  %aptr = getelementptr <2 x i32>, <2 x i32> addrspace(5)* %bufv2, i32 1
+  %a = load volatile <2 x i32>, <2 x i32> addrspace(5)* %aptr
+
+  ; Force %a to spill.
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  ; Ensure the alloca sticks around.
+  %bptr = getelementptr i32, i32 addrspace(5)* %bufv1, i32 1
+  %b = load volatile i32, i32 addrspace(5)* %bptr
+
+  ; Ensure the spill is of the full super-reg.
+  call void asm sideeffect "; $0", "r"(<2 x i32> %a)
+
+  ret void
+}
+
+; CHECK-LABEL: test_inst_offset_subregs_kernel
+define amdgpu_kernel void @test_inst_offset_subregs_kernel() {
+entry:
+  ; Occupy 4092 bytes of scratch, so that the spill of the last subreg of %a
+  ; does not fit below offset 4096 (4092 + 8 - 4 = 4096), and has to live
+  ; in the SGPR offset.
+  %alloca = alloca i8, i32 4088, align 4, addrspace(5)
+  %bufv1 = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+  %bufv2 = bitcast i8 addrspace(5)* %alloca to <2 x i32> addrspace(5)*
+
+  ; 0x3ff00 / 64 = 4092 (for wave64)
+  ; CHECK: s_add_u32 s7, s7, 0x3ff00
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s7 ; 4-byte Folded Spill
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s7 offset:4 ; 4-byte Folded Spill
+  ; CHECK: s_sub_u32 s7, s7, 0x3ff00
+  %aptr = getelementptr <2 x i32>, <2 x i32> addrspace(5)* %bufv2, i32 1
+  %a = load volatile <2 x i32>, <2 x i32> addrspace(5)* %aptr
+
+  ; Force %a to spill.
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  ; Ensure the alloca sticks around.
+  %bptr = getelementptr i32, i32 addrspace(5)* %bufv1, i32 1
+  %b = load volatile i32, i32 addrspace(5)* %bptr
+
+  ; Ensure the spill is of the full super-reg.
+  call void asm sideeffect "; $0", "r"(<2 x i32> %a)
+
+  ret void
+}
+
+; CHECK-LABEL: test_inst_offset_function
+define void @test_inst_offset_function() {
+entry:
+  ; Occupy 4092 bytes of scratch, so the offset of the spill of %a just fits in
+  ; the instruction offset field.
+  %alloca = alloca i8, i32 4088, align 4, addrspace(5)
+  %buf = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+
+  %aptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s{{[0-9]+}} offset:4092 ; 4-byte Folded Spill
+  %a = load volatile i32, i32 addrspace(5)* %aptr
+
+  ; Force %a to spill.
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  %outptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  store volatile i32 %a, i32 addrspace(5)* %outptr
+
+  ret void
+}
+
+; CHECK-LABEL: test_sgpr_offset_function
+define void @test_sgpr_offset_function() {
+entry:
+  ; Occupy 4096 bytes of scratch, so the offset of the spill of %a does not
+  ; fit in the instruction, and has to live in the SGPR offset.
+  %alloca = alloca i8, i32 4092, align 4, addrspace(5)
+  %buf = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+
+  %aptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  ; 0x40000 / 64 = 4096 (for wave64)
+  ; CHECK: s_add_u32 s5, s5, 0x40000
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s5 ; 4-byte Folded Spill
+  ; CHECK: s_sub_u32 s5, s5, 0x40000
+  %a = load volatile i32, i32 addrspace(5)* %aptr
+
+  ; Force %a to spill
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  %outptr = getelementptr i32, i32 addrspace(5)* %buf, i32 1
+  store volatile i32 %a, i32 addrspace(5)* %outptr
+
+  ret void
+}
+
+; CHECK-LABEL: test_sgpr_offset_subregs_function
+define void @test_sgpr_offset_subregs_function() {
+entry:
+  ; Occupy 4088 bytes of scratch, so that the spill of the last subreg of %a
+  ; still fits below offset 4096 (4088 + 8 - 4 = 4092), and can be placed in
+  ; the instruction offset field.
+  %alloca = alloca i8, i32 4084, align 4, addrspace(5)
+  %bufv1 = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+  %bufv2 = bitcast i8 addrspace(5)* %alloca to <2 x i32> addrspace(5)*
+
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s{{[0-9]+}} offset:4088 ; 4-byte Folded Spill
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s{{[0-9]+}} offset:4092 ; 4-byte Folded Spill
+  %aptr = getelementptr <2 x i32>, <2 x i32> addrspace(5)* %bufv2, i32 1
+  %a = load volatile <2 x i32>, <2 x i32> addrspace(5)* %aptr
+
+  ; Force %a to spill.
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  ; Ensure the alloca sticks around.
+  %bptr = getelementptr i32, i32 addrspace(5)* %bufv1, i32 1
+  %b = load volatile i32, i32 addrspace(5)* %bptr
+
+  ; Ensure the spill is of the full super-reg.
+  call void asm sideeffect "; $0", "r"(<2 x i32> %a)
+
+  ret void
+}
+
+; CHECK-LABEL: test_inst_offset_subregs_function
+define void @test_inst_offset_subregs_function() {
+entry:
+  ; Occupy 4092 bytes of scratch, so that the spill of the last subreg of %a
+  ; does not fit below offset 4096 (4092 + 8 - 4 = 4096), and has to live
+  ; in the SGPR offset.
+  %alloca = alloca i8, i32 4088, align 4, addrspace(5)
+  %bufv1 = bitcast i8 addrspace(5)* %alloca to i32 addrspace(5)*
+  %bufv2 = bitcast i8 addrspace(5)* %alloca to <2 x i32> addrspace(5)*
+
+  ; 0x3ff00 / 64 = 4092 (for wave64)
+  ; CHECK: s_add_u32 s5, s5, 0x3ff00
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s5 ; 4-byte Folded Spill
+  ; CHECK: buffer_store_dword v{{[0-9]+}}, off, s[{{[0-9]+:[0-9]+}}], s5 offset:4 ; 4-byte Folded Spill
+  ; CHECK: s_sub_u32 s5, s5, 0x3ff00
+  %aptr = getelementptr <2 x i32>, <2 x i32> addrspace(5)* %bufv2, i32 1
+  %a = load volatile <2 x i32>, <2 x i32> addrspace(5)* %aptr
+
+  ; Force %a to spill.
+  call void asm sideeffect "", "~{v0},~{v1},~{v2},~{v3},~{v4},~{v5},~{v6},~{v7}" ()
+
+  ; Ensure the alloca sticks around.
+  %bptr = getelementptr i32, i32 addrspace(5)* %bufv1, i32 1
+  %b = load volatile i32, i32 addrspace(5)* %bptr
+
+  ; Ensure the spill is of the full super-reg.
+  call void asm sideeffect "; $0", "r"(<2 x i32> %a)
+
+  ret void
+}