1//===- XRayInstrumentation.cpp - Adds XRay instrumentation to functions. --===//
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// This file implements a MachineFunctionPass that inserts the appropriate
10// XRay instrumentation instructions. We look for XRay-specific attributes
11// on the function to determine whether we should insert the replacement
12// operations.
13//
14//===---------------------------------------------------------------------===//
15
16#include "llvm/CodeGen/XRayInstrumentation.h"
17#include "llvm/ADT/STLExtras.h"
18#include "llvm/ADT/SmallVector.h"
19#include "llvm/CodeGen/MachineBasicBlock.h"
20#include "llvm/CodeGen/MachineDominators.h"
21#include "llvm/CodeGen/MachineFunction.h"
22#include "llvm/CodeGen/MachineFunctionAnalysis.h"
23#include "llvm/CodeGen/MachineFunctionPass.h"
24#include "llvm/CodeGen/MachineInstrBuilder.h"
25#include "llvm/CodeGen/MachineLoopInfo.h"
26#include "llvm/CodeGen/MachinePassManager.h"
27#include "llvm/CodeGen/TargetInstrInfo.h"
28#include "llvm/CodeGen/TargetSubtargetInfo.h"
29#include "llvm/IR/Attributes.h"
30#include "llvm/IR/DiagnosticInfo.h"
31#include "llvm/IR/Function.h"
32#include "llvm/InitializePasses.h"
33#include "llvm/Pass.h"
34#include "llvm/Target/TargetMachine.h"
35#include "llvm/TargetParser/Triple.h"
36
37using namespace llvm;
38
39namespace {
40
41struct InstrumentationOptions {
42 // Whether to emit PATCHABLE_TAIL_CALL.
43 bool HandleTailcall;
44
45 // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
46 // return, e.g. conditional return.
47 bool HandleAllReturns;
48};
49
50struct XRayInstrumentationLegacy : public MachineFunctionPass {
51 static char ID;
52
53 XRayInstrumentationLegacy() : MachineFunctionPass(ID) {
54 initializeXRayInstrumentationLegacyPass(*PassRegistry::getPassRegistry());
55 }
56
57 void getAnalysisUsage(AnalysisUsage &AU) const override {
58 AU.setPreservesCFG();
59 AU.addPreserved<MachineLoopInfoWrapperPass>();
60 AU.addPreserved<MachineDominatorTreeWrapperPass>();
61 MachineFunctionPass::getAnalysisUsage(AU);
62 }
63
64 bool runOnMachineFunction(MachineFunction &MF) override;
65};
66
67struct 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(Kind: "function-instrument");
77 return InstrAttr.isStringAttribute() &&
78 InstrAttr.getValueAsString() == "xray-always";
79 }
80
81 static bool needMDTAndMLIAnalyses(Function &F) {
82 auto IgnoreLoopsAttr = F.getFnAttribute(Kind: "xray-ignore-loops");
83 auto AlwaysInstrument = XRayInstrumentation::alwaysInstrument(F);
84 return !AlwaysInstrument && !IgnoreLoopsAttr.isValid();
85 }
86
87private:
88 // Replace the original RET instruction with the exit sled code ("patchable
89 // ret" pseudo-instruction), so that at runtime XRay can replace the sled
90 // with a code jumping to XRay trampoline, which calls the tracing handler
91 // and, in the end, issues the RET instruction.
92 // This is the approach to go on CPUs which have a single RET instruction,
93 // like x86/x86_64.
94 void replaceRetWithPatchableRet(MachineFunction &MF,
95 const TargetInstrInfo *TII,
96 InstrumentationOptions);
97
98 // Prepend the original return instruction with the exit sled code ("patchable
99 // function exit" pseudo-instruction), preserving the original return
100 // instruction just after the exit sled code.
101 // This is the approach to go on CPUs which have multiple options for the
102 // return instruction, like ARM. For such CPUs we can't just jump into the
103 // XRay trampoline and issue a single return instruction there. We rather
104 // have to call the trampoline and return from it to the original return
105 // instruction of the function being instrumented.
106 void prependRetWithPatchableExit(MachineFunction &MF,
107 const TargetInstrInfo *TII,
108 InstrumentationOptions);
109
110 MachineDominatorTree *MDT;
111 MachineLoopInfo *MLI;
112};
113
114} // end anonymous namespace
115
116void XRayInstrumentation::replaceRetWithPatchableRet(
117 MachineFunction &MF, const TargetInstrInfo *TII,
118 InstrumentationOptions op) {
119 // We look for *all* terminators and returns, then replace those with
120 // PATCHABLE_RET instructions.
121 SmallVector<MachineInstr *, 4> Terminators;
122 for (auto &MBB : MF) {
123 for (auto &T : MBB.terminators()) {
124 unsigned Opc = 0;
125 if (T.isReturn() &&
126 (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
127 // Replace return instructions with:
128 // PATCHABLE_RET <Opcode>, <Operand>...
129 Opc = TargetOpcode::PATCHABLE_RET;
130 }
131 if (TII->isTailCall(Inst: T) && op.HandleTailcall) {
132 // Treat the tail call as a return instruction, which has a
133 // different-looking sled than the normal return case.
134 Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
135 }
136 if (Opc != 0) {
137 auto MIB = BuildMI(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc))
138 .addImm(Val: T.getOpcode());
139 for (auto &MO : T.operands())
140 MIB.add(MO);
141 Terminators.push_back(Elt: &T);
142 if (T.shouldUpdateAdditionalCallInfo())
143 MF.eraseAdditionalCallInfo(MI: &T);
144 }
145 }
146 }
147
148 for (auto &I : Terminators)
149 I->eraseFromParent();
150}
151
152void XRayInstrumentation::prependRetWithPatchableExit(
153 MachineFunction &MF, const TargetInstrInfo *TII,
154 InstrumentationOptions op) {
155 for (auto &MBB : MF)
156 for (auto &T : MBB.terminators()) {
157 unsigned Opc = 0;
158 if (T.isReturn() &&
159 (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
160 Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
161 }
162 if (TII->isTailCall(Inst: T) && op.HandleTailcall) {
163 Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
164 }
165 if (Opc != 0) {
166 // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
167 // PATCHABLE_TAIL_CALL .
168 BuildMI(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc));
169 }
170 }
171}
172
173PreservedAnalyses
174XRayInstrumentationPass::run(MachineFunction &MF,
175 MachineFunctionAnalysisManager &MFAM) {
176 MachineDominatorTree *MDT = nullptr;
177 MachineLoopInfo *MLI = nullptr;
178
179 if (XRayInstrumentation::needMDTAndMLIAnalyses(F&: MF.getFunction())) {
180 MDT = MFAM.getCachedResult<MachineDominatorTreeAnalysis>(IR&: MF);
181 MLI = MFAM.getCachedResult<MachineLoopAnalysis>(IR&: 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
192bool XRayInstrumentationLegacy::runOnMachineFunction(MachineFunction &MF) {
193 MachineDominatorTree *MDT = nullptr;
194 MachineLoopInfo *MLI = nullptr;
195 if (XRayInstrumentation::needMDTAndMLIAnalyses(F&: 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
205bool XRayInstrumentation::run(MachineFunction &MF) {
206 auto &F = MF.getFunction();
207 auto InstrAttr = F.getFnAttribute(Kind: "function-instrument");
208 bool AlwaysInstrument = alwaysInstrument(F);
209 bool NeverInstrument = InstrAttr.isStringAttribute() &&
210 InstrAttr.getValueAsString() == "xray-never";
211 if (NeverInstrument && !AlwaysInstrument)
212 return false;
213 auto IgnoreLoopsAttr = F.getFnAttribute(Kind: "xray-ignore-loops");
214
215 uint64_t XRayThreshold = 0;
216 if (!AlwaysInstrument) {
217 bool IgnoreLoops = IgnoreLoopsAttr.isValid();
218 XRayThreshold = F.getFnAttributeAsParsedInteger(
219 Kind: "xray-instruction-threshold", Default: std::numeric_limits<uint64_t>::max());
220 if (XRayThreshold == std::numeric_limits<uint64_t>::max())
221 return false;
222
223 // Count the number of MachineInstr`s in MachineFunction
224 uint64_t MICount = 0;
225 for (const auto &MBB : MF)
226 MICount += MBB.size();
227
228 bool TooFewInstrs = MICount < XRayThreshold;
229
230 if (!IgnoreLoops) {
231 // Get MachineDominatorTree or compute it on the fly if it's unavailable
232 MachineDominatorTree ComputedMDT;
233 if (!MDT) {
234 ComputedMDT.recalculate(Func&: MF);
235 MDT = &ComputedMDT;
236 }
237
238 // Get MachineLoopInfo or compute it on the fly if it's unavailable
239 MachineLoopInfo ComputedMLI;
240 if (!MLI) {
241 ComputedMLI.analyze(DomTree: *MDT);
242 MLI = &ComputedMLI;
243 }
244
245 // Check if we have a loop.
246 // FIXME: Maybe make this smarter, and see whether the loops are dependent
247 // on inputs or side-effects?
248 if (MLI->empty() && TooFewInstrs)
249 return false; // Function is too small and has no loops.
250 } else if (TooFewInstrs) {
251 // Function is too small
252 return false;
253 }
254 }
255
256 // We look for the first non-empty MachineBasicBlock, so that we can insert
257 // the function instrumentation in the appropriate place.
258 auto MBI = llvm::find_if(
259 Range&: MF, P: [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
260 if (MBI == MF.end())
261 return false; // The function is empty.
262
263 auto *TII = MF.getSubtarget().getInstrInfo();
264 auto &FirstMBB = *MBI;
265 auto &FirstMI = *FirstMBB.begin();
266
267 if (!MF.getSubtarget().isXRaySupported()) {
268
269 const Function &Fn = FirstMBB.getParent()->getFunction();
270 Fn.getContext().diagnose(DI: DiagnosticInfoUnsupported(
271 Fn, "An attempt to perform XRay instrumentation for an"
272 " unsupported target."));
273
274 return false;
275 }
276
277 if (!F.hasFnAttribute(Kind: "xray-skip-entry")) {
278 // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
279 // MachineFunction.
280 BuildMI(BB&: FirstMBB, I&: FirstMI, MIMD: FirstMI.getDebugLoc(),
281 MCID: TII->get(Opcode: TargetOpcode::PATCHABLE_FUNCTION_ENTER));
282 }
283
284 if (!F.hasFnAttribute(Kind: "xray-skip-exit")) {
285 switch (MF.getTarget().getTargetTriple().getArch()) {
286 case Triple::ArchType::arm:
287 case Triple::ArchType::thumb:
288 case Triple::ArchType::aarch64:
289 case Triple::ArchType::hexagon:
290 case Triple::ArchType::loongarch64:
291 case Triple::ArchType::mips:
292 case Triple::ArchType::mipsel:
293 case Triple::ArchType::mips64:
294 case Triple::ArchType::mips64el:
295 case Triple::ArchType::riscv32:
296 case Triple::ArchType::riscv64: {
297 // For the architectures which don't have a single return instruction
298 InstrumentationOptions op;
299 // AArch64 and RISC-V support patching tail calls.
300 op.HandleTailcall = MF.getTarget().getTargetTriple().isAArch64() ||
301 MF.getTarget().getTargetTriple().isRISCV();
302 op.HandleAllReturns = true;
303 prependRetWithPatchableExit(MF, TII, op);
304 break;
305 }
306 case Triple::ArchType::ppc64le:
307 case Triple::ArchType::systemz: {
308 // PPC has conditional returns. Turn them into branch and plain returns.
309 InstrumentationOptions op;
310 op.HandleTailcall = false;
311 op.HandleAllReturns = true;
312 replaceRetWithPatchableRet(MF, TII, op);
313 break;
314 }
315 default: {
316 // For the architectures that have a single return instruction (such as
317 // RETQ on x86_64).
318 InstrumentationOptions op;
319 op.HandleTailcall = true;
320 op.HandleAllReturns = false;
321 replaceRetWithPatchableRet(MF, TII, op);
322 break;
323 }
324 }
325 }
326 return true;
327}
328
329char XRayInstrumentationLegacy::ID = 0;
330char &llvm::XRayInstrumentationID = XRayInstrumentationLegacy::ID;
331INITIALIZE_PASS_BEGIN(XRayInstrumentationLegacy, "xray-instrumentation",
332 "Insert XRay ops", false, false)
333INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
334INITIALIZE_PASS_END(XRayInstrumentationLegacy, "xray-instrumentation",
335 "Insert XRay ops", false, false)
336