[X86] Add AMXProgModel to YAML serialization (#94988)

This allows tested passes to depend on the AMX model in the function
info. Preparatory work for to adopt #94358 for other AMX passes.

Author: aengelke

llvm/lib/Target/X86/X86MachineFunctionInfo.cpp

@@ -13,13 +13,26 @@
 
 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)
     const {
   return DestMF.cloneInfo<X86MachineFunctionInfo>(*this);
 }
 
+void X86MachineFunctionInfo::initializeBaseYamlFields(
+    const yaml::X86MachineFunctionInfo &YamlMFI) {
+  AMXProgModel = YamlMFI.AMXProgModel;
+}
+
 void X86MachineFunctionInfo::anchor() { }
 
 void X86MachineFunctionInfo::setRestoreBasePointer(const MachineFunction *MF) {

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 @@ class X86MachineFunctionInfo : public MachineFunctionInfo {
         const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &Src2DstMBB)
       const override;
 
+  void initializeBaseYamlFields(const yaml::X86MachineFunctionInfo &YamlMFI);
+
   bool getForceFramePointer() const { return ForceFramePointer;}
   void setForceFramePointer(bool forceFP) { ForceFramePointer = forceFP; }
 

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!");

llvm/lib/Target/X86/X86TargetMachine.h

@@ -58,6 +58,14 @@ class X86TargetMachine final : public LLVMTargetMachine {
   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;
 

llvm/test/CodeGen/X86/AMX/amx-fastconfig-phi.mir

@@ -87,7 +87,8 @@ liveins:
   - { reg: '$rsi', virtual-reg: '%14' }
 frameInfo:
   maxAlignment:    1
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   ; CHECK-LABEL: name: foo
   ; CHECK: bb.0.entry:

llvm/test/CodeGen/X86/AMX/amx-fastconfig-phi2.mir

@@ -34,7 +34,8 @@ liveins:
   - { reg: '$edi', virtual-reg: '%12' }
 frameInfo:
   maxAlignment:    1
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   ; CHECK-LABEL: name: foo
   ; CHECK: bb.0.entry:

llvm/test/CodeGen/X86/AMX/amx-fastconfig-phi4.mir

@@ -35,7 +35,8 @@ liveins:
   - { reg: '$edi', virtual-reg: '%12' }
 frameInfo:
   maxAlignment:    1
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   ; CHECK-LABEL: name: foo
   ; CHECK: bb.0.entry:

llvm/test/CodeGen/X86/AMX/amx-fastconfig-spill.mir

@@ -23,7 +23,8 @@ frameInfo:
 stack:
   - { id: 0, size: 1024, alignment: 16 }
   - { id: 1, size: 64, alignment: 4 }
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   ; CHECK-LABEL: name: foo
   ; CHECK: bb.0.entry:
@@ -100,7 +101,8 @@ frameInfo:
 stack:
   - { id: 0, size: 1024, alignment: 16 }
   - { id: 1, size: 64, alignment: 4 }
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   ; CHECK-LABEL: name: copy
   ; CHECK: bb.0.entry:

llvm/test/CodeGen/X86/AMX/amx-fastconfig.mir

@@ -77,7 +77,8 @@ liveins:
   - { reg: '$edx', virtual-reg: '%11' }
 frameInfo:
   maxAlignment:    1
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   ; CHECK-LABEL: name: test_api
   ; CHECK: bb.0.entry:

llvm/test/CodeGen/X86/AMX/amx-fastpreconfig.mir

@@ -23,7 +23,8 @@ frameInfo:
 stack:
   - { id: 0, size: 1024, alignment: 16 }
   - { id: 1, size: 64, alignment: 4 }
-machineFunctionInfo: {}
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: main
@@ -79,6 +80,8 @@ registers:
 liveins:
   - { reg: '$rdi', virtual-reg: '' }
   - { reg: '$rsi', virtual-reg: '' }
+machineFunctionInfo:
+  amxProgModel: ManagedRA
 body:             |
   bb.1.entry:
     liveins: $rdi, $rsi