Skip to content

Support line_num in TraceMeEncode by making TraceMeEncode a macro that calls TraceMeEncodeWithLocation #3579

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Support line_num in TraceMeEncode by making TraceMeEncode a macro that calls TraceMeEncodeWithLocation #3579

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 32 additions & 21 deletions tsl/profiler/lib/connected_traceme.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename NameT>
explicit TraceMeProducer(NameT&& name,
ContextType context_type = ContextType::kGeneric,
std::optional<uint64_t> 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<uint64_t> context_id = std::nullopt,
int level = tsl::profiler::TraceMeLevel::kCritical)
: TraceMe(std::forward<NameT>(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);
});
}

Expand All @@ -96,26 +97,36 @@ class TraceMeProducer : public TraceMe {
uint64_t context_id_;
};

class TraceMeConsumer : public TraceMe {
class TraceMeConsumerImpl : public TraceMe {
public:
template <typename NameT>
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<uint64_t> context_id = std::nullopt,
int level = tsl::profiler::TraceMeLevel::kCritical)
: TraceMe(std::forward<NameT>(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 <typename NameT>
TraceMeConsumer(NameT&& name, uint64_t context_id, int level = 2,
const char* source_loc = TRACEME_DEFAULT_FILE)
: TraceMeConsumer(std::forward<NameT>(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<NameT>(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

Expand Down
64 changes: 25 additions & 39 deletions tsl/profiler/lib/traceme_encode.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<TraceMeArg> 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<TraceMeArg> args) {
return AppendArgs(std::move(name), args, source_loc);
}
TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode(
absl::string_view name, std::initializer_list<TraceMeArg> 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<TraceMeArg> args) {
return AppendArgs(std::string(name), args, source_loc);
}
TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeEncode(
const char* name, std::initializer_list<TraceMeArg> 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<TraceMeArg> 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<TraceMeArg> 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<TraceMeArg> 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(
Expand Down
19 changes: 5 additions & 14 deletions tsl/profiler/lib/traceme_encode_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,15 +36,15 @@ 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) {
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_TRUE(absl::StrContains(encoded, "traceme_encode_test.cc:"));
}

TEST(TraceMeEncodeTest, ThreeArgsTest) {
Expand All @@ -54,15 +54,15 @@ 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)
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

Expand All @@ -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

Expand Down
Loading