From d59a218656565ebaa1cdb0f38adf2566530723ef Mon Sep 17 00:00:00 2001
From: Ethan Luis McDonough <ethanluismcdonough@gmail.com>
Date: Mon, 24 Feb 2025 15:22:08 -0600
Subject: [PATCH] [Instrumentor] Add target configuration options

---
 .../Transforms/Instrumentation/Instrumentor.h | 12 +++++++++
 .../Instrumentation/Instrumentor.cpp          | 26 +++++++++++++++++++
 .../Instrumentor/default_config.json          |  8 +++++-
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Transforms/Instrumentation/Instrumentor.h b/llvm/include/llvm/Transforms/Instrumentation/Instrumentor.h
index 59fe9f8bf4b9a..aa4cb36fb4a80 100644
--- a/llvm/include/llvm/Transforms/Instrumentation/Instrumentor.h
+++ b/llvm/include/llvm/Transforms/Instrumentation/Instrumentor.h
@@ -405,6 +405,15 @@ struct InstrumentationConfig {
     DemangleFunctionNames = BaseConfigurationOpportunity::getBoolOption(
         *this, "demangle_function_names",
         "Demangle functions names passed to the runtime.", true);
+    TargetRegex = BaseConfigurationOpportunity::getStringOption(
+        *this, "target_regex",
+        "Regular expression to be matched against the module target. "
+        "Only targets that match this regex will be instrumented",
+        "");
+    HostEnabled = BaseConfigurationOpportunity::getBoolOption(
+        *this, "host_enabled", "Instrument non-GPU targets", true);
+    GPUEnabled = BaseConfigurationOpportunity::getBoolOption(
+        *this, "gpu_enabled", "Instrument GPU targets", true);
   }
 
   bool ReadConfig = true;
@@ -425,6 +434,9 @@ struct InstrumentationConfig {
   BaseConfigurationOpportunity *RuntimePrefix;
   BaseConfigurationOpportunity *RuntimeStubsFile;
   BaseConfigurationOpportunity *DemangleFunctionNames;
+  BaseConfigurationOpportunity *TargetRegex;
+  BaseConfigurationOpportunity *HostEnabled;
+  BaseConfigurationOpportunity *GPUEnabled;
 
   EnumeratedArray<StringMap<InstrumentationOpportunity *>,
                   InstrumentationLocation::KindTy>
diff --git a/llvm/lib/Transforms/Instrumentation/Instrumentor.cpp b/llvm/lib/Transforms/Instrumentation/Instrumentor.cpp
index 04857919660df..469873a8ee8a4 100644
--- a/llvm/lib/Transforms/Instrumentation/Instrumentor.cpp
+++ b/llvm/lib/Transforms/Instrumentation/Instrumentor.cpp
@@ -52,6 +52,7 @@
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Regex.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
@@ -367,6 +368,7 @@ class InstrumentorImpl final {
   bool instrument();
 
 private:
+  bool shouldInstrumentTarget();
   bool shouldInstrumentFunction(Function &Fn);
   bool shouldInstrumentGlobalVariable(GlobalVariable &GV);
 
@@ -434,6 +436,28 @@ class InstrumentorImpl final {
 
 } // end anonymous namespace
 
+bool InstrumentorImpl::shouldInstrumentTarget() {
+  const auto &TripleStr = M.getTargetTriple();
+  Triple T(TripleStr);
+  const bool IsGPU = T.isAMDGPU() || T.isNVPTX();
+
+  bool RegexMatches = true;
+  const auto TargetRegexStr = IConf.TargetRegex->getString();
+  if (!TargetRegexStr.empty()) {
+    llvm::Regex TargetRegex(TargetRegexStr);
+    std::string ErrMsg;
+    if (!TargetRegex.isValid(ErrMsg)) {
+      errs() << "WARNING: failed to parse target regex: " << ErrMsg << "\n";
+      return false;
+    }
+    RegexMatches = TargetRegex.match(TripleStr);
+  }
+
+  return ((IsGPU && IConf.GPUEnabled->getBool()) ||
+          (!IsGPU && IConf.HostEnabled->getBool())) &&
+         RegexMatches;
+}
+
 bool InstrumentorImpl::shouldInstrumentFunction(Function &Fn) {
   if (Fn.isDeclaration())
     return false;
@@ -567,6 +591,8 @@ bool InstrumentorImpl::instrumentModule() {
 
 bool InstrumentorImpl::instrument() {
   bool Changed = false;
+  if (!shouldInstrumentTarget())
+    return Changed;
 
   for (auto &ChoiceIt :
        IConf.IChoices[InstrumentationLocation::INSTRUCTION_PRE])
diff --git a/llvm/test/Instrumentation/Instrumentor/default_config.json b/llvm/test/Instrumentation/Instrumentor/default_config.json
index 6f680fa14aa8e..841c60fbc79fb 100644
--- a/llvm/test/Instrumentation/Instrumentor/default_config.json
+++ b/llvm/test/Instrumentation/Instrumentor/default_config.json
@@ -5,7 +5,13 @@
     "runtime_stubs_file": "test.c",
     "runtime_stubs_file.description": "The file into which runtime stubs should be written.",
     "demangle_function_names": true,
-    "demangle_function_names.description": "Demangle functions names passed to the runtime."
+    "demangle_function_names.description": "Demangle functions names passed to the runtime.",
+    "target_regex": "",
+    "target_regex.description": "Regular expression to be matched against the module target. Only targets that match this regex will be instrumented",
+    "host_enabled": true,
+    "host_enabled.description": "Instrument non-GPU targets",
+    "gpu_enabled": true,
+    "gpu_enabled.description": "Instrument GPU targets"
   },
   "module_pre": {
     "module": {