From 4a081383975e520b1d9e59e2417525cf26ed21d1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 3 Feb 2026 19:21:23 -0800
Subject: [PATCH] Always run the name generators passed to TraceMe on debug
 builds to help catch bugs

Today, the name generators are run only when there is an active tracing session, which means that none of the code inside the name generators are tested at all. Since users sometimes perform non-trivial work (such as accessing pointers that may be dangling) inside the name generators, it is useful to test such code by always invoking the name generators in debug builds where the extra overhead is n...

PiperOrigin-RevId: 865154634
---
 tsl/profiler/lib/traceme.h | 34 ++++++----------------------------
 1 file changed, 6 insertions(+), 28 deletions(-)

diff --git a/tsl/profiler/lib/traceme.h b/tsl/profiler/lib/traceme.h
index b58c74b31..7c9a1f4ac 100644
--- a/tsl/profiler/lib/traceme.h
+++ b/tsl/profiler/lib/traceme.h
@@ -155,12 +155,6 @@ class TraceMe {
                          TraceMeRecorder::CheckFilter(filter_mask))) {
       name_.Emplace(std::forward<NameGeneratorT>(name_generator)());
       start_time_ = GetCurrentTimeNanos();
-    } else {
-#ifndef NDEBUG
-      // Always invoke the name generator on debug builds to catch bugs in name
-      // generators that would otherwise only be caught when tracing is enabled.
-      std::forward<NameGeneratorT>(name_generator)();
-#endif
     }
 #endif
   }
@@ -218,16 +212,12 @@ class TraceMe {
       std::enable_if_t<std::is_invocable_v<MetadataGeneratorT>, bool> = true>
   void AppendMetadata(MetadataGeneratorT&& metadata_generator) {
 #if !defined(IS_MOBILE_PLATFORM)
-    if (TF_PREDICT_FALSE(start_time_ != kUntracedActivity) &&
-        TF_PREDICT_TRUE(TraceMeRecorder::Active())) {
-      traceme_internal::AppendMetadata(
-          &name_.value, std::forward<MetadataGeneratorT>(metadata_generator)());
-    } else {
-#ifndef NDEBUG
-      // Always invoke the name generator on debug builds to catch bugs in name
-      // generators that would otherwise only be caught when tracing is enabled.
-      std::forward<MetadataGeneratorT>(metadata_generator)();
-#endif
+    if (TF_PREDICT_FALSE(start_time_ != kUntracedActivity)) {
+      if (TF_PREDICT_TRUE(TraceMeRecorder::Active())) {
+        traceme_internal::AppendMetadata(
+            &name_.value,
+            std::forward<MetadataGeneratorT>(metadata_generator)());
+      }
     }
 #endif
   }
@@ -249,12 +239,6 @@ class TraceMe {
       TraceMeRecorder::Record({std::forward<NameGeneratorT>(name_generator)(),
                                GetCurrentTimeNanos(), -activity_id});
       return activity_id;
-    } else {
-#ifndef NDEBUG
-      // Always invoke the name generator on debug builds to catch bugs in name
-      // generators that would otherwise only be caught when tracing is enabled.
-      std::forward<NameGeneratorT>(name_generator)();
-#endif
     }
 #endif
     return kUntracedActivity;
@@ -316,12 +300,6 @@ class TraceMe {
       int64_t now = GetCurrentTimeNanos();
       TraceMeRecorder::Record({std::forward<NameGeneratorT>(name_generator)(),
                                /*start_time=*/now, /*end_time=*/now});
-    } else {
-#ifndef NDEBUG
-      // Always invoke the name generator on debug builds to catch bugs in name
-      // generators that would otherwise only be caught when tracing is enabled.
-      std::forward<NameGeneratorT>(name_generator)();
-#endif
     }
 #endif
   }