xref: /freebsd/contrib/llvm-project/llvm/lib/CodeGen/XRayInstrumentation.cpp (revision 770cf0a5f02dc8983a89c6568d741fbc25baa999)
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 
37 using namespace llvm;
38 
39 namespace {
40 
41 struct 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 
50 struct 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 
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   }
86 
87 private:
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 
116 void 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(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(MBB, T, T.getDebugLoc(), TII->get(Opc))
138                        .addImm(T.getOpcode());
139         for (auto &MO : T.operands())
140           MIB.add(MO);
141         Terminators.push_back(&T);
142         if (T.shouldUpdateAdditionalCallInfo())
143           MF.eraseAdditionalCallInfo(&T);
144       }
145     }
146   }
147 
148   for (auto &I : Terminators)
149     I->eraseFromParent();
150 }
151 
152 void 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(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(MBB, T, T.getDebugLoc(), TII->get(Opc));
169       }
170     }
171 }
172 
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) {
206   auto &F = MF.getFunction();
207   auto InstrAttr = F.getFnAttribute("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("xray-ignore-loops");
214 
215   uint64_t XRayThreshold = 0;
216   if (!AlwaysInstrument) {
217     bool IgnoreLoops = IgnoreLoopsAttr.isValid();
218     XRayThreshold = F.getFnAttributeAsParsedInteger(
219         "xray-instruction-threshold", 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(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(*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       MF, [&](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(DiagnosticInfoUnsupported(
271         Fn, "An attempt to perform XRay instrumentation for an"
272             " unsupported target."));
273 
274     return false;
275   }
276 
277   if (!F.hasFnAttribute("xray-skip-entry")) {
278     // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
279     // MachineFunction.
280     BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
281             TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
282   }
283 
284   if (!F.hasFnAttribute("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 
329 char XRayInstrumentationLegacy::ID = 0;
330 char &llvm::XRayInstrumentationID = XRayInstrumentationLegacy::ID;
331 INITIALIZE_PASS_BEGIN(XRayInstrumentationLegacy, "xray-instrumentation",
332                       "Insert XRay ops", false, false)
333 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
334 INITIALIZE_PASS_END(XRayInstrumentationLegacy, "xray-instrumentation",
335                     "Insert XRay ops", false, false)
336