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