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