From a53d97cd4955ee95bba328b12c4f029ce7c45d56 Mon Sep 17 00:00:00 2001 From: Surbhi Jain Date: Tue, 30 Jan 2024 16:49:48 -0800 Subject: [PATCH] Add prefix to lines containing traces from Libtpu in Trace Viewer PiperOrigin-RevId: 602878733 --- xla/backends/profiler/plugin/BUILD | 1 + xla/backends/profiler/plugin/plugin_tracer.cc | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/xla/backends/profiler/plugin/BUILD b/xla/backends/profiler/plugin/BUILD index 53fd864e1ad9c..9de1a9da531bc 100644 --- a/xla/backends/profiler/plugin/BUILD +++ b/xla/backends/profiler/plugin/BUILD @@ -36,6 +36,7 @@ cc_library( "@tsl//tsl/profiler/lib:profiler_interface", "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc", "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + "@tsl//tsl/profiler/utils:xplane_schema", ], alwayslink = True, ) diff --git a/xla/backends/profiler/plugin/plugin_tracer.cc b/xla/backends/profiler/plugin/plugin_tracer.cc index 8176ac55f553e..078187beeddcf 100644 --- a/xla/backends/profiler/plugin/plugin_tracer.cc +++ b/xla/backends/profiler/plugin/plugin_tracer.cc @@ -21,15 +21,18 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "xla/backends/profiler/plugin/profiler_c_api.h" #include "xla/status.h" #include "tsl/platform/logging.h" #include "tsl/profiler/protobuf/xplane.pb.h" +#include "tsl/profiler/utils/xplane_schema.h" namespace xla { namespace profiler { +using tensorflow::profiler::XLine; using tensorflow::profiler::XPlane; using tensorflow::profiler::XSpace; @@ -169,6 +172,11 @@ Status PluginTracer::CollectData(XSpace* space) { xspace.ParseFromArray(args.buffer, args.buffer_size_in_bytes); for (XPlane& tpu_plane : *xspace.mutable_planes()) { XPlane* plane = space->add_planes(); + if (tpu_plane.name() == tsl::profiler::kHostThreadsPlaneName) { + for (XLine& xline : *tpu_plane.mutable_lines()) { + xline.set_display_name(absl::StrCat("libtpu:", xline.name())); + } + } plane->Swap(&tpu_plane); } }