contrib/llvm/lib/CodeGen/XRayInstrumentation.cpp

   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/SmallVector.h"
  18 #include "llvm/ADT/STLExtras.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       unsigned Opc = 0;
 115       if (T.isReturn()) {
 116         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
 117       }
 118       if (TII->isTailCall(T)) {
 119         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
 120       }
 121       if (Opc != 0) {
 122         // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
 123         //   PATCHABLE_TAIL_CALL .
 124         BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
 125       }
 126     }
 127   }
 128 }
 129
 130 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
 131   auto &F = *MF.getFunction();
 132   auto InstrAttr = F.getFnAttribute("function-instrument");
 133   bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
 134                           InstrAttr.isStringAttribute() &&
 135                           InstrAttr.getValueAsString() == "xray-always";
 136   Attribute Attr = F.getFnAttribute("xray-instruction-threshold");
 137   unsigned XRayThreshold = 0;
 138   if (!AlwaysInstrument) {
 139     if (Attr.hasAttribute(Attribute::None) || !Attr.isStringAttribute())
 140       return false; // XRay threshold attribute not found.
 141     if (Attr.getValueAsString().getAsInteger(10, XRayThreshold))
 142       return false; // Invalid value for threshold.
 143
 144     // Count the number of MachineInstr`s in MachineFunction
 145     int64_t MICount = 0;
 146     for (const auto& MBB : MF)
 147       MICount += MBB.size();
 148
 149     // Check if we have a loop.
 150     // FIXME: Maybe make this smarter, and see whether the loops are dependent
 151     // on inputs or side-effects?
 152     MachineLoopInfo &MLI = getAnalysis<MachineLoopInfo>();
 153     if (MLI.empty() && MICount < XRayThreshold)
 154       return false; // Function is too small and has no loops.
 155   }
 156
 157   // We look for the first non-empty MachineBasicBlock, so that we can insert
 158   // the function instrumentation in the appropriate place.
 159   auto MBI = llvm::find_if(
 160       MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
 161   if (MBI == MF.end())
 162     return false; // The function is empty.
 163
 164   auto *TII = MF.getSubtarget().getInstrInfo();
 165   auto &FirstMBB = *MBI;
 166   auto &FirstMI = *FirstMBB.begin();
 167
 168   if (!MF.getSubtarget().isXRaySupported()) {
 169     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
 170                       " unsupported target.");
 171     return false;
 172   }
 173
 174   // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
 175   // MachineFunction.
 176   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
 177           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 178
 179   switch (MF.getTarget().getTargetTriple().getArch()) {
 180   case Triple::ArchType::arm:
 181   case Triple::ArchType::thumb:
 182   case Triple::ArchType::aarch64:
 183   case Triple::ArchType::ppc64le:
 184   case Triple::ArchType::mips:
 185   case Triple::ArchType::mipsel:
 186   case Triple::ArchType::mips64:
 187   case Triple::ArchType::mips64el:
 188     // For the architectures which don't have a single return instruction
 189     prependRetWithPatchableExit(MF, TII);
 190     break;
 191   default:
 192     // For the architectures that have a single return instruction (such as
 193     //   RETQ on x86_64).
 194     replaceRetWithPatchableRet(MF, TII);
 195     break;
 196   }
 197   return true;
 198 }
 199
 200 char XRayInstrumentation::ID = 0;
 201 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
 202 INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
 203                       "Insert XRay ops", false, false)
 204 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 205 INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
 206                     "Insert XRay ops", false, false)