[llvm-branch-commits] [llvm] [CodeGen][NPM] Port XRayInstrumentation to NPM (PR #129865)

Akshat Oke via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Wed Mar 5 02:43:53 PST 2025


https://github.com/optimisan created https://github.com/llvm/llvm-project/pull/129865

None

>From 32a8bd59f64750fb3c2e72a7e26ba7a81ff86210 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 5 Mar 2025 10:11:27 +0000
Subject: [PATCH] [CodeGen][NPM] Port XRayInstrumentation to NPM

---
 .../llvm/CodeGen/XRayInstrumentation.h        | 24 +++++
 llvm/include/llvm/InitializePasses.h          |  2 +-
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |  1 +
 .../llvm/Passes/MachinePassRegistry.def       |  2 +-
 llvm/lib/CodeGen/CodeGen.cpp                  |  2 +-
 llvm/lib/CodeGen/FEntryInserter.cpp           |  8 +-
 llvm/lib/CodeGen/XRayInstrumentation.cpp      | 90 +++++++++++++++----
 llvm/lib/Passes/PassBuilder.cpp               |  1 +
 llvm/test/CodeGen/X86/xray-empty-firstmbb.mir |  1 +
 .../X86/xray-multiplerets-in-blocks.mir       |  1 +
 10 files changed, 110 insertions(+), 22 deletions(-)
 create mode 100644 llvm/include/llvm/CodeGen/XRayInstrumentation.h

diff --git a/llvm/include/llvm/CodeGen/XRayInstrumentation.h b/llvm/include/llvm/CodeGen/XRayInstrumentation.h
new file mode 100644
index 0000000000000..ed39a7f3c1654
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/XRayInstrumentation.h
@@ -0,0 +1,24 @@
+//===- 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);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_XRAYINSTRUMENTATION_H
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 63917b2b7f729..dcfd9fc6a86b9 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/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
 
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index bab475d740467..426dc6c7eacfd 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -85,6 +85,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"
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index 667a7352930ea..8b1373c0ffefd 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -188,6 +188,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
@@ -296,5 +297,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
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index effb556e63435..b299983503232 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -144,5 +144,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeVirtRegRewriterPass(Registry);
   initializeWasmEHPreparePass(Registry);
   initializeWinEHPreparePass(Registry);
-  initializeXRayInstrumentationPass(Registry);
+  initializeXRayInstrumentationLegacyPass(Registry);
 }
diff --git a/llvm/lib/CodeGen/FEntryInserter.cpp b/llvm/lib/CodeGen/FEntryInserter.cpp
index 4f1bd7df6a204..79949dac51448 100644
--- a/llvm/lib/CodeGen/FEntryInserter.cpp
+++ b/llvm/lib/CodeGen/FEntryInserter.cpp
@@ -59,7 +59,7 @@ bool FEntryInserter::run(MachineFunction &MF) {
   return true;
 }
 
-char FEntryInserter::ID = 0;
-char &llvm::FEntryInserterID = FEntryInserter::ID;
-INITIALIZE_PASS(FEntryInserter, "fentry-insert", "Insert fentry calls", false,
-                false)
+char FEntryInserterLegacy::ID = 0;
+char &llvm::FEntryInserterID = FEntryInserterLegacy::ID;
+INITIALIZE_PASS(FEntryInserterLegacy, "fentry-insert", "Insert fentry calls",
+                false, false)
diff --git a/llvm/lib/CodeGen/XRayInstrumentation.cpp b/llvm/lib/CodeGen/XRayInstrumentation.cpp
index 0873d9956356e..00aa26c5f369a 100644
--- a/llvm/lib/CodeGen/XRayInstrumentation.cpp
+++ b/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,43 @@ 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();
+
+  return getMachineFunctionPassPreservedAnalyses()
+      .preserve<MachineDominatorTreeAnalysis>()
+      .preserve<MachineLoopAnalysis>()
+      .preserveSet<CFGAnalyses>();
+}
+
+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,9 +230,9 @@ 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;
+      // auto *MDTWrapper =
+      //     getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
+      // auto *MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
       MachineDominatorTree ComputedMDT;
       if (!MDT) {
         ComputedMDT.recalculate(MF);
@@ -181,8 +240,9 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
       }
 
       // Get MachineLoopInfo or compute it on the fly if it's unavailable
-      auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
-      auto *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
+      // auto *MLIWrapper =
+      // getAnalysisIfAvailable<MachineLoopInfoWrapperPass>(); auto *MLI =
+      // MLIWrapper ? &MLIWrapper->getLI() : nullptr;
       MachineLoopInfo ComputedMLI;
       if (!MLI) {
         ComputedMLI.analyze(*MDT);
@@ -272,10 +332,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)
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index b9fb16844f6bf..a7e2c6b09a70c 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -159,6 +159,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"
diff --git a/llvm/test/CodeGen/X86/xray-empty-firstmbb.mir b/llvm/test/CodeGen/X86/xray-empty-firstmbb.mir
index df5dc7b28ec1a..cd8b04b96ba2b 100644
--- a/llvm/test/CodeGen/X86/xray-empty-firstmbb.mir
+++ b/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.
 
diff --git a/llvm/test/CodeGen/X86/xray-multiplerets-in-blocks.mir b/llvm/test/CodeGen/X86/xray-multiplerets-in-blocks.mir
index 60a33b95f1412..0ddd5037e4265 100644
--- a/llvm/test/CodeGen/X86/xray-multiplerets-in-blocks.mir
+++ b/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.



More information about the llvm-branch-commits mailing list