[NewPM][CodeGen] Add MachineFunctionAnalysis

Author: paperchalice

llvm/include/llvm/CodeGen/FreeMachineFunction.h

@@ -13,6 +13,7 @@
 
 namespace llvm {
 
+// TODO: Convert it to function pass.
 class FreeMachineFunctionPass : public PassInfoMixin<FreeMachineFunctionPass> {
 public:
   PreservedAnalyses run(MachineFunction &MF,

llvm/include/llvm/CodeGen/MIRParser/MIRParser.h

@@ -34,6 +34,9 @@ class MachineModuleInfo;
 class SMDiagnostic;
 class StringRef;
 
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
+using ModuleAnalysisManager = AnalysisManager<Module>;
+
 typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
@@ -60,6 +63,15 @@ class MIRParser {
   ///
   /// \returns true if an error occurred.
   bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI);
+
+  /// Parses MachineFunctions in the MIR file and add them as the result
+  /// of MachineFunctionAnalysis in ModulePassManager \p MAM.
+  /// User should register at least MachineFunctionAnalysis,
+  /// MachineModuleAnalysis, FunctionAnalysisManagerModuleProxy and
+  /// PassInstrumentationAnalysis in \p MAM before parsing MIR.
+  ///
+  /// \returns true if an error occurred.
+  bool parseMachineFunctions(Module &M, ModuleAnalysisManager &MAM);
 };
 
 /// This function is the main interface to the MIR serialization format parser.

llvm/include/llvm/CodeGen/MachineFunctionAnalysis.h

@@ -0,0 +1,48 @@
+//===- llvm/CodeGen/MachineFunctionAnalysis.h -------------------*- C++ -*-===//
+//
+// 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 declares the MachineFunctionAnalysis class.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_MACHINEFUNCTIONANALYSIS
+#define LLVM_CODEGEN_MACHINEFUNCTIONANALYSIS
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class MachineFunction;
+class LLVMTargetMachine;
+
+class MachineFunctionAnalysis
+    : public AnalysisInfoMixin<MachineFunctionAnalysis> {
+  friend AnalysisInfoMixin<MachineFunctionAnalysis>;
+
+  static AnalysisKey Key;
+
+  const LLVMTargetMachine *TM;
+
+public:
+  class Result {
+    std::unique_ptr<MachineFunction> MF;
+
+  public:
+    Result(std::unique_ptr<MachineFunction> MF) : MF(std::move(MF)) {}
+    MachineFunction &getMF() { return *MF; };
+    bool invalidate(Function &, const PreservedAnalyses &PA,
+                    FunctionAnalysisManager::Invalidator &);
+  };
+
+  MachineFunctionAnalysis(const LLVMTargetMachine *TM) : TM(TM){};
+  Result run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_MACHINEFUNCTIONANALYSIS

llvm/include/llvm/CodeGen/MachineModuleInfo.h

@@ -147,10 +147,14 @@ class MachineModuleInfo {
 
   /// Returns the MachineFunction constructed for the IR function \p F.
   /// Creates a new MachineFunction if none exists yet.
+  /// NOTE: New pass manager clients shall not use this method to get
+  /// the `MachineFunction`, use `MachineFunctionAnalysis` instead.
   MachineFunction &getOrCreateMachineFunction(Function &F);
 
   /// \brief Returns the MachineFunction associated to IR function \p F if there
   /// is one, otherwise nullptr.
+  /// NOTE: New pass manager clients shall not use this method to get
+  /// the `MachineFunction`, use `MachineFunctionAnalysis` instead.
   MachineFunction *getMachineFunction(const Function &F) const;
 
   /// Delete the MachineFunction \p MF and reset the link in the IR Function to

llvm/include/llvm/CodeGen/MachinePassManager.h

@@ -76,6 +76,10 @@ struct MachinePassModel
 #endif
 
     auto PA = this->Pass.run(IR, AM);
+    // Machine function passes are not allowed to modify the LLVM
+    // representation, therefore we should preserve all IR analyses.
+    PA.template preserveSet<AllAnalysesOn<Module>>();
+    PA.template preserveSet<AllAnalysesOn<Function>>();
 
     if constexpr (is_detected<has_get_set_properties_t, PassT>::value)
       IR.getProperties().set(PassT::getSetProperties());

llvm/lib/CodeGen/CMakeLists.txt

@@ -120,6 +120,7 @@ add_llvm_component_library(LLVMCodeGen
   MachineDominators.cpp
   MachineFrameInfo.cpp
   MachineFunction.cpp
+  MachineFunctionAnalysis.cpp
   MachineFunctionPass.cpp
   MachineFunctionPrinterPass.cpp
   MachineFunctionSplitter.cpp

llvm/lib/CodeGen/FreeMachineFunction.cpp

@@ -8,15 +8,15 @@
 
 #include "llvm/CodeGen/FreeMachineFunction.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 
 using namespace llvm;
 
 PreservedAnalyses
 FreeMachineFunctionPass::run(MachineFunction &MF,
                              MachineFunctionAnalysisManager &MFAM) {
-  auto &MMI = MF.getMMI();
-  MFAM.invalidate(MF, PreservedAnalyses::none());
-  MMI.deleteMachineFunctionFor(MF.getFunction()); // MF is dangling now.
-  return PreservedAnalyses::none();
+  PreservedAnalyses PA = PreservedAnalyses::none();
+  PA.abandon<MachineFunctionAnalysis>();
+  return PA;
 }

llvm/lib/CodeGen/MIRParser/MIRParser.cpp

@@ -21,6 +21,7 @@
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
@@ -97,13 +98,15 @@ class MIRParserImpl {
   /// Create an empty function with the given name.
   Function *createDummyFunction(StringRef Name, Module &M);
 
-  bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI);
+  bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI,
+                             ModuleAnalysisManager *FAM = nullptr);
 
   /// Parse the machine function in the current YAML document.
   ///
   ///
   /// Return true if an error occurred.
-  bool parseMachineFunction(Module &M, MachineModuleInfo &MMI);
+  bool parseMachineFunction(Module &M, MachineModuleInfo &MMI,
+                            ModuleAnalysisManager *FAM);
 
   /// Initialize the machine function to the state that's described in the MIR
   /// file.
@@ -275,13 +278,14 @@ MIRParserImpl::parseIRModule(DataLayoutCallbackTy DataLayoutCallback) {
   return M;
 }
 
-bool MIRParserImpl::parseMachineFunctions(Module &M, MachineModuleInfo &MMI) {
+bool MIRParserImpl::parseMachineFunctions(Module &M, MachineModuleInfo &MMI,
+                                          ModuleAnalysisManager *MAM) {
   if (NoMIRDocuments)
     return false;
 
   // Parse the machine functions.
   do {
-    if (parseMachineFunction(M, MMI))
+    if (parseMachineFunction(M, MMI, MAM))
       return true;
     In.nextDocument();
   } while (In.setCurrentDocument());
@@ -303,7 +307,8 @@ Function *MIRParserImpl::createDummyFunction(StringRef Name, Module &M) {
   return F;
 }
 
-bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI) {
+bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI,
+                                         ModuleAnalysisManager *MAM) {
   // Parse the yaml.
   yaml::MachineFunction YamlMF;
   yaml::EmptyContext Ctx;
@@ -327,14 +332,28 @@ bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI) {
                    "' isn't defined in the provided LLVM IR");
     }
   }
-  if (MMI.getMachineFunction(*F) != nullptr)
-    return error(Twine("redefinition of machine function '") + FunctionName +
-                 "'");
 
-  // Create the MachineFunction.
-  MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
-  if (initializeMachineFunction(YamlMF, MF))
-    return true;
+  if (!MAM) {
+    if (MMI.getMachineFunction(*F) != nullptr)
+      return error(Twine("redefinition of machine function '") + FunctionName +
+                   "'");
+
+    // Create the MachineFunction.
+    MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
+    if (initializeMachineFunction(YamlMF, MF))
+      return true;
+  } else {
+    auto &FAM =
+        MAM->getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+    if (FAM.getCachedResult<MachineFunctionAnalysis>(*F))
+      return error(Twine("redefinition of machine function '") + FunctionName +
+                   "'");
+
+    // Create the MachineFunction.
+    MachineFunction &MF = FAM.getResult<MachineFunctionAnalysis>(*F).getMF();
+    if (initializeMachineFunction(YamlMF, MF))
+      return true;
+  }
 
   return false;
 }
@@ -1101,6 +1120,11 @@ bool MIRParser::parseMachineFunctions(Module &M, MachineModuleInfo &MMI) {
   return Impl->parseMachineFunctions(M, MMI);
 }
 
+bool MIRParser::parseMachineFunctions(Module &M, ModuleAnalysisManager &MAM) {
+  auto &MMI = MAM.getResult<MachineModuleAnalysis>(M).getMMI();
+  return Impl->parseMachineFunctions(M, MMI, &MAM);
+}
+
 std::unique_ptr<MIRParser> llvm::createMIRParserFromFile(
     StringRef Filename, SMDiagnostic &Error, LLVMContext &Context,
     std::function<void(Function &)> ProcessIRFunction) {

llvm/lib/CodeGen/MachineFunctionAnalysis.cpp

@@ -0,0 +1,46 @@
+//===- MachineFunctionAnalysis.cpp ----------------------------------------===//
+//
+// 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 contains the definitions of the MachineFunctionAnalysis members.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+AnalysisKey MachineFunctionAnalysis::Key;
+
+bool MachineFunctionAnalysis::Result::invalidate(
+    Function &, const PreservedAnalyses &PA,
+    FunctionAnalysisManager::Invalidator &) {
+  // Unless it is invalidated explicitly, it should remain preserved.
+  auto PAC = PA.getChecker<MachineFunctionAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
+MachineFunctionAnalysis::Result
+MachineFunctionAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
+  /// Next unique number available for a MachineFunction.
+  static unsigned NextFnNum = 0;
+
+  const TargetSubtargetInfo &STI = *TM->getSubtargetImpl(F);
+  auto &MMI = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F)
+                  .getCachedResult<MachineModuleAnalysis>(*F.getParent())
+                  ->getMMI();
+  auto MF = std::make_unique<MachineFunction>(F, *TM, STI, NextFnNum++, MMI);
+  MF->initTargetMachineFunctionInfo(STI);
+
+  // MRI callback for target specific initializations.
+  TM->registerMachineRegisterInfoCallback(*MF);
+
+  return Result(std::move(MF));
+}

llvm/lib/CodeGen/MachinePassManager.cpp

@@ -11,7 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/CodeGen/FreeMachineFunction.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/PassManagerImpl.h"
 
@@ -71,7 +73,9 @@ bool MachineFunctionAnalysisManagerModuleProxy::Result::invalidate(
 
 PreservedAnalyses
 ModuleToMachineFunctionPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
-  auto &MMI = AM.getResult<MachineModuleAnalysis>(M).getMMI();
+  // Ensure we have a MachineModuleInfo
+  AM.getResult<MachineModuleAnalysis>(M).getMMI();
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   MachineFunctionAnalysisManager &MFAM =
       AM.getResult<MachineFunctionAnalysisManagerModuleProxy>(M).getManager();
   PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(M);
@@ -82,19 +86,20 @@ ModuleToMachineFunctionPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
     if (F.isDeclaration() || F.hasAvailableExternallyLinkage())
       continue;
 
-    MachineFunction &MF = MMI.getOrCreateMachineFunction(F);
+    MachineFunction &MF = FAM.getResult<MachineFunctionAnalysis>(F).getMF();
 
     if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
       continue;
     PreservedAnalyses PassPA = Pass->run(MF, MFAM);
-    if (MMI.getMachineFunction(F)) {
-      MFAM.invalidate(MF, PassPA);
+    MFAM.invalidate(MF, PassPA);
+    if (Pass->name() != FreeMachineFunctionPass::name()) {
       PI.runAfterPass(*Pass, MF, PassPA);
+      PA.intersect(std::move(PassPA));
     } else {
-      MFAM.clear(MF, F.getName());
-      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
+      PA.intersect(std::move(PassPA));
+      FAM.invalidate(F, PA);
+      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PA);
     }
-    PA.intersect(std::move(PassPA));
   }
 
   return PA;
@@ -112,25 +117,24 @@ PreservedAnalyses
 PassManager<MachineFunction>::run(MachineFunction &MF,
                                   AnalysisManager<MachineFunction> &MFAM) {
   PassInstrumentation PI = MFAM.getResult<PassInstrumentationAnalysis>(MF);
-  Function &F = MF.getFunction();
-  MachineModuleInfo &MMI =
-      MFAM.getResult<ModuleAnalysisManagerMachineFunctionProxy>(MF)
-          .getCachedResult<MachineModuleAnalysis>(*F.getParent())
-          ->getMMI();
+  FunctionAnalysisManager &FAM =
+      MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
+          .getManager();
   PreservedAnalyses PA = PreservedAnalyses::all();
   for (auto &Pass : Passes) {
     if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
       continue;
 
     PreservedAnalyses PassPA = Pass->run(MF, MFAM);
-    if (MMI.getMachineFunction(F)) {
-      MFAM.invalidate(MF, PassPA);
+    MFAM.invalidate(MF, PassPA);
+    if (Pass->name() != FreeMachineFunctionPass::name()) {
       PI.runAfterPass(*Pass, MF, PassPA);
+      PA.intersect(std::move(PassPA));
     } else {
-      MFAM.clear(MF, F.getName());
-      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
+      PA.intersect(std::move(PassPA));
+      FAM.invalidate(MF.getFunction(), PA);
+      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PA);
     }
-    PA.intersect(std::move(PassPA));
   }
   return PA;
 }

llvm/lib/Passes/PassBuilder.cpp

@@ -92,6 +92,7 @@
 #include "llvm/CodeGen/JMCInstrumenter.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
 #include "llvm/CodeGen/MIRPrinter.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"

llvm/lib/Passes/PassRegistry.def

@@ -258,6 +258,9 @@ FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis())
 FUNCTION_ANALYSIS("inliner-size-estimator", InlineSizeEstimatorAnalysis())
 FUNCTION_ANALYSIS("lazy-value-info", LazyValueAnalysis())
 FUNCTION_ANALYSIS("loops", LoopAnalysis())
+FUNCTION_ANALYSIS(
+    "machine-function-info",
+    MachineFunctionAnalysis(static_cast<const LLVMTargetMachine *>(TM)))
 FUNCTION_ANALYSIS("memdep", MemoryDependenceAnalysis())
 FUNCTION_ANALYSIS("memoryssa", MemorySSAAnalysis())
 FUNCTION_ANALYSIS("no-op-function", NoOpFunctionAnalysis())