Skip to content

Support source_info with file name in host-level TraceMes #3553

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
+50 −131
Open

Support source_info with file name in host-level TraceMes #3553

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
3 changes: 0 additions & 3 deletions tsl/profiler/lib/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
)

Expand Down
17 changes: 6 additions & 11 deletions tsl/profiler/lib/connected_traceme.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,14 +79,12 @@ class TraceMeProducer : public TraceMe {
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)
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 TraceMeEncode({{"_pt", context_type}, {"_p", context_id_}});
});
}

Expand All @@ -100,20 +98,17 @@ class TraceMeConsumer : 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)
int level = tsl::profiler::TraceMeLevel::kCritical)
: TraceMe(std::forward<NameT>(name), level) {
AppendMetadata([&] {
return TraceMeEncode({{"_ct", context_type}, {"_c", context_id}},
source_loc);
return TraceMeEncode({{"_ct", context_type}, {"_c", context_id}});
});
}

template <typename NameT>
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<NameT>(name), ContextType::kGeneric,
context_id, level, source_loc) {}
context_id, level) {}
};

} // namespace profiler
Expand Down
109 changes: 27 additions & 82 deletions tsl/profiler/lib/traceme_encode.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,32 +19,13 @@ limitations under the License.

#include <initializer_list>
#include <string>
#include <utility>

#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 {
Expand Down Expand Up @@ -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<TraceMeArg> 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<TraceMeArg> 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;
}

Expand All @@ -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<TraceMeArg> args,
const char* source_loc = TRACEME_DEFAULT_FILE) {
return traceme_internal::AppendArgs(std::move(name), args, source_loc);
std::string name, std::initializer_list<TraceMeArg> args) {
return traceme_internal::AppendArgs(std::move(name), args);
}
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);
absl::string_view name, std::initializer_list<TraceMeArg> args) {
return traceme_internal::AppendArgs(std::string(name), args);
}
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);
const char* name, std::initializer_list<TraceMeArg> args) {
return traceme_internal::AppendArgs(std::string(name), args);
}

// Encodes arguments into TraceMe metadata.
Expand All @@ -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<TraceMeArg> args,
const char* source_loc = TRACEME_DEFAULT_FILE) {
return traceme_internal::AppendArgs(std::string(), args, source_loc);
std::initializer_list<TraceMeArg> args) {
return traceme_internal::AppendArgs(std::string(), args);
}

// Concatenates op_name and op_type.
Expand Down
52 changes: 17 additions & 35 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 @@ -16,7 +16,6 @@ limitations under the License.

#include <string>

#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "xla/tsl/platform/test.h"
Expand All @@ -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

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