[AMDGPU] Enable serializing of argument info.

[hliao]

Summary

• Support serialization of all arguments in machine function info. This enables fabricating MIR tests depending on argument info.

[hliao]

llvm/test/CodeGen/AMDGPU/hazard-hidden-bundle.mir

6–7

As non-empty arg info has '}' as well, need to match the basic block label to ensure '}' matching the end of the bundle.

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

1082

Default mask value in ArgDescriptor is not 0, it is ~0.

[hliao]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

1082

Optional<unsigned> only return true when is assigned with values. It's not equivalent to the checking of 0.

[rampitec]

LGTM, but let's wait for Matt.

[arsenm]

Can you avoid having an explicit isReg field? It should be able to infer if it's a register or stack offset

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

129–132

braces

[hliao]

patch is enhanced to remove 'isReg'

[hliao]

minor style fixing

[arsenm]

LGTM

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

if (MFI->FrameOffsetReg != AMDGPU::FP_REG &&
     return diagnoseRegisterClass(YamlMFI.FrameOffsetReg);
   }
 
   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) {
+    // Skip parsing if it's not present.
+    if (!A)
+      return false;
+
+    if (A->IsRegister) {
+      unsigned Reg;
+      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());
+
+    return false;
+  };
+
+  if (YamlMFI.ArgInfo &&
+      (parseAndCheckArgument(YamlMFI.ArgInfo->PrivateSegmentBuffer,
+                             AMDGPU::SReg_128RegClass,
+                             MFI->ArgInfo.PrivateSegmentBuffer) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->DispatchPtr,
+                             AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.DispatchPtr) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->QueuePtr, AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.QueuePtr) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->KernargSegmentPtr,
+                             AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.KernargSegmentPtr) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->DispatchID,
+                             AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.DispatchID) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->FlatScratchInit,
+                             AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.FlatScratchInit) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->PrivateSegmentSize,
+                             AMDGPU::SGPR_32RegClass,
+                             MFI->ArgInfo.PrivateSegmentSize) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupIDX,
+                             AMDGPU::SGPR_32RegClass,
+                             MFI->ArgInfo.WorkGroupIDX) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupIDY,
+                             AMDGPU::SGPR_32RegClass,
+                             MFI->ArgInfo.WorkGroupIDY) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupIDZ,
+                             AMDGPU::SGPR_32RegClass,
+                             MFI->ArgInfo.WorkGroupIDZ) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->WorkGroupInfo,
+                             AMDGPU::SGPR_32RegClass,
+                             MFI->ArgInfo.WorkGroupInfo) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->PrivateSegmentWaveByteOffset,
+                             AMDGPU::SGPR_32RegClass,
+                             MFI->ArgInfo.PrivateSegmentWaveByteOffset) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->ImplicitArgPtr,
+                             AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.ImplicitArgPtr) ||
+       parseAndCheckArgument(YamlMFI.ArgInfo->ImplicitBufferPtr,
+                             AMDGPU::SReg_64RegClass,
+                             MFI->ArgInfo.ImplicitBufferPtr) ||
+       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;
 }

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h

public:
 
   void printCustom(raw_ostream &OS) const override {
     OS << "GWSResource";
   }
 };
 
 namespace yaml {
 
+struct SIArgument {
+  bool IsRegister;
+  union {
+    StringValue RegisterName;
+    unsigned StackOffset;
+  };
+  Optional<unsigned> Mask;
+
+  // Default constructor, which creates a stack argument.
+  SIArgument() : IsRegister(false), StackOffset(0) {}
+  SIArgument(const SIArgument &Other) {
+    IsRegister = Other.IsRegister;
+    if (IsRegister)
+      ::new ((void *)std::addressof(RegisterName))
+          StringValue(Other.RegisterName);
+    else
+      StackOffset = Other.StackOffset;
+    Mask = Other.Mask;
+  }
+  SIArgument &operator=(const SIArgument &Other) {
+    IsRegister = Other.IsRegister;
+    if (IsRegister)
+      ::new ((void *)std::addressof(RegisterName))
+          StringValue(Other.RegisterName);
+    else
+      StackOffset = Other.StackOffset;
+    Mask = Other.Mask;
+    return *this;
+  }
+  ~SIArgument() {
+    if (IsRegister)
+      RegisterName.~StringValue();
+  }
+
+  // Helper to create a register or stack argument.
+  static inline SIArgument createArgument(bool IsReg) {
+    if (IsReg)
+      return SIArgument(IsReg);
+    return SIArgument();
+  }
+
+private:
+  // Construct a register argument.
+  SIArgument(bool) : IsRegister(true), RegisterName() {}
+};
+
+template <> struct MappingTraits<SIArgument> {
+  static void mapping(IO &YamlIO, SIArgument &A) {
+    if (YamlIO.outputting()) {
+      if (A.IsRegister)
+        YamlIO.mapRequired("reg", A.RegisterName);
+      else
+        YamlIO.mapRequired("offset", A.StackOffset);
+    } else {
+      auto Keys = YamlIO.keys();
+      if (is_contained(Keys, "reg")) {
+        A = SIArgument::createArgument(true);
+        YamlIO.mapRequired("reg", A.RegisterName);
+      } else if (is_contained(Keys, "offset")) {
+        YamlIO.mapRequired("offset", A.StackOffset);
+      } else {
+        YamlIO.setError("missing required key 'reg' or 'offset'");
+      }
+    }
+    YamlIO.mapOptional("mask", A.Mask);
+  }
+  static const bool flow = true;
+};
+
+struct SIArgumentInfo {
+  Optional<SIArgument> PrivateSegmentBuffer;
+  Optional<SIArgument> DispatchPtr;
+  Optional<SIArgument> QueuePtr;
+  Optional<SIArgument> KernargSegmentPtr;
+  Optional<SIArgument> DispatchID;
+  Optional<SIArgument> FlatScratchInit;
+  Optional<SIArgument> PrivateSegmentSize;
+
+  Optional<SIArgument> WorkGroupIDX;
+  Optional<SIArgument> WorkGroupIDY;
+  Optional<SIArgument> WorkGroupIDZ;
+  Optional<SIArgument> WorkGroupInfo;
+  Optional<SIArgument> PrivateSegmentWaveByteOffset;
+
+  Optional<SIArgument> ImplicitArgPtr;
+  Optional<SIArgument> ImplicitBufferPtr;
+
+  Optional<SIArgument> WorkItemIDX;
+  Optional<SIArgument> WorkItemIDY;
+  Optional<SIArgument> WorkItemIDZ;
+};
+
+template <> struct MappingTraits<SIArgumentInfo> {
+  static void mapping(IO &YamlIO, SIArgumentInfo &AI) {
+    YamlIO.mapOptional("privateSegmentBuffer", AI.PrivateSegmentBuffer);
+    YamlIO.mapOptional("dispatchPtr", AI.DispatchPtr);
+    YamlIO.mapOptional("queuePtr", AI.QueuePtr);
+    YamlIO.mapOptional("kernargSegmentPtr", AI.KernargSegmentPtr);
+    YamlIO.mapOptional("dispatchID", AI.DispatchID);
+    YamlIO.mapOptional("flatScratchInit", AI.FlatScratchInit);
+    YamlIO.mapOptional("privateSegmentSize", AI.PrivateSegmentSize);
+
+    YamlIO.mapOptional("workGroupIDX", AI.WorkGroupIDX);
+    YamlIO.mapOptional("workGroupIDY", AI.WorkGroupIDY);
+    YamlIO.mapOptional("workGroupIDZ", AI.WorkGroupIDZ);
+    YamlIO.mapOptional("workGroupInfo", AI.WorkGroupInfo);
+    YamlIO.mapOptional("privateSegmentWaveByteOffset",
+                       AI.PrivateSegmentWaveByteOffset);
+
+    YamlIO.mapOptional("implicitArgPtr", AI.ImplicitArgPtr);
+    YamlIO.mapOptional("implicitBufferPtr", AI.ImplicitBufferPtr);
+
+    YamlIO.mapOptional("workItemIDX", AI.WorkItemIDX);
+    YamlIO.mapOptional("workItemIDY", AI.WorkItemIDY);
+    YamlIO.mapOptional("workItemIDZ", AI.WorkItemIDZ);
+  }
+};
+
 struct SIMachineFunctionInfo final : public yaml::MachineFunctionInfo {
   uint64_t ExplicitKernArgSize = 0;
   unsigned MaxKernArgAlign = 0;
   unsigned LDSSize = 0;
   bool IsEntryFunction = false;
   bool NoSignedZerosFPMath = false;
   bool MemoryBound = false;
   bool WaveLimiter = false;
 
   StringValue ScratchRSrcReg = "$private_rsrc_reg";
   StringValue ScratchWaveOffsetReg = "$scratch_wave_offset_reg";
   StringValue FrameOffsetReg = "$fp_reg";
   StringValue StackPtrOffsetReg = "$sp_reg";
 
+  Optional<SIArgumentInfo> ArgInfo;
+
   SIMachineFunctionInfo() = default;
   SIMachineFunctionInfo(const llvm::SIMachineFunctionInfo &,
                         const TargetRegisterInfo &TRI);
 
   void mappingImpl(yaml::IO &YamlIO) override;
   ~SIMachineFunctionInfo() = default;
 };
 
static void mapping(IO &YamlIO, SIMachineFunctionInfo &MFI) {
     YamlIO.mapOptional("scratchRSrcReg", MFI.ScratchRSrcReg,
                        StringValue("$private_rsrc_reg"));
     YamlIO.mapOptional("scratchWaveOffsetReg", MFI.ScratchWaveOffsetReg,
                        StringValue("$scratch_wave_offset_reg"));
     YamlIO.mapOptional("frameOffsetReg", MFI.FrameOffsetReg,
                        StringValue("$fp_reg"));
     YamlIO.mapOptional("stackPtrOffsetReg", MFI.StackPtrOffsetReg,
                        StringValue("$sp_reg"));
+    YamlIO.mapOptional("argumentInfo", MFI.ArgInfo);
   }
 };
 
 } // end namespace yaml
 
 /// This class keeps track of the SPI_SP_INPUT_ADDR config register, which
 /// tells the hardware which interpolation parameters to load.
 class SIMachineFunctionInfo final : public AMDGPUMachineFunction {

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

static yaml::StringValue regToString(unsigned Reg,
   yaml::StringValue Dest;
   {
     raw_string_ostream OS(Dest.Value);
     OS << printReg(Reg, &TRI);
   }
   return Dest;
 }
 
+static Optional<yaml::SIArgumentInfo>
+convertArgumentInfo(const AMDGPUFunctionArgInfo &ArgInfo,
+                    const TargetRegisterInfo &TRI) {
+  yaml::SIArgumentInfo AI;
+
+  auto convertArg = [&](Optional<yaml::SIArgument> &A,
+                        const ArgDescriptor &Arg) {
+    if (!Arg)
+      return false;
+
+    // Create a register or stack argument.
+    yaml::SIArgument SA = yaml::SIArgument::createArgument(Arg.isRegister());
+    if (Arg.isRegister()) {
+      raw_string_ostream OS(SA.RegisterName.Value);
+      OS << printReg(Arg.getRegister(), &TRI);
+    } else
+      SA.StackOffset = Arg.getStackOffset();
+    // Check and update the optional mask.
+    if (Arg.isMasked())
+      SA.Mask = Arg.getMask();
+
+    A = SA;
+    return true;
+  };
+
+  bool Any = false;
+  Any |= convertArg(AI.PrivateSegmentBuffer, ArgInfo.PrivateSegmentBuffer);
+  Any |= convertArg(AI.DispatchPtr, ArgInfo.DispatchPtr);
+  Any |= convertArg(AI.QueuePtr, ArgInfo.QueuePtr);
+  Any |= convertArg(AI.KernargSegmentPtr, ArgInfo.KernargSegmentPtr);
+  Any |= convertArg(AI.DispatchID, ArgInfo.DispatchID);
+  Any |= convertArg(AI.FlatScratchInit, ArgInfo.FlatScratchInit);
+  Any |= convertArg(AI.PrivateSegmentSize, ArgInfo.PrivateSegmentSize);
+  Any |= convertArg(AI.WorkGroupIDX, ArgInfo.WorkGroupIDX);
+  Any |= convertArg(AI.WorkGroupIDY, ArgInfo.WorkGroupIDY);
+  Any |= convertArg(AI.WorkGroupIDZ, ArgInfo.WorkGroupIDZ);
+  Any |= convertArg(AI.WorkGroupInfo, ArgInfo.WorkGroupInfo);
+  Any |= convertArg(AI.PrivateSegmentWaveByteOffset,
+                    ArgInfo.PrivateSegmentWaveByteOffset);
+  Any |= convertArg(AI.ImplicitArgPtr, ArgInfo.ImplicitArgPtr);
+  Any |= convertArg(AI.ImplicitBufferPtr, ArgInfo.ImplicitBufferPtr);
+  Any |= convertArg(AI.WorkItemIDX, ArgInfo.WorkItemIDX);
+  Any |= convertArg(AI.WorkItemIDY, ArgInfo.WorkItemIDY);
+  Any |= convertArg(AI.WorkItemIDZ, ArgInfo.WorkItemIDZ);
+
+  if (Any)
+    return AI;
+
+  return None;
+}
+
 yaml::SIMachineFunctionInfo::SIMachineFunctionInfo(
   const llvm::SIMachineFunctionInfo& MFI,
   const TargetRegisterInfo &TRI)
   : ExplicitKernArgSize(MFI.getExplicitKernArgSize()),
     MaxKernArgAlign(MFI.getMaxKernArgAlign()),
     LDSSize(MFI.getLDSSize()),
     IsEntryFunction(MFI.isEntryFunction()),
     NoSignedZerosFPMath(MFI.hasNoSignedZerosFPMath()),
     MemoryBound(MFI.isMemoryBound()),
     WaveLimiter(MFI.needsWaveLimiter()),
     ScratchRSrcReg(regToString(MFI.getScratchRSrcReg(), TRI)),
     ScratchWaveOffsetReg(regToString(MFI.getScratchWaveOffsetReg(), TRI)),
     FrameOffsetReg(regToString(MFI.getFrameOffsetReg(), TRI)),
-    StackPtrOffsetReg(regToString(MFI.getStackPtrOffsetReg(), TRI)) {}
+    StackPtrOffsetReg(regToString(MFI.getStackPtrOffsetReg(), TRI)),
+    ArgInfo(convertArgumentInfo(MFI.getArgInfo(), TRI)) {}
 
 void yaml::SIMachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
   MappingTraits<SIMachineFunctionInfo>::mapping(YamlIO, *this);
 }
 
 bool SIMachineFunctionInfo::initializeBaseYamlFields(
   const yaml::SIMachineFunctionInfo &YamlMFI) {
   ExplicitKernArgSize = YamlMFI.ExplicitKernArgSize;
   MaxKernArgAlign = YamlMFI.MaxKernArgAlign;
   LDSSize = YamlMFI.LDSSize;
   IsEntryFunction = YamlMFI.IsEntryFunction;
   NoSignedZerosFPMath = YamlMFI.NoSignedZerosFPMath;
   MemoryBound = YamlMFI.MemoryBound;
   WaveLimiter = YamlMFI.WaveLimiter;
   return false;
 }

llvm/test/CodeGen/AMDGPU/hazard-hidden-bundle.mir

 # RUN: llc -march=amdgcn -mcpu=gfx902 -verify-machineinstrs -run-pass  post-RA-hazard-rec %s -o - | FileCheck -check-prefixes=GCN,XNACK,GCX9 %s
 # RUN: llc -march=amdgcn -mcpu=gfx900 -verify-machineinstrs -run-pass  post-RA-hazard-rec %s -o - | FileCheck -check-prefixes=GCN,NOXNACK,GFX9 %s
 # RUN: llc -march=amdgcn -mcpu=gfx1010 -mattr=-WavefrontSize32,+WavefrontSize64 -verify-machineinstrs -run-pass  post-RA-hazard-rec %s -o - | FileCheck -check-prefixes=GCN,NOXNACK,GFX10 %s
 
 # GCN-LABEL: name: break_smem_clause_simple_load_smrd8_ptr_hidden_bundle
+# GCN: bb.0:
 # GCN:          }
 # XNACK-NEXT:   S_NOP
 # NOXNACK-NOT:  S_NOP
 # GCN:          S_LOAD_DWORDX2_IMM
 ---
 name: break_smem_clause_simple_load_smrd8_ptr_hidden_bundle
 body: |
   bb.0:
     BUNDLE implicit-def $sgpr6_sgpr7 {

llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir

 # FULL-NEXT: isEntryFunction: true
 # FULL-NEXT: noSignedZerosFPMath: false
 # FULL-NEXT: memoryBound:     true
 # FULL-NEXT: waveLimiter:     true
 # FULL-NEXT: scratchRSrcReg:  '$sgpr8_sgpr9_sgpr10_sgpr11'
 # FULL-NEXT: scratchWaveOffsetReg:  '$sgpr12'
 # FULL-NEXT: frameOffsetReg:  '$sgpr12'
 # FULL-NEXT: stackPtrOffsetReg:  '$sgpr13'
+# FULL-NEXT: argumentInfo:
+# FULL-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# FULL-NEXT: kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+# FULL-NEXT: workGroupIDX: { reg: '$sgpr6' }
+# FULL-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+# FULL-NEXT: workItemIDX: { reg: '$vgpr0' }
 # FULL-NEXT: body:
 
 # SIMPLE: machineFunctionInfo:
 # SIMPLE-NEXT: explicitKernArgSize: 128
 # SIMPLE-NEXT: maxKernArgAlign: 64
 # SIMPLE-NEXT: ldsSize: 2048
 # SIMPLE-NEXT: isEntryFunction: true
 # SIMPLE-NEXT: memoryBound: true
 # SIMPLE-NEXT: waveLimiter: true
 # SIMPLE-NEXT: scratchRSrcReg: '$sgpr8_sgpr9_sgpr10_sgpr11'
 # SIMPLE-NEXT: scratchWaveOffsetReg:  '$sgpr12'
 # SIMPLE-NEXT: frameOffsetReg:  '$sgpr12'
 # SIMPLE-NEXT: stackPtrOffsetReg:  '$sgpr13'
+# SIMPLE-NEXT: argumentInfo:
+# SIMPLE-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# SIMPLE-NEXT: kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+# SIMPLE-NEXT: workGroupIDX: { reg: '$sgpr6' }
+# SIMPLE-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+# SIMPLE-NEXT: workItemIDX: { reg: '$vgpr0' }
 # SIMPLE-NEXT: body:
 name: kernel0
 machineFunctionInfo:
   explicitKernArgSize: 128
   maxKernArgAlign: 64
   ldsSize: 2048
   isEntryFunction: true
   noSignedZerosFPMath: false
   memoryBound:     true
   waveLimiter:     true
   scratchRSrcReg:  '$sgpr8_sgpr9_sgpr10_sgpr11'
   scratchWaveOffsetReg:  '$sgpr12'
   frameOffsetReg: '$sgpr12'
   stackPtrOffsetReg:  '$sgpr13'
+  argumentInfo:
+    privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+    kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+    workGroupIDX:    { reg: '$sgpr6' }
+    privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+    workItemIDX:     { reg: '$vgpr0' }
 body:             |
   bb.0:
     S_ENDPGM 0
 
 ...
 
 # FIXME: Should be able to not print section for simple
 ---
 # ALL-LABEL: name: no_mfi
 # FULL: machineFunctionInfo:
 # FULL-NEXT: explicitKernArgSize: 0
 # FULL-NEXT: maxKernArgAlign: 0
 # FULL-NEXT: ldsSize: 0
 # FULL-NEXT: isEntryFunction: false
 # FULL-NEXT: noSignedZerosFPMath: false
 # FULL-NEXT: memoryBound:     false
 # FULL-NEXT: waveLimiter:     false
 # FULL-NEXT: scratchRSrcReg:  '$private_rsrc_reg'
 # FULL-NEXT: scratchWaveOffsetReg:  '$scratch_wave_offset_reg'
 # FULL-NEXT: frameOffsetReg:  '$fp_reg'
 # FULL-NEXT: stackPtrOffsetReg:  '$sp_reg'
+# FULL-NEXT: argumentInfo:
+# FULL-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# FULL-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
+# FULL-NEXT: body:
 
 # SIMPLE: machineFunctionInfo:
+# SIMPLE-NEXT: argumentInfo:
+# SIMPLE-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# SIMPLE-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
 # SIMPLE-NEXT: body:
 
 name: no_mfi
 body:             |
   bb.0:
     S_ENDPGM 0
 
 ...
 
 ---
 # ALL-LABEL: name: empty_mfi
 # FULL: machineFunctionInfo:
 # FULL-NEXT: explicitKernArgSize: 0
 # FULL-NEXT: maxKernArgAlign: 0
 # FULL-NEXT: ldsSize: 0
 # FULL-NEXT: isEntryFunction: false
 # FULL-NEXT: noSignedZerosFPMath: false
 # FULL-NEXT: memoryBound:     false
 # FULL-NEXT: waveLimiter:     false
 # FULL-NEXT: scratchRSrcReg:  '$private_rsrc_reg'
 # FULL-NEXT: scratchWaveOffsetReg:  '$scratch_wave_offset_reg'
 # FULL-NEXT: frameOffsetReg:  '$fp_reg'
 # FULL-NEXT: stackPtrOffsetReg:  '$sp_reg'
+# FULL-NEXT: argumentInfo:
+# FULL-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# FULL-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
+# FULL-NEXT: body:
 
 # SIMPLE: machineFunctionInfo:
+# SIMPLE-NEXT: argumentInfo:
+# SIMPLE-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# SIMPLE-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
 # SIMPLE-NEXT: body:
 
 name: empty_mfi
 machineFunctionInfo:
 body:             |
   bb.0:
     S_ENDPGM 0
 
 ...
 
 ---
 # ALL-LABEL: name: empty_mfi_entry_func
 # FULL: machineFunctionInfo:
 # FULL-NEXT: explicitKernArgSize: 0
 # FULL-NEXT: maxKernArgAlign: 0
 # FULL-NEXT: ldsSize: 0
 # FULL-NEXT: isEntryFunction: true
 # FULL-NEXT: noSignedZerosFPMath: false
 # FULL-NEXT: memoryBound:     false
 # FULL-NEXT: waveLimiter:     false
 # FULL-NEXT: scratchRSrcReg:  '$private_rsrc_reg'
 # FULL-NEXT: scratchWaveOffsetReg:  '$scratch_wave_offset_reg'
 # FULL-NEXT: frameOffsetReg:  '$fp_reg'
 # FULL-NEXT: stackPtrOffsetReg:  '$sp_reg'
+# FULL-NEXT: argumentInfo:
+# FULL-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# FULL-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
+# FULL-NEXT: body:
 
 # SIMPLE: machineFunctionInfo:
 # SIMPLE-NEXT: isEntryFunction: true
+# SIMPLE-NEXT: argumentInfo:
+# SIMPLE-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# SIMPLE-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
 # SIMPLE-NEXT: body:
 
 name: empty_mfi_entry_func
 machineFunctionInfo:
   isEntryFunction: true
 body:             |
   bb.0:
     S_ENDPGM 0

 machineFunctionInfo:
   scratchRSrcReg:  '$private_rsrc_reg'
 
 body:             |
   bb.0:
     S_ENDPGM 0
 
 ...
+
+---
+# ALL-LABEL: name: fake_stack_arginfo
+
+# FULL: argumentInfo:
+# FULL-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# FULL-NEXT: flatScratchInit: { offset: 4 }
+# FULL-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
+# FULL-NEXT: workItemIDY: { reg: '$vgpr0', mask: 65280 }
+
+# SIMPLE: argumentInfo:
+# SIMPLE-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+# SIMPLE-NEXT: flatScratchInit: { offset: 4 }
+# SIMPLE-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
+# SIMPLE-NEXT: workItemIDY: { reg: '$vgpr0', mask: 65280 }
+name: fake_stack_arginfo
+machineFunctionInfo:
+  argumentInfo:
+    flatScratchInit: { offset: 4 }
+    workItemIDY:     { reg: '$vgpr0' , mask: 0xff00 }
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

llvm/test/CodeGen/MIR/AMDGPU/machine-function-info.ll

 ; CHECK-NEXT: isEntryFunction: true
 ; CHECK-NEXT: noSignedZerosFPMath: false
 ; CHECK-NEXT: memoryBound: false
 ; CHECK-NEXT: waveLimiter: false
 ; CHECK-NEXT: scratchRSrcReg:  '$sgpr96_sgpr97_sgpr98_sgpr99'
 ; CHECK-NEXT: scratchWaveOffsetReg: '$sgpr101'
 ; CHECK-NEXT: frameOffsetReg:  '$sgpr101'
 ; CHECK-NEXT: stackPtrOffsetReg: '$sgpr101'
+; CHECK-NEXT: argumentInfo:
+; CHECK-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+; CHECK-NEXT: kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+; CHECK-NEXT: workGroupIDX: { reg: '$sgpr6' }
+; CHECK-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+; CHECK-NEXT: workItemIDX: { reg: '$vgpr0' }
 ; CHECK-NEXT: body:
 define amdgpu_kernel void @kernel(i32 %arg0, i64 %arg1, <16 x i32> %arg2) {
   %gep = getelementptr inbounds [512 x float], [512 x float] addrspace(3)* @lds, i32 0, i32 %arg0
   store float 0.0, float addrspace(3)* %gep, align 4
   ret void
 }
 
 ; CHECK-LABEL: {{^}}name: ps_shader
 ; CHECK: machineFunctionInfo:
 ; CHECK-NEXT: explicitKernArgSize: 0
 ; CHECK-NEXT: maxKernArgAlign: 0
 ; CHECK-NEXT: ldsSize: 0
 ; CHECK-NEXT: isEntryFunction: true
 ; CHECK-NEXT: noSignedZerosFPMath: false
 ; CHECK-NEXT: memoryBound: false
 ; CHECK-NEXT: waveLimiter: false
 ; CHECK-NEXT: scratchRSrcReg:  '$sgpr96_sgpr97_sgpr98_sgpr99'
 ; CHECK-NEXT: scratchWaveOffsetReg: '$sgpr101'
 ; CHECK-NEXT: frameOffsetReg:  '$sgpr101'
 ; CHECK-NEXT: stackPtrOffsetReg: '$sgpr101'
+; CHECK-NEXT: argumentInfo:
+; CHECK-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr3' }
+; CHECK-NEXT: implicitBufferPtr: { reg: '$sgpr0_sgpr1' }
 ; CHECK-NEXT: body:
 define amdgpu_ps void @ps_shader(i32 %arg0, i32 inreg %arg1) {
   ret void
 }
 
 ; CHECK-LABEL: {{^}}name: function
 ; CHECK: machineFunctionInfo:
 ; CHECK-NEXT: explicitKernArgSize: 0
 ; CHECK-NEXT: maxKernArgAlign: 0
 ; CHECK-NEXT: ldsSize: 0
 ; CHECK-NEXT: isEntryFunction: false
 ; CHECK-NEXT: noSignedZerosFPMath: false
 ; CHECK-NEXT: memoryBound: false
 ; CHECK-NEXT: waveLimiter: false
 ; CHECK-NEXT: scratchRSrcReg: '$sgpr0_sgpr1_sgpr2_sgpr3'
 ; CHECK-NEXT: scratchWaveOffsetReg: '$sgpr33'
 ; CHECK-NEXT: frameOffsetReg: '$sgpr5'
 ; CHECK-NEXT: stackPtrOffsetReg: '$sgpr32'
+; CHECK-NEXT: argumentInfo:
+; CHECK-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+; CHECK-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
 ; CHECK-NEXT: body:
 define void @function() {
   ret void
 }
 
 ; CHECK-LABEL: {{^}}name: function_nsz
 ; CHECK: machineFunctionInfo:
 ; CHECK-NEXT: explicitKernArgSize: 0
 ; CHECK-NEXT: maxKernArgAlign: 0
 ; CHECK-NEXT: ldsSize: 0
 ; CHECK-NEXT: isEntryFunction: false
 ; CHECK-NEXT: noSignedZerosFPMath: true
 ; CHECK-NEXT: memoryBound: false
 ; CHECK-NEXT: waveLimiter: false
 ; CHECK-NEXT: scratchRSrcReg: '$sgpr0_sgpr1_sgpr2_sgpr3'
 ; CHECK-NEXT: scratchWaveOffsetReg: '$sgpr33'
 ; CHECK-NEXT: frameOffsetReg: '$sgpr5'
 ; CHECK-NEXT: stackPtrOffsetReg: '$sgpr32'
+; CHECK-NEXT: argumentInfo:
+; CHECK-NEXT: privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+; CHECK-NEXT: privateSegmentWaveByteOffset: { reg: '$sgpr33' }
 ; CHECK-NEXT: body:
 define void @function_nsz() #0 {
   ret void
 }
 
 attributes #0 = { "no-signed-zeros-fp-math" = "true" }