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/ADT/STLExtras.h"
17#include "llvm/ADT/SmallVector.h"
18#include "llvm/CodeGen/MachineBasicBlock.h"
19#include "llvm/CodeGen/MachineDominators.h"
20#include "llvm/CodeGen/MachineFunction.h"
21#include "llvm/CodeGen/MachineFunctionPass.h"
22#include "llvm/CodeGen/MachineInstrBuilder.h"
23#include "llvm/CodeGen/MachineLoopInfo.h"
24#include "llvm/CodeGen/TargetInstrInfo.h"
25#include "llvm/CodeGen/TargetSubtargetInfo.h"
26#include "llvm/IR/Attributes.h"
27#include "llvm/IR/Function.h"
28#include "llvm/InitializePasses.h"
29#include "llvm/Pass.h"
30#include "llvm/Target/TargetMachine.h"
31#include "llvm/TargetParser/Triple.h"
32
33using namespace llvm;
34
35namespace {
36
37struct InstrumentationOptions {
38 // Whether to emit PATCHABLE_TAIL_CALL.
39 bool HandleTailcall;
40
41 // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
42 // return, e.g. conditional return.
43 bool HandleAllReturns;
44};
45
46struct XRayInstrumentation : public MachineFunctionPass {
47 static char ID;
48
49 XRayInstrumentation() : MachineFunctionPass(ID) {
50 initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
51 }
52
53 void getAnalysisUsage(AnalysisUsage &AU) const override {
54 AU.setPreservesCFG();
55 AU.addPreserved<MachineLoopInfoWrapperPass>();
56 AU.addPreserved<MachineDominatorTreeWrapperPass>();
57 MachineFunctionPass::getAnalysisUsage(AU);
58 }
59
60 bool runOnMachineFunction(MachineFunction &MF) override;
61
62private:
63 // Replace the original RET instruction with the exit sled code ("patchable
64 // ret" pseudo-instruction), so that at runtime XRay can replace the sled
65 // with a code jumping to XRay trampoline, which calls the tracing handler
66 // and, in the end, issues the RET instruction.
67 // This is the approach to go on CPUs which have a single RET instruction,
68 // like x86/x86_64.
69 void replaceRetWithPatchableRet(MachineFunction &MF,
70 const TargetInstrInfo *TII,
71 InstrumentationOptions);
72
73 // Prepend the original return instruction with the exit sled code ("patchable
74 // function exit" pseudo-instruction), preserving the original return
75 // instruction just after the exit sled code.
76 // This is the approach to go on CPUs which have multiple options for the
77 // return instruction, like ARM. For such CPUs we can't just jump into the
78 // XRay trampoline and issue a single return instruction there. We rather
79 // have to call the trampoline and return from it to the original return
80 // instruction of the function being instrumented.
81 void prependRetWithPatchableExit(MachineFunction &MF,
82 const TargetInstrInfo *TII,
83 InstrumentationOptions);
84};
85
86} // end anonymous namespace
87
88void XRayInstrumentation::replaceRetWithPatchableRet(
89 MachineFunction &MF, const TargetInstrInfo *TII,
90 InstrumentationOptions op) {
91 // We look for *all* terminators and returns, then replace those with
92 // PATCHABLE_RET instructions.
93 SmallVector<MachineInstr *, 4> Terminators;
94 for (auto &MBB : MF) {
95 for (auto &T : MBB.terminators()) {
96 unsigned Opc = 0;
97 if (T.isReturn() &&
98 (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
99 // Replace return instructions with:
100 // PATCHABLE_RET <Opcode>, <Operand>...
101 Opc = TargetOpcode::PATCHABLE_RET;
102 }
103 if (TII->isTailCall(Inst: T) && op.HandleTailcall) {
104 // Treat the tail call as a return instruction, which has a
105 // different-looking sled than the normal return case.
106 Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
107 }
108 if (Opc != 0) {
109 auto MIB = BuildMI(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc))
110 .addImm(Val: T.getOpcode());
111 for (auto &MO : T.operands())
112 MIB.add(MO);
113 Terminators.push_back(Elt: &T);
114 if (T.shouldUpdateCallSiteInfo())
115 MF.eraseCallSiteInfo(MI: &T);
116 }
117 }
118 }
119
120 for (auto &I : Terminators)
121 I->eraseFromParent();
122}
123
124void XRayInstrumentation::prependRetWithPatchableExit(
125 MachineFunction &MF, const TargetInstrInfo *TII,
126 InstrumentationOptions op) {
127 for (auto &MBB : MF)
128 for (auto &T : MBB.terminators()) {
129 unsigned Opc = 0;
130 if (T.isReturn() &&
131 (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
132 Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
133 }
134 if (TII->isTailCall(Inst: T) && op.HandleTailcall) {
135 Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
136 }
137 if (Opc != 0) {
138 // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
139 // PATCHABLE_TAIL_CALL .
140 BuildMI(BB&: MBB, I&: T, MIMD: T.getDebugLoc(), MCID: TII->get(Opcode: Opc));
141 }
142 }
143}
144
145bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
146 auto &F = MF.getFunction();
147 auto InstrAttr = F.getFnAttribute(Kind: "function-instrument");
148 bool AlwaysInstrument = InstrAttr.isStringAttribute() &&
149 InstrAttr.getValueAsString() == "xray-always";
150 bool NeverInstrument = InstrAttr.isStringAttribute() &&
151 InstrAttr.getValueAsString() == "xray-never";
152 if (NeverInstrument && !AlwaysInstrument)
153 return false;
154 auto IgnoreLoopsAttr = F.getFnAttribute(Kind: "xray-ignore-loops");
155
156 uint64_t XRayThreshold = 0;
157 if (!AlwaysInstrument) {
158 bool IgnoreLoops = IgnoreLoopsAttr.isValid();
159 XRayThreshold = F.getFnAttributeAsParsedInteger(
160 Kind: "xray-instruction-threshold", Default: std::numeric_limits<uint64_t>::max());
161 if (XRayThreshold == std::numeric_limits<uint64_t>::max())
162 return false;
163
164 // Count the number of MachineInstr`s in MachineFunction
165 uint64_t MICount = 0;
166 for (const auto &MBB : MF)
167 MICount += MBB.size();
168
169 bool TooFewInstrs = MICount < XRayThreshold;
170
171 if (!IgnoreLoops) {
172 // Get MachineDominatorTree or compute it on the fly if it's unavailable
173 auto *MDTWrapper =
174 getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
175 auto *MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
176 MachineDominatorTree ComputedMDT;
177 if (!MDT) {
178 ComputedMDT.getBase().recalculate(Func&: MF);
179 MDT = &ComputedMDT;
180 }
181
182 // Get MachineLoopInfo or compute it on the fly if it's unavailable
183 auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
184 auto *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
185 MachineLoopInfo ComputedMLI;
186 if (!MLI) {
187 ComputedMLI.analyze(DomTree: MDT->getBase());
188 MLI = &ComputedMLI;
189 }
190
191 // Check if we have a loop.
192 // FIXME: Maybe make this smarter, and see whether the loops are dependent
193 // on inputs or side-effects?
194 if (MLI->empty() && TooFewInstrs)
195 return false; // Function is too small and has no loops.
196 } else if (TooFewInstrs) {
197 // Function is too small
198 return false;
199 }
200 }
201
202 // We look for the first non-empty MachineBasicBlock, so that we can insert
203 // the function instrumentation in the appropriate place.
204 auto MBI = llvm::find_if(
205 Range&: MF, P: [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
206 if (MBI == MF.end())
207 return false; // The function is empty.
208
209 auto *TII = MF.getSubtarget().getInstrInfo();
210 auto &FirstMBB = *MBI;
211 auto &FirstMI = *FirstMBB.begin();
212
213 if (!MF.getSubtarget().isXRaySupported()) {
214 FirstMI.emitError(Msg: "An attempt to perform XRay instrumentation for an"
215 " unsupported target.");
216 return false;
217 }
218
219 if (!F.hasFnAttribute(Kind: "xray-skip-entry")) {
220 // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
221 // MachineFunction.
222 BuildMI(BB&: FirstMBB, I&: FirstMI, MIMD: FirstMI.getDebugLoc(),
223 MCID: TII->get(Opcode: TargetOpcode::PATCHABLE_FUNCTION_ENTER));
224 }
225
226 if (!F.hasFnAttribute(Kind: "xray-skip-exit")) {
227 switch (MF.getTarget().getTargetTriple().getArch()) {
228 case Triple::ArchType::arm:
229 case Triple::ArchType::thumb:
230 case Triple::ArchType::aarch64:
231 case Triple::ArchType::hexagon:
232 case Triple::ArchType::loongarch64:
233 case Triple::ArchType::mips:
234 case Triple::ArchType::mipsel:
235 case Triple::ArchType::mips64:
236 case Triple::ArchType::mips64el: {
237 // For the architectures which don't have a single return instruction
238 InstrumentationOptions op;
239 op.HandleTailcall = false;
240 op.HandleAllReturns = true;
241 prependRetWithPatchableExit(MF, TII, op);
242 break;
243 }
244 case Triple::ArchType::ppc64le: {
245 // PPC has conditional returns. Turn them into branch and plain returns.
246 InstrumentationOptions op;
247 op.HandleTailcall = false;
248 op.HandleAllReturns = true;
249 replaceRetWithPatchableRet(MF, TII, op);
250 break;
251 }
252 default: {
253 // For the architectures that have a single return instruction (such as
254 // RETQ on x86_64).
255 InstrumentationOptions op;
256 op.HandleTailcall = true;
257 op.HandleAllReturns = false;
258 replaceRetWithPatchableRet(MF, TII, op);
259 break;
260 }
261 }
262 }
263 return true;
264}
265
266char XRayInstrumentation::ID = 0;
267char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
268INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
269 "Insert XRay ops", false, false)
270INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
271INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
272 "Insert XRay ops", false, false)
273