[XRay] support conditional return on PPC.

Summary: Conditional returns were not taken into consideration at all. Implement them by turning them into jumps and normal returns. This means there is a slightly higher performance penalty for conditional returns, but this is the best we can do, and it still disturbs little of the rest.

Reviewers: dberris, echristo

Subscribers: sanjoy, nemanjai, hiraditya, kbarton, llvm-commits

Differential Revision: https://reviews.llvm.org/D38102

llvm-svn: 314005

4 files changed, 164 insertions, 72 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;
 }
 
diff --git a/llvm/lib/Target/PowerPC/PPC.h b/llvm/lib/Target/PowerPC/PPC.h
index 40790011f9b..4b9f0e1b1b7 100644
--- a/llvm/lib/Target/PowerPC/PPC.h
+++ b/llvm/lib/Target/PowerPC/PPC.h
@@ -26,8 +26,10 @@ namespace llvm {
   class PassRegistry;
   class FunctionPass;
   class MachineInstr;
+  class MachineOperand;
   class AsmPrinter;
   class MCInst;
+  class MCOperand;
 
   FunctionPass *createPPCCTRLoops();
 #ifndef NDEBUG
@@ -49,6 +51,9 @@ namespace llvm {
   FunctionPass *createPPCExpandISELPass();
   void LowerPPCMachineInstrToMCInst(const MachineInstr *MI, MCInst &OutMI,
                                     AsmPrinter &AP, bool isDarwin);
+  bool LowerPPCMachineOperandToMCOperand(const MachineOperand &MO,
+                                         MCOperand &OutMO, AsmPrinter &AP,
+                                         bool isDarwin);
 
   void initializePPCVSXFMAMutatePass(PassRegistry&);
   void initializePPCBoolRetToIntPass(PassRegistry&);
diff --git a/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp b/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp
index 67c849b37bd..7fee5ff1bf8 100644
--- a/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp
+++ b/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp
@@ -19,6 +19,7 @@
 #include "InstPrinter/PPCInstPrinter.h"
 #include "MCTargetDesc/PPCMCExpr.h"
 #include "MCTargetDesc/PPCMCTargetDesc.h"
+#include "MCTargetDesc/PPCPredicates.h"
 #include "PPC.h"
 #include "PPCInstrInfo.h"
 #include "PPCMachineFunctionInfo.h"
@@ -1089,7 +1090,61 @@ void PPCLinuxAsmPrinter::EmitInstruction(const MachineInstr *MI) {
     recordSled(BeginOfSled, *MI, SledKind::FUNCTION_ENTER);
     break;
   }
-  case TargetOpcode::PATCHABLE_FUNCTION_EXIT: {
+  case TargetOpcode::PATCHABLE_RET: {
+    unsigned RetOpcode = MI->getOperand(0).getImm();
+    MCInst RetInst;
+    RetInst.setOpcode(RetOpcode);
+    for (const auto &MO :
+         make_range(std::next(MI->operands_begin()), MI->operands_end())) {
+      MCOperand MCOp;
+      if (LowerPPCMachineOperandToMCOperand(MO, MCOp, *this, false))
+        RetInst.addOperand(MCOp);
+    }
+
+    bool IsConditional;
+    if (RetOpcode == PPC::BCCLR) {
+      IsConditional = true;
+    } else if (RetOpcode == PPC::TCRETURNdi8 || RetOpcode == PPC::TCRETURNri8 ||
+               RetOpcode == PPC::TCRETURNai8) {
+      break;
+    } else if (RetOpcode == PPC::BLR8 || RetOpcode == PPC::TAILB8) {
+      IsConditional = false;
+    } else {
+      EmitToStreamer(*OutStreamer, RetInst);
+      break;
+    }
+
+    MCSymbol *FallthroughLabel;
+    if (IsConditional) {
+      // Before:
+      //   bgtlr cr0
+      //
+      // After:
+      //   ble cr0, .end
+      // .p2align 3
+      // .begin:
+      //   blr    # lis 0, FuncId[16..32]
+      //   nop    # li  0, FuncId[0..15]
+      //   std 0, -8(1)
+      //   mflr 0
+      //   bl __xray_FunctionExit
+      //   mtlr 0
+      //   blr
+      // .end:
+      //
+      // Update compiler-rt/lib/xray/xray_powerpc64.cc accordingly when number
+      // of instructions change.
+      FallthroughLabel = OutContext.createTempSymbol();
+      EmitToStreamer(
+          *OutStreamer,
+          MCInstBuilder(PPC::BCC)
+              .addImm(PPC::InvertPredicate(
+                  static_cast<PPC::Predicate>(MI->getOperand(1).getImm())))
+              .addReg(MI->getOperand(2).getReg())
+              .addExpr(MCSymbolRefExpr::create(FallthroughLabel, OutContext)));
+      RetInst = MCInst();
+      RetInst.setOpcode(PPC::BLR8);
+    }
     // .p2align 3
     // .begin:
     //   b(lr)? # lis 0, FuncId[16..32]
@@ -1098,24 +1153,14 @@ void PPCLinuxAsmPrinter::EmitInstruction(const MachineInstr *MI) {
     //   mflr 0
     //   bl __xray_FunctionExit
     //   mtlr 0
-    // .end:
     //   b(lr)?
     //
     // Update compiler-rt/lib/xray/xray_powerpc64.cc accordingly when number
     // of instructions change.
-    const MachineInstr *Next = [&] {
-      MachineBasicBlock::const_iterator It(MI);
-      assert(It != MI->getParent()->end());
-      ++It;
-      assert(It->isReturn());
-      return &*It;
-    }();
     OutStreamer->EmitCodeAlignment(8);
     MCSymbol *BeginOfSled = OutContext.createTempSymbol();
     OutStreamer->EmitLabel(BeginOfSled);
-    MCInst TmpInst;
-    LowerPPCMachineInstrToMCInst(Next, TmpInst, *this, false);
-    EmitToStreamer(*OutStreamer, TmpInst);
+    EmitToStreamer(*OutStreamer, RetInst);
     EmitToStreamer(*OutStreamer, MCInstBuilder(PPC::NOP));
     EmitToStreamer(
         *OutStreamer,
@@ -1127,17 +1172,18 @@ void PPCLinuxAsmPrinter::EmitInstruction(const MachineInstr *MI) {
                            OutContext.getOrCreateSymbol("__xray_FunctionExit"),
                            OutContext)));
     EmitToStreamer(*OutStreamer, MCInstBuilder(PPC::MTLR8).addReg(PPC::X0));
+    EmitToStreamer(*OutStreamer, RetInst);
+    if (IsConditional)
+      OutStreamer->EmitLabel(FallthroughLabel);
     recordSled(BeginOfSled, *MI, SledKind::FUNCTION_EXIT);
     break;
   }
+  case TargetOpcode::PATCHABLE_FUNCTION_EXIT:
+    llvm_unreachable("PATCHABLE_FUNCTION_EXIT should never be emitted");
   case TargetOpcode::PATCHABLE_TAIL_CALL:
     // TODO: Define a trampoline `__xray_FunctionTailExit` and differentiate a
     // normal function exit from a tail exit.
-  case TargetOpcode::PATCHABLE_RET:
-    // PPC's tail call instruction, e.g. PPC::TCRETURNdi8, doesn't really
-    // lower to a PPC::B instruction. The PPC::B instruction is generated
-    // before it, and handled by the normal case.
-    llvm_unreachable("Tail call is handled in the normal case. See comments"
+    llvm_unreachable("Tail call is handled in the normal case. See comments "
                      "around this assert.");
   }
 }
diff --git a/llvm/lib/Target/PowerPC/PPCMCInstLower.cpp b/llvm/lib/Target/PowerPC/PPCMCInstLower.cpp
index b310493587a..cde511cb2e6 100644
--- a/llvm/lib/Target/PowerPC/PPCMCInstLower.cpp
+++ b/llvm/lib/Target/PowerPC/PPCMCInstLower.cpp
@@ -143,45 +143,48 @@ void llvm::LowerPPCMachineInstrToMCInst(const MachineInstr *MI, MCInst &OutMI,
   OutMI.setOpcode(MI->getOpcode());
   
   for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
-    const MachineOperand &MO = MI->getOperand(i);
-    
     MCOperand MCOp;
-    switch (MO.getType()) {
-    default:
-      MI->print(errs());
-      llvm_unreachable("unknown operand type");
-    case MachineOperand::MO_Register:
-      assert(!MO.getSubReg() && "Subregs should be eliminated!");
-      assert(MO.getReg() > PPC::NoRegister &&
-             MO.getReg() < PPC::NUM_TARGET_REGS &&
-             "Invalid register for this target!");
-      MCOp = MCOperand::createReg(MO.getReg());
-      break;
-    case MachineOperand::MO_Immediate:
-      MCOp = MCOperand::createImm(MO.getImm());
-      break;
-    case MachineOperand::MO_MachineBasicBlock:
-      MCOp = MCOperand::createExpr(MCSymbolRefExpr::create(
-                                      MO.getMBB()->getSymbol(), AP.OutContext));
-      break;
-    case MachineOperand::MO_GlobalAddress:
-    case MachineOperand::MO_ExternalSymbol:
-      MCOp = GetSymbolRef(MO, GetSymbolFromOperand(MO, AP), AP, isDarwin);
-      break;
-    case MachineOperand::MO_JumpTableIndex:
-      MCOp = GetSymbolRef(MO, AP.GetJTISymbol(MO.getIndex()), AP, isDarwin);
-      break;
-    case MachineOperand::MO_ConstantPoolIndex:
-      MCOp = GetSymbolRef(MO, AP.GetCPISymbol(MO.getIndex()), AP, isDarwin);
-      break;
-    case MachineOperand::MO_BlockAddress:
-      MCOp = GetSymbolRef(MO,AP.GetBlockAddressSymbol(MO.getBlockAddress()),AP,
-                          isDarwin);
-      break;
-    case MachineOperand::MO_RegisterMask:
-      continue;
-    }
-    
-    OutMI.addOperand(MCOp);
+    if (LowerPPCMachineOperandToMCOperand(MI->getOperand(i), MCOp, AP,
+                                          isDarwin))
+      OutMI.addOperand(MCOp);
+  }
+}
+
+bool llvm::LowerPPCMachineOperandToMCOperand(const MachineOperand &MO,
+                                             MCOperand &OutMO, AsmPrinter &AP,
+                                             bool isDarwin) {
+  switch (MO.getType()) {
+  default:
+    llvm_unreachable("unknown operand type");
+  case MachineOperand::MO_Register:
+    assert(!MO.getSubReg() && "Subregs should be eliminated!");
+    assert(MO.getReg() > PPC::NoRegister &&
+           MO.getReg() < PPC::NUM_TARGET_REGS &&
+           "Invalid register for this target!");
+    OutMO = MCOperand::createReg(MO.getReg());
+    return true;
+  case MachineOperand::MO_Immediate:
+    OutMO = MCOperand::createImm(MO.getImm());
+    return true;
+  case MachineOperand::MO_MachineBasicBlock:
+    OutMO = MCOperand::createExpr(
+        MCSymbolRefExpr::create(MO.getMBB()->getSymbol(), AP.OutContext));
+    return true;
+  case MachineOperand::MO_GlobalAddress:
+  case MachineOperand::MO_ExternalSymbol:
+    OutMO = GetSymbolRef(MO, GetSymbolFromOperand(MO, AP), AP, isDarwin);
+    return true;
+  case MachineOperand::MO_JumpTableIndex:
+    OutMO = GetSymbolRef(MO, AP.GetJTISymbol(MO.getIndex()), AP, isDarwin);
+    return true;
+  case MachineOperand::MO_ConstantPoolIndex:
+    OutMO = GetSymbolRef(MO, AP.GetCPISymbol(MO.getIndex()), AP, isDarwin);
+    return true;
+  case MachineOperand::MO_BlockAddress:
+    OutMO = GetSymbolRef(MO, AP.GetBlockAddressSymbol(MO.getBlockAddress()), AP,
+                         isDarwin);
+    return true;
+  case MachineOperand::MO_RegisterMask:
+    return false;
   }
 }