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