Rename xplane_visitor module to profile_data.

PiperOrigin-RevId: 687672929

xla/python/BUILD

@@ -958,7 +958,7 @@ cc_library(
         "//xla/pjrt:status_casters",
         "//xla/pjrt/c:pjrt_c_api_hdrs",
         "//xla/pjrt/c:pjrt_c_api_profiler_extension_hdrs",
-        "//xla/python/profiler:xplane_visitor_lib",
+        "//xla/python/profiler:profile_data_lib",
         "//xla/tsl/profiler/rpc:profiler_server_impl",
         "//xla/tsl/profiler/rpc/client:capture_profile",
         "//xla/tsl/profiler/rpc/client:profiler_client_impl",

xla/python/profiler.cc

@@ -37,7 +37,7 @@ limitations under the License.
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/python/aggregate_profile.h"
-#include "xla/python/profiler/xplane_visitor.h"
+#include "xla/python/profiler/profile_data.h"
 #include "xla/python/profiler_utils.h"
 #include "xla/python/xplane_to_profile_instructions.h"
 #include "xla/tsl/profiler/rpc/client/capture_profile.h"
@@ -201,11 +201,11 @@ void BuildProfilerSubmodule(nb::module_& m) {
            })
       .def("stop_and_get_profile_data",
            [](ProfilerSessionWrapper* sess)
-               -> tensorflow::profiler::python::XSpaceVisitor {
+               -> tensorflow::profiler::python::ProfileData {
              std::shared_ptr<tensorflow::profiler::XSpace> xspace;
              // Disables the ProfilerSession
              xla::ThrowIfError(sess->session->CollectData(xspace.get()));
-             return tensorflow::profiler::python::XSpaceVisitor(xspace);
+             return tensorflow::profiler::python::ProfileData(xspace);
            })
       .def("export",
            [](ProfilerSessionWrapper* sess, nb::bytes xspace,

xla/python/profiler/BUILD

@@ -4,9 +4,9 @@ load("//xla/tsl:tsl.default.bzl", "tsl_pybind_extension")
 # copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
 
 cc_library(
-    name = "xplane_visitor_lib",
-    srcs = ["xplane_visitor.cc"],
-    hdrs = ["xplane_visitor.h"],
+    name = "profile_data_lib",
+    srcs = ["profile_data.cc"],
+    hdrs = ["profile_data.h"],
     copts = [
         "-fexceptions",
         "-fno-strict-aliasing",
@@ -28,16 +28,16 @@ cc_library(
 )
 
 tsl_pybind_extension(
-    name = "xplane_visitor",
-    srcs = ["py_xplane_visitor.cc"],
+    name = "profile_data",
+    srcs = ["py_profile_data.cc"],
     copts = [
         "-fexceptions",
         "-fno-strict-aliasing",
     ],
-    pytype_srcs = ["xplane_visitor.pyi"],
+    pytype_srcs = ["profile_data.pyi"],
     visibility = ["//visibility:public"],
     deps = [
-        ":xplane_visitor_lib",
+        ":profile_data_lib",
         "@nanobind",
     ],
 )

xla/python/profiler/xplane_visitor.cc → xla/python/profiler/profile_data.cc

@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/python/profiler/xplane_visitor.h"
+#include "xla/python/profiler/profile_data.h"
 
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/string.h>  // IWYU pragma: keep. For automatic conversion of std::string to Python string.
@@ -74,9 +74,9 @@ nb::tuple stats_to_tuple(const XStat& stat, const XPlane* plane) {
   return nb::make_tuple(nb::none(), nb::none());
 }
 
-XEventVisitor::XEventVisitor(const XEvent* event, int64_t line_timestamp_ns,
-                             const XPlane* plane,
-                             std::shared_ptr<const XSpace> xspace)
+ProfileEvent::ProfileEvent(const XEvent* event, int64_t line_timestamp_ns,
+                           const XPlane* plane,
+                           std::shared_ptr<const XSpace> xspace)
     : event_(event),
       plane_(plane),
       line_timestamp_ns_(line_timestamp_ns),
@@ -86,120 +86,118 @@ XEventVisitor::XEventVisitor(const XEvent* event, int64_t line_timestamp_ns,
   CHECK_NOTNULL(xspace_);
 }
 
-double XEventVisitor::start_ns() const {
+double ProfileEvent::start_ns() const {
   return event_->offset_ps() / 1000 + line_timestamp_ns_;
 }
 
-double XEventVisitor::duration_ns() const {
+double ProfileEvent::duration_ns() const {
   return event_->duration_ps() / 1000;
 }
 
-double XEventVisitor::end_ns() const { return start_ns() + duration_ns(); }
+double ProfileEvent::end_ns() const { return start_ns() + duration_ns(); }
 
-std::string XEventVisitor::name() const {
+std::string ProfileEvent::name() const {
   if (plane_->event_metadata().contains(event_->metadata_id())) {
     return plane_->event_metadata().at(event_->metadata_id()).name();
   }
   return "";
 }
 
-VisitorIterator<nb::tuple, XStat> XEventVisitor::stats_begin() {
+VisitorIterator<nb::tuple, XStat> ProfileEvent::stats_begin() {
   return VisitorIterator<nb::tuple, XStat>(
       &event_->stats(),
       [this](const XStat& stat) { return stats_to_tuple(stat, plane_); });
 }
-VisitorIterator<nb::tuple, XStat> XEventVisitor::stats_end() {
+VisitorIterator<nb::tuple, XStat> ProfileEvent::stats_end() {
   return VisitorIterator<nb::tuple, XStat>(
       &event_->stats(),
       [this](const XStat& stat) { return stats_to_tuple(stat, plane_); },
       event_->stats().size());
 }
 
-XLineVisitor::XLineVisitor(const XLine* line, const XPlane* plane,
-                           std::shared_ptr<const XSpace> xspace)
+ProfileLine::ProfileLine(const XLine* line, const XPlane* plane,
+                         std::shared_ptr<const XSpace> xspace)
     : line_(line), plane_(plane), xspace_(xspace) {
   CHECK_NOTNULL(line_);
   CHECK_NOTNULL(plane_);
   CHECK_NOTNULL(xspace_);
 }
 
-const std::string& XLineVisitor::name() const { return line_->name(); }
+const std::string& ProfileLine::name() const { return line_->name(); }
 
-VisitorIterator<XEventVisitor, XEvent> XLineVisitor::events_begin() {
-  return VisitorIterator<XEventVisitor, XEvent>(
+VisitorIterator<ProfileEvent, XEvent> ProfileLine::events_begin() {
+  return VisitorIterator<ProfileEvent, XEvent>(
       &line_->events(), [this](const XEvent& event) {
-        return XEventVisitor(&event, line_->timestamp_ns(), plane_, xspace_);
+        return ProfileEvent(&event, line_->timestamp_ns(), plane_, xspace_);
       });
 }
 
-VisitorIterator<XEventVisitor, XEvent> XLineVisitor::events_end() {
-  return VisitorIterator<XEventVisitor, XEvent>(
+VisitorIterator<ProfileEvent, XEvent> ProfileLine::events_end() {
+  return VisitorIterator<ProfileEvent, XEvent>(
       &line_->events(),
       [this](const XEvent& event) {
-        return XEventVisitor(&event, line_->timestamp_ns(), plane_, xspace_);
+        return ProfileEvent(&event, line_->timestamp_ns(), plane_, xspace_);
       },
       line_->events().size());
 }
 
-XPlaneVisitor::XPlaneVisitor(const XPlane* plane,
-                             std::shared_ptr<const XSpace> xspace)
+ProfilePlane::ProfilePlane(const XPlane* plane,
+                           std::shared_ptr<const XSpace> xspace)
     : plane_(plane), xspace_(xspace) {
   CHECK_NOTNULL(plane_);
   CHECK_NOTNULL(xspace_);
 }
 
-const std::string& XPlaneVisitor::name() const { return plane_->name(); }
+const std::string& ProfilePlane::name() const { return plane_->name(); }
 
-VisitorIterator<XLineVisitor, XLine> XPlaneVisitor::lines_begin() {
-  return VisitorIterator<XLineVisitor, XLine>(
+VisitorIterator<ProfileLine, XLine> ProfilePlane::lines_begin() {
+  return VisitorIterator<ProfileLine, XLine>(
       &plane_->lines(), [this](const XLine& line) {
-        return XLineVisitor(&line, plane_, xspace_);
+        return ProfileLine(&line, plane_, xspace_);
       });
 }
-VisitorIterator<XLineVisitor, XLine> XPlaneVisitor::lines_end() {
-  return VisitorIterator<XLineVisitor, XLine>(
+VisitorIterator<ProfileLine, XLine> ProfilePlane::lines_end() {
+  return VisitorIterator<ProfileLine, XLine>(
       &plane_->lines(),
-      [this](const XLine& line) {
-        return XLineVisitor(&line, plane_, xspace_);
-      },
+      [this](const XLine& line) { return ProfileLine(&line, plane_, xspace_); },
       plane_->lines().size());
 }
 
-VisitorIterator<nb::tuple, XStat> XPlaneVisitor::stats_begin() {
+VisitorIterator<nb::tuple, XStat> ProfilePlane::stats_begin() {
   return VisitorIterator<nb::tuple, XStat>(
       &plane_->stats(),
       [this](const XStat& stat) { return stats_to_tuple(stat, plane_); });
 }
 
-VisitorIterator<nb::tuple, XStat> XPlaneVisitor::stats_end() {
+VisitorIterator<nb::tuple, XStat> ProfilePlane::stats_end() {
   return VisitorIterator<nb::tuple, XStat>(
       &plane_->stats(),
       [this](const XStat& stat) { return stats_to_tuple(stat, plane_); },
       plane_->stats().size());
 }
 
-/*static*/ XSpaceVisitor XSpaceVisitor::from_serialized_xspace(
+/*static*/ ProfileData ProfileData::from_serialized_xspace(
     const nb::bytes& serialized_xspace) {
-  return XSpaceVisitor(serialized_xspace);
+  return ProfileData(serialized_xspace);
 }
 
-/*static*/ XSpaceVisitor XSpaceVisitor::from_file(
+/*static*/ ProfileData ProfileData::from_file(
     const std::string& proto_file_path) {
   std::string serialized_xspace;
   TF_CHECK_OK(tsl::ReadFileToString(tsl::Env::Default(), proto_file_path,
                                     &serialized_xspace));
-  return XSpaceVisitor(serialized_xspace.c_str(), serialized_xspace.size());
+  return ProfileData(serialized_xspace.c_str(), serialized_xspace.size());
 }
 
-/*static*/ XSpaceVisitor XSpaceVisitor::from_raw_cpp_ptr(nb::capsule capsule) {
+/*static*/ ProfileData ProfileData::from_raw_cpp_ptr(nb::capsule capsule) {
   auto raw_ptr = static_cast<XSpace*>(capsule.data());
   auto proto_ptr = std::shared_ptr<XSpace>(raw_ptr);
 
-  return XSpaceVisitor(proto_ptr);
+  return ProfileData(proto_ptr);
 }
 
-XSpaceVisitor::XSpaceVisitor(const char* serialized_xspace_ptr,
-                             size_t serialized_xspace_size) {
+ProfileData::ProfileData(const char* serialized_xspace_ptr,
+                         size_t serialized_xspace_size) {
   CHECK_NOTNULL(serialized_xspace_ptr);
 
   if (!xspace_) {
@@ -208,36 +206,35 @@ XSpaceVisitor::XSpaceVisitor(const char* serialized_xspace_ptr,
   CHECK(xspace_->ParseFromArray(serialized_xspace_ptr, serialized_xspace_size));
 }
 
-/*explicit*/ XSpaceVisitor::XSpaceVisitor(std::shared_ptr<XSpace> xspace_ptr) {
+/*explicit*/ ProfileData::ProfileData(std::shared_ptr<XSpace> xspace_ptr) {
   xspace_ = xspace_ptr;
 }
 
-/*explicit*/ XSpaceVisitor::XSpaceVisitor(const nb::bytes& serialized_xspace) {
+/*explicit*/ ProfileData::ProfileData(const nb::bytes& serialized_xspace) {
   if (!xspace_) {
     xspace_ = std::make_shared<XSpace>();
   }
   CHECK(xspace_->ParseFromArray(serialized_xspace.data(),
                                 serialized_xspace.size()));
 }
 
-VisitorIterator<XPlaneVisitor, XPlane> XSpaceVisitor::planes_begin() {
-  return VisitorIterator<XPlaneVisitor, XPlane>(
+VisitorIterator<ProfilePlane, XPlane> ProfileData::planes_begin() {
+  return VisitorIterator<ProfilePlane, XPlane>(
       &xspace_->planes(),
-      [this](const XPlane& plane) { return XPlaneVisitor(&plane, xspace_); });
+      [this](const XPlane& plane) { return ProfilePlane(&plane, xspace_); });
 }
 
-VisitorIterator<XPlaneVisitor, XPlane> XSpaceVisitor::planes_end() {
-  return VisitorIterator<XPlaneVisitor, XPlane>(
+VisitorIterator<ProfilePlane, XPlane> ProfileData::planes_end() {
+  return VisitorIterator<ProfilePlane, XPlane>(
       &xspace_->planes(),
-      [this](const XPlane& plane) { return XPlaneVisitor(&plane, xspace_); },
+      [this](const XPlane& plane) { return ProfilePlane(&plane, xspace_); },
       xspace_->planes().size());
 }
 
-XPlaneVisitor* XSpaceVisitor::find_plane_with_name(
-    const std::string& name) const {
+ProfilePlane* ProfileData::find_plane_with_name(const std::string& name) const {
   for (const auto& plane : xspace_->planes()) {
     if (plane.name() == name) {
-      return new XPlaneVisitor(&plane, xspace_);
+      return new ProfilePlane(&plane, xspace_);
     }
   }
   return nullptr;