10b57cec5SDimitry Andric //===- XRayInstrumentation.cpp - Adds XRay instrumentation to functions. --===//
20b57cec5SDimitry Andric //
30b57cec5SDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
40b57cec5SDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
50b57cec5SDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
60b57cec5SDimitry Andric //
70b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
80b57cec5SDimitry Andric //
90b57cec5SDimitry Andric // This file implements a MachineFunctionPass that inserts the appropriate
100b57cec5SDimitry Andric // XRay instrumentation instructions. We look for XRay-specific attributes
110b57cec5SDimitry Andric // on the function to determine whether we should insert the replacement
120b57cec5SDimitry Andric // operations.
130b57cec5SDimitry Andric //
140b57cec5SDimitry Andric //===---------------------------------------------------------------------===//
150b57cec5SDimitry Andric
160b57cec5SDimitry Andric #include "llvm/ADT/STLExtras.h"
170b57cec5SDimitry Andric #include "llvm/ADT/SmallVector.h"
180b57cec5SDimitry Andric #include "llvm/CodeGen/MachineBasicBlock.h"
190b57cec5SDimitry Andric #include "llvm/CodeGen/MachineDominators.h"
200b57cec5SDimitry Andric #include "llvm/CodeGen/MachineFunction.h"
210b57cec5SDimitry Andric #include "llvm/CodeGen/MachineFunctionPass.h"
220b57cec5SDimitry Andric #include "llvm/CodeGen/MachineInstrBuilder.h"
230b57cec5SDimitry Andric #include "llvm/CodeGen/MachineLoopInfo.h"
240b57cec5SDimitry Andric #include "llvm/CodeGen/TargetInstrInfo.h"
250b57cec5SDimitry Andric #include "llvm/CodeGen/TargetSubtargetInfo.h"
260b57cec5SDimitry Andric #include "llvm/IR/Attributes.h"
270b57cec5SDimitry Andric #include "llvm/IR/Function.h"
28480093f4SDimitry Andric #include "llvm/InitializePasses.h"
290b57cec5SDimitry Andric #include "llvm/Pass.h"
300b57cec5SDimitry Andric #include "llvm/Target/TargetMachine.h"
3106c3fb27SDimitry Andric #include "llvm/TargetParser/Triple.h"
320b57cec5SDimitry Andric
330b57cec5SDimitry Andric using namespace llvm;
340b57cec5SDimitry Andric
350b57cec5SDimitry Andric namespace {
360b57cec5SDimitry Andric
370b57cec5SDimitry Andric struct InstrumentationOptions {
380b57cec5SDimitry Andric // Whether to emit PATCHABLE_TAIL_CALL.
390b57cec5SDimitry Andric bool HandleTailcall;
400b57cec5SDimitry Andric
410b57cec5SDimitry Andric // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
420b57cec5SDimitry Andric // return, e.g. conditional return.
430b57cec5SDimitry Andric bool HandleAllReturns;
440b57cec5SDimitry Andric };
450b57cec5SDimitry Andric
460b57cec5SDimitry Andric struct XRayInstrumentation : public MachineFunctionPass {
470b57cec5SDimitry Andric static char ID;
480b57cec5SDimitry Andric
XRayInstrumentation__anonc74285340111::XRayInstrumentation490b57cec5SDimitry Andric XRayInstrumentation() : MachineFunctionPass(ID) {
500b57cec5SDimitry Andric initializeXRayInstrumentationPass(*PassRegistry::getPassRegistry());
510b57cec5SDimitry Andric }
520b57cec5SDimitry Andric
getAnalysisUsage__anonc74285340111::XRayInstrumentation530b57cec5SDimitry Andric void getAnalysisUsage(AnalysisUsage &AU) const override {
540b57cec5SDimitry Andric AU.setPreservesCFG();
550b57cec5SDimitry Andric AU.addPreserved<MachineLoopInfo>();
560b57cec5SDimitry Andric AU.addPreserved<MachineDominatorTree>();
570b57cec5SDimitry Andric MachineFunctionPass::getAnalysisUsage(AU);
580b57cec5SDimitry Andric }
590b57cec5SDimitry Andric
600b57cec5SDimitry Andric bool runOnMachineFunction(MachineFunction &MF) override;
610b57cec5SDimitry Andric
620b57cec5SDimitry Andric private:
630b57cec5SDimitry Andric // Replace the original RET instruction with the exit sled code ("patchable
640b57cec5SDimitry Andric // ret" pseudo-instruction), so that at runtime XRay can replace the sled
650b57cec5SDimitry Andric // with a code jumping to XRay trampoline, which calls the tracing handler
660b57cec5SDimitry Andric // and, in the end, issues the RET instruction.
670b57cec5SDimitry Andric // This is the approach to go on CPUs which have a single RET instruction,
680b57cec5SDimitry Andric // like x86/x86_64.
690b57cec5SDimitry Andric void replaceRetWithPatchableRet(MachineFunction &MF,
700b57cec5SDimitry Andric const TargetInstrInfo *TII,
710b57cec5SDimitry Andric InstrumentationOptions);
720b57cec5SDimitry Andric
730b57cec5SDimitry Andric // Prepend the original return instruction with the exit sled code ("patchable
740b57cec5SDimitry Andric // function exit" pseudo-instruction), preserving the original return
750b57cec5SDimitry Andric // instruction just after the exit sled code.
760b57cec5SDimitry Andric // This is the approach to go on CPUs which have multiple options for the
770b57cec5SDimitry Andric // return instruction, like ARM. For such CPUs we can't just jump into the
780b57cec5SDimitry Andric // XRay trampoline and issue a single return instruction there. We rather
790b57cec5SDimitry Andric // have to call the trampoline and return from it to the original return
800b57cec5SDimitry Andric // instruction of the function being instrumented.
810b57cec5SDimitry Andric void prependRetWithPatchableExit(MachineFunction &MF,
820b57cec5SDimitry Andric const TargetInstrInfo *TII,
830b57cec5SDimitry Andric InstrumentationOptions);
840b57cec5SDimitry Andric };
850b57cec5SDimitry Andric
860b57cec5SDimitry Andric } // end anonymous namespace
870b57cec5SDimitry Andric
replaceRetWithPatchableRet(MachineFunction & MF,const TargetInstrInfo * TII,InstrumentationOptions op)880b57cec5SDimitry Andric void XRayInstrumentation::replaceRetWithPatchableRet(
890b57cec5SDimitry Andric MachineFunction &MF, const TargetInstrInfo *TII,
900b57cec5SDimitry Andric InstrumentationOptions op) {
910b57cec5SDimitry Andric // We look for *all* terminators and returns, then replace those with
920b57cec5SDimitry Andric // PATCHABLE_RET instructions.
930b57cec5SDimitry Andric SmallVector<MachineInstr *, 4> Terminators;
940b57cec5SDimitry Andric for (auto &MBB : MF) {
950b57cec5SDimitry Andric for (auto &T : MBB.terminators()) {
960b57cec5SDimitry Andric unsigned Opc = 0;
970b57cec5SDimitry Andric if (T.isReturn() &&
980b57cec5SDimitry Andric (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
990b57cec5SDimitry Andric // Replace return instructions with:
1000b57cec5SDimitry Andric // PATCHABLE_RET <Opcode>, <Operand>...
1010b57cec5SDimitry Andric Opc = TargetOpcode::PATCHABLE_RET;
1020b57cec5SDimitry Andric }
1030b57cec5SDimitry Andric if (TII->isTailCall(T) && op.HandleTailcall) {
1040b57cec5SDimitry Andric // Treat the tail call as a return instruction, which has a
1050b57cec5SDimitry Andric // different-looking sled than the normal return case.
1060b57cec5SDimitry Andric Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
1070b57cec5SDimitry Andric }
1080b57cec5SDimitry Andric if (Opc != 0) {
1090b57cec5SDimitry Andric auto MIB = BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc))
1100b57cec5SDimitry Andric .addImm(T.getOpcode());
1110b57cec5SDimitry Andric for (auto &MO : T.operands())
1120b57cec5SDimitry Andric MIB.add(MO);
1130b57cec5SDimitry Andric Terminators.push_back(&T);
1145ffd83dbSDimitry Andric if (T.shouldUpdateCallSiteInfo())
1158bcb0991SDimitry Andric MF.eraseCallSiteInfo(&T);
1160b57cec5SDimitry Andric }
1170b57cec5SDimitry Andric }
1180b57cec5SDimitry Andric }
1190b57cec5SDimitry Andric
1200b57cec5SDimitry Andric for (auto &I : Terminators)
1210b57cec5SDimitry Andric I->eraseFromParent();
1220b57cec5SDimitry Andric }
1230b57cec5SDimitry Andric
prependRetWithPatchableExit(MachineFunction & MF,const TargetInstrInfo * TII,InstrumentationOptions op)1240b57cec5SDimitry Andric void XRayInstrumentation::prependRetWithPatchableExit(
1250b57cec5SDimitry Andric MachineFunction &MF, const TargetInstrInfo *TII,
1260b57cec5SDimitry Andric InstrumentationOptions op) {
1270b57cec5SDimitry Andric for (auto &MBB : MF)
1280b57cec5SDimitry Andric for (auto &T : MBB.terminators()) {
1290b57cec5SDimitry Andric unsigned Opc = 0;
1300b57cec5SDimitry Andric if (T.isReturn() &&
1310b57cec5SDimitry Andric (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
1320b57cec5SDimitry Andric Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
1330b57cec5SDimitry Andric }
1340b57cec5SDimitry Andric if (TII->isTailCall(T) && op.HandleTailcall) {
1350b57cec5SDimitry Andric Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
1360b57cec5SDimitry Andric }
1370b57cec5SDimitry Andric if (Opc != 0) {
1380b57cec5SDimitry Andric // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
1390b57cec5SDimitry Andric // PATCHABLE_TAIL_CALL .
1400b57cec5SDimitry Andric BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
1410b57cec5SDimitry Andric }
1420b57cec5SDimitry Andric }
1430b57cec5SDimitry Andric }
1440b57cec5SDimitry Andric
runOnMachineFunction(MachineFunction & MF)1450b57cec5SDimitry Andric bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
1460b57cec5SDimitry Andric auto &F = MF.getFunction();
1470b57cec5SDimitry Andric auto InstrAttr = F.getFnAttribute("function-instrument");
148e8d8bef9SDimitry Andric bool AlwaysInstrument = InstrAttr.isStringAttribute() &&
1490b57cec5SDimitry Andric InstrAttr.getValueAsString() == "xray-always";
150e8d8bef9SDimitry Andric bool NeverInstrument = InstrAttr.isStringAttribute() &&
151e8d8bef9SDimitry Andric InstrAttr.getValueAsString() == "xray-never";
152e8d8bef9SDimitry Andric if (NeverInstrument && !AlwaysInstrument)
153e8d8bef9SDimitry Andric return false;
1545ffd83dbSDimitry Andric auto IgnoreLoopsAttr = F.getFnAttribute("xray-ignore-loops");
1550b57cec5SDimitry Andric
156bdd1243dSDimitry Andric uint64_t XRayThreshold = 0;
157bdd1243dSDimitry Andric if (!AlwaysInstrument) {
158e8d8bef9SDimitry Andric bool IgnoreLoops = IgnoreLoopsAttr.isValid();
159bdd1243dSDimitry Andric XRayThreshold = F.getFnAttributeAsParsedInteger(
160bdd1243dSDimitry Andric "xray-instruction-threshold", std::numeric_limits<uint64_t>::max());
161bdd1243dSDimitry Andric if (XRayThreshold == std::numeric_limits<uint64_t>::max())
162bdd1243dSDimitry Andric return false;
1635ffd83dbSDimitry Andric
1640b57cec5SDimitry Andric // Count the number of MachineInstr`s in MachineFunction
165bdd1243dSDimitry Andric uint64_t MICount = 0;
1660b57cec5SDimitry Andric for (const auto &MBB : MF)
1670b57cec5SDimitry Andric MICount += MBB.size();
1680b57cec5SDimitry Andric
1695ffd83dbSDimitry Andric bool TooFewInstrs = MICount < XRayThreshold;
1705ffd83dbSDimitry Andric
1715ffd83dbSDimitry Andric if (!IgnoreLoops) {
1720b57cec5SDimitry Andric // Get MachineDominatorTree or compute it on the fly if it's unavailable
1730b57cec5SDimitry Andric auto *MDT = getAnalysisIfAvailable<MachineDominatorTree>();
1740b57cec5SDimitry Andric MachineDominatorTree ComputedMDT;
1750b57cec5SDimitry Andric if (!MDT) {
1760b57cec5SDimitry Andric ComputedMDT.getBase().recalculate(MF);
1770b57cec5SDimitry Andric MDT = &ComputedMDT;
1780b57cec5SDimitry Andric }
1790b57cec5SDimitry Andric
1800b57cec5SDimitry Andric // Get MachineLoopInfo or compute it on the fly if it's unavailable
1810b57cec5SDimitry Andric auto *MLI = getAnalysisIfAvailable<MachineLoopInfo>();
1820b57cec5SDimitry Andric MachineLoopInfo ComputedMLI;
1830b57cec5SDimitry Andric if (!MLI) {
1840b57cec5SDimitry Andric ComputedMLI.getBase().analyze(MDT->getBase());
1850b57cec5SDimitry Andric MLI = &ComputedMLI;
1860b57cec5SDimitry Andric }
1870b57cec5SDimitry Andric
1880b57cec5SDimitry Andric // Check if we have a loop.
1890b57cec5SDimitry Andric // FIXME: Maybe make this smarter, and see whether the loops are dependent
1900b57cec5SDimitry Andric // on inputs or side-effects?
1915ffd83dbSDimitry Andric if (MLI->empty() && TooFewInstrs)
1920b57cec5SDimitry Andric return false; // Function is too small and has no loops.
1935ffd83dbSDimitry Andric } else if (TooFewInstrs) {
1945ffd83dbSDimitry Andric // Function is too small
1955ffd83dbSDimitry Andric return false;
1965ffd83dbSDimitry Andric }
1970b57cec5SDimitry Andric }
1980b57cec5SDimitry Andric
1990b57cec5SDimitry Andric // We look for the first non-empty MachineBasicBlock, so that we can insert
2000b57cec5SDimitry Andric // the function instrumentation in the appropriate place.
2010b57cec5SDimitry Andric auto MBI = llvm::find_if(
2020b57cec5SDimitry Andric MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
2030b57cec5SDimitry Andric if (MBI == MF.end())
2040b57cec5SDimitry Andric return false; // The function is empty.
2050b57cec5SDimitry Andric
2060b57cec5SDimitry Andric auto *TII = MF.getSubtarget().getInstrInfo();
2070b57cec5SDimitry Andric auto &FirstMBB = *MBI;
2080b57cec5SDimitry Andric auto &FirstMI = *FirstMBB.begin();
2090b57cec5SDimitry Andric
2100b57cec5SDimitry Andric if (!MF.getSubtarget().isXRaySupported()) {
2110b57cec5SDimitry Andric FirstMI.emitError("An attempt to perform XRay instrumentation for an"
2120b57cec5SDimitry Andric " unsupported target.");
2130b57cec5SDimitry Andric return false;
2140b57cec5SDimitry Andric }
2150b57cec5SDimitry Andric
2165ffd83dbSDimitry Andric if (!F.hasFnAttribute("xray-skip-entry")) {
2170b57cec5SDimitry Andric // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
2180b57cec5SDimitry Andric // MachineFunction.
2190b57cec5SDimitry Andric BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
2200b57cec5SDimitry Andric TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
2215ffd83dbSDimitry Andric }
2220b57cec5SDimitry Andric
2235ffd83dbSDimitry Andric if (!F.hasFnAttribute("xray-skip-exit")) {
2240b57cec5SDimitry Andric switch (MF.getTarget().getTargetTriple().getArch()) {
2250b57cec5SDimitry Andric case Triple::ArchType::arm:
2260b57cec5SDimitry Andric case Triple::ArchType::thumb:
2270b57cec5SDimitry Andric case Triple::ArchType::aarch64:
2280eae32dcSDimitry Andric case Triple::ArchType::hexagon:
22906c3fb27SDimitry Andric case Triple::ArchType::loongarch64:
2300b57cec5SDimitry Andric case Triple::ArchType::mips:
2310b57cec5SDimitry Andric case Triple::ArchType::mipsel:
2320b57cec5SDimitry Andric case Triple::ArchType::mips64:
2330b57cec5SDimitry Andric case Triple::ArchType::mips64el: {
2340b57cec5SDimitry Andric // For the architectures which don't have a single return instruction
2350b57cec5SDimitry Andric InstrumentationOptions op;
2360b57cec5SDimitry Andric op.HandleTailcall = false;
2370b57cec5SDimitry Andric op.HandleAllReturns = true;
2380b57cec5SDimitry Andric prependRetWithPatchableExit(MF, TII, op);
2390b57cec5SDimitry Andric break;
2400b57cec5SDimitry Andric }
2410b57cec5SDimitry Andric case Triple::ArchType::ppc64le: {
2420b57cec5SDimitry Andric // PPC has conditional returns. Turn them into branch and plain returns.
2430b57cec5SDimitry Andric InstrumentationOptions op;
2440b57cec5SDimitry Andric op.HandleTailcall = false;
2450b57cec5SDimitry Andric op.HandleAllReturns = true;
2460b57cec5SDimitry Andric replaceRetWithPatchableRet(MF, TII, op);
2470b57cec5SDimitry Andric break;
2480b57cec5SDimitry Andric }
2490b57cec5SDimitry Andric default: {
2500b57cec5SDimitry Andric // For the architectures that have a single return instruction (such as
2510b57cec5SDimitry Andric // RETQ on x86_64).
2520b57cec5SDimitry Andric InstrumentationOptions op;
2530b57cec5SDimitry Andric op.HandleTailcall = true;
2540b57cec5SDimitry Andric op.HandleAllReturns = false;
2550b57cec5SDimitry Andric replaceRetWithPatchableRet(MF, TII, op);
2560b57cec5SDimitry Andric break;
2570b57cec5SDimitry Andric }
2580b57cec5SDimitry Andric }
2595ffd83dbSDimitry Andric }
2600b57cec5SDimitry Andric return true;
2610b57cec5SDimitry Andric }
2620b57cec5SDimitry Andric
2630b57cec5SDimitry Andric char XRayInstrumentation::ID = 0;
2640b57cec5SDimitry Andric char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
2650b57cec5SDimitry Andric INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
2660b57cec5SDimitry Andric "Insert XRay ops", false, false)
2670b57cec5SDimitry Andric INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
2680b57cec5SDimitry Andric INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
2690b57cec5SDimitry Andric "Insert XRay ops", false, false)
270