xref: /freebsd/contrib/llvm-project/llvm/lib/CodeGen/XRayInstrumentation.cpp (revision 56b17de1e8360fe131d425de20b5e75ff3ea897c)
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 
33 using namespace llvm;
34 
35 namespace {
36 
37 struct 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 
46 struct 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 
62 private:
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 
88 void 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(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(MBB, T, T.getDebugLoc(), TII->get(Opc))
110                        .addImm(T.getOpcode());
111         for (auto &MO : T.operands())
112           MIB.add(MO);
113         Terminators.push_back(&T);
114         if (T.shouldUpdateCallSiteInfo())
115           MF.eraseCallSiteInfo(&T);
116       }
117     }
118   }
119 
120   for (auto &I : Terminators)
121     I->eraseFromParent();
122 }
123 
124 void 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(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(MBB, T, T.getDebugLoc(), TII->get(Opc));
141       }
142     }
143 }
144 
145 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
146   auto &F = MF.getFunction();
147   auto InstrAttr = F.getFnAttribute("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("xray-ignore-loops");
155 
156   uint64_t XRayThreshold = 0;
157   if (!AlwaysInstrument) {
158     bool IgnoreLoops = IgnoreLoopsAttr.isValid();
159     XRayThreshold = F.getFnAttributeAsParsedInteger(
160         "xray-instruction-threshold", 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(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(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       MF, [&](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("An attempt to perform XRay instrumentation for an"
215                       " unsupported target.");
216     return false;
217   }
218 
219   if (!F.hasFnAttribute("xray-skip-entry")) {
220     // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
221     // MachineFunction.
222     BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
223             TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
224   }
225 
226   if (!F.hasFnAttribute("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 
266 char XRayInstrumentation::ID = 0;
267 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
268 INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
269                       "Insert XRay ops", false, false)
270 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
271 INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
272                     "Insert XRay ops", false, false)
273