[AMDGPU] Correct the setting of `FlatScratchInit`.

[hliao]

Summary

• That flag setting should skip spilling stack slot.

[hliao]

FlatScratchInit is set by checking whether there are any stack objects in the early stage of codegen. As we could load MIR in any stage, even after RA, it may set that flag incorrectly by counting spilling stack slot.

[hliao]

llvm/test/CodeGen/AMDGPU/undefined-physreg-sgpr-spill.mir

44

this test case needs revising as the max-sgprs is incorrectly calculated when FlatScratchInit is set wrongly. That results in $sgpr10 and sgpr11 being reserved. After fixing FlatScratchInit setting, we need to add them into livein registers.

[arsenm]

LGTM

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
   } else if (ST.isMesaGfxShader(F)) {
     ImplicitBufferPtr = true;
   }
 
   if (F.hasFnAttribute("amdgpu-kernarg-segment-ptr"))
     KernargSegmentPtr = true;
 
   if (ST.hasFlatAddressSpace() && isEntryFunction() && isAmdHsaOrMesa) {
+    auto hasNonSpillStackObjects = [&]() {
+      // Avoid expensive checking if there's no stack objects.
+      if (!HasStackObjects)
+        return false;
+      for (auto OI = FrameInfo.getObjectIndexBegin(),
+                OE = FrameInfo.getObjectIndexEnd(); OI != OE; ++OI)
+        if (!FrameInfo.isSpillSlotObjectIndex(OI))
+          return true;
+      // All stack objects are spill slots.
+      return false;
+    };
     // 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"))
+    if (hasNonSpillStackObjects() || 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);
 

llvm/test/CodeGen/AMDGPU/undefined-physreg-sgpr-spill.mir

stack:
   - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 4,
       stack-id: sgpr-spill, callee-saved-register: '', callee-saved-restored: true,
       debug-info-variable: '', debug-info-expression: '',
       debug-info-location: '' }
 constants:
 body:             |
   bb.0:
     successors: %bb.1, %bb.2
-    liveins: $vgpr0, $sgpr4_sgpr5
+    liveins: $vgpr0, $sgpr4_sgpr5, $sgpr8_sgpr9_sgpr10_sgpr11, $sgpr13
 
     $vgpr1_vgpr2 = COPY killed $sgpr4_sgpr5, implicit $exec
     $vgpr1 = GLOBAL_LOAD_UBYTE killed $vgpr1_vgpr2, 0, 0, 0, 0, implicit $exec :: (non-temporal dereferenceable invariant load 1 from `i1 addrspace(4)* undef`)
     $vcc = V_CMP_NE_U32_e64 0, $vgpr0, implicit $exec
     $sgpr0_sgpr1 = V_CMP_EQ_U32_e64 1, killed $vgpr1, implicit $exec
     $vgpr1 = V_CNDMASK_B32_e64 0, 0, 0, -1, killed $sgpr0_sgpr1, implicit $exec
     $sgpr0_sgpr1 = COPY $exec, implicit-def $exec
     SI_SPILL_S64_SAVE $sgpr0_sgpr1, %stack.0, implicit $exec, implicit $sgpr8_sgpr9_sgpr10_sgpr11, implicit $sgpr13, implicit-def dead $m0 :: (store 8 into %stack.0, align 4, addrspace 5)
     $sgpr2_sgpr3 = S_AND_B64 killed $sgpr0_sgpr1, killed $vcc, implicit-def dead $scc
     $exec = S_MOV_B64_term killed $sgpr2_sgpr3
     SI_MASK_BRANCH %bb.2, implicit $exec
     S_BRANCH %bb.1
 
   bb.1:
     successors: %bb.3(0x80000000)
     liveins: $vgpr0, $vgpr1
 
     $sgpr2_sgpr3 = S_MOV_B64 0
     $vgpr2 = V_MOV_B32_e32 0, implicit $exec
     $sgpr4_sgpr5 = IMPLICIT_DEF
     S_BRANCH %bb.3
 
   bb.2:
     successors:
+    liveins: $sgpr8_sgpr9_sgpr10_sgpr11, $sgpr13
 
     $sgpr0_sgpr1 = SI_SPILL_S64_RESTORE %stack.0, implicit $exec, implicit $sgpr8_sgpr9_sgpr10_sgpr11, implicit $sgpr13, implicit-def dead $m0 :: (load 8 from %stack.0, align 4, addrspace 5)
     $exec = S_OR_B64 $exec, killed $sgpr0_sgpr1, implicit-def $scc
 
   bb.3:
     liveins: $vgpr0, $vgpr1, $vgpr2, $sgpr2_sgpr3, $sgpr4_sgpr5
 
     $vcc = COPY $vgpr1
stack:
   - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 4,
       stack-id: sgpr-spill, callee-saved-register: '', callee-saved-restored: true,
       debug-info-variable: '', debug-info-expression: '',
       debug-info-location: '' }
 constants:
 body:             |
   bb.0:
     successors: %bb.1, %bb.2
-    liveins: $vgpr0, $sgpr4_sgpr5
+    liveins: $vgpr0, $sgpr4_sgpr5, $sgpr8_sgpr9_sgpr10_sgpr11, $sgpr13
 
     $vgpr1_vgpr2 = COPY killed $sgpr4_sgpr5, implicit $exec
     $vgpr1 = GLOBAL_LOAD_UBYTE killed $vgpr1_vgpr2, 0, 0, 0, 0, implicit $exec :: (non-temporal dereferenceable invariant load 1 from `i1 addrspace(4)* undef`)
     $vcc = V_CMP_NE_U32_e64 0, $vgpr0, implicit $exec
     $sgpr0_sgpr1 = V_CMP_EQ_U32_e64 1, killed $vgpr1, implicit $exec
     $vgpr1 = V_CNDMASK_B32_e64 0, 0, 0, -1, killed $sgpr0_sgpr1, implicit $exec
     $sgpr0_sgpr1 = COPY $exec, implicit-def $exec
     $sgpr2_sgpr3 = S_AND_B64 $sgpr0_sgpr1, killed $vcc, implicit-def dead $scc
     SI_SPILL_S64_SAVE killed $sgpr0_sgpr1, %stack.0, implicit $exec, implicit $sgpr8_sgpr9_sgpr10_sgpr11, implicit $sgpr13, implicit-def dead $m0 :: (store 8 into %stack.0, align 4, addrspace 5)
     $exec = S_MOV_B64_term killed $sgpr2_sgpr3
     SI_MASK_BRANCH %bb.2, implicit $exec
     S_BRANCH %bb.1
 
   bb.1:
     successors: %bb.3(0x80000000)
     liveins: $vgpr0, $vgpr1
 
     $sgpr2_sgpr3 = S_MOV_B64 0
     $vgpr2 = V_MOV_B32_e32 0, implicit $exec
     $sgpr4_sgpr5 = IMPLICIT_DEF
     S_BRANCH %bb.3
 
   bb.2:
     successors:
+    liveins: $sgpr8_sgpr9_sgpr10_sgpr11, $sgpr13
 
     $sgpr0_sgpr1 = SI_SPILL_S64_RESTORE %stack.0, implicit $exec, implicit $sgpr8_sgpr9_sgpr10_sgpr11, implicit $sgpr13, implicit-def dead $m0 :: (load 8 from %stack.0, align 4, addrspace 5)
     $exec = S_OR_B64 $exec, killed $sgpr0_sgpr1, implicit-def $scc
 
   bb.3:
     liveins: $vgpr0, $vgpr1, $vgpr2, $sgpr2_sgpr3, $sgpr4_sgpr5
 
     $vcc = COPY $vgpr1
     S_ENDPGM 0
 
 ...