[CodeGen][NPM] Port XRayInstrumentation to NPM #129865

optimisan · 2025-03-05T10:43:46Z

No description provided.

optimisan · 2025-03-05T10:44:05Z

This stack of pull requests is managed by Graphite. Learn more about stacking.

llvmbot · 2025-03-11T08:35:34Z

@llvm/pr-subscribers-backend-x86

Author: Akshat Oke (optimisan)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/129865.diff

9 Files Affected:

(added) llvm/include/llvm/CodeGen/XRayInstrumentation.h (+25)
(modified) llvm/include/llvm/InitializePasses.h (+1-1)
(modified) llvm/include/llvm/Passes/CodeGenPassBuilder.h (+1)
(modified) llvm/include/llvm/Passes/MachinePassRegistry.def (+1-1)
(modified) llvm/lib/CodeGen/CodeGen.cpp (+1-1)
(modified) llvm/lib/CodeGen/XRayInstrumentation.cpp (+69-15)
(modified) llvm/lib/Passes/PassBuilder.cpp (+1)
(modified) llvm/test/CodeGen/X86/xray-empty-firstmbb.mir (+1)
(modified) llvm/test/CodeGen/X86/xray-multiplerets-in-blocks.mir (+1)

diff --git a/llvm/include/llvm/CodeGen/XRayInstrumentation.h b/llvm/include/llvm/CodeGen/XRayInstrumentation.h
new file mode 100644
index 0000000000000..cc65d61627fc0
--- /dev/null
+++ b/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
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 6c5def57d9f05..75992aea07a73 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 35d9b124bd859..39df80a6b13ed 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -86,6 +86,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 416e5506dfd0d..3e2ead944dd7e 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -189,6 +189,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 4a69e06d84a9a..d2b70539a95ef 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/XRayInstrumentation.cpp b/llvm/lib/CodeGen/XRayInstrumentation.cpp
index 0873d9956356e..e5020d6d19412 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,6 @@ 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);
@@ -181,8 +237,6 @@ 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;
       MachineLoopInfo ComputedMLI;
       if (!MLI) {
         ComputedMLI.analyze(*MDT);
@@ -272,10 +326,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 4e623834e5e5d..7ea6951df47d6 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -160,6 +160,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.

llvm/lib/CodeGen/XRayInstrumentation.cpp

paperchalice · 2025-03-11T09:52:14Z

llvm/lib/CodeGen/XRayInstrumentation.cpp

+  return getMachineFunctionPassPreservedAnalyses()
+      .preserve<MachineDominatorTreeAnalysis>()
+      .preserve<MachineLoopAnalysis>()
+      .preserveSet<CFGAnalyses>();


I think preserveSet<CFGAnalyses>() is enough here.

This was referenced Mar 5, 2025

[CodeGen][NPM] Port MachineBlockPlacement to NPM #129828

Merged

[CodeGen][NPM] Port MachineBlockPlacementStats to NPM #129853

Merged

optimisan force-pushed the users/optimisan/postbb/05-03-port-fentryinserter branch from 0ba6ca6 to 4b18d66 Compare March 7, 2025 10:09

optimisan force-pushed the users/optimisan/postbb/05-03-port-xray branch from 32a8bd5 to 20188b3 Compare March 7, 2025 10:10

optimisan mentioned this pull request Mar 10, 2025

[AMDGPU][NFC] Format GCNCreateVOPD.cpp #130548

Merged

optimisan force-pushed the users/optimisan/postbb/05-03-port-fentryinserter branch 3 times, most recently from 0712e45 to 192d2f2 Compare March 11, 2025 08:29

optimisan force-pushed the users/optimisan/postbb/05-03-port-xray branch 2 times, most recently from 103b199 to c00160b Compare March 11, 2025 08:34

optimisan marked this pull request as ready for review March 11, 2025 08:35

llvmbot added the backend:X86 label Mar 11, 2025

optimisan requested review from arsenm and paperchalice March 11, 2025 08:35

optimisan requested review from cdevadas and vikramRH March 11, 2025 08:35

paperchalice reviewed Mar 11, 2025

View reviewed changes

llvm/lib/CodeGen/XRayInstrumentation.cpp Show resolved Hide resolved

paperchalice reviewed Mar 11, 2025

View reviewed changes

optimisan force-pushed the users/optimisan/postbb/05-03-port-fentryinserter branch from 192d2f2 to c029a58 Compare March 12, 2025 05:44

optimisan added 2 commits March 12, 2025 05:55

[CodeGen][NPM] Port XRayInstrumentation to NPM

d6256a4

only preserve CFG

b2be4e3

optimisan force-pushed the users/optimisan/postbb/05-03-port-xray branch from c00160b to b2be4e3 Compare March 12, 2025 05:56

optimisan changed the base branch from users/optimisan/postbb/05-03-port-fentryinserter to main March 12, 2025 05:57

paperchalice approved these changes Mar 12, 2025

View reviewed changes

arsenm approved these changes Mar 17, 2025

View reviewed changes

arsenm requested a review from aeubanks March 17, 2025 06:30

optimisan merged commit 4a68702 into main Apr 1, 2025
11 checks passed

optimisan deleted the users/optimisan/postbb/05-03-port-xray branch April 1, 2025 10:08

Ankur-0429 pushed a commit to Ankur-0429/llvm-project that referenced this pull request Apr 2, 2025

[CodeGen][NPM] Port XRayInstrumentation to NPM (llvm#129865)

55f4132

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[CodeGen][NPM] Port XRayInstrumentation to NPM #129865

[CodeGen][NPM] Port XRayInstrumentation to NPM #129865

Uh oh!

optimisan commented Mar 5, 2025

Uh oh!

optimisan commented Mar 5, 2025 •

edited

Loading

Uh oh!

llvmbot commented Mar 11, 2025

Uh oh!

Uh oh!

paperchalice Mar 11, 2025

Uh oh!

Uh oh!

Uh oh!

[CodeGen][NPM] Port XRayInstrumentation to NPM #129865

[CodeGen][NPM] Port XRayInstrumentation to NPM #129865

Uh oh!

Conversation

optimisan commented Mar 5, 2025

Uh oh!

optimisan commented Mar 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Mar 11, 2025

Uh oh!

Uh oh!

paperchalice Mar 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

optimisan commented Mar 5, 2025 •

edited

Loading