Add plane name to xplane_schema
PiperOrigin-RevId: 285316051
Change-Id: I61f925ba6c168ea78facbd36fa18d17fd24cecd8
diff --git a/tensorflow/core/profiler/internal/cpu/BUILD b/tensorflow/core/profiler/internal/cpu/BUILD
index a3048df..75240b5 100644
--- a/tensorflow/core/profiler/internal/cpu/BUILD
+++ b/tensorflow/core/profiler/internal/cpu/BUILD
@@ -33,6 +33,7 @@
"//tensorflow/core/profiler/internal:profiler_interface",
"//tensorflow/core/profiler/internal:traceme_recorder",
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
+ "//tensorflow/core/profiler/utils:xplane_schema",
"@com_google_absl//absl/strings",
],
alwayslink = True,
@@ -50,6 +51,7 @@
"//tensorflow/core/profiler/internal:profiler_interface",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
+ "//tensorflow/core/profiler/utils:xplane_schema",
"@com_google_absl//absl/types:optional",
"@com_google_googletest//:gtest_main",
],
diff --git a/tensorflow/core/profiler/internal/cpu/host_tracer.cc b/tensorflow/core/profiler/internal/cpu/host_tracer.cc
index 79fe7cf..9e1d2fb 100644
--- a/tensorflow/core/profiler/internal/cpu/host_tracer.cc
+++ b/tensorflow/core/profiler/internal/cpu/host_tracer.cc
@@ -24,6 +24,7 @@
#include "tensorflow/core/profiler/internal/profiler_interface.h"
#include "tensorflow/core/profiler/internal/traceme_recorder.h"
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
+#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/util/env_var.h"
@@ -140,8 +141,9 @@
return errors::Internal("TraceMeRecorder not stopped");
}
MakeCompleteEvents(&events_);
- ConvertCompleteEventsToXPlane(start_timestamp_ns_, events_,
- space->add_planes());
+ XPlane* plane = space->add_planes();
+ plane->set_name(string(kHostThreads));
+ ConvertCompleteEventsToXPlane(start_timestamp_ns_, events_, plane);
events_.clear();
return Status::OK();
}
diff --git a/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc b/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
index ffe702a..f98912a 100644
--- a/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
+++ b/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
@@ -24,6 +24,7 @@
#include "tensorflow/core/profiler/internal/profiler_interface.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
+#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
@@ -126,6 +127,7 @@
ASSERT_EQ(space.planes_size(), 1);
const auto& plane = space.planes(0);
+ ASSERT_EQ(plane.name(), kHostThreads);
ASSERT_EQ(plane.lines_size(), 1);
ASSERT_EQ(plane.event_metadata_size(), 6);
ASSERT_EQ(plane.stat_metadata_size(), 2);
diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index cab8b69..511594f 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -18,6 +18,8 @@
namespace tensorflow {
namespace profiler {
+const absl::string_view kHostThreads = "Host Threads";
+
const int kNumStatTypes = static_cast<int>(StatType::kHloModule) + 1;
static constexpr absl::string_view kStatTypeStrMap[kNumStatTypes] = {
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index 5211406..2797785 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -23,6 +23,9 @@
namespace tensorflow {
namespace profiler {
+// Name of XPlane that contains TraceMe events.
+ABSL_CONST_INIT extern const absl::string_view kHostThreads;
+
enum StatType {
kUnknown = 0,
// TraceMe arguments.
