Skip to content

Commit 4a68702

Browse files
optimisanoptimisan
authored
[CodeGen][NPM] Port XRayInstrumentation to NPM (#129865)
1 parent 337bad3 commit 4a68702

File tree

9 files changed

+100
-18
lines changed

9 files changed

+100
-18
lines changed

‎llvm/include/llvm/CodeGen/XRayInstrumentation.h‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
//===- llvm/CodeGen/XRayInstrumentation.h -----------------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef LLVM_CODEGEN_XRAYINSTRUMENTATION_H
10+
#define LLVM_CODEGEN_XRAYINSTRUMENTATION_H
11+
12+
#include "llvm/CodeGen/MachinePassManager.h"
13+
14+
namespace llvm {
15+
16+
class XRayInstrumentationPass : public PassInfoMixin<XRayInstrumentationPass> {
17+
public:
18+
PreservedAnalyses run(MachineFunction &MF,
19+
MachineFunctionAnalysisManager &MFAM);
20+
static bool isRequired() { return true; }
21+
};
22+
23+
} // namespace llvm
24+
25+
#endif // LLVM_CODEGEN_XRAYINSTRUMENTATION_H

‎llvm/include/llvm/InitializePasses.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,7 @@ void initializeVirtRegRewriterPass(PassRegistry &);
322322
void initializeWasmEHPreparePass(PassRegistry &);
323323
void initializeWinEHPreparePass(PassRegistry &);
324324
void initializeWriteBitcodePassPass(PassRegistry &);
325-
void initializeXRayInstrumentationPass(PassRegistry &);
325+
void initializeXRayInstrumentationLegacyPass(PassRegistry &);
326326

327327
} // end namespace llvm
328328

‎llvm/include/llvm/Passes/CodeGenPassBuilder.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
#include "llvm/CodeGen/UnreachableBlockElim.h"
9191
#include "llvm/CodeGen/WasmEHPrepare.h"
9292
#include "llvm/CodeGen/WinEHPrepare.h"
93+
#include "llvm/CodeGen/XRayInstrumentation.h"
9394
#include "llvm/IR/PassManager.h"
9495
#include "llvm/IR/Verifier.h"
9596
#include "llvm/IRPrinter/IRPrintingPasses.h"

‎llvm/include/llvm/Passes/MachinePassRegistry.def‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,7 @@ MACHINE_FUNCTION_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
190190
MACHINE_FUNCTION_PASS("two-address-instruction", TwoAddressInstructionPass())
191191
MACHINE_FUNCTION_PASS("verify", MachineVerifierPass())
192192
MACHINE_FUNCTION_PASS("verify<machine-trace-metrics>", MachineTraceMetricsVerifierPass())
193+
MACHINE_FUNCTION_PASS("xray-instrumentation", XRayInstrumentationPass())
193194
#undef MACHINE_FUNCTION_PASS
194195

195196
#ifndef MACHINE_FUNCTION_PASS_WITH_PARAMS
@@ -315,5 +316,4 @@ DUMMY_MACHINE_FUNCTION_PASS("stack-frame-layout", StackFrameLayoutAnalysisPass)
315316
DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
316317
DUMMY_MACHINE_FUNCTION_PASS("unpack-mi-bundles", UnpackMachineBundlesPass)
317318
DUMMY_MACHINE_FUNCTION_PASS("virtregrewriter", VirtRegRewriterPass)
318-
DUMMY_MACHINE_FUNCTION_PASS("xray-instrumentation", XRayInstrumentationPass)
319319
#undef DUMMY_MACHINE_FUNCTION_PASS

‎llvm/lib/CodeGen/CodeGen.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,5 +145,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
145145
initializeVirtRegRewriterPass(Registry);
146146
initializeWasmEHPreparePass(Registry);
147147
initializeWinEHPreparePass(Registry);
148-
initializeXRayInstrumentationPass(Registry);
148+
initializeXRayInstrumentationLegacyPass(Registry);
149149
}

‎llvm/lib/CodeGen/XRayInstrumentation.cpp‎

Lines changed: 68 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,17 @@
1313
//
1414
//===---------------------------------------------------------------------===//
1515

16+
#include "llvm/CodeGen/XRayInstrumentation.h"
1617
#include "llvm/ADT/STLExtras.h"
1718
#include "llvm/ADT/SmallVector.h"
1819
#include "llvm/CodeGen/MachineBasicBlock.h"
1920
#include "llvm/CodeGen/MachineDominators.h"
2021
#include "llvm/CodeGen/MachineFunction.h"
22+
#include "llvm/CodeGen/MachineFunctionAnalysis.h"
2123
#include "llvm/CodeGen/MachineFunctionPass.h"
2224
#include "llvm/CodeGen/MachineInstrBuilder.h"
2325
#include "llvm/CodeGen/MachineLoopInfo.h"
26+
#include "llvm/CodeGen/MachinePassManager.h"
2427
#include "llvm/CodeGen/TargetInstrInfo.h"
2528
#include "llvm/CodeGen/TargetSubtargetInfo.h"
2629
#include "llvm/IR/Attributes.h"
@@ -44,11 +47,11 @@ struct InstrumentationOptions {
4447
bool HandleAllReturns;
4548
};
4649

47-
struct XRayInstrumentation : public MachineFunctionPass {
50+
struct XRayInstrumentationLegacy : public MachineFunctionPass {
4851
static char ID;
4952

50-
XRayInstrumentation() : MachineFunctionPass(ID) {
51-
initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
53+
XRayInstrumentationLegacy() : MachineFunctionPass(ID) {
54+
initializeXRayInstrumentationLegacyPass(*PassRegistry::getPassRegistry());
5255
}
5356

5457
void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -59,6 +62,27 @@ struct XRayInstrumentation : public MachineFunctionPass {
5962
}
6063

6164
bool runOnMachineFunction(MachineFunction &MF) override;
65+
};
66+
67+
struct XRayInstrumentation {
68+
XRayInstrumentation(MachineDominatorTree *MDT, MachineLoopInfo *MLI)
69+
: MDT(MDT), MLI(MLI) {}
70+
71+
bool run(MachineFunction &MF);
72+
73+
// Methods for use in the NPM and legacy passes, can be removed once migration
74+
// is complete.
75+
static bool alwaysInstrument(Function &F) {
76+
auto InstrAttr = F.getFnAttribute("function-instrument");
77+
return InstrAttr.isStringAttribute() &&
78+
InstrAttr.getValueAsString() == "xray-always";
79+
}
80+
81+
static bool needMDTAndMLIAnalyses(Function &F) {
82+
auto IgnoreLoopsAttr = F.getFnAttribute("xray-ignore-loops");
83+
auto AlwaysInstrument = XRayInstrumentation::alwaysInstrument(F);
84+
return !AlwaysInstrument && !IgnoreLoopsAttr.isValid();
85+
}
6286

6387
private:
6488
// Replace the original RET instruction with the exit sled code ("patchable
@@ -82,6 +106,9 @@ struct XRayInstrumentation : public MachineFunctionPass {
82106
void prependRetWithPatchableExit(MachineFunction &MF,
83107
const TargetInstrInfo *TII,
84108
InstrumentationOptions);
109+
110+
MachineDominatorTree *MDT;
111+
MachineLoopInfo *MLI;
85112
};
86113

87114
} // end anonymous namespace
@@ -143,11 +170,42 @@ void XRayInstrumentation::prependRetWithPatchableExit(
143170
}
144171
}
145172

146-
bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
173+
PreservedAnalyses
174+
XRayInstrumentationPass::run(MachineFunction &MF,
175+
MachineFunctionAnalysisManager &MFAM) {
176+
MachineDominatorTree *MDT = nullptr;
177+
MachineLoopInfo *MLI = nullptr;
178+
179+
if (XRayInstrumentation::needMDTAndMLIAnalyses(MF.getFunction())) {
180+
MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(MF);
181+
MLI = MFAM.getCachedResult<MachineLoopAnalysis>(MF);
182+
}
183+
184+
if (!XRayInstrumentation(MDT, MLI).run(MF))
185+
return PreservedAnalyses::all();
186+
187+
auto PA = getMachineFunctionPassPreservedAnalyses();
188+
PA.preserveSet<CFGAnalyses>();
189+
return PA;
190+
}
191+
192+
bool XRayInstrumentationLegacy::runOnMachineFunction(MachineFunction &MF) {
193+
MachineDominatorTree *MDT = nullptr;
194+
MachineLoopInfo *MLI = nullptr;
195+
if (XRayInstrumentation::needMDTAndMLIAnalyses(MF.getFunction())) {
196+
auto *MDTWrapper =
197+
getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
198+
MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
199+
auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
200+
MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
201+
}
202+
return XRayInstrumentation(MDT, MLI).run(MF);
203+
}
204+
205+
bool XRayInstrumentation::run(MachineFunction &MF) {
147206
auto &F = MF.getFunction();
148207
auto InstrAttr = F.getFnAttribute("function-instrument");
149-
bool AlwaysInstrument = InstrAttr.isStringAttribute() &&
150-
InstrAttr.getValueAsString() == "xray-always";
208+
bool AlwaysInstrument = alwaysInstrument(F);
151209
bool NeverInstrument = InstrAttr.isStringAttribute() &&
152210
InstrAttr.getValueAsString() == "xray-never";
153211
if (NeverInstrument && !AlwaysInstrument)
@@ -171,18 +229,13 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
171229

172230
if (!IgnoreLoops) {
173231
// Get MachineDominatorTree or compute it on the fly if it's unavailable
174-
auto *MDTWrapper =
175-
getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
176-
auto *MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
177232
MachineDominatorTree ComputedMDT;
178233
if (!MDT) {
179234
ComputedMDT.recalculate(MF);
180235
MDT = &ComputedMDT;
181236
}
182237

183238
// Get MachineLoopInfo or compute it on the fly if it's unavailable
184-
auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
185-
auto *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
186239
MachineLoopInfo ComputedMLI;
187240
if (!MLI) {
188241
ComputedMLI.analyze(*MDT);
@@ -272,10 +325,10 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
272325
return true;
273326
}
274327

275-
char XRayInstrumentation::ID = 0;
276-
char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
277-
INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
328+
char XRayInstrumentationLegacy::ID = 0;
329+
char &llvm::XRayInstrumentationID = XRayInstrumentationLegacy::ID;
330+
INITIALIZE_PASS_BEGIN(XRayInstrumentationLegacy, "xray-instrumentation",
278331
"Insert XRay ops", false, false)
279332
INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
280-
INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
333+
INITIALIZE_PASS_END(XRayInstrumentationLegacy, "xray-instrumentation",
281334
"Insert XRay ops", false, false)

‎llvm/lib/Passes/PassBuilder.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,7 @@
165165
#include "llvm/CodeGen/VirtRegMap.h"
166166
#include "llvm/CodeGen/WasmEHPrepare.h"
167167
#include "llvm/CodeGen/WinEHPrepare.h"
168+
#include "llvm/CodeGen/XRayInstrumentation.h"
168169
#include "llvm/IR/DebugInfo.h"
169170
#include "llvm/IR/Dominators.h"
170171
#include "llvm/IR/PassManager.h"

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# RUN: llc -run-pass=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
2+
# RUN: llc -passes=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
23
#
34
# Make sure we can handle empty first basic blocks.
45

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# RUN: llc -run-pass=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
2+
# RUN: llc -passes=xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
23
#
34
# Make sure we can handle multiple ret instructions in a single basic block for
45
# XRay.

0 commit comments

Comments
 (0)