[CodeGen][NPM] Port XRayInstrumentation to NPM

llvm/include/llvm/CodeGen/XRayInstrumentation.h

@@ -0,0 +1,25 @@
+//===- llvm/CodeGen/XRayInstrumentation.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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_XRAYINSTRUMENTATION_H
+#define LLVM_CODEGEN_XRAYINSTRUMENTATION_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+class XRayInstrumentationPass : public PassInfoMixin<XRayInstrumentationPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+  static bool isRequired() { return true; }
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_XRAYINSTRUMENTATION_H

llvm/include/llvm/InitializePasses.h

@@ -320,7 +320,7 @@ void initializeVirtRegRewriterPass(PassRegistry &);
 void initializeWasmEHPreparePass(PassRegistry &);
 void initializeWinEHPreparePass(PassRegistry &);
 void initializeWriteBitcodePassPass(PassRegistry &);
-void initializeXRayInstrumentationPass(PassRegistry &);
+void initializeXRayInstrumentationLegacyPass(PassRegistry &);
 
 } // end namespace llvm
 

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -84,6 +84,7 @@
 #include "llvm/CodeGen/UnreachableBlockElim.h"
 #include "llvm/CodeGen/WasmEHPrepare.h"
 #include "llvm/CodeGen/WinEHPrepare.h"
+#include "llvm/CodeGen/XRayInstrumentation.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/IRPrinter/IRPrintingPasses.h"

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -187,6 +187,7 @@ MACHINE_FUNCTION_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
 MACHINE_FUNCTION_PASS("two-address-instruction", TwoAddressInstructionPass())
 MACHINE_FUNCTION_PASS("verify", MachineVerifierPass())
 MACHINE_FUNCTION_PASS("verify<machine-trace-metrics>", MachineTraceMetricsVerifierPass())
+MACHINE_FUNCTION_PASS("xray-instrumentation", XRayInstrumentationPass())
 #undef MACHINE_FUNCTION_PASS
 
 #ifndef MACHINE_FUNCTION_PASS_WITH_PARAMS
@@ -288,5 +289,4 @@ DUMMY_MACHINE_FUNCTION_PASS("stack-frame-layout", StackFrameLayoutAnalysisPass)
 DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
 DUMMY_MACHINE_FUNCTION_PASS("unpack-mi-bundles", UnpackMachineBundlesPass)
 DUMMY_MACHINE_FUNCTION_PASS("virtregrewriter", VirtRegRewriterPass)
-DUMMY_MACHINE_FUNCTION_PASS("xray-instrumentation", XRayInstrumentationPass)
 #undef DUMMY_MACHINE_FUNCTION_PASS

llvm/lib/CodeGen/CodeGen.cpp

@@ -144,5 +144,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeVirtRegRewriterPass(Registry);
   initializeWasmEHPreparePass(Registry);
   initializeWinEHPreparePass(Registry);
-  initializeXRayInstrumentationPass(Registry);
+  initializeXRayInstrumentationLegacyPass(Registry);
 }

llvm/lib/CodeGen/XRayInstrumentation.cpp

@@ -13,14 +13,17 @@
 //
 //===---------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/XRayInstrumentation.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/Attributes.h"
@@ -44,11 +47,11 @@ struct InstrumentationOptions {
   bool HandleAllReturns;
 };
 
-struct XRayInstrumentation : public MachineFunctionPass {
+struct XRayInstrumentationLegacy : public MachineFunctionPass {
   static char ID;
 
-  XRayInstrumentation() : MachineFunctionPass(ID) {
-    initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
+  XRayInstrumentationLegacy() : MachineFunctionPass(ID) {
+    initializeXRayInstrumentationLegacyPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -59,6 +62,27 @@ struct XRayInstrumentation : public MachineFunctionPass {
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
+};
+
+struct XRayInstrumentation {
+  XRayInstrumentation(MachineDominatorTree *MDT, MachineLoopInfo *MLI)
+      : MDT(MDT), MLI(MLI) {}
+
+  bool run(MachineFunction &MF);
+
+  // Methods for use in the NPM and legacy passes, can be removed once migration
+  // is complete.
+  static bool alwaysInstrument(Function &F) {
+    auto InstrAttr = F.getFnAttribute("function-instrument");
+    return InstrAttr.isStringAttribute() &&
+           InstrAttr.getValueAsString() == "xray-always";
+  }
+
+  static bool needMDTAndMLIAnalyses(Function &F) {
+    auto IgnoreLoopsAttr = F.getFnAttribute("xray-ignore-loops");
+    auto AlwaysInstrument = XRayInstrumentation::alwaysInstrument(F);
+    return !AlwaysInstrument && !IgnoreLoopsAttr.isValid();
+  }
 
 private:
   // Replace the original RET instruction with the exit sled code ("patchable
@@ -82,6 +106,9 @@ struct XRayInstrumentation : public MachineFunctionPass {
   void prependRetWithPatchableExit(MachineFunction &MF,
                                    const TargetInstrInfo *TII,
                                    InstrumentationOptions);
+
+  MachineDominatorTree *MDT;
+  MachineLoopInfo *MLI;
 };
 
 } // end anonymous namespace
@@ -143,11 +170,42 @@ void XRayInstrumentation::prependRetWithPatchableExit(
     }
 }
 
-bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
+PreservedAnalyses
+XRayInstrumentationPass::run(MachineFunction &MF,
+                             MachineFunctionAnalysisManager &MFAM) {
+  MachineDominatorTree *MDT = nullptr;
+  MachineLoopInfo *MLI = nullptr;
+
+  if (XRayInstrumentation::needMDTAndMLIAnalyses(MF.getFunction())) {
+    MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(MF);
+    MLI = MFAM.getCachedResult<MachineLoopAnalysis>(MF);
+  }
+
+  if (!XRayInstrumentation(MDT, MLI).run(MF))
+    return PreservedAnalyses::all();
+
+  auto PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}
+
+bool XRayInstrumentationLegacy::runOnMachineFunction(MachineFunction &MF) {
+  MachineDominatorTree *MDT = nullptr;
+  MachineLoopInfo *MLI = nullptr;
+  if (XRayInstrumentation::needMDTAndMLIAnalyses(MF.getFunction())) {
+    auto *MDTWrapper =
+        getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
+    MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
+    auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
+    MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
+  }
+  return XRayInstrumentation(MDT, MLI).run(MF);
+}
+
+bool XRayInstrumentation::run(MachineFunction &MF) {
   auto &F = MF.getFunction();
   auto InstrAttr = F.getFnAttribute("function-instrument");
-  bool AlwaysInstrument = InstrAttr.isStringAttribute() &&
-                          InstrAttr.getValueAsString() == "xray-always";
+  bool AlwaysInstrument = alwaysInstrument(F);
   bool NeverInstrument = InstrAttr.isStringAttribute() &&
                          InstrAttr.getValueAsString() == "xray-never";
   if (NeverInstrument && !AlwaysInstrument)
@@ -171,18 +229,13 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
 
     if (!IgnoreLoops) {
       // Get MachineDominatorTree or compute it on the fly if it's unavailable
-      auto *MDTWrapper =
-          getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
-      auto *MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
       MachineDominatorTree ComputedMDT;
       if (!MDT) {
         ComputedMDT.recalculate(MF);
         MDT = &ComputedMDT;
       }
 
       // Get MachineLoopInfo or compute it on the fly if it's unavailable
-      auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
-      auto *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
       MachineLoopInfo ComputedMLI;
       if (!MLI) {
         ComputedMLI.analyze(*MDT);
@@ -272,10 +325,10 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   return true;
 }
 
-char XRayInstrumentation::ID = 0;
-char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
-INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
+char XRayInstrumentationLegacy::ID = 0;
+char &llvm::XRayInstrumentationID = XRayInstrumentationLegacy::ID;
+INITIALIZE_PASS_BEGIN(XRayInstrumentationLegacy, "xray-instrumentation",
                       "Insert XRay ops", false, false)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
-INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
+INITIALIZE_PASS_END(XRayInstrumentationLegacy, "xray-instrumentation",
                     "Insert XRay ops", false, false)

llvm/lib/Passes/PassBuilder.cpp

@@ -158,6 +158,7 @@
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/CodeGen/WasmEHPrepare.h"
 #include "llvm/CodeGen/WinEHPrepare.h"
+#include "llvm/CodeGen/XRayInstrumentation.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/PassManager.h"

llvm/test/CodeGen/X86/xray-empty-firstmbb.mir

@@ -1,4 +1,5 @@
 # RUN: llc -run-pass=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+# RUN: llc -passes=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
 #
 # Make sure we can handle empty first basic blocks.
 

llvm/test/CodeGen/X86/xray-multiplerets-in-blocks.mir

@@ -1,4 +1,5 @@
 # RUN: llc -run-pass=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+# RUN: llc -passes=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
 #
 # Make sure we can handle multiple ret instructions in a single basic block for
 # XRay.