1 files changed, 54 insertions, 16 deletions

diff --git a/llvm/lib/CodeGen/XRayInstrumentation.cpp b/llvm/lib/CodeGen/XRayInstrumentation.cpp
index e6b65825f5b..2063ab11a74 100644
--- a/llvm/lib/CodeGen/XRayInstrumentation.cpp
+++ b/llvm/lib/CodeGen/XRayInstrumentation.cpp
@@ -34,6 +34,15 @@ 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,24 @@ void XRayInstrumentation::replaceRetWithPatchableRet(
 }
 
 void XRayInstrumentation::prependRetWithPatchableExit(
-    MachineFunction &MF, const TargetInstrInfo *TII) {
+    MachineFunction &MF, const TargetInstrInfo *TII,
+    InstrumentationOptions op) {
   for (auto &MBB : MF)
-    for (auto &T : MBB.terminators())
-      if (T.isReturn()) {
-        // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT.
-        BuildMI(MBB, T, T.getDebugLoc(),
-                TII->get(TargetOpcode::PATCHABLE_FUNCTION_EXIT));
+    for (auto &T : MBB.terminators()) {
+      unsigned Opc = 0;
+      if (T.isReturn() &&
+          (op.HandleAllReturns || T.getOpcode() == TII->getReturnOpcode())) {
+        Opc = TargetOpcode::PATCHABLE_FUNCTION_EXIT;
+      }
+      if (TII->isTailCall(T) && op.HandleTailcall) {
+        Opc = TargetOpcode::PATCHABLE_TAIL_CALL;
       }
+      if (Opc != 0) {
+        // Prepend the return instruction with PATCHABLE_FUNCTION_EXIT or
+        //   PATCHABLE_TAIL_CALL .
+        BuildMI(MBB, T, T.getDebugLoc(), TII->get(Opc));
+      }
+    }
 }
 
 bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
@@ -171,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;
-  default:
+  }
+  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: {
     // 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;
 }