1 //===- XRayInstrumentation.cpp - Adds XRay instrumentation to functions. --===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file implements a MachineFunctionPass that inserts the appropriate
11 // XRay instrumentation instructions. We look for XRay-specific attributes
12 // on the function to determine whether we should insert the replacement
13 // operations.
14 //
15 //===---------------------------------------------------------------------===//
16 
17 #include "llvm/ADT/STLExtras.h"
18 #include "llvm/ADT/SmallVector.h"
19 #include "llvm/ADT/Triple.h"
20 #include "llvm/CodeGen/MachineBasicBlock.h"
21 #include "llvm/CodeGen/MachineDominators.h"
22 #include "llvm/CodeGen/MachineFunction.h"
23 #include "llvm/CodeGen/MachineFunctionPass.h"
24 #include "llvm/CodeGen/MachineInstrBuilder.h"
25 #include "llvm/CodeGen/MachineLoopInfo.h"
26 #include "llvm/CodeGen/TargetInstrInfo.h"
27 #include "llvm/CodeGen/TargetSubtargetInfo.h"
28 #include "llvm/IR/Attributes.h"
29 #include "llvm/IR/Function.h"
30 #include "llvm/Pass.h"
31 #include "llvm/Target/TargetMachine.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.addRequired<MachineLoopInfo>();
56     AU.addPreserved<MachineLoopInfo>();
57     AU.addPreserved<MachineDominatorTree>();
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       }
116     }
117   }
118 
119   for (auto &I : Terminators)
120     I->eraseFromParent();
121 }
122 
123 void XRayInstrumentation::prependRetWithPatchableExit(
124     MachineFunction &MF, const TargetInstrInfo *TII,
125     InstrumentationOptions op) {
126   for (auto &MBB : MF)
127     for (auto &T : MBB.terminators()) {
128       unsigned Opc = 0;
129       if (T.isReturn() &&
130           (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
131         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
132       }
133       if (TII->isTailCall(T) && op.HandleTailcall) {
134         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
135       }
136       if (Opc != 0) {
137         // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
138         //   PATCHABLE_TAIL_CALL .
139         BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
140       }
141     }
142 }
143 
144 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
145   auto &F = *MF.getFunction();
146   auto InstrAttr = F.getFnAttribute("function-instrument");
147   bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
148                           InstrAttr.isStringAttribute() &&
149                           InstrAttr.getValueAsString() == "xray-always";
150   Attribute Attr = F.getFnAttribute("xray-instruction-threshold");
151   unsigned XRayThreshold = 0;
152   if (!AlwaysInstrument) {
153     if (Attr.hasAttribute(Attribute::None) || !Attr.isStringAttribute())
154       return false; // XRay threshold attribute not found.
155     if (Attr.getValueAsString().getAsInteger(10, XRayThreshold))
156       return false; // Invalid value for threshold.
157 
158     // Count the number of MachineInstr`s in MachineFunction
159     int64_t MICount = 0;
160     for (const auto &MBB : MF)
161       MICount += MBB.size();
162 
163     // Check if we have a loop.
164     // FIXME: Maybe make this smarter, and see whether the loops are dependent
165     // on inputs or side-effects?
166     MachineLoopInfo &MLI = getAnalysis<MachineLoopInfo>();
167     if (MLI.empty() && MICount < XRayThreshold)
168       return false; // Function is too small and has no loops.
169   }
170 
171   // We look for the first non-empty MachineBasicBlock, so that we can insert
172   // the function instrumentation in the appropriate place.
173   auto MBI = llvm::find_if(
174       MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
175   if (MBI == MF.end())
176     return false; // The function is empty.
177 
178   auto *TII = MF.getSubtarget().getInstrInfo();
179   auto &FirstMBB = *MBI;
180   auto &FirstMI = *FirstMBB.begin();
181 
182   if (!MF.getSubtarget().isXRaySupported()) {
183     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
184                       " unsupported target.");
185     return false;
186   }
187 
188   // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
189   // MachineFunction.
190   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
191           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
192 
193   switch (MF.getTarget().getTargetTriple().getArch()) {
194   case Triple::ArchType::arm:
195   case Triple::ArchType::thumb:
196   case Triple::ArchType::aarch64:
197   case Triple::ArchType::mips:
198   case Triple::ArchType::mipsel:
199   case Triple::ArchType::mips64:
200   case Triple::ArchType::mips64el: {
201     // For the architectures which don't have a single return instruction
202     InstrumentationOptions op;
203     op.HandleTailcall = false;
204     op.HandleAllReturns = true;
205     prependRetWithPatchableExit(MF, TII, op);
206     break;
207   }
208   case Triple::ArchType::ppc64le: {
209     // PPC has conditional returns. Turn them into branch and plain returns.
210     InstrumentationOptions op;
211     op.HandleTailcall = false;
212     op.HandleAllReturns = true;
213     replaceRetWithPatchableRet(MF, TII, op);
214     break;
215   }
216   default: {
217     // For the architectures that have a single return instruction (such as
218     //   RETQ on x86_64).
219     InstrumentationOptions op;
220     op.HandleTailcall = true;
221     op.HandleAllReturns = false;
222     replaceRetWithPatchableRet(MF, TII, op);
223     break;
224   }
225   }
226   return true;
227 }
228 
229 char XRayInstrumentation::ID = 0;
230 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
231 INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
232                       "Insert XRay ops", false, false)
233 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
234 INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
235                     "Insert XRay ops", false, false)
236