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/CodeGen/Analysis.h"
18 #include "llvm/CodeGen/MachineFunction.h"
19 #include "llvm/CodeGen/MachineFunctionPass.h"
20 #include "llvm/CodeGen/MachineInstrBuilder.h"
21 #include "llvm/CodeGen/Passes.h"
22 #include "llvm/Support/TargetRegistry.h"
23 #include "llvm/Target/TargetInstrInfo.h"
24 #include "llvm/Target/TargetSubtargetInfo.h"
25 
26 using namespace llvm;
27 
28 namespace {
29 struct XRayInstrumentation : public MachineFunctionPass {
30   static char ID;
31 
32   XRayInstrumentation() : MachineFunctionPass(ID) {
33     initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
34   }
35 
36   bool runOnMachineFunction(MachineFunction &MF) override;
37 
38 private:
39   // Replace the original RET instruction with the exit sled code ("patchable
40   //   ret" pseudo-instruction), so that at runtime XRay can replace the sled
41   //   with a code jumping to XRay trampoline, which calls the tracing handler
42   //   and, in the end, issues the RET instruction.
43   // This is the approach to go on CPUs which have a single RET instruction,
44   //   like x86/x86_64.
45   void replaceRetWithPatchableRet(MachineFunction &MF,
46     const TargetInstrInfo *TII);
47 
48   // Prepend the original return instruction with the exit sled code ("patchable
49   //   function exit" pseudo-instruction), preserving the original return
50   //   instruction just after the exit sled code.
51   // This is the approach to go on CPUs which have multiple options for the
52   //   return instruction, like ARM. For such CPUs we can't just jump into the
53   //   XRay trampoline and issue a single return instruction there. We rather
54   //   have to call the trampoline and return from it to the original return
55   //   instruction of the function being instrumented.
56   void prependRetWithPatchableExit(MachineFunction &MF,
57     const TargetInstrInfo *TII);
58 };
59 } // anonymous namespace
60 
61 void XRayInstrumentation::replaceRetWithPatchableRet(MachineFunction &MF,
62   const TargetInstrInfo *TII)
63 {
64   // We look for *all* terminators and returns, then replace those with
65   // PATCHABLE_RET instructions.
66   SmallVector<MachineInstr *, 4> Terminators;
67   for (auto &MBB : MF) {
68     for (auto &T : MBB.terminators()) {
69       unsigned Opc = 0;
70       if (T.isReturn() && T.getOpcode() == TII->getReturnOpcode()) {
71         // Replace return instructions with:
72         //   PATCHABLE_RET <Opcode>, <Operand>...
73         Opc = TargetOpcode::PATCHABLE_RET;
74       }
75       if (TII->isTailCall(T)) {
76         // Treat the tail call as a return instruction, which has a
77         // different-looking sled than the normal return case.
78         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
79       }
80       if (Opc != 0) {
81         auto MIB = BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc))
82                        .addImm(T.getOpcode());
83         for (auto &MO : T.operands())
84           MIB.addOperand(MO);
85         Terminators.push_back(&T);
86       }
87     }
88   }
89 
90   for (auto &I : Terminators)
91     I->eraseFromParent();
92 }
93 
94 void XRayInstrumentation::prependRetWithPatchableExit(MachineFunction &MF,
95   const TargetInstrInfo *TII)
96 {
97   for (auto &MBB : MF) {
98     for (auto &T : MBB.terminators()) {
99       unsigned Opc = 0;
100       if (T.isReturn()) {
101         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
102       }
103       if (TII->isTailCall(T)) {
104         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
105       }
106       if (Opc != 0) {
107         // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
108         //   PATCHABLE_TAIL_CALL .
109         BuildMI(MBB, T, T.getDebugLoc(),TII->get(Opc));
110       }
111     }
112   }
113 }
114 
115 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
116   auto &F = *MF.getFunction();
117   auto InstrAttr = F.getFnAttribute("function-instrument");
118   bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
119                           InstrAttr.isStringAttribute() &&
120                           InstrAttr.getValueAsString() == "xray-always";
121   Attribute Attr = F.getFnAttribute("xray-instruction-threshold");
122   unsigned XRayThreshold = 0;
123   if (!AlwaysInstrument) {
124     if (Attr.hasAttribute(Attribute::None) || !Attr.isStringAttribute())
125       return false; // XRay threshold attribute not found.
126     if (Attr.getValueAsString().getAsInteger(10, XRayThreshold))
127       return false; // Invalid value for threshold.
128     if (F.size() < XRayThreshold)
129       return false; // Function is too small.
130   }
131 
132   auto &FirstMBB = *MF.begin();
133   auto &FirstMI = *FirstMBB.begin();
134 
135   if (!MF.getSubtarget().isXRaySupported()) {
136     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
137       " unsupported target.");
138     return false;
139   }
140 
141   // FIXME: Do the loop triviality analysis here or in an earlier pass.
142 
143   // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
144   // MachineFunction.
145   auto *TII = MF.getSubtarget().getInstrInfo();
146   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
147           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
148 
149   switch (MF.getTarget().getTargetTriple().getArch()) {
150   case Triple::ArchType::arm:
151   case Triple::ArchType::thumb:
152   case Triple::ArchType::aarch64:
153     // For the architectures which don't have a single return instruction
154     prependRetWithPatchableExit(MF, TII);
155     break;
156   default:
157     // For the architectures that have a single return instruction (such as
158     //   RETQ on x86_64).
159     replaceRetWithPatchableRet(MF, TII);
160     break;
161   }
162   return true;
163 }
164 
165 char XRayInstrumentation::ID = 0;
166 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
167 INITIALIZE_PASS(XRayInstrumentation, "xray-instrumentation", "Insert XRay ops",
168                 false, false)
169