llvm/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/CodeGen/MachineBasicBlock.h"
  19 #include "llvm/CodeGen/MachineDominators.h"
  20 #include "llvm/CodeGen/MachineFunction.h"
  21 #include "llvm/CodeGen/MachineFunctionPass.h"
  22 #include "llvm/CodeGen/MachineInstrBuilder.h"
  23 #include "llvm/CodeGen/MachineLoopInfo.h"
  24 #include "llvm/CodeGen/TargetInstrInfo.h"
  25 #include "llvm/CodeGen/TargetSubtargetInfo.h"
  26 #include "llvm/IR/Attributes.h"
  27 #include "llvm/IR/Function.h"
  28 #include "llvm/InitializePasses.h"
  29 #include "llvm/Pass.h"
  30 #include "llvm/Target/TargetMachine.h"
  31 #include "llvm/TargetParser/Triple.h"
  32
  33 using namespace llvm;
  34
  35 namespace {
  36
  37 struct InstrumentationOptions {
  38   // Whether to emit PATCHABLE_TAIL_CALL.
  39   bool HandleTailcall;
  40
  41   // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
  42   // return, e.g. conditional return.
  43   bool HandleAllReturns;
  44 };
  45
  46 struct XRayInstrumentation : public MachineFunctionPass {
  47   static char ID;
  48
  49   XRayInstrumentation() : MachineFunctionPass(ID) {
  50     initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
  51   }
  52
  53   void getAnalysisUsage(AnalysisUsage &AU) const override {
  54     AU.setPreservesCFG();
  55     AU.addPreserved<MachineLoopInfoWrapperPass>();
  56     AU.addPreserved<MachineDominatorTreeWrapperPass>();
  57     MachineFunctionPass::getAnalysisUsage(AU);
  58   }
  59
  60   bool runOnMachineFunction(MachineFunction &MF) override;
  61
  62 private:
  63   // Replace the original RET instruction with the exit sled code ("patchable
  64   //   ret" pseudo-instruction), so that at runtime XRay can replace the sled
  65   //   with a code jumping to XRay trampoline, which calls the tracing handler
  66   //   and, in the end, issues the RET instruction.
  67   // This is the approach to go on CPUs which have a single RET instruction,
  68   //   like x86/x86_64.
  69   void replaceRetWithPatchableRet(MachineFunction &MF,
  70                                   const TargetInstrInfo *TII,
  71                                   InstrumentationOptions);
  72
  73   // Prepend the original return instruction with the exit sled code ("patchable
  74   //   function exit" pseudo-instruction), preserving the original return
  75   //   instruction just after the exit sled code.
  76   // This is the approach to go on CPUs which have multiple options for the
  77   //   return instruction, like ARM. For such CPUs we can't just jump into the
  78   //   XRay trampoline and issue a single return instruction there. We rather
  79   //   have to call the trampoline and return from it to the original return
  80   //   instruction of the function being instrumented.
  81   void prependRetWithPatchableExit(MachineFunction &MF,
  82                                    const TargetInstrInfo *TII,
  83                                    InstrumentationOptions);
  84 };
  85
  86 } // end anonymous namespace
  87
  88 void XRayInstrumentation::replaceRetWithPatchableRet(
  89     MachineFunction &MF, const TargetInstrInfo *TII,
  90     InstrumentationOptions op) {
  91   // We look for *all* terminators and returns, then replace those with
  92   // PATCHABLE_RET instructions.
  93   SmallVector<MachineInstr *, 4> Terminators;
  94   for (auto &MBB : MF) {
  95     for (auto &T : MBB.terminators()) {
  96       unsigned Opc = 0;
  97       if (T.isReturn() &&
  98           (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
  99         // Replace return instructions with:
 100         //   PATCHABLE_RET <Opcode>, <Operand>...
 101         Opc = TargetOpcode::PATCHABLE_RET;
 102       }
 103       if (TII->isTailCall(T) && op.HandleTailcall) {
 104         // Treat the tail call as a return instruction, which has a
 105         // different-looking sled than the normal return case.
 106         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
 107       }
 108       if (Opc != 0) {
 109         auto MIB = BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc))
 110                        .addImm(T.getOpcode());
 111         for (auto &MO : T.operands())
 112           MIB.add(MO);
 113         Terminators.push_back(&T);
 114         if (T.shouldUpdateCallSiteInfo())
 115           MF.eraseCallSiteInfo(&T);
 116       }
 117     }
 118   }
 119
 120   for (auto &I : Terminators)
 121     I->eraseFromParent();
 122 }
 123
 124 void XRayInstrumentation::prependRetWithPatchableExit(
 125     MachineFunction &MF, const TargetInstrInfo *TII,
 126     InstrumentationOptions op) {
 127   for (auto &MBB : MF)
 128     for (auto &T : MBB.terminators()) {
 129       unsigned Opc = 0;
 130       if (T.isReturn() &&
 131           (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
 132         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
 133       }
 134       if (TII->isTailCall(T) && op.HandleTailcall) {
 135         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
 136       }
 137       if (Opc != 0) {
 138         // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
 139         //   PATCHABLE_TAIL_CALL .
 140         BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
 141       }
 142     }
 143 }
 144
 145 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
 146   auto &F = MF.getFunction();
 147   auto InstrAttr = F.getFnAttribute("function-instrument");
 148   bool AlwaysInstrument = InstrAttr.isStringAttribute() &&
 149                           InstrAttr.getValueAsString() == "xray-always";
 150   bool NeverInstrument = InstrAttr.isStringAttribute() &&
 151                          InstrAttr.getValueAsString() == "xray-never";
 152   if (NeverInstrument && !AlwaysInstrument)
 153     return false;
 154   auto IgnoreLoopsAttr = F.getFnAttribute("xray-ignore-loops");
 155
 156   uint64_t XRayThreshold = 0;
 157   if (!AlwaysInstrument) {
 158     bool IgnoreLoops = IgnoreLoopsAttr.isValid();
 159     XRayThreshold = F.getFnAttributeAsParsedInteger(
 160         "xray-instruction-threshold", std::numeric_limits<uint64_t>::max());
 161     if (XRayThreshold == std::numeric_limits<uint64_t>::max())
 162       return false;
 163
 164     // Count the number of MachineInstr`s in MachineFunction
 165     uint64_t MICount = 0;
 166     for (const auto &MBB : MF)
 167       MICount += MBB.size();
 168
 169     bool TooFewInstrs = MICount < XRayThreshold;
 170
 171     if (!IgnoreLoops) {
 172       // Get MachineDominatorTree or compute it on the fly if it's unavailable
 173       auto *MDTWrapper =
 174           getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
 175       auto *MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
 176       MachineDominatorTree ComputedMDT;
 177       if (!MDT) {
 178         ComputedMDT.getBase().recalculate(MF);
 179         MDT = &ComputedMDT;
 180       }
 181
 182       // Get MachineLoopInfo or compute it on the fly if it's unavailable
 183       auto *MLIWrapper = getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
 184       auto *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
 185       MachineLoopInfo ComputedMLI;
 186       if (!MLI) {
 187         ComputedMLI.analyze(MDT->getBase());
 188         MLI = &ComputedMLI;
 189       }
 190
 191       // Check if we have a loop.
 192       // FIXME: Maybe make this smarter, and see whether the loops are dependent
 193       // on inputs or side-effects?
 194       if (MLI->empty() && TooFewInstrs)
 195         return false; // Function is too small and has no loops.
 196     } else if (TooFewInstrs) {
 197       // Function is too small
 198       return false;
 199     }
 200   }
 201
 202   // We look for the first non-empty MachineBasicBlock, so that we can insert
 203   // the function instrumentation in the appropriate place.
 204   auto MBI = llvm::find_if(
 205       MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
 206   if (MBI == MF.end())
 207     return false; // The function is empty.
 208
 209   auto *TII = MF.getSubtarget().getInstrInfo();
 210   auto &FirstMBB = *MBI;
 211   auto &FirstMI = *FirstMBB.begin();
 212
 213   if (!MF.getSubtarget().isXRaySupported()) {
 214     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
 215                       " unsupported target.");
 216     return false;
 217   }
 218
 219   if (!F.hasFnAttribute("xray-skip-entry")) {
 220     // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
 221     // MachineFunction.
 222     BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
 223             TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 224   }
 225
 226   if (!F.hasFnAttribute("xray-skip-exit")) {
 227     switch (MF.getTarget().getTargetTriple().getArch()) {
 228     case Triple::ArchType::arm:
 229     case Triple::ArchType::thumb:
 230     case Triple::ArchType::aarch64:
 231     case Triple::ArchType::hexagon:
 232     case Triple::ArchType::loongarch64:
 233     case Triple::ArchType::mips:
 234     case Triple::ArchType::mipsel:
 235     case Triple::ArchType::mips64:
 236     case Triple::ArchType::mips64el: {
 237       // For the architectures which don't have a single return instruction
 238       InstrumentationOptions op;
 239       op.HandleTailcall = false;
 240       op.HandleAllReturns = true;
 241       prependRetWithPatchableExit(MF, TII, op);
 242       break;
 243     }
 244     case Triple::ArchType::ppc64le: {
 245       // PPC has conditional returns. Turn them into branch and plain returns.
 246       InstrumentationOptions op;
 247       op.HandleTailcall = false;
 248       op.HandleAllReturns = true;
 249       replaceRetWithPatchableRet(MF, TII, op);
 250       break;
 251     }
 252     default: {
 253       // For the architectures that have a single return instruction (such as
 254       //   RETQ on x86_64).
 255       InstrumentationOptions op;
 256       op.HandleTailcall = true;
 257       op.HandleAllReturns = false;
 258       replaceRetWithPatchableRet(MF, TII, op);
 259       break;
 260     }
 261     }
 262   }
 263   return true;
 264 }
 265
 266 char XRayInstrumentation::ID = 0;
 267 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
 268 INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
 269                       "Insert XRay ops", false, false)
 270 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
 271 INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
 272                     "Insert XRay ops", false, false)