[xray] Allow instrumenting only function entry and/or only function exit

Extend -fxray-instrumentation-bundle to split function-entry and
function-exit into two separate options, so that it is possible to
instrument only function entry or only function exit.  For use cases
that only care about one or the other this will save significant overhead
and code size.

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

clang/include/clang/Basic/XRayInstr.h

@@ -28,17 +28,19 @@ namespace XRayInstrKind {
 
 // TODO: Auto-generate these as we add more instrumentation kinds.
 enum XRayInstrOrdinal : XRayInstrMask {
-  XRIO_Function,
+  XRIO_FunctionEntry,
+  XRIO_FunctionExit,
   XRIO_Custom,
   XRIO_Typed,
   XRIO_Count
 };
 
 constexpr XRayInstrMask None = 0;
-constexpr XRayInstrMask Function = 1U << XRIO_Function;
+constexpr XRayInstrMask FunctionEntry = 1U << XRIO_FunctionEntry;
+constexpr XRayInstrMask FunctionExit = 1U << XRIO_FunctionExit;
 constexpr XRayInstrMask Custom = 1U << XRIO_Custom;
 constexpr XRayInstrMask Typed = 1U << XRIO_Typed;
-constexpr XRayInstrMask All = Function | Custom | Typed;
+constexpr XRayInstrMask All = FunctionEntry | FunctionExit | Custom | Typed;
 
 } // namespace XRayInstrKind
 
@@ -51,7 +53,6 @@ struct XRayInstrSet {
   bool hasOneOf(XRayInstrMask K) const { return Mask & K; }
 
   void set(XRayInstrMask K, bool Value) {
-    assert(llvm::isPowerOf2_32(K));
     Mask = Value ? (Mask | K) : (Mask & ~K);
   }
 

clang/include/clang/Driver/Options.td

@@ -1308,7 +1308,7 @@ def fnoxray_link_deps : Flag<["-"], "fnoxray-link-deps">, Group<f_Group>,
 def fxray_instrumentation_bundle :
   JoinedOrSeparate<["-"], "fxray-instrumentation-bundle=">,
   Group<f_Group>, Flags<[CC1Option]>,
-  HelpText<"Select which XRay instrumentation points to emit. Options: all, none, function, custom. Default is 'all'.">;
+  HelpText<"Select which XRay instrumentation points to emit. Options: all, none, function-entry, function-exit, function, custom. Default is 'all'.  'function' includes both 'function-entry' and 'function-exit'.">;
 
 def ffine_grained_bitfield_accesses : Flag<["-"],
   "ffine-grained-bitfield-accesses">, Group<f_clang_Group>, Flags<[CC1Option]>,

clang/lib/Basic/XRayInstr.cpp

@@ -16,13 +16,17 @@
 namespace clang {
 
 XRayInstrMask parseXRayInstrValue(StringRef Value) {
-  XRayInstrMask ParsedKind = llvm::StringSwitch<XRayInstrMask>(Value)
+  XRayInstrMask ParsedKind =
-                                 .Case("all", XRayInstrKind::All)
+      llvm::StringSwitch<XRayInstrMask>(Value)
-                                 .Case("custom", XRayInstrKind::Custom)
+          .Case("all", XRayInstrKind::All)
-                                 .Case("function", XRayInstrKind::Function)
+          .Case("custom", XRayInstrKind::Custom)
-                                 .Case("typed", XRayInstrKind::Typed)
+          .Case("function",
-                                 .Case("none", XRayInstrKind::None)
+                XRayInstrKind::FunctionEntry | XRayInstrKind::FunctionExit)
-                                 .Default(XRayInstrKind::None);
+          .Case("function-entry", XRayInstrKind::FunctionEntry)
+          .Case("function-exit", XRayInstrKind::FunctionExit)
+          .Case("typed", XRayInstrKind::Typed)
+          .Case("none", XRayInstrKind::None)
+          .Default(XRayInstrKind::None);
   return ParsedKind;
 }
 

clang/lib/CodeGen/CodeGenFunction.cpp

@@ -803,7 +803,9 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
     // Apply xray attributes to the function (as a string, for now)
     if (const auto *XRayAttr = D->getAttr<XRayInstrumentAttr>()) {
       if (CGM.getCodeGenOpts().XRayInstrumentationBundle.has(
-              XRayInstrKind::Function)) {
+              XRayInstrKind::FunctionEntry) ||
+          CGM.getCodeGenOpts().XRayInstrumentationBundle.has(
+              XRayInstrKind::FunctionExit)) {
         if (XRayAttr->alwaysXRayInstrument() && ShouldXRayInstrumentFunction())
           Fn->addFnAttr("function-instrument", "xray-always");
         if (XRayAttr->neverXRayInstrument())
@@ -812,6 +814,14 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
           if (ShouldXRayInstrumentFunction())
             Fn->addFnAttr("xray-log-args",
                           llvm::utostr(LogArgs->getArgumentCount()));
+        if (!CGM.getCodeGenOpts().XRayInstrumentationBundle.has(
+                XRayInstrKind::FunctionExit)) {
+          Fn->addFnAttr("xray-skip-exit");
+        }
+        if (!CGM.getCodeGenOpts().XRayInstrumentationBundle.has(
+                XRayInstrKind::FunctionEntry)) {
+          Fn->addFnAttr("xray-skip-entry");
+        }
       }
     } else {
       if (ShouldXRayInstrumentFunction() && !CGM.imbueXRayAttrs(Fn, Loc))

clang/lib/Driver/XRayArgs.cpp

@@ -113,7 +113,8 @@ XRayArgs::XRayArgs(const ToolChain &TC, const ArgList &Args) {
         for (const auto &P : BundleParts) {
           // TODO: Automate the generation of the string case table.
           auto Valid = llvm::StringSwitch<bool>(P)
-                           .Cases("none", "all", "function", "custom", true)
+                           .Cases("none", "all", "function", "function-entry",
+                                  "function-exit", "custom", true)
                            .Default(false);
 
           if (!Valid) {
@@ -237,8 +238,14 @@ void XRayArgs::addArgs(const ToolChain &TC, const ArgList &Args,
   } else if (InstrumentationBundle.empty()) {
     Bundle += "none";
   } else {
-    if (InstrumentationBundle.has(XRayInstrKind::Function))
+    if (InstrumentationBundle.has(XRayInstrKind::FunctionEntry) &&
+        InstrumentationBundle.has(XRayInstrKind::FunctionExit))
       Bundle += "function";
+    else if (InstrumentationBundle.has(XRayInstrKind::FunctionEntry))
+      Bundle += "function-entry";
+    else if (InstrumentationBundle.has(XRayInstrKind::FunctionExit))
+      Bundle += "function-exit";
+
     if (InstrumentationBundle.has(XRayInstrKind::Custom))
       Bundle += "custom";
     if (InstrumentationBundle.has(XRayInstrKind::Typed))

clang/test/CodeGen/xray-instrumentation-bundles.cpp

@@ -34,6 +34,18 @@
 // RUN:     -fxray-instrumentation-bundle=typed -x c++ \
 // RUN:     -std=c++11 -triple x86_64-unknown-unknown -emit-llvm -o - %s \
 // RUN:     | FileCheck --check-prefixes CHECK,FUNCTION,CUSTOM,TYPED %s
+// RUN: %clang_cc1 -fxray-instrument \
+// RUN:     -fxray-instrumentation-bundle=function-entry -x c++ \
+// RUN:     -std=c++11 -triple x86_64-unknown-unknown -emit-llvm -o - %s \
+// RUN:     | FileCheck --check-prefixes CHECK,NOCUSTOM,NOTYPED,SKIPEXIT %s
+// RUN: %clang_cc1 -fxray-instrument \
+// RUN:     -fxray-instrumentation-bundle=function-exit -x c++ \
+// RUN:     -std=c++11 -triple x86_64-unknown-unknown -emit-llvm -o - %s \
+// RUN:     | FileCheck --check-prefixes CHECK,NOCUSTOM,NOTYPED,SKIPENTRY %s
+// RUN: %clang_cc1 -fxray-instrument \
+// RUN:     -fxray-instrumentation-bundle=function-entry,function-exit -x c++ \
+// RUN:     -std=c++11 -triple x86_64-unknown-unknown -emit-llvm -o - %s \
+// RUN:     | FileCheck --check-prefixes CHECK,FUNCTION,NOCUSTOM,NOTYPED %s
 
 // CHECK: define void @_Z16alwaysInstrumentv() #[[ALWAYSATTR:[0-9]+]] {
 [[clang::xray_always_instrument]] void alwaysInstrument() {
@@ -48,3 +60,6 @@
 
 // FUNCTION: attributes #[[ALWAYSATTR]] = {{.*}} "function-instrument"="xray-always" {{.*}}
 // NOFUNCTION-NOT: attributes #[[ALWAYSATTR]] = {{.*}} "function-instrument"="xray-always" {{.*}}
+
+// SKIPENTRY: attributes #[[ALWAYSATTR]] = {{.*}} "function-instrument"="xray-always" {{.*}} "xray-skip-entry" {{.*}}
+// SKIPEXIT: attributes #[[ALWAYSATTR]] = {{.*}} "function-instrument"="xray-always" {{.*}} "xray-skip-exit" {{.*}}

llvm/lib/CodeGen/XRayInstrumentation.cpp

@@ -212,43 +212,47 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
     return false;
   }
 
-  // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
+  if (!F.hasFnAttribute("xray-skip-entry")) {
-  // MachineFunction.
+    // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
-  BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
+    // MachineFunction.
-          TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
+    BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
+            TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
+  }
 
-  switch (MF.getTarget().getTargetTriple().getArch()) {
+  if (!F.hasFnAttribute("xray-skip-exit")) {
-  case Triple::ArchType::arm:
+    switch (MF.getTarget().getTargetTriple().getArch()) {
-  case Triple::ArchType::thumb:
+    case Triple::ArchType::arm:
-  case Triple::ArchType::aarch64:
+    case Triple::ArchType::thumb:
-  case Triple::ArchType::mips:
+    case Triple::ArchType::aarch64:
-  case Triple::ArchType::mipsel:
+    case Triple::ArchType::mips:
-  case Triple::ArchType::mips64:
+    case Triple::ArchType::mipsel:
-  case Triple::ArchType::mips64el: {
+    case Triple::ArchType::mips64:
-    // For the architectures which don't have a single return instruction
+    case Triple::ArchType::mips64el: {
-    InstrumentationOptions op;
+      // For the architectures which don't have a single return instruction
-    op.HandleTailcall = false;
+      InstrumentationOptions op;
-    op.HandleAllReturns = true;
+      op.HandleTailcall = false;
-    prependRetWithPatchableExit(MF, TII, op);
+      op.HandleAllReturns = true;
-    break;
+      prependRetWithPatchableExit(MF, TII, op);
-  }
+      break;
-  case Triple::ArchType::ppc64le: {
+    }
-    // PPC has conditional returns. Turn them into branch and plain returns.
+    case Triple::ArchType::ppc64le: {
-    InstrumentationOptions op;
+      // PPC has conditional returns. Turn them into branch and plain returns.
-    op.HandleTailcall = false;
+      InstrumentationOptions op;
-    op.HandleAllReturns = true;
+      op.HandleTailcall = false;
-    replaceRetWithPatchableRet(MF, TII, op);
+      op.HandleAllReturns = true;
-    break;
+      replaceRetWithPatchableRet(MF, TII, op);
-  }
+      break;
-  default: {
+    }
-    // For the architectures that have a single return instruction (such as
+    default: {
-    //   RETQ on x86_64).
+      // For the architectures that have a single return instruction (such as
-    InstrumentationOptions op;
+      //   RETQ on x86_64).
-    op.HandleTailcall = true;
+      InstrumentationOptions op;
-    op.HandleAllReturns = false;
+      op.HandleTailcall = true;
-    replaceRetWithPatchableRet(MF, TII, op);
+      op.HandleAllReturns = false;
-    break;
+      replaceRetWithPatchableRet(MF, TII, op);
-  }
+      break;
+    }
+    }
   }
   return true;
 }

llvm/test/CodeGen/AArch64/xray-partial-instrumentation-skip-entry.ll

@@ -0,0 +1,21 @@
+; RUN: llc -filetype=asm -o - -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
+
+define i32 @foo() nounwind noinline uwtable "function-instrument"="xray-always" "xray-skip-entry" {
+; CHECK-NOT: Lxray_sled_0:
+  ret i32 0
+; CHECK-LABEL: Lxray_sled_0:
+; CHECK-NEXT:  b  #32
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-LABEL: Ltmp0:
+; CHECK-NEXT:  ret
+}
+; CHECK-LABEL: xray_instr_map
+; CHECK-LABEL: Lxray_sleds_start0
+; CHECK:       .xword .Lxray_sled_0
+; CHECK-LABEL: Lxray_sleds_end0

llvm/test/CodeGen/AArch64/xray-partial-instrumentation-skip-exit.ll

@@ -0,0 +1,21 @@
+; RUN: llc -filetype=asm -o - -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
+
+define i32 @foo() nounwind noinline uwtable "function-instrument"="xray-always" "xray-skip-exit" {
+; CHECK-LABEL: Lxray_sled_0:
+; CHECK-NEXT:  b  #32
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-NEXT:  nop
+; CHECK-LABEL: Ltmp0:
+  ret i32 0
+; CHECK-NOT: Lxray_sled_1:
+; CHECK:  ret
+}
+; CHECK-LABEL: xray_instr_map
+; CHECK-LABEL: Lxray_sleds_start0
+; CHECK:       .xword .Lxray_sled_0
+; CHECK-LABEL: Lxray_sleds_end0

llvm/test/CodeGen/X86/xray-partial-instrumentation-skip-entry.ll

@@ -0,0 +1,50 @@
+; RUN: llc -verify-machineinstrs -filetype=asm -o - -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: llc -verify-machineinstrs -filetype=asm -o - \
+; RUN:     -mtriple=x86_64-unknown-linux-gnu -relocation-model=pic < %s | FileCheck %s
+; RUN: llc -verify-machineinstrs -filetype=asm -o - -mtriple=x86_64-darwin-unknown    < %s | FileCheck %s
+
+define i32 @foo() nounwind noinline uwtable "function-instrument"="xray-always" "xray-skip-entry" {
+; CHECK-NOT: Lxray_sled_0:
+  ret i32 0
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_0:
+; CHECK:       retq
+; CHECK-NEXT:  nopw %cs:512(%rax,%rax)
+}
+; CHECK-LABEL: xray_instr_map
+; CHECK-LABEL: Lxray_sleds_start0:
+; CHECK:       .quad {{.*}}xray_sled_0
+; CHECK-LABEL: Lxray_sleds_end0:
+; CHECK-LABEL: xray_fn_idx
+; CHECK:       .quad {{.*}}xray_sleds_start0
+; CHECK-NEXT:  .quad {{.*}}xray_sleds_end0
+
+
+; We test multiple returns in a single function to make sure we're getting all
+; of them with XRay instrumentation.
+define i32 @bar(i32 %i) nounwind noinline uwtable "function-instrument"="xray-always" "xray-skip-entry" {
+; CHECK-NOT: Lxray_sled_1:
+Test:
+  %cond = icmp eq i32 %i, 0
+  br i1 %cond, label %IsEqual, label %NotEqual
+IsEqual:
+  ret i32 0
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_1:
+; CHECK:       retq
+; CHECK-NEXT:  nopw %cs:512(%rax,%rax)
+NotEqual:
+  ret i32 1
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_2:
+; CHECK:       retq
+; CHECK-NEXT:  nopw %cs:512(%rax,%rax)
+}
+; CHECK-LABEL: xray_instr_map
+; CHECK-LABEL: Lxray_sleds_start1:
+; CHECK:       .quad {{.*}}xray_sled_1
+; CHECK:       .quad {{.*}}xray_sled_2
+; CHECK-LABEL: Lxray_sleds_end1:
+; CHECK-LABEL: xray_fn_idx
+; CHECK:       .quad {{.*}}xray_sleds_start1
+; CHECK-NEXT:  .quad {{.*}}xray_sleds_end1

llvm/test/CodeGen/X86/xray-partial-instrumentation-skip-exit.ll

@@ -0,0 +1,49 @@
+; RUN: llc -verify-machineinstrs -filetype=asm -o - -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: llc -verify-machineinstrs -filetype=asm -o - \
+; RUN:     -mtriple=x86_64-unknown-linux-gnu -relocation-model=pic < %s | FileCheck %s
+; RUN: llc -verify-machineinstrs -filetype=asm -o - -mtriple=x86_64-darwin-unknown    < %s | FileCheck %s
+
+define i32 @foo() nounwind noinline uwtable "function-instrument"="xray-always" "xray-skip-exit" {
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_0:
+; CHECK:       .ascii "\353\t"
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+  ret i32 0
+; CHECK-NOT: Lxray_sled_1:
+; CHECK:       retq
+}
+; CHECK-LABEL: xray_instr_map
+; CHECK-LABEL: Lxray_sleds_start0:
+; CHECK:       .quad {{.*}}xray_sled_0
+; CHECK-LABEL: Lxray_sleds_end0:
+; CHECK-LABEL: xray_fn_idx
+; CHECK:       .quad {{.*}}xray_sleds_start0
+; CHECK-NEXT:  .quad {{.*}}xray_sleds_end0
+
+
+; We test multiple returns in a single function to make sure we're skipping all
+; of them with XRay instrumentation.
+define i32 @bar(i32 %i) nounwind noinline uwtable "function-instrument"="xray-always" "xray-skip-exit" {
+; CHECK:       .p2align 1, 0x90
+; CHECK-LABEL: Lxray_sled_1:
+; CHECK:       .ascii "\353\t"
+; CHECK-NEXT:  nopw 512(%rax,%rax)
+Test:
+  %cond = icmp eq i32 %i, 0
+  br i1 %cond, label %IsEqual, label %NotEqual
+IsEqual:
+  ret i32 0
+; CHECK-NOT: Lxray_sled_{{.*}}:
+; CHECK:       retq
+NotEqual:
+  ret i32 1
+; CHECK-NOT: Lxray_sled_{{.*}}:
+; CHECK:       retq
+}
+; CHECK-LABEL: xray_instr_map
+; CHECK-LABEL: Lxray_sleds_start1:
+; CHECK:       .quad {{.*}}xray_sled_1
+; CHECK-LABEL: Lxray_sleds_end1:
+; CHECK-LABEL: xray_fn_idx
+; CHECK:       .quad {{.*}}xray_sleds_start1
+; CHECK-NEXT:  .quad {{.*}}xray_sleds_end1