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

[ianlevesque]

Summary

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.

[dberris]

LGTM

Do you mind adding x86 tests too?

[ianlevesque]

Add x86 tests

clang/include/clang/Basic/XRayInstr.h

 namespace clang {
 
 using XRayInstrMask = uint32_t;
 
 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
 
 struct XRayInstrSet {
   bool has(XRayInstrMask K) const {
     assert(llvm::isPowerOf2_32(K));
     return Mask & K;
   }
 
   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);
   }
 
   void clear(XRayInstrMask K = XRayInstrKind::All) { Mask &= ~K; }
 
   bool empty() const { return Mask == 0; }
 
   bool full() const { return Mask == XRayInstrKind::All; }

clang/include/clang/Driver/Options.td

def fxray_link_deps : Flag<["-"], "fxray-link-deps">, Group<f_Group>,
   Flags<[CC1Option]>,
   HelpText<"Tells clang to add the link dependencies for XRay.">;
 def fnoxray_link_deps : Flag<["-"], "fnoxray-link-deps">, Group<f_Group>,
   Flags<[CC1Option]>;
 
 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]>,
   HelpText<"Use separate accesses for consecutive bitfield runs with legal widths and alignments.">;
 def fno_fine_grained_bitfield_accesses : Flag<["-"],
   "fno-fine-grained-bitfield-accesses">, Group<f_clang_Group>, Flags<[CC1Option]>,
   HelpText<"Use large-integer access for consecutive bitfield runs.">;
 

clang/lib/Basic/XRayInstr.cpp

 //===----------------------------------------------------------------------===//
 
 #include "clang/Basic/XRayInstr.h"
 #include "llvm/ADT/StringSwitch.h"
 
 namespace clang {
 
 XRayInstrMask parseXRayInstrValue(StringRef Value) {
-  XRayInstrMask ParsedKind = llvm::StringSwitch<XRayInstrMask>(Value)
-                                 .Case("all", XRayInstrKind::All)
-                                 .Case("custom", XRayInstrKind::Custom)
-                                 .Case("function", XRayInstrKind::Function)
-                                 .Case("typed", XRayInstrKind::Typed)
-                                 .Case("none", XRayInstrKind::None)
-                                 .Default(XRayInstrKind::None);
+  XRayInstrMask ParsedKind =
+      llvm::StringSwitch<XRayInstrMask>(Value)
+          .Case("function",
+                XRayInstrKind::FunctionEntry | XRayInstrKind::FunctionExit)
+          .Case("function-entry", XRayInstrKind::FunctionEntry)
+          .Case("function-exit", XRayInstrKind::FunctionExit)
   return ParsedKind;
 }
 
 } // namespace clang

clang/lib/CodeGen/CodeGenFunction.cpp

if (const auto *FD = dyn_cast<FunctionDecl>(D))
       if (FD->getBody() &&
           FD->getBody()->getStmtClass() == Stmt::CoroutineBodyStmtClass)
         SanOpts.Mask &= ~SanitizerKind::Null;
 
   if (D) {
     // 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())
           Fn->addFnAttr("function-instrument", "xray-never");
         if (const auto *LogArgs = D->getAttr<XRayLogArgsAttr>())
           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))
         Fn->addFnAttr(
             "xray-instruction-threshold",
             llvm::itostr(CGM.getCodeGenOpts().XRayInstructionThreshold));
       if (CGM.getCodeGenOpts().XRayIgnoreLoops) {
         Fn->addFnAttr("xray-ignore-loops");

clang/lib/Driver/XRayArgs.cpp

if (Bundles.empty())
       InstrumentationBundle.Mask = XRayInstrKind::All;
     else
       for (const auto &B : Bundles) {
         llvm::SmallVector<StringRef, 2> BundleParts;
         llvm::SplitString(B, BundleParts, ",");
         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) {
             D.Diag(clang::diag::err_drv_invalid_value)
                 << "-fxray-instrumentation-bundle=" << P;
             continue;
           }
 
void XRayArgs::addArgs(const ToolChain &TC, const ArgList &Args,
   }
 
   SmallString<64> Bundle("-fxray-instrumentation-bundle=");
   if (InstrumentationBundle.full()) {
     Bundle += "all";
   } 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))
       Bundle += "typed";
   }
   CmdArgs.push_back(Args.MakeArgString(Bundle));
 }

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

 // 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 \
 // RUN:     -fxray-instrumentation-bundle=custom \
 // 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() {
   static constexpr char kPhase[] = "always";
   __xray_customevent(kPhase, 6);
   __xray_typedevent(1, kPhase, 6);
   // CUSTOM: call void @llvm.xray.customevent(i8*{{.*}}, i32 6)
   // NOCUSTOM-NOT: call void @llvm.xray.customevent(i8*{{.*}}, i32 6)
   // TYPED: call void @llvm.xray.typedevent(i16 {{.*}}, i8*{{.*}}, i32 6)
   // NOTYPED-NOT: call void @llvm.xray.typedevent(i16 {{.*}}, i8*{{.*}}, i32 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

bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
   auto &FirstMI = *FirstMBB.begin();
 
   if (!MF.getSubtarget().isXRaySupported()) {
     FirstMI.emitError("An attempt to perform XRay instrumentation for an"
                       " unsupported target.");
     return false;
   }
 
+  if (!F.hasFnAttribute("xray-skip-entry")) {
     // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
     // MachineFunction.
     BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
             TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
+  }
 
+  if (!F.hasFnAttribute("xray-skip-exit")) {
     switch (MF.getTarget().getTargetTriple().getArch()) {
     case Triple::ArchType::arm:
     case Triple::ArchType::thumb:
     case Triple::ArchType::aarch64:
     case Triple::ArchType::mips:
     case Triple::ArchType::mipsel:
     case Triple::ArchType::mips64:
     case Triple::ArchType::mips64el: {
       // For the architectures which don't have a single return instruction
       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: {
       // For the architectures that have a single return instruction (such as
       //   RETQ on x86_64).
       InstrumentationOptions op;
       op.HandleTailcall = true;
       op.HandleAllReturns = false;
       replaceRetWithPatchableRet(MF, TII, op);
       break;
     }
     }
+  }
   return true;
 }
 
 char XRayInstrumentation::ID = 0;
 char &llvm::XRayInstrumentationID = XRayInstrumentation::ID;
 INITIALIZE_PASS_BEGIN(XRayInstrumentation, "xray-instrumentation",
                       "Insert XRay ops", false, false)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_END(XRayInstrumentation, "xray-instrumentation",
                     "Insert XRay ops", false, false)

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

+; 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

+; 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

+; 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

+; 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