lib/CodeGen/XRayInstrumentation.cpp

   1 //===- XRayInstrumentation.cpp - Adds XRay instrumentation to functions. --===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements a MachineFunctionPass that inserts the appropriate
  10 // XRay instrumentation instructions. We look for XRay-specific attributes
  11 // on the function to determine whether we should insert the replacement
  12 // operations.
  13 //
  14 //===---------------------------------------------------------------------===//
  15
  16 #include "llvm/ADT/STLExtras.h"
  17 #include "llvm/ADT/SmallVector.h"
  18 #include "llvm/ADT/Triple.h"
  19 #include "llvm/CodeGen/MachineBasicBlock.h"
  20 #include "llvm/CodeGen/MachineDominators.h"
  21 #include "llvm/CodeGen/MachineFunction.h"
  22 #include "llvm/CodeGen/MachineFunctionPass.h"
  23 #include "llvm/CodeGen/MachineInstrBuilder.h"
  24 #include "llvm/CodeGen/MachineLoopInfo.h"
  25 #include "llvm/CodeGen/TargetInstrInfo.h"
  26 #include "llvm/CodeGen/TargetSubtargetInfo.h"
  27 #include "llvm/IR/Attributes.h"
  28 #include "llvm/IR/Function.h"
  29 #include "llvm/Pass.h"
  30 #include "llvm/Target/TargetMachine.h"
  31
  32 using namespace llvm;
  33
  34 namespace {
  35
  36 struct InstrumentationOptions {
  37   // Whether to emit PATCHABLE_TAIL_CALL.
  38   bool HandleTailcall;
  39
  40   // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
  41   // return, e.g. conditional return.
  42   bool HandleAllReturns;
  43 };
  44
  45 struct XRayInstrumentation : public MachineFunctionPass {
  46   static char ID;
  47
  48   XRayInstrumentation() : MachineFunctionPass(ID) {
  49     initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
  50   }
  51
  52   void getAnalysisUsage(AnalysisUsage &AU) const override {
  53     AU.setPreservesCFG();
  54     AU.addPreserved<MachineLoopInfo>();
  55     AU.addPreserved<MachineDominatorTree>();
  56     MachineFunctionPass::getAnalysisUsage(AU);
  57   }
  58
  59   bool runOnMachineFunction(MachineFunction &MF) override;
  60
  61 private:
  62   // Replace the original RET instruction with the exit sled code ("patchable
  63   //   ret" pseudo-instruction), so that at runtime XRay can replace the sled
  64   //   with a code jumping to XRay trampoline, which calls the tracing handler
  65   //   and, in the end, issues the RET instruction.
  66   // This is the approach to go on CPUs which have a single RET instruction,
  67   //   like x86/x86_64.
  68   void replaceRetWithPatchableRet(MachineFunction &MF,
  69                                   const TargetInstrInfo *TII,
  70                                   InstrumentationOptions);
  71
  72   // Prepend the original return instruction with the exit sled code ("patchable
  73   //   function exit" pseudo-instruction), preserving the original return
  74   //   instruction just after the exit sled code.
  75   // This is the approach to go on CPUs which have multiple options for the
  76   //   return instruction, like ARM. For such CPUs we can't just jump into the
  77   //   XRay trampoline and issue a single return instruction there. We rather
  78   //   have to call the trampoline and return from it to the original return
  79   //   instruction of the function being instrumented.
  80   void prependRetWithPatchableExit(MachineFunction &MF,
  81                                    const TargetInstrInfo *TII,
  82                                    InstrumentationOptions);
  83 };
  84
  85 } // end anonymous namespace
  86
  87 void XRayInstrumentation::replaceRetWithPatchableRet(
  88     MachineFunction &MF, const TargetInstrInfo *TII,
  89     InstrumentationOptions op) {
  90   // We look for *all* terminators and returns, then replace those with
  91   // PATCHABLE_RET instructions.
  92   SmallVector<MachineInstr *, 4> Terminators;
  93   for (auto &MBB : MF) {
  94     for (auto &T : MBB.terminators()) {
  95       unsigned Opc = 0;
  96       if (T.isReturn() &&
  97           (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
  98         // Replace return instructions with:
  99         //   PATCHABLE_RET <Opcode>, <Operand>...
 100         Opc = TargetOpcode::PATCHABLE_RET;
 101       }
 102       if (TII->isTailCall(T) && op.HandleTailcall) {
 103         // Treat the tail call as a return instruction, which has a
 104         // different-looking sled than the normal return case.
 105         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
 106       }
 107       if (Opc != 0) {
 108         auto MIB = BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc))
 109                        .addImm(T.getOpcode());
 110         for (auto &MO : T.operands())
 111           MIB.add(MO);
 112         Terminators.push_back(&T);
 113         if (T.isCall())
 114           MF.updateCallSiteInfo(&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     // Get MachineDominatorTree or compute it on the fly if it's unavailable
 164     auto *MDT = getAnalysisIfAvailable<MachineDominatorTree>();
 165     MachineDominatorTree ComputedMDT;
 166     if (!MDT) {
 167       ComputedMDT.getBase().recalculate(MF);
 168       MDT = &ComputedMDT;
 169     }
 170
 171     // Get MachineLoopInfo or compute it on the fly if it's unavailable
 172     auto *MLI = getAnalysisIfAvailable<MachineLoopInfo>();
 173     MachineLoopInfo ComputedMLI;
 174     if (!MLI) {
 175       ComputedMLI.getBase().analyze(MDT->getBase());
 176       MLI = &ComputedMLI;
 177     }
 178
 179     // Check if we have a loop.
 180     // FIXME: Maybe make this smarter, and see whether the loops are dependent
 181     // on inputs or side-effects?
 182     if (MLI->empty() && MICount < XRayThreshold)
 183       return false; // Function is too small and has no loops.
 184   }
 185
 186   // We look for the first non-empty MachineBasicBlock, so that we can insert
 187   // the function instrumentation in the appropriate place.
 188   auto MBI = llvm::find_if(
 189       MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
 190   if (MBI == MF.end())
 191     return false; // The function is empty.
 192
 193   auto *TII = MF.getSubtarget().getInstrInfo();
 194   auto &FirstMBB = *MBI;
 195   auto &FirstMI = *FirstMBB.begin();
 196
 197   if (!MF.getSubtarget().isXRaySupported()) {
 198     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
 199                       " unsupported target.");
 200     return false;
 201   }
 202
 203   // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
 204   // MachineFunction.
 205   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
 206           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 207
 208   switch (MF.getTarget().getTargetTriple().getArch()) {
 209   case Triple::ArchType::arm:
 210   case Triple::ArchType::thumb:
 211   case Triple::ArchType::aarch64:
 212   case Triple::ArchType::mips:
 213   case Triple::ArchType::mipsel:
 214   case Triple::ArchType::mips64:
 215   case Triple::ArchType::mips64el: {
 216     // For the architectures which don't have a single return instruction
 217     InstrumentationOptions op;
 218     op.HandleTailcall = false;
 219     op.HandleAllReturns = true;
 220     prependRetWithPatchableExit(MF, TII, op);
 221     break;
 222   }
 223   case Triple::ArchType::ppc64le: {
 224     // PPC has conditional returns. Turn them into branch and plain returns.
 225     InstrumentationOptions op;
 226     op.HandleTailcall = false;
 227     op.HandleAllReturns = true;
 228     replaceRetWithPatchableRet(MF, TII, op);
 229     break;
 230   }
 231   default: {
 232     // For the architectures that have a single return instruction (such as
 233     //   RETQ on x86_64).
 234     InstrumentationOptions op;
 235     op.HandleTailcall = true;
 236     op.HandleAllReturns = false;
 237     replaceRetWithPatchableRet(MF, TII, op);
 238     break;
 239   }
 240   }
 241   return true;
 242 }
 243
 244 char XRayInstrumentation::ID = 0;
 245 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
 246 INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
 247                       "Insert XRay ops", false, false)
 248 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 249 INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
 250                     "Insert XRay ops", false, false)