1 files changed, 49 insertions, 20 deletions

diff --git a/lib/CodeGen/XRayInstrumentation.cpp b/lib/CodeGen/XRayInstrumentation.cpp
index 0b4c6e551667..3d83afcf1fc5 100644
--- a/lib/CodeGen/XRayInstrumentation.cpp
+++ b/lib/CodeGen/XRayInstrumentation.cpp
@@ -14,8 +14,8 @@
 //
 //===---------------------------------------------------------------------===//
 
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominators.h"
@@ -23,17 +23,26 @@
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/Pass.h"
-#include "llvm/Target/TargetInstrInfo.h"
 #include "llvm/Target/TargetMachine.h"
-#include "llvm/Target/TargetSubtargetInfo.h"
 
 using namespace llvm;
 
 namespace {
 
+struct InstrumentationOptions {
+  // Whether to emit PATCHABLE_TAIL_CALL.
+  bool HandleTailcall;
+
+  // Whether to emit PATCHABLE_RET/PATCHABLE_FUNCTION_EXIT for all forms of
+  // return, e.g. conditional return.
+  bool HandleAllReturns;
+};
+
 struct XRayInstrumentation : public MachineFunctionPass {
   static char ID;
 
@@ -59,7 +68,8 @@ private:
   // This is the approach to go on CPUs which have a single RET instruction,
   //   like x86/x86_64.
   void replaceRetWithPatchableRet(MachineFunction &MF,
-                                  const TargetInstrInfo *TII);
+                                  const TargetInstrInfo *TII,
+                                  InstrumentationOptions);
 
   // Prepend the original return instruction with the exit sled code ("patchable
   //   function exit" pseudo-instruction), preserving the original return
@@ -70,25 +80,28 @@ private:
   //   have to call the trampoline and return from it to the original return
   //   instruction of the function being instrumented.
   void prependRetWithPatchableExit(MachineFunction &MF,
-                                   const TargetInstrInfo *TII);
+                                   const TargetInstrInfo *TII,
+                                   InstrumentationOptions);
 };
 
 } // end anonymous namespace
 
 void XRayInstrumentation::replaceRetWithPatchableRet(
-    MachineFunction &MF, const TargetInstrInfo *TII) {
+    MachineFunction &MF, const TargetInstrInfo *TII,
+    InstrumentationOptions op) {
   // We look for *all* terminators and returns, then replace those with
   // PATCHABLE_RET instructions.
   SmallVector<MachineInstr *, 4> Terminators;
   for (auto &MBB : MF) {
     for (auto &T : MBB.terminators()) {
       unsigned Opc = 0;
-      if (T.isReturn() && T.getOpcode() == TII->getReturnOpcode()) {
+      if (T.isReturn() &&
+          (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
         // Replace return instructions with:
         //   PATCHABLE_RET <Opcode>, <Operand>...
         Opc = TargetOpcode::PATCHABLE_RET;
       }
-      if (TII->isTailCall(T)) {
+      if (TII->isTailCall(T) && op.HandleTailcall) {
         // Treat the tail call as a return instruction, which has a
         // different-looking sled than the normal return case.
         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
@@ -108,14 +121,16 @@ void XRayInstrumentation::replaceRetWithPatchableRet(
 }
 
 void XRayInstrumentation::prependRetWithPatchableExit(
-    MachineFunction &MF, const TargetInstrInfo *TII) {
-  for (auto &MBB : MF) {
+    MachineFunction &MF, const TargetInstrInfo *TII,
+    InstrumentationOptions op) {
+  for (auto &MBB : MF)
     for (auto &T : MBB.terminators()) {
       unsigned Opc = 0;
-      if (T.isReturn()) {
+      if (T.isReturn() &&
+          (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
         Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
       }
-      if (TII->isTailCall(T)) {
+      if (TII->isTailCall(T) && op.HandleTailcall) {
         Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
       }
       if (Opc != 0) {
@@ -124,11 +139,10 @@ void XRayInstrumentation::prependRetWithPatchableExit(
         BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
       }
     }
-  }
 }
 
 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
-  auto &F = *MF.getFunction();
+  auto &F = MF.getFunction();
   auto InstrAttr = F.getFnAttribute("function-instrument");
   bool AlwaysInstrument = !InstrAttr.hasAttribute(Attribute::None) &&
                           InstrAttr.isStringAttribute() &&
@@ -143,7 +157,7 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
 
     // Count the number of MachineInstr`s in MachineFunction
     int64_t MICount = 0;
-    for (const auto& MBB : MF)
+    for (const auto &MBB : MF)
       MICount += MBB.size();
 
     // Check if we have a loop.
@@ -180,20 +194,35 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   case Triple::ArchType::arm:
   case Triple::ArchType::thumb:
   case Triple::ArchType::aarch64:
-  case Triple::ArchType::ppc64le:
   case Triple::ArchType::mips:
   case Triple::ArchType::mipsel:
   case Triple::ArchType::mips64:
-  case Triple::ArchType::mips64el:
+  case Triple::ArchType::mips64el: {
     // For the architectures which don't have a single return instruction
-    prependRetWithPatchableExit(MF, TII);
+    InstrumentationOptions op;
+    op.HandleTailcall = false;
+    op.HandleAllReturns = true;
+    prependRetWithPatchableExit(MF, TII, op);
+    break;
+  }
+  case Triple::ArchType::ppc64le: {
+    // PPC has conditional returns. Turn them into branch and plain returns.
+    InstrumentationOptions op;
+    op.HandleTailcall = false;
+    op.HandleAllReturns = true;
+    replaceRetWithPatchableRet(MF, TII, op);
     break;
-  default:
+  }
+  default: {
     // For the architectures that have a single return instruction (such as
     //   RETQ on x86_64).
-    replaceRetWithPatchableRet(MF, TII);
+    InstrumentationOptions op;
+    op.HandleTailcall = true;
+    op.HandleAllReturns = false;
+    replaceRetWithPatchableRet(MF, TII, op);
     break;
   }
+  }
   return true;
 }