From e4e09248f2de35a5d95d29273ccd01f51e18843a Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 11 May 2026 19:32:23 -0700 Subject: [PATCH] Support source_info with file name in host-level TraceMes PiperOrigin-RevId: 914005742 --- tsl/profiler/lib/BUILD | 3 - tsl/profiler/lib/connected_traceme.h | 17 ++-- tsl/profiler/lib/traceme_encode.h | 109 ++++++------------------ tsl/profiler/lib/traceme_encode_test.cc | 52 ++++------- 4 files changed, 50 insertions(+), 131 deletions(-) diff --git a/tsl/profiler/lib/BUILD b/tsl/profiler/lib/BUILD index 752c17bfc..1071f08e2 100644 --- a/tsl/profiler/lib/BUILD +++ b/tsl/profiler/lib/BUILD @@ -245,13 +245,10 @@ cc_library( hdrs = ["traceme_encode.h"], visibility = ["//visibility:public"], deps = [ - "//tsl/platform", - "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/strings", "@xla//xla/tsl/platform:logging", "@xla//xla/tsl/platform:macros", - "@xla//xla/tsl/profiler/utils:traceme_global_flags", ], ) diff --git a/tsl/profiler/lib/connected_traceme.h b/tsl/profiler/lib/connected_traceme.h index af95b994d..40195d33a 100644 --- a/tsl/profiler/lib/connected_traceme.h +++ b/tsl/profiler/lib/connected_traceme.h @@ -79,14 +79,12 @@ class TraceMeProducer : public TraceMe { explicit TraceMeProducer(NameT&& name, ContextType context_type = ContextType::kGeneric, std::optional context_id = std::nullopt, - int level = tsl::profiler::TraceMeLevel::kCritical, - const char* source_loc = TRACEME_DEFAULT_FILE) + int level = tsl::profiler::TraceMeLevel::kCritical) : TraceMe(std::forward(name), level), context_id_(context_id.has_value() ? context_id.value() : TraceMe::NewActivityId()) { AppendMetadata([&] { - return TraceMeEncode({{"_pt", context_type}, {"_p", context_id_}}, - source_loc); + return TraceMeEncode({{"_pt", context_type}, {"_p", context_id_}}); }); } @@ -100,20 +98,17 @@ class TraceMeConsumer : public TraceMe { public: template TraceMeConsumer(NameT&& name, ContextType context_type, uint64_t context_id, - int level = tsl::profiler::TraceMeLevel::kCritical, - const char* source_loc = TRACEME_DEFAULT_FILE) + int level = tsl::profiler::TraceMeLevel::kCritical) : TraceMe(std::forward(name), level) { AppendMetadata([&] { - return TraceMeEncode({{"_ct", context_type}, {"_c", context_id}}, - source_loc); + return TraceMeEncode({{"_ct", context_type}, {"_c", context_id}}); }); } template - TraceMeConsumer(NameT&& name, uint64_t context_id, int level = 2, - const char* source_loc = TRACEME_DEFAULT_FILE) + TraceMeConsumer(NameT&& name, uint64_t context_id, int level = 2) : TraceMeConsumer(std::forward(name), ContextType::kGeneric, - context_id, level, source_loc) {} + context_id, level) {} }; } // namespace profiler diff --git a/tsl/profiler/lib/traceme_encode.h b/tsl/profiler/lib/traceme_encode.h index e22b5a86b..69f12dd08 100644 --- a/tsl/profiler/lib/traceme_encode.h +++ b/tsl/profiler/lib/traceme_encode.h @@ -19,32 +19,13 @@ limitations under the License. #include #include -#include #include "absl/base/attributes.h" -#include "absl/base/config.h" // IWYU pragma: keep #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/macros.h" -#include "tsl/platform/platform.h" // IWYU pragma: keep -#if !defined(IS_MOBILE_PLATFORM) -#include "xla/tsl/profiler/utils/traceme_global_flags.h" -#endif - -#define TRACEME_ENCODE_STRINGIFY(x) #x -#define TRACEME_ENCODE_TOSTRING(x) TRACEME_ENCODE_STRINGIFY(x) -#define TRACEME_FILE_AND_LINE __FILE__ ":" TRACEME_ENCODE_TOSTRING(__LINE__) - -#if !defined(LIBTPU_ON_GCE) && ABSL_HAVE_BUILTIN(__builtin_FILE) -// TODO(b/507077868): Switch to absl::SourceLocation after XLA upgrades to the -// next absl version. For more details, see -// https://gist.github.com/youchunni/24ee88f9daa9566312f055d71513dbea -#define TRACEME_DEFAULT_FILE __builtin_FILE() -#else -#define TRACEME_DEFAULT_FILE "" -#endif namespace tsl { namespace profiler { @@ -86,60 +67,28 @@ TF_ATTRIBUTE_ALWAYS_INLINE inline char* Append(char* out, return out; } -// Appends arguments encoded as TraceMe metadata to `name`. -// -// The resulting string format is: -// name#key1=value1,key2=value2,_src=file.cc# -// -// Performance implementation details: -// To minimize overhead on the critical path, this function pre-calculates the -// exact string length for a single resize, then appends the arguments using -// raw buffer manipulation pointed by `out`. +// Appends args encoded as TraceMe metadata to name. TF_ATTRIBUTE_ALWAYS_INLINE inline std::string AppendArgs( - std::string name, std::initializer_list args, - const char* source_loc_arg) { - absl::string_view source_loc(source_loc_arg); -#if !defined(IS_MOBILE_PLATFORM) - if (TF_PREDICT_FALSE(!TraceMeGlobalFlags::IsSourceLocationEnabled())) { - source_loc = ""; - } -#else - source_loc = ""; -#endif - if (args.size() == 0 && source_loc.empty()) { - return name; - } - const auto old_size = name.size(); - // `args.size() * 2`: Accounts for '=' and ',' added for each arg. - // `+ 2`: Accounts for the initial '#' and the final '#'. - auto new_size = old_size + args.size() * 2 + 2; - for (const auto& arg : args) { - new_size += arg.key.size() + arg.value.size(); - } - - if (!source_loc.empty()) { - new_size += source_loc.size() + 5; // `+ 5`: Accounts for '_src=' - } else { - new_size -= 1; // Minus one ',' between the last arg and '_src='. - } - name.resize(new_size); - char* const begin = &name[0]; - char* out = begin + old_size; - *out++ = '#'; - for (const auto& arg : args) { - out = Append(out, arg.key); - *out++ = '='; - out = Append(out, arg.value); - *out++ = ','; - } - if (!source_loc.empty()) { - out = Append(out, "_src="); - out = Append(out, source_loc); + std::string name, std::initializer_list args) { + if (TF_PREDICT_TRUE(args.size() > 0)) { + const auto old_size = name.size(); + auto new_size = old_size + args.size() * 2 + 1; + for (const auto& arg : args) { + new_size += arg.key.size() + arg.value.size(); + } + name.resize(new_size); + char* const begin = &name[0]; + char* out = begin + old_size; *out++ = '#'; - } else { - *(out - 1) = '#'; // Replace the last ',' with '#'. + for (const auto& arg : args) { + out = Append(out, arg.key); + *out++ = '='; + out = Append(out, arg.value); + *out++ = ','; + } + *(out - 1) = '#'; + DCHECK_EQ(out, begin + new_size); } - DCHECK_EQ(out, begin + new_size); return name; } @@ -166,19 +115,16 @@ TF_ATTRIBUTE_ALWAYS_INLINE inline void AppendMetadata( // return TraceMeEncode("my_trace", {{"key1", value1}, {"key2", 42}}); // }); TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode( - std::string name, std::initializer_list args, - const char* source_loc = TRACEME_DEFAULT_FILE) { - return traceme_internal::AppendArgs(std::move(name), args, source_loc); + std::string name, std::initializer_list args) { + return traceme_internal::AppendArgs(std::move(name), args); } TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode( - absl::string_view name, std::initializer_list args, - const char* source_loc = TRACEME_DEFAULT_FILE) { - return traceme_internal::AppendArgs(std::string(name), args, source_loc); + absl::string_view name, std::initializer_list args) { + return traceme_internal::AppendArgs(std::string(name), args); } TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode( - const char* name, std::initializer_list args, - const char* source_loc = TRACEME_DEFAULT_FILE) { - return traceme_internal::AppendArgs(name, args, source_loc); + const char* name, std::initializer_list args) { + return traceme_internal::AppendArgs(std::string(name), args); } // Encodes arguments into TraceMe metadata. @@ -190,9 +136,8 @@ TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode( // return TraceMeEncode({{"key1", value1}, {"key2", 42}}); // }); TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode( - std::initializer_list args, - const char* source_loc = TRACEME_DEFAULT_FILE) { - return traceme_internal::AppendArgs(std::string(), args, source_loc); + std::initializer_list args) { + return traceme_internal::AppendArgs(std::string(), args); } // Concatenates op_name and op_type. diff --git a/tsl/profiler/lib/traceme_encode_test.cc b/tsl/profiler/lib/traceme_encode_test.cc index cb3be0bad..f8dc39196 100644 --- a/tsl/profiler/lib/traceme_encode_test.cc +++ b/tsl/profiler/lib/traceme_encode_test.cc @@ -16,7 +16,6 @@ limitations under the License. #include -#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "xla/tsl/platform/test.h" @@ -28,41 +27,31 @@ namespace profiler { namespace { TEST(TraceMeEncodeTest, NoArgTest) { - std::string encoded = TraceMeEncode("Hello!", {}); - EXPECT_TRUE(absl::StrContains(encoded, "Hello!#_src=")); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc")); + EXPECT_EQ(TraceMeEncode("Hello!", {}), "Hello!"); } TEST(TraceMeEncodeTest, OneArgTest) { - std::string encoded = TraceMeEncode("Hello", {{"context", "World"}}); - EXPECT_TRUE(absl::StrContains(encoded, "Hello#context=World,_src=")); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc")); + EXPECT_EQ(TraceMeEncode("Hello", {{"context", "World"}}), + "Hello#context=World#"); } TEST(TraceMeEncodeTest, TwoArgsTest) { - std::string encoded = - TraceMeEncode("Hello", {{"context", "World"}, {"request_id", 42}}); - EXPECT_TRUE( - absl::StrContains(encoded, "Hello#context=World,request_id=42,_src=")); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc")); + EXPECT_EQ(TraceMeEncode("Hello", {{"context", "World"}, {"request_id", 42}}), + "Hello#context=World,request_id=42#"); } TEST(TraceMeEncodeTest, ThreeArgsTest) { - std::string encoded = - TraceMeEncode("Hello", {{"context", "World"}, - {"request_id", 42}, - {"addr", absl::Hex(0xdeadbeef)}}); - EXPECT_TRUE(absl::StrContains( - encoded, "Hello#context=World,request_id=42,addr=deadbeef,_src=")); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc")); + EXPECT_EQ(TraceMeEncode("Hello", {{"context", "World"}, + {"request_id", 42}, + {"addr", absl::Hex(0xdeadbeef)}}), + "Hello#context=World,request_id=42,addr=deadbeef#"); } #if !defined(PLATFORM_WINDOWS) TEST(TraceMeEncodeTest, TemporaryStringTest) { - std::string encoded = - TraceMeEncode("Hello", {{"context", absl::StrCat("World:", 2020)}}); - EXPECT_TRUE(absl::StrContains(encoded, "Hello#context=World:2020,_src=")); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc")); + EXPECT_EQ(TraceMeEncode("Hello", {{std::string("context"), + absl::StrCat("World:", 2020)}}), + "Hello#context=World:2020#"); } #endif @@ -81,22 +70,15 @@ struct Point { }; TEST(TraceMeEncodeTest, AbslStringifyTest) { - std::string encoded = TraceMeEncode("Plot", {{"point", Point{10, 20}}}); - EXPECT_TRUE(absl::StrContains(encoded, "Plot#point=(10, 20),_src=")); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc")); + EXPECT_EQ(TraceMeEncode("Plot", {{"point", Point{10, 20}}}), + "Plot#point=(10, 20)#"); } #endif -TEST(TraceMeEncodeTest, AppendLineNumberTest) { - std::string encoded = - TraceMeEncode("Hello", {{"context", "World"}}, TRACEME_FILE_AND_LINE); - EXPECT_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:")); -} - -TEST(TraceMeEncodeTest, EmptySourceLocTest) { - std::string encoded = TraceMeEncode("Hello", {{"context", "World"}}, ""); - EXPECT_EQ(encoded, "Hello#context=World#"); +TEST(TraceMeEncodeTest, NoNameTest) { + EXPECT_EQ(TraceMeEncode({{"context", "World"}, {"request_id", 42}}), + "#context=World,request_id=42#"); } } // namespace