[llvm] [BOLT] Add --instrument-funcs-file option to scope instrumentation (PR #197226)

[YongKang Zhu via llvm-commits]

https://github.com/yozhu created https://github.com/llvm/llvm-project/pull/197226

Adds a new option that takes a file containing one function name per
line (exact match) and restricts instrumentation to those functions.
Composes with --instrument-hot-only as an AND filter.

>From 5a806f7ec9959f7f0dce38110b90f38cec463281 Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at fb.com>
Date: Mon, 11 May 2026 10:38:48 -0700
Subject: [PATCH] [BOLT] Add --instrument-funcs-file option to scope
 instrumentation

Adds a new option that takes a file containing one function name per
line (exact match) and restricts instrumentation to those functions.
Composes with --instrument-hot-only as an AND filter.
---
 bolt/lib/Passes/Instrumentation.cpp   | 41 +++++++++++++-
 bolt/test/X86/instrument-funcs-file.s | 82 +++++++++++++++++++++++++++
 2 files changed, 121 insertions(+), 2 deletions(-)
 create mode 100644 bolt/test/X86/instrument-funcs-file.s

diff --git a/bolt/lib/Passes/Instrumentation.cpp b/bolt/lib/Passes/Instrumentation.cpp
index ffcfb26415971..ae9faf4e9cc2d 100644
--- a/bolt/lib/Passes/Instrumentation.cpp
+++ b/bolt/lib/Passes/Instrumentation.cpp
@@ -15,8 +15,10 @@
 #include "bolt/RuntimeLibs/InstrumentationRuntimeLibrary.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "bolt/Utils/Utils.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/RWMutex.h"
+#include <fstream>
 #include <queue>
 #include <stack>
 
@@ -76,6 +78,13 @@ cl::opt<bool> InstrumentationWaitForks(
              "(use with instrumentation-sleep-time option)"),
     cl::init(false), cl::Optional, cl::cat(BoltInstrCategory));
 
+cl::opt<std::string> InstrumentFuncsFile(
+    "instrument-funcs-file",
+    cl::desc("file with list of function names (one per line) to instrument; "
+             "only functions whose name exactly matches a line in this file "
+             "will be instrumented"),
+    cl::Optional, cl::cat(BoltInstrCategory));
+
 cl::opt<bool>
     InstrumentHotOnly("instrument-hot-only",
                       cl::desc("only insert instrumentation on hot functions "
@@ -643,9 +652,37 @@ Error Instrumentation::runOnFunctions(BinaryContext &BC) {
 
   createAuxiliaryFunctions(BC);
 
+  const bool HasInstrumentFuncsFilter = !opts::InstrumentFuncsFile.empty();
+  StringSet<> InstrumentFuncsSet;
+  if (HasInstrumentFuncsFilter) {
+    std::ifstream FuncsFile(opts::InstrumentFuncsFile, std::ios::in);
+    if (!FuncsFile)
+      return createFatalBOLTError(Twine("instrument-funcs-file \"") +
+                                  Twine(opts::InstrumentFuncsFile) +
+                                  Twine("\" can't be opened."));
+    std::string FuncName;
+    while (std::getline(FuncsFile, FuncName))
+      if (!FuncName.empty())
+        InstrumentFuncsSet.insert(FuncName);
+  }
+
   ParallelUtilities::PredicateTy SkipPredicate = [&](const BinaryFunction &BF) {
-    return (!BF.isSimple() || BF.isIgnored() ||
-            (opts::InstrumentHotOnly && !BF.getKnownExecutionCount()));
+    if (!BF.isSimple() || BF.isIgnored())
+      return true;
+    if (opts::InstrumentHotOnly && !BF.getKnownExecutionCount())
+      return true;
+    if (HasInstrumentFuncsFilter) {
+      bool Found = false;
+      for (const StringRef Name : BF.getNames()) {
+        if (InstrumentFuncsSet.contains(Name)) {
+          Found = true;
+          break;
+        }
+      }
+      if (!Found)
+        return true;
+    }
+    return false;
   };
 
   ParallelUtilities::WorkFuncWithAllocTy WorkFun =
diff --git a/bolt/test/X86/instrument-funcs-file.s b/bolt/test/X86/instrument-funcs-file.s
new file mode 100644
index 0000000000000..18aefd8cce0ed
--- /dev/null
+++ b/bolt/test/X86/instrument-funcs-file.s
@@ -0,0 +1,82 @@
+# Test --instrument-funcs-file, alone and combined with --instrument-hot-only.
+#
+# The binary defines three functions (foo, bar, baz). We attach a profile that
+# only marks foo as hot. With --instrument-funcs-file listing foo and bar, only
+# those two are instrumented. Adding --instrument-hot-only further restricts
+# instrumentation to foo (the only function that is both listed and hot).
+
+# REQUIRES: system-linux,bolt-runtime
+
+# RUN: %clang %cflags %s -o %t.exe -Wl,-q
+# RUN: link_fdata %s %t.exe %t.fdata
+
+# Funcs file lists foo and bar (baz is intentionally omitted).
+# RUN: echo "foo" > %t.funcs
+# RUN: echo "bar" >> %t.funcs
+
+# Test A: only --instrument-funcs-file. Both foo and bar get instrumented.
+# RUN: llvm-bolt --instrument --instrument-funcs-file=%t.funcs \
+# RUN:     -o %t.a.out %t.exe 2>&1 | FileCheck %s --check-prefix=CHECK-A
+
+# Test B: --instrument-funcs-file combined with --instrument-hot-only. Profile
+# marks only foo as hot, so bar is filtered out by --instrument-hot-only.
+# RUN: llvm-bolt --instrument --instrument-funcs-file=%t.funcs \
+# RUN:     --instrument-hot-only --data %t.fdata \
+# RUN:     -o %t.b.out %t.exe 2>&1 | FileCheck %s --check-prefix=CHECK-B
+
+# Test C: empty file means "no functions match", so nothing is instrumented.
+# RUN: rm -f %t.empty && touch %t.empty
+# RUN: llvm-bolt --instrument --instrument-funcs-file=%t.empty \
+# RUN:     -o %t.c.out %t.exe 2>&1 | FileCheck %s --check-prefix=CHECK-C
+
+# Test D: missing file produces a fatal error.
+# RUN: not llvm-bolt --instrument --instrument-funcs-file=%t.missing \
+# RUN:     -o %t.d.out %t.exe 2>&1 | FileCheck %s --check-prefix=CHECK-D
+
+# CHECK-A: BOLT-INSTRUMENTER: Number of function descriptors: 2
+# CHECK-B: BOLT-INSTRUMENTER: Number of function descriptors: 1
+# CHECK-C: BOLT-INSTRUMENTER: Number of function descriptors: 0
+# CHECK-D: instrument-funcs-file {{.*}}.missing{{.*}} can't be opened
+
+    .text
+    .globl _start
+    .type _start, %function
+_start:
+    call foo
+    call bar
+    call baz
+    retq
+    .size _start, .-_start
+
+    .globl foo
+    .type foo, %function
+foo:
+# FDATA: 0 [unknown] 0 1 foo 0 0 100
+    retq
+    .size foo, .-foo
+
+    .globl bar
+    .type bar, %function
+bar:
+    retq
+    .size bar, .-bar
+
+    .globl baz
+    .type baz, %function
+baz:
+    retq
+    .size baz, .-baz
+
+    .globl _init
+    .type _init, %function
+    # Force DT_INIT to be created (needed for instrumentation).
+_init:
+    retq
+    .size _init, .-_init
+
+    .globl _fini
+    .type _fini, %function
+    # Force DT_FINI to be created (needed for instrumentation).
+_fini:
+    retq
+    .size _fini, .-_fini