MachineFunctionInfo for AArch64 in MIR

Starting with hasRedZone adding MachineFunctionInfo to be put in the YAML for MIR files.

Split out of: D78062

Based on implementation for MachineFunctionInfo for WebAssembly

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

Patch by Andrew Litteken! (AndrewLitteken)

lib/Target/AArch64/AArch64MachineFunctionInfo.cpp

@@ -0,0 +1,32 @@
+//=- AArch64MachineFunctionInfo.cpp - AArch64 Machine Function Info ---------=//
+
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file implements AArch64-specific per-machine-function
+/// information.
+///
+//===----------------------------------------------------------------------===//
+
+#include "AArch64MachineFunctionInfo.h"
+
+using namespace llvm;
+
+yaml::AArch64FunctionInfo::AArch64FunctionInfo(
+    const llvm::AArch64FunctionInfo &MFI)
+    : HasRedZone(MFI.hasRedZone()) {}
+
+void yaml::AArch64FunctionInfo::mappingImpl(yaml::IO &YamlIO) {
+  MappingTraits<AArch64FunctionInfo>::mapping(YamlIO, *this);
+}
+
+void AArch64FunctionInfo::initializeBaseYamlFields(
+    const yaml::AArch64FunctionInfo &YamlMFI) {
+  if (YamlMFI.HasRedZone.hasValue())
+    HasRedZone = YamlMFI.HasRedZone;
+}

lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/CallingConvLower.h"
+#include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/IR/Function.h"
@@ -26,6 +27,10 @@
 
 namespace llvm {
 
+namespace yaml {
+struct AArch64FunctionInfo;
+} // end namespace yaml
+
 class MachineInstr;
 
 /// AArch64FunctionInfo - This class is derived from MachineFunctionInfo and
@@ -137,6 +142,7 @@ public:
     if (MF.getFunction().hasFnAttribute(Attribute::NoRedZone))
       HasRedZone = false;
   }
+  void initializeBaseYamlFields(const yaml::AArch64FunctionInfo &YamlMFI);
 
   unsigned getBytesInStackArgArea() const { return BytesInStackArgArea; }
   void setBytesInStackArgArea(unsigned bytes) { BytesInStackArgArea = bytes; }
@@ -333,6 +339,25 @@ private:
   DenseMap<int, std::pair<unsigned, MCSymbol *>> JumpTableEntryInfo;
 };
 
+namespace yaml {
+struct AArch64FunctionInfo final : public yaml::MachineFunctionInfo {
+  Optional<bool> HasRedZone;
+
+  AArch64FunctionInfo() = default;
+  AArch64FunctionInfo(const llvm::AArch64FunctionInfo &MFI);
+
+  void mappingImpl(yaml::IO &YamlIO) override;
+  ~AArch64FunctionInfo() = default;
+};
+
+template <> struct MappingTraits<AArch64FunctionInfo> {
+  static void mapping(IO &YamlIO, AArch64FunctionInfo &MFI) {
+    YamlIO.mapOptional("hasRedZone", MFI.HasRedZone);
+  }
+};
+
+} // end namespace yaml
+
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H

lib/Target/AArch64/AArch64TargetMachine.cpp

@@ -11,6 +11,7 @@
 
 #include "AArch64TargetMachine.h"
 #include "AArch64.h"
+#include "AArch64MachineFunctionInfo.h"
 #include "AArch64MacroFusion.h"
 #include "AArch64Subtarget.h"
 #include "AArch64TargetObjectFile.h"
@@ -26,6 +27,7 @@
 #include "llvm/CodeGen/GlobalISel/Legalizer.h"
 #include "llvm/CodeGen/GlobalISel/Localizer.h"
 #include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
+#include "llvm/CodeGen/MIRParser/MIParser.h"
 #include "llvm/CodeGen/MachineScheduler.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
@@ -661,3 +663,24 @@ void AArch64PassConfig::addPreEmitPass() {
   // SVE bundles move prefixes with destructive operations.
   addPass(createUnpackMachineBundles(nullptr));
 }
+
+yaml::MachineFunctionInfo *
+AArch64TargetMachine::createDefaultFuncInfoYAML() const {
+  return new yaml::AArch64FunctionInfo();
+}
+
+yaml::MachineFunctionInfo *
+AArch64TargetMachine::convertFuncInfoToYAML(const MachineFunction &MF) const {
+  const auto *MFI = MF.getInfo<AArch64FunctionInfo>();
+  return new yaml::AArch64FunctionInfo(*MFI);
+}
+
+bool AArch64TargetMachine::parseMachineFunctionInfo(
+    const yaml::MachineFunctionInfo &MFI, PerFunctionMIParsingState &PFS,
+    SMDiagnostic &Error, SMRange &SourceRange) const {
+  const auto &YamlMFI =
+      reinterpret_cast<const yaml::AArch64FunctionInfo &>(MFI);
+  MachineFunction &MF = PFS.MF;
+  MF.getInfo<AArch64FunctionInfo>()->initializeBaseYamlFields(YamlMFI);
+  return false;
+}

lib/Target/AArch64/AArch64TargetMachine.h

@@ -49,6 +49,14 @@ public:
     return TLOF.get();
   }
 
+  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;
+
 private:
   bool isLittle;
 };

lib/Target/AArch64/CMakeLists.txt

@@ -48,6 +48,7 @@ add_llvm_target(AArch64CodeGen
   AArch64InstructionSelector.cpp
   AArch64LegalizerInfo.cpp
   AArch64LoadStoreOptimizer.cpp
+  AArch64MachineFunctionInfo.cpp
   AArch64MacroFusion.cpp
   AArch64MCInstLower.cpp
   AArch64PreLegalizerCombiner.cpp

test/CodeGen/AArch64/function-info-noredzone-present.ll

@@ -0,0 +1,21 @@
+; RUN: llc -mtriple=aarch64---  --verify-machineinstrs -stop-before=finalize-isel -simplify-mir -o - < %s | FileCheck %s
+
+; Here we check thatt the noredzone attribute is carried through the machine
+; IR generation and is put in MachineFunctionInfo
+
+define void @baz() {
+  entry:
+    ; CHECK-LABEL: name:            baz
+    ; CHECK: machineFunctionInfo: {}
+    ret void
+}
+
+define void @bar() #0 {
+  entry:
+    ; CHECK-LABEL: name:            bar
+    ; CHECK: machineFunctionInfo:
+    ; CHECK-NEXT: hasRedZone:      false
+    ret void
+}
+
+attributes #0 = { noredzone }

test/CodeGen/MIR/AArch64/function-info-noredzone-present.mir

@@ -0,0 +1,59 @@
+# RUN: llc -mtriple=aarch64--- --verify-machineinstrs -simplify-mir -run-pass none %s -o - | FileCheck %s
+
+# This test checks for persistence of the hasRedZone attribute through a
+# llc transformation that shouldn't do anything
+
+# CHECK-NAME: name:            foo
+# CHECK-LABEL: machineFunctionInfo: {}
+
+---
+name:            foo
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0.entry:
+    RET_ReallyLR
+
+...
+
+# CHECK-LABEL: name:            bar
+# CHECK: machineFunctionInfo:
+# CHECK-NEXT: hasRedZone:      false
+
+---
+name:            bar
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0:
+    RET_ReallyLR
+
+...
+
+# CHECK-LABEL: name:            baz
+# CHECK: machineFunctionInfo:
+# CHECK-NEXT: hasRedZone:      true
+
+---
+name:            baz
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      true
+body:             |
+  bb.0.entry:
+    RET_ReallyLR
+
+...