From 3cc3ccaf045c395433ab09dff25910a2eac6d693 Mon Sep 17 00:00:00 2001 From: Youchun Ni Date: Fri, 5 Jun 2026 18:29:19 -0700 Subject: [PATCH] Support line_num in TraceMeEncode by making TraceMeEncode a macro that calls TraceMeEncodeWithLocation PiperOrigin-RevId: 927571824 --- tsl/profiler/lib/connected_traceme.h | 53 ++++++++++++-------- tsl/profiler/lib/traceme_encode.h | 64 ++++++++++--------------- tsl/profiler/lib/traceme_encode_test.cc | 19 ++------ 3 files changed, 62 insertions(+), 74 deletions(-) diff --git a/tsl/profiler/lib/connected_traceme.h b/tsl/profiler/lib/connected_traceme.h index af95b994d..f8144c465 100644 --- a/tsl/profiler/lib/connected_traceme.h +++ b/tsl/profiler/lib/connected_traceme.h @@ -73,20 +73,21 @@ namespace profiler { * // context_id is passed from the producer thread. * TraceMeConsumer consumer([&] { return "op_execute"; }, context_id); */ -class TraceMeProducer : public TraceMe { +class TraceMeProducerImpl : public TraceMe { public: template - 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) + explicit TraceMeProducerImpl( + const char* source_loc, NameT&& name, + ContextType context_type = ContextType::kGeneric, + std::optional context_id = std::nullopt, + 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 traceme_internal::AppendArgs( + std::string(), {{"_pt", context_type}, {"_p", context_id_}}, + source_loc); }); } @@ -96,26 +97,36 @@ class TraceMeProducer : public TraceMe { uint64_t context_id_; }; -class TraceMeConsumer : public TraceMe { +class TraceMeConsumerImpl : 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) + explicit TraceMeConsumerImpl( + const char* source_loc, NameT&& name, + ContextType context_type = ContextType::kGeneric, + std::optional context_id = std::nullopt, + int level = tsl::profiler::TraceMeLevel::kCritical) : TraceMe(std::forward(name), level) { - AppendMetadata([&] { - return TraceMeEncode({{"_ct", context_type}, {"_c", context_id}}, - source_loc); - }); + if (context_id.has_value()) { + AppendMetadata([&] { + return traceme_internal::AppendArgs( + std::string(), {{"_ct", context_type}, {"_c", context_id.value()}}, + source_loc); + }); + } } - template - TraceMeConsumer(NameT&& name, uint64_t context_id, int level = 2, - const char* source_loc = TRACEME_DEFAULT_FILE) - : TraceMeConsumer(std::forward(name), ContextType::kGeneric, - context_id, level, source_loc) {} + TraceMeConsumerImpl(const char* source_loc, NameT&& name, uint64_t context_id, + int level = 2) + : TraceMeConsumerImpl(source_loc, std::forward(name), + ContextType::kGeneric, context_id, level) {} }; +#define TraceMeProducer(name, ...) \ + ::tsl::profiler::TraceMeProducerImpl(CURRENT_LOCATION, name, ##__VA_ARGS__) + +#define TraceMeConsumer(name, ...) \ + ::tsl::profiler::TraceMeConsumerImpl(CURRENT_LOCATION, name, ##__VA_ARGS__) + } // namespace profiler } // namespace tsl diff --git a/tsl/profiler/lib/traceme_encode.h b/tsl/profiler/lib/traceme_encode.h index e22b5a86b..c9aad6e86 100644 --- a/tsl/profiler/lib/traceme_encode.h +++ b/tsl/profiler/lib/traceme_encode.h @@ -33,18 +33,10 @@ limitations under the License. #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 +#define STRINGIZE_DETAIL(x) #x +#define STRINGIZE(x) STRINGIZE_DETAIL(x) +#define __LINE_STR__ STRINGIZE(__LINE__) +#define CURRENT_LOCATION __FILE__ ":" __LINE_STR__ namespace tsl { namespace profiler { @@ -157,43 +149,37 @@ TF_ATTRIBUTE_ALWAYS_INLINE inline void AppendMetadata( } } -} // namespace traceme_internal - // Encodes an event name and arguments into TraceMe metadata. // Use within a lambda to avoid expensive operations when tracing is disabled. // Example Usage: // TraceMe trace_me([value1]() { // 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); + +TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncodeWithLocation( + const char* source_loc, std::string name, + std::initializer_list args) { + return AppendArgs(std::move(name), args, source_loc); } -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); +TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncodeWithLocation( + const char* source_loc, absl::string_view name, + std::initializer_list args) { + return AppendArgs(std::string(name), args, source_loc); } -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); +TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncodeWithLocation( + const char* source_loc, const char* name, + std::initializer_list args) { + return AppendArgs(name, args, source_loc); } - -// Encodes arguments into TraceMe metadata. -// Use within a lambda to avoid expensive operations when tracing is disabled. -// Example Usage: -// TraceMe trace_me("my_trace"); -// ... -// trace_me.AppendMetadata([value1]() { -// 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); +TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncodeWithLocation( + const char* source_loc, std::initializer_list args) { + return AppendArgs(std::string(), args, source_loc); } +} // namespace traceme_internal + +#define TraceMeEncode(...) \ + ::tsl::profiler::traceme_internal::TraceMeEncodeWithLocation( \ + CURRENT_LOCATION, __VA_ARGS__) // Concatenates op_name and op_type. TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp( diff --git a/tsl/profiler/lib/traceme_encode_test.cc b/tsl/profiler/lib/traceme_encode_test.cc index cb3be0bad..c086861c7 100644 --- a/tsl/profiler/lib/traceme_encode_test.cc +++ b/tsl/profiler/lib/traceme_encode_test.cc @@ -36,7 +36,7 @@ TEST(TraceMeEncodeTest, NoArgTest) { 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_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:")); } TEST(TraceMeEncodeTest, TwoArgsTest) { @@ -44,7 +44,7 @@ TEST(TraceMeEncodeTest, TwoArgsTest) { 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_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:")); } TEST(TraceMeEncodeTest, ThreeArgsTest) { @@ -54,7 +54,7 @@ TEST(TraceMeEncodeTest, ThreeArgsTest) { {"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_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:")); } #if !defined(PLATFORM_WINDOWS) @@ -62,7 +62,7 @@ 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_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:")); } #endif @@ -83,21 +83,12 @@ 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_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:")); } #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#"); -} } // namespace