[AMDGPU] Serialize MFInfo::ScavengeFI

Serialize ScavengeFI from SIMachineFunctionInfo into yaml.

ScavengeFI is not used outside of the PrologEpilogInserter,
so this shouldn't change anything.

Differential Revision: https://reviews.llvm.org/D101367

include/llvm/CodeGen/MIRYamlMapping.h

@@ -381,6 +381,50 @@ template <> struct MappingTraits<FixedMachineStackObject> {
   static const bool flow = true;
 };
 
+/// A serializaable representation of a reference to a stack object or fixed
+/// stack object.
+struct FrameIndex {
+  // The frame index as printed. This is always a positive number, even for
+  // fixed objects. To obtain the real index,
+  // MachineFrameInfo::getObjectIndexBegin has to be added.
+  int FI;
+  bool IsFixed;
+  SMRange SourceRange;
+
+  FrameIndex() {}
+  FrameIndex(int FI, const llvm::MachineFrameInfo &MFI);
+
+  Expected<int> getFI(const llvm::MachineFrameInfo &MFI) const;
+};
+
+template <> struct ScalarTraits<FrameIndex> {
+  static void output(const FrameIndex &FI, void *, raw_ostream &OS) {
+    MachineOperand::printStackObjectReference(OS, FI.FI, FI.IsFixed, "");
+  }
+
+  static StringRef input(StringRef Scalar, void *Ctx, FrameIndex &FI) {
+    FI.IsFixed = false;
+    StringRef Num;
+    if (Scalar.startswith("%stack.")) {
+      Num = Scalar.substr(7);
+    } else if (Scalar.startswith("%fixed-stack.")) {
+      Num = Scalar.substr(13);
+      FI.IsFixed = true;
+    } else {
+      return "Invalid frame index, needs to start with %stack. or "
+             "%fixed-stack.";
+    }
+    if (Num.consumeInteger(10, FI.FI))
+      return "Invalid frame index, not a valid number";
+
+    if (const auto *Node =
+            reinterpret_cast<yaml::Input *>(Ctx)->getCurrentNode())
+      FI.SourceRange = Node->getSourceRange();
+    return StringRef();
+  }
+
+  static QuotingType mustQuote(StringRef S) { return needsQuotes(S); }
+};
 
 /// Serializable representation of CallSiteInfo.
 struct CallSiteInfo {

lib/CodeGen/CMakeLists.txt

@@ -108,6 +108,7 @@ add_llvm_component_library(LLVMCodeGen
   MachineStripDebug.cpp
   MachineTraceMetrics.cpp
   MachineVerifier.cpp
+  MIRYamlMapping.cpp
   ModuloSchedule.cpp
   MultiHazardRecognizer.cpp
   PatchableFunction.cpp

lib/CodeGen/MIRYamlMapping.cpp

@@ -0,0 +1,43 @@
+//===- MIRYamlMapping.cpp - Describe mapping between MIR and YAML ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the mapping between various MIR data structures and
+// their corresponding YAML representation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MIRYamlMapping.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace llvm;
+using namespace llvm::yaml;
+
+FrameIndex::FrameIndex(int FI, const llvm::MachineFrameInfo &MFI) {
+  IsFixed = MFI.isFixedObjectIndex(FI);
+  if (IsFixed)
+    FI -= MFI.getObjectIndexBegin();
+  this->FI = FI;
+}
+
+// Returns the value and if the frame index is fixed or not.
+Expected<int> FrameIndex::getFI(const llvm::MachineFrameInfo &MFI) const {
+  int FI = this->FI;
+  if (IsFixed) {
+    if (unsigned(FI) >= MFI.getNumFixedObjects())
+      return make_error<StringError>(
+          formatv("invalid fixed frame index {0}", FI).str(),
+          inconvertibleErrorCode());
+    FI += MFI.getObjectIndexBegin();
+  }
+  if (unsigned(FI + MFI.getNumFixedObjects()) >= MFI.getNumObjects())
+    return make_error<StringError>(formatv("invalid frame index {0}", FI).str(),
+                                   inconvertibleErrorCode());
+  return FI;
+}

lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -1237,8 +1237,8 @@ yaml::MachineFunctionInfo *GCNTargetMachine::createDefaultFuncInfoYAML() const {
 yaml::MachineFunctionInfo *
 GCNTargetMachine::convertFuncInfoToYAML(const MachineFunction &MF) const {
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
-  return new yaml::SIMachineFunctionInfo(*MFI,
-                                         *MF.getSubtarget().getRegisterInfo());
+  return new yaml::SIMachineFunctionInfo(
+      *MFI, *MF.getSubtarget().getRegisterInfo(), MF);
 }
 
 bool GCNTargetMachine::parseMachineFunctionInfo(
@@ -1249,7 +1249,8 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
   MachineFunction &MF = PFS.MF;
   SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
 
-  MFI->initializeBaseYamlFields(YamlMFI);
+  if (MFI->initializeBaseYamlFields(YamlMFI, MF, PFS, Error, SourceRange))
+    return true;
 
   if (MFI->Occupancy == 0) {
     // Fixup the subtarget dependent default value.

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

@@ -8,6 +8,7 @@
 
 #include "SIMachineFunctionInfo.h"
 #include "AMDGPUTargetMachine.h"
+#include "llvm/CodeGen/MIRParser/MIParser.h"
 
 #define MAX_LANES 64
 
@@ -547,7 +548,8 @@ convertArgumentInfo(const AMDGPUFunctionArgInfo &ArgInfo,
 }
 
 yaml::SIMachineFunctionInfo::SIMachineFunctionInfo(
-    const llvm::SIMachineFunctionInfo &MFI, const TargetRegisterInfo &TRI)
+    const llvm::SIMachineFunctionInfo &MFI, const TargetRegisterInfo &TRI,
+    const llvm::MachineFunction &MF)
     : ExplicitKernArgSize(MFI.getExplicitKernArgSize()),
       MaxKernArgAlign(MFI.getMaxKernArgAlign()), LDSSize(MFI.getLDSSize()),
       DynLDSAlign(MFI.getDynLDSAlign()), IsEntryFunction(MFI.isEntryFunction()),
@@ -561,6 +563,9 @@ yaml::SIMachineFunctionInfo::SIMachineFunctionInfo(
       FrameOffsetReg(regToString(MFI.getFrameOffsetReg(), TRI)),
       StackPtrOffsetReg(regToString(MFI.getStackPtrOffsetReg(), TRI)),
       ArgInfo(convertArgumentInfo(MFI.getArgInfo(), TRI)), Mode(MFI.getMode()) {
+  auto SFI = MFI.getOptionalScavengeFI();
+  if (SFI)
+    ScavengeFI = yaml::FrameIndex(*SFI, MF.getFrameInfo());
 }
 
 void yaml::SIMachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
@@ -568,7 +573,8 @@ void yaml::SIMachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
 }
 
 bool SIMachineFunctionInfo::initializeBaseYamlFields(
-  const yaml::SIMachineFunctionInfo &YamlMFI) {
+    const yaml::SIMachineFunctionInfo &YamlMFI, const MachineFunction &MF,
+    PerFunctionMIParsingState &PFS, SMDiagnostic &Error, SMRange &SourceRange) {
   ExplicitKernArgSize = YamlMFI.ExplicitKernArgSize;
   MaxKernArgAlign = assumeAligned(YamlMFI.MaxKernArgAlign);
   LDSSize = YamlMFI.LDSSize;
@@ -581,6 +587,24 @@ bool SIMachineFunctionInfo::initializeBaseYamlFields(
   WaveLimiter = YamlMFI.WaveLimiter;
   HasSpilledSGPRs = YamlMFI.HasSpilledSGPRs;
   HasSpilledVGPRs = YamlMFI.HasSpilledVGPRs;
+
+  if (YamlMFI.ScavengeFI) {
+    auto FIOrErr = YamlMFI.ScavengeFI->getFI(MF.getFrameInfo());
+    if (!FIOrErr) {
+      // Create a diagnostic for a the frame index.
+      const MemoryBuffer &Buffer =
+          *PFS.SM->getMemoryBuffer(PFS.SM->getMainFileID());
+
+      Error = SMDiagnostic(*PFS.SM, SMLoc(), Buffer.getBufferIdentifier(), 1, 1,
+                           SourceMgr::DK_Error, toString(FIOrErr.takeError()),
+                           "", None, None);
+      SourceRange = YamlMFI.ScavengeFI->SourceRange;
+      return true;
+    }
+    ScavengeFI = *FIOrErr;
+  } else {
+    ScavengeFI = None;
+  }
   return false;
 }
 

lib/Target/AMDGPU/SIMachineFunctionInfo.h

@@ -289,10 +289,12 @@ struct SIMachineFunctionInfo final : public yaml::MachineFunctionInfo {
 
   Optional<SIArgumentInfo> ArgInfo;
   SIMode Mode;
+  Optional<FrameIndex> ScavengeFI;
 
   SIMachineFunctionInfo() = default;
   SIMachineFunctionInfo(const llvm::SIMachineFunctionInfo &,
-                        const TargetRegisterInfo &TRI);
+                        const TargetRegisterInfo &TRI,
+                        const llvm::MachineFunction &MF);
 
   void mappingImpl(yaml::IO &YamlIO) override;
   ~SIMachineFunctionInfo() = default;
@@ -322,6 +324,7 @@ template <> struct MappingTraits<SIMachineFunctionInfo> {
     YamlIO.mapOptional("highBitsOf32BitAddress",
                        MFI.HighBitsOf32BitAddress, 0u);
     YamlIO.mapOptional("occupancy", MFI.Occupancy, 0);
+    YamlIO.mapOptional("scavengeFI", MFI.ScavengeFI);
   }
 };
 
@@ -502,7 +505,10 @@ public: // FIXME
 public:
   SIMachineFunctionInfo(const MachineFunction &MF);
 
-  bool initializeBaseYamlFields(const yaml::SIMachineFunctionInfo &YamlMFI);
+  bool initializeBaseYamlFields(const yaml::SIMachineFunctionInfo &YamlMFI,
+                                const MachineFunction &MF,
+                                PerFunctionMIParsingState &PFS,
+                                SMDiagnostic &Error, SMRange &SourceRange);
 
   void reserveWWMRegister(Register Reg, Optional<int> FI) {
     WWMReservedRegs.insert(std::make_pair(Reg, FI));
@@ -546,6 +552,7 @@ public:
   void removeDeadFrameIndices(MachineFrameInfo &MFI);
 
   int getScavengeFI(MachineFrameInfo &MFI, const SIRegisterInfo &TRI);
+  Optional<int> getOptionalScavengeFI() const { return ScavengeFI; }
 
   bool hasCalculatedTID() const { return TIDReg != 0; };
   Register getTIDReg() const { return TIDReg; };

test/CodeGen/MIR/AMDGPU/invalid-frame-index-invalid-fixed-stack.mir

@@ -0,0 +1,17 @@
+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck %s
+
+---
+name: invalid_scavenge_fi
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 16, alignment: 4, stack-id: noalloc }
+  - { id: 1, name: '', type: spill-slot, offset: 16, size: 16, alignment: 4, stack-id: noalloc }
+  - { id: 2, name: '', type: spill-slot, offset: 32, size: 16, alignment: 4, stack-id: noalloc }
+machineFunctionInfo:
+  # CHECK: [[@LINE+1]]:17: invalid fixed frame index 2
+  scavengeFI: '%fixed-stack.2'
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

test/CodeGen/MIR/AMDGPU/invalid-frame-index-invalid-stack.mir

@@ -0,0 +1,17 @@
+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck %s
+
+---
+name: invalid_scavenge_fi
+fixedStack:
+  - { id: 0, type: spill-slot, offset: 0, size: 16, alignment: 4, stack-id: noalloc }
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 16, alignment: 4, stack-id: noalloc }
+machineFunctionInfo:
+  # CHECK: [[@LINE+1]]:17: invalid frame index 2
+  scavengeFI: '%stack.2'
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

test/CodeGen/MIR/AMDGPU/invalid-frame-index-no-stack.mir

@@ -0,0 +1,13 @@
+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck %s
+
+---
+name: invalid_scavenge_fi
+machineFunctionInfo:
+  # CHECK: [[@LINE+1]]:17: invalid frame index 0
+  scavengeFI: '%stack.0'
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

test/CodeGen/MIR/AMDGPU/invalid-frame-index.mir

@@ -0,0 +1,13 @@
+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck %s
+
+---
+name: invalid_scavenge_fi
+machineFunctionInfo:
+  # CHECK: [[@LINE+1]]:15: Invalid frame index, needs to start with %stack. or %fixed-stack
+  scavengeFI: 0
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

test/CodeGen/MIR/AMDGPU/invalid-frame-index2.mir

@@ -0,0 +1,13 @@
+# RUN: not llc -mtriple=amdgcn-amd-amdhsa -run-pass=none -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck %s
+
+---
+name: invalid_scavenge_fi
+machineFunctionInfo:
+  # CHECK: [[@LINE+1]]:15: Invalid frame index, not a valid number
+  scavengeFI: '%stack.abc'
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...

test/CodeGen/MIR/AMDGPU/machine-function-info-after-pei.ll

@@ -0,0 +1,50 @@
+; RUN: llc -mtriple=amdgcn-mesa-mesa3d -mcpu=tahiti -amdgpu-spill-sgpr-to-vgpr=0 -stop-after prologepilog -verify-machineinstrs %s -o - | FileCheck -check-prefix=AFTER-PEI %s
+
+; Test that the ScavengeFI is serialized in the SIMachineFunctionInfo.
+
+; AFTER-PEI-LABEL: {{^}}name: scavenge_fi
+; AFTER-PEI: machineFunctionInfo:
+; AFTER-PEI-NEXT: explicitKernArgSize: 12
+; AFTER-PEI-NEXT: maxKernArgAlign: 8
+; AFTER-PEI-NEXT: ldsSize:         0
+; AFTER-PEI-NEXT: dynLDSAlign:     1
+; AFTER-PEI-NEXT: isEntryFunction: true
+; AFTER-PEI-NEXT: noSignedZerosFPMath: false
+; AFTER-PEI-NEXT: memoryBound:     false
+; AFTER-PEI-NEXT: waveLimiter:     false
+; AFTER-PEI-NEXT: hasSpilledSGPRs: true
+; AFTER-PEI-NEXT: hasSpilledVGPRs: false
+; AFTER-PEI-NEXT: scratchRSrcReg:  '$sgpr68_sgpr69_sgpr70_sgpr71'
+; AFTER-PEI-NEXT: frameOffsetReg:  '$fp_reg'
+; AFTER-PEI-NEXT: stackPtrOffsetReg: '$sgpr32'
+; AFTER-PEI-NEXT: argumentInfo:
+; AFTER-PEI-NEXT:   privateSegmentBuffer: { reg: '$sgpr0_sgpr1_sgpr2_sgpr3' }
+; AFTER-PEI-NEXT:   kernargSegmentPtr: { reg: '$sgpr4_sgpr5' }
+; AFTER-PEI-NEXT:   workGroupIDX:    { reg: '$sgpr6' }
+; AFTER-PEI-NEXT:   privateSegmentWaveByteOffset: { reg: '$sgpr7' }
+; AFTER-PEI-NEXT:   workItemIDX:     { reg: '$vgpr0' }
+; AFTER-PEI-NEXT: mode:
+; AFTER-PEI-NEXT:   ieee:            true
+; AFTER-PEI-NEXT:   dx10-clamp:      true
+; AFTER-PEI-NEXT:   fp32-input-denormals: true
+; AFTER-PEI-NEXT:   fp32-output-denormals: true
+; AFTER-PEI-NEXT:   fp64-fp16-input-denormals: true
+; AFTER-PEI-NEXT:   fp64-fp16-output-denormals: true
+; AFTER-PEI-NEXT: highBitsOf32BitAddress: 0
+; AFTER-PEI-NEXT: occupancy: 5
+; AFTER-PEI-NEXT: scavengeFI: '%fixed-stack.0'
+; AFTER-PEI-NEXT: body:
+define amdgpu_kernel void @scavenge_fi(i32 addrspace(1)* %out, i32 %in) #0 {
+  %wide.sgpr0 = call <32 x i32>  asm sideeffect "; def $0", "=s" () #0
+  %wide.sgpr1 = call <32 x i32>  asm sideeffect "; def $0", "=s" () #0
+  %wide.sgpr2 = call <32 x i32>  asm sideeffect "; def $0", "=s" () #0
+  %wide.sgpr3 = call <32 x i32>  asm sideeffect "; def $0", "=s" () #0
+
+  call void asm sideeffect "; use $0", "s"(<32 x i32> %wide.sgpr0) #0
+  call void asm sideeffect "; use $0", "s"(<32 x i32> %wide.sgpr1) #0
+  call void asm sideeffect "; use $0", "s"(<32 x i32> %wide.sgpr2) #0
+  call void asm sideeffect "; use $0", "s"(<32 x i32> %wide.sgpr3) #0
+  ret void
+}
+
+attributes #0 = { nounwind }

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

@@ -340,3 +340,18 @@ body:             |
     S_ENDPGM 0
 
 ...
+
+---
+# ALL-LABEL: name: scavenge_fi
+# ALL: scavengeFI: '%stack.0'
+name: scavenge_fi
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 4, alignment: 4 }
+machineFunctionInfo:
+  scavengeFI: '%stack.0'
+
+body:             |
+  bb.0:
+    S_ENDPGM 0
+
+...