[AMDGPU] Update NumUserSGPRs and NumSystemSGPRs in MIR parsing.

[hliao]

NumUserSGPRs and NumSystemSGPRs need updating during the parsing argument info.

[arsenm]

Test? I’m also note sure we get much value by tracking these here. I don’t think they are accurate for shaders

[hliao]

In D64144#1569828, @arsenm wrote:

Test? I’m also note sure we get much value by tracking these here. I don’t think they are accurate for shaders

There are several places behaving slightly differently by checking getNumPreloadedSGPRs, which is the sum of NumUserSGPRs and NumSystemSGRPs. The test in https://reviews.llvm.org/D64145 relies on that to avoid assigning a preloaded register again.

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

bool GCNTargetMachine::parseMachineFunctionInfo(
 
   if (MFI->StackPtrOffsetReg != AMDGPU::SP_REG &&
       !AMDGPU::SGPR_32RegClass.contains(MFI->StackPtrOffsetReg)) {
     return diagnoseRegisterClass(YamlMFI.StackPtrOffsetReg);
   }
 
   auto parseAndCheckArgument = [&](const Optional<yaml::SIArgument> &A,
                                    const TargetRegisterClass &RC,
-                                   ArgDescriptor &Arg) {
+                                   ArgDescriptor &Arg, unsigned UserSGPRs = 0,
+                                   unsigned SystemSGPRs = 0) {
     // Skip parsing if it's not present.
     if (!A)
       return false;
 
     if (A->IsRegister) {
       unsigned Reg;
-      if (parseNamedRegisterReference(PFS, Reg, A->RegisterName.Value,
-                                      Error)) {
+      if (parseNamedRegisterReference(PFS, Reg, A->RegisterName.Value, Error)) {
         SourceRange = A->RegisterName.SourceRange;
         return true;
       }
       if (!RC.contains(Reg))
         return diagnoseRegisterClass(A->RegisterName);
       Arg = ArgDescriptor::createRegister(Reg);
     } else
       Arg = ArgDescriptor::createStack(A->StackOffset);
     // Check and apply the optional mask.
     if (A->Mask)
       Arg = ArgDescriptor::createArg(Arg, A->Mask.getValue());
 
+    MFI->NumUserSGPRs += UserSGPRs;
+    MFI->NumSystemSGPRs += SystemSGPRs;
     return false;
   };
 
   if (YamlMFI.ArgInfo &&
       (parseAndCheckArgument(YamlMFI.ArgInfo->PrivateSegmentBuffer,
                              AMDGPU::SReg_128RegClass,
-                             MFI->ArgInfo.PrivateSegmentBuffer) ||
+                             MFI->ArgInfo.PrivateSegmentBuffer, 4, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->DispatchPtr,
-                             AMDGPU::SReg_64RegClass,
-                             MFI->ArgInfo.DispatchPtr) ||
+                             AMDGPU::SReg_64RegClass, MFI->ArgInfo.DispatchPtr,
+                             2, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->QueuePtr, AMDGPU::SReg_64RegClass,
-                             MFI->ArgInfo.QueuePtr) ||
+                             MFI->ArgInfo.QueuePtr, 2, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->KernargSegmentPtr,
                              AMDGPU::SReg_64RegClass,
-                             MFI->ArgInfo.KernargSegmentPtr) ||
+                             MFI->ArgInfo.KernargSegmentPtr, 2, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->DispatchID,
-                             AMDGPU::SReg_64RegClass,
-                             MFI->ArgInfo.DispatchID) ||
+                             AMDGPU::SReg_64RegClass, MFI->ArgInfo.DispatchID,
+                             2, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->FlatScratchInit,
                              AMDGPU::SReg_64RegClass,
-                             MFI->ArgInfo.FlatScratchInit) ||
+                             MFI->ArgInfo.FlatScratchInit, 2, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->PrivateSegmentSize,
                              AMDGPU::SGPR_32RegClass,
                              MFI->ArgInfo.PrivateSegmentSize) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupIDX,
-                             AMDGPU::SGPR_32RegClass,
-                             MFI->ArgInfo.WorkGroupIDX) ||
+                             AMDGPU::SGPR_32RegClass, MFI->ArgInfo.WorkGroupIDX,
+                             0, 1) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupIDY,
-                             AMDGPU::SGPR_32RegClass,
-                             MFI->ArgInfo.WorkGroupIDY) ||
+                             AMDGPU::SGPR_32RegClass, MFI->ArgInfo.WorkGroupIDY,
+                             0, 1) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupIDZ,
-                             AMDGPU::SGPR_32RegClass,
-                             MFI->ArgInfo.WorkGroupIDZ) ||
+                             AMDGPU::SGPR_32RegClass, MFI->ArgInfo.WorkGroupIDZ,
+                             0, 1) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupInfo,
                              AMDGPU::SGPR_32RegClass,
-                             MFI->ArgInfo.WorkGroupInfo) ||
+                             MFI->ArgInfo.WorkGroupInfo, 0, 1) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->PrivateSegmentWaveByteOffset,
                              AMDGPU::SGPR_32RegClass,
-                             MFI->ArgInfo.PrivateSegmentWaveByteOffset) ||
+                             MFI->ArgInfo.PrivateSegmentWaveByteOffset, 0, 1) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->ImplicitArgPtr,
                              AMDGPU::SReg_64RegClass,
                              MFI->ArgInfo.ImplicitArgPtr) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->ImplicitBufferPtr,
                              AMDGPU::SReg_64RegClass,
-                             MFI->ArgInfo.ImplicitBufferPtr) ||
+                             MFI->ArgInfo.ImplicitBufferPtr, 2, 0) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkItemIDX,
                              AMDGPU::VGPR_32RegClass,
                              MFI->ArgInfo.WorkItemIDX) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkItemIDY,
                              AMDGPU::VGPR_32RegClass,
                              MFI->ArgInfo.WorkItemIDY) ||
        parseAndCheckArgument(YamlMFI.ArgInfo->WorkItemIDZ,
                              AMDGPU::VGPR_32RegClass,
                              MFI->ArgInfo.WorkItemIDZ)))
     return true;
 
   return false;
 }