Add prefix to lines containing traces from Libtpu in Trace Viewer

PiperOrigin-RevId: 602878733
openxla · Jan 31, 2024 · a53d97c · a53d97c
1 parent 5c6462d
commit a53d97c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git 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
@@ -21,15 +21,18 @@ limitations under the License.
 #include <vector>
 
 #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);
     }
   }