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