4 files changed, 73 insertions, 0 deletions
diff --git a/llvm/lib/Target/X86/X86MachineFunctionInfo.cpp b/llvm/lib/Target/X86/X86MachineFunctionInfo.cpp
index 2e88e01..7b57f7c 100644
--- a/llvm/lib/Target/X86/X86MachineFunctionInfo.cpp
+++ b/llvm/lib/Target/X86/X86MachineFunctionInfo.cpp
@@ -13,6 +13,14 @@
 
 using namespace llvm;
 
+yaml::X86MachineFunctionInfo::X86MachineFunctionInfo(
+    const llvm::X86MachineFunctionInfo &MFI)
+    : AMXProgModel(MFI.getAMXProgModel()) {}
+
+void yaml::X86MachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
+  MappingTraits<X86MachineFunctionInfo>::mapping(YamlIO, *this);
+}
+
 MachineFunctionInfo *X86MachineFunctionInfo::clone(
     BumpPtrAllocator &Allocator, MachineFunction &DestMF,
     const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &Src2DstMBB)
@@ -20,6 +28,11 @@ MachineFunctionInfo *X86MachineFunctionInfo::clone(
   return DestMF.cloneInfo<X86MachineFunctionInfo>(*this);
 }
 
+void X86MachineFunctionInfo::initializeBaseYamlFields(
+    const yaml::X86MachineFunctionInfo &YamlMFI) {
+  AMXProgModel = YamlMFI.AMXProgModel;
+}
+
 void X86MachineFunctionInfo::anchor() { }
 
 void X86MachineFunctionInfo::setRestoreBasePointer(const MachineFunction *MF) {
diff --git a/llvm/lib/Target/X86/X86MachineFunctionInfo.h b/llvm/lib/Target/X86/X86MachineFunctionInfo.h
index 8aaa499..af2de2e 100644
--- a/llvm/lib/Target/X86/X86MachineFunctionInfo.h
+++ b/llvm/lib/Target/X86/X86MachineFunctionInfo.h
@@ -16,13 +16,43 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/CallingConvLower.h"
+#include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/Support/YAMLTraits.h"
 #include <set>
 
 namespace llvm {
 
 enum AMXProgModelEnum { None = 0, DirectReg = 1, ManagedRA = 2 };
 
+class X86MachineFunctionInfo;
+
+namespace yaml {
+template <> struct ScalarEnumerationTraits<AMXProgModelEnum> {
+  static void enumeration(IO &YamlIO, AMXProgModelEnum &Value) {
+    YamlIO.enumCase(Value, "None", AMXProgModelEnum::None);
+    YamlIO.enumCase(Value, "DirectReg", AMXProgModelEnum::DirectReg);
+    YamlIO.enumCase(Value, "ManagedRA", AMXProgModelEnum::ManagedRA);
+  }
+};
+
+struct X86MachineFunctionInfo final : public yaml::MachineFunctionInfo {
+  AMXProgModelEnum AMXProgModel;
+
+  X86MachineFunctionInfo() = default;
+  X86MachineFunctionInfo(const llvm::X86MachineFunctionInfo &MFI);
+
+  void mappingImpl(yaml::IO &YamlIO) override;
+  ~X86MachineFunctionInfo() = default;
+};
+
+template <> struct MappingTraits<X86MachineFunctionInfo> {
+  static void mapping(IO &YamlIO, X86MachineFunctionInfo &MFI) {
+    YamlIO.mapOptional("amxProgModel", MFI.AMXProgModel);
+  }
+};
+} // end namespace yaml
+
 /// X86MachineFunctionInfo - This class is derived from MachineFunction and
 /// contains private X86 target-specific information for each MachineFunction.
 class X86MachineFunctionInfo : public MachineFunctionInfo {
@@ -160,6 +190,8 @@ public:
         const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &Src2DstMBB)
       const override;
 
+  void initializeBaseYamlFields(const yaml::X86MachineFunctionInfo &YamlMFI);
+
   bool getForceFramePointer() const { return ForceFramePointer;}
   void setForceFramePointer(bool forceFP) { ForceFramePointer = forceFP; }
 
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 27542e5..d4e642c 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -31,6 +31,8 @@
 #include "llvm/CodeGen/GlobalISel/InstructionSelector.h"
 #include "llvm/CodeGen/GlobalISel/Legalizer.h"
 #include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
+#include "llvm/CodeGen/MIRParser/MIParser.h"
+#include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineScheduler.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/RegAllocRegistry.h"
@@ -344,6 +346,24 @@ X86TargetMachine::getSubtargetImpl(const Function &F) const {
   return I.get();
 }
 
+yaml::MachineFunctionInfo *X86TargetMachine::createDefaultFuncInfoYAML() const {
+  return new yaml::X86MachineFunctionInfo();
+}
+
+yaml::MachineFunctionInfo *
+X86TargetMachine::convertFuncInfoToYAML(const MachineFunction &MF) const {
+  const auto *MFI = MF.getInfo<X86MachineFunctionInfo>();
+  return new yaml::X86MachineFunctionInfo(*MFI);
+}
+
+bool X86TargetMachine::parseMachineFunctionInfo(
+    const yaml::MachineFunctionInfo &MFI, PerFunctionMIParsingState &PFS,
+    SMDiagnostic &Error, SMRange &SourceRange) const {
+  const auto &YamlMFI = static_cast<const yaml::X86MachineFunctionInfo &>(MFI);
+  PFS.MF.getInfo<X86MachineFunctionInfo>()->initializeBaseYamlFields(YamlMFI);
+  return false;
+}
+
 bool X86TargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
                                            unsigned DestAS) const {
   assert(SrcAS != DestAS && "Expected different address spaces!");
diff --git a/llvm/lib/Target/X86/X86TargetMachine.h b/llvm/lib/Target/X86/X86TargetMachine.h
index 4a5f20f..916445c 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.h
+++ b/llvm/lib/Target/X86/X86TargetMachine.h
@@ -58,6 +58,14 @@ public:
   createMachineFunctionInfo(BumpPtrAllocator &Allocator, const Function &F,
                             const TargetSubtargetInfo *STI) const override;
 
+  yaml::MachineFunctionInfo *createDefaultFuncInfoYAML() const override;
+  yaml::MachineFunctionInfo *
+  convertFuncInfoToYAML(const MachineFunction &MF) const override;
+  bool parseMachineFunctionInfo(const yaml::MachineFunctionInfo &,
+                                PerFunctionMIParsingState &PFS,
+                                SMDiagnostic &Error,
+                                SMRange &SourceRange) const override;
+
   void registerPassBuilderCallbacks(PassBuilder &PB,
                                     bool PopulateClassToPassNames) override;