[xprof] Add support for tracemark_lower and tracemark_upper profiler options

PiperOrigin-RevId: 816343394

Author: Matt-Hurd

third_party/tsl/tsl/profiler/protobuf/profiler_options.proto

@@ -75,12 +75,14 @@ message ProfileOptions {
   TraceOptions trace_options = 11;
 
   // AdvancedConfigValue represents the configuration value, it can be one of
-  // the following types: string, bool, int64, depending upon the config type.
+  // the following types: string, bool, int64, int32, depending upon the config
+  // type.
   message AdvancedConfigValue {
     oneof value {
       string string_value = 1;
       bool bool_value = 2;
       int64 int64_value = 3;
+      int32 int32_value = 4;
     }
   }
 

xla/tsl/profiler/utils/profiler_options_util.cc

@@ -23,7 +23,7 @@ limitations under the License.
 namespace tsl {
 namespace profiler {
 
-std::optional<std::variant<std::string, bool, int64_t>> GetConfigValue(
+std::optional<std::variant<std::string, bool, int64_t, int32_t>> GetConfigValue(
     const tensorflow::ProfileOptions& options, const std::string& key) {
   auto config = options.advanced_configuration().find(key);
 
@@ -36,6 +36,8 @@ std::optional<std::variant<std::string, bool, int64_t>> GetConfigValue(
       return config_value.bool_value();
     } else if (config_value.has_int64_value()) {
       return config_value.int64_value();
+    } else if (config_value.has_int32_value()) {
+      return config_value.int32_value();
     }
   }
 

xla/tsl/profiler/utils/profiler_options_util.h

@@ -32,7 +32,7 @@ namespace tsl {
 namespace profiler {
 // Get config value from the profiler options, if the key is not found, return
 // std::nullopt.
-std::optional<std::variant<std::string, bool, int64_t>> GetConfigValue(
+std::optional<std::variant<std::string, bool, int64_t, int32_t>> GetConfigValue(
     const tensorflow::ProfileOptions& options, const std::string& key);
 
 template <typename T>

xla/tsl/profiler/utils/profiler_options_util_test.cc

@@ -40,7 +40,7 @@ TEST(ProfilerOptionsUtilTest, GetConfigValueString) {
   config_value.set_string_value("test_value");
   advanced_config["test_key"] = config_value;
 
-  std::optional<std::variant<std::string, bool, int64_t>> result =
+  std::optional<std::variant<std::string, bool, int64_t, int32_t>> result =
       GetConfigValue(options, "test_key");
   EXPECT_THAT(result, Optional(VariantWith<std::string>("test_value")));
 }
@@ -52,7 +52,7 @@ TEST(ProfilerOptionsUtilTest, GetConfigValueBool) {
   config_value.set_bool_value(true);
   advanced_config["test_key"] = config_value;
 
-  std::optional<std::variant<std::string, bool, int64_t>> result =
+  std::optional<std::variant<std::string, bool, int64_t, int32_t>> result =
       GetConfigValue(options, "test_key");
   EXPECT_THAT(result, Optional(VariantWith<bool>(true)));
 }
@@ -64,15 +64,15 @@ TEST(ProfilerOptionsUtilTest, GetConfigValueInt64) {
   config_value.set_int64_value(12345);
   advanced_config["test_key"] = config_value;
 
-  std::optional<std::variant<std::string, bool, int64_t>> result =
+  std::optional<std::variant<std::string, bool, int64_t, int32_t>> result =
       GetConfigValue(options, "test_key");
   EXPECT_THAT(result, Optional(VariantWith<int64_t>(12345)));
 }
 
 TEST(ProfilerOptionsUtilTest, GetConfigValueNotFound) {
   tensorflow::ProfileOptions options;
 
-  std::optional<std::variant<std::string, bool, int64_t>> result =
+  std::optional<std::variant<std::string, bool, int64_t, int32_t>> result =
       GetConfigValue(options, "test_key");
   EXPECT_FALSE(result.has_value());
   EXPECT_THAT(result, Eq(std::nullopt));

xla/tsl/profiler/utils/session_manager.cc

@@ -143,6 +143,30 @@ RemoteProfilerSessionManagerOptions GetRemoteSessionManagerOptionsLocked(
             options.mutable_profiler_options()->set_session_id(value);
           },
           nullptr);
+    } else if (key == "tracemark_lower") {
+      SetOption<int>(
+          key, kw.second,
+          [](tensorflow::ProfileOptions* options, int value) {
+            options->mutable_advanced_configuration()->insert(
+                {"tracemark_lower",
+                 tensorflow::ProfileOptions::AdvancedConfigValue()});
+            options->mutable_advanced_configuration()
+                ->at("tracemark_lower")
+                .set_int32_value(value);
+          },
+          options.mutable_profiler_options());
+    } else if (key == "tracemark_upper") {
+      SetOption<int>(
+          key, kw.second,
+          [](tensorflow::ProfileOptions* options, int value) {
+            options->mutable_advanced_configuration()->insert(
+                {"tracemark_upper",
+                 tensorflow::ProfileOptions::AdvancedConfigValue()});
+            options->mutable_advanced_configuration()
+                ->at("tracemark_upper")
+                .set_int32_value(value);
+          },
+          options.mutable_profiler_options());
     } else {
       LOG(WARNING) << "Unrecognised key: " << key;
     }