From d0d716fb63a0433ecf75c4c298f968083dacfc6f Mon Sep 17 00:00:00 2001 From: xla authors Date: Sat, 19 Oct 2024 13:16:11 -0700 Subject: [PATCH] Rename xplane_visitor module to profile_data. PiperOrigin-RevId: 687672929 --- xla/python/BUILD | 2 +- xla/python/profiler.cc | 6 +- xla/python/profiler/BUILD | 14 +-- .../{xplane_visitor.cc => profile_data.cc} | 99 +++++++++---------- .../{xplane_visitor.h => profile_data.h} | 88 ++++++++--------- .../{xplane_visitor.pyi => profile_data.pyi} | 32 +++--- ...y_xplane_visitor.cc => py_profile_data.cc} | 53 +++++----- 7 files changed, 144 insertions(+), 150 deletions(-) rename xla/python/profiler/{xplane_visitor.cc => profile_data.cc} (64%) rename xla/python/profiler/{xplane_visitor.h => profile_data.h} (60%) rename xla/python/profiler/{xplane_visitor.pyi => profile_data.pyi} (75%) rename xla/python/profiler/{py_xplane_visitor.cc => py_profile_data.cc} (59%) diff --git a/xla/python/BUILD b/xla/python/BUILD index 4dcee866f00f5..7aa8908b0c26a 100644 --- a/xla/python/BUILD +++ b/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", diff --git a/xla/python/profiler.cc b/xla/python/profiler.cc index 70c437b1baaa6..f26ef1b046438 100644 --- a/xla/python/profiler.cc +++ b/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 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, diff --git a/xla/python/profiler/BUILD b/xla/python/profiler/BUILD index 1898f3ff68731..5b3373d896677 100644 --- a/xla/python/profiler/BUILD +++ b/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", ], ) diff --git a/xla/python/profiler/xplane_visitor.cc b/xla/python/profiler/profile_data.cc similarity index 64% rename from xla/python/profiler/xplane_visitor.cc rename to xla/python/profiler/profile_data.cc index d98daea40f874..91c9e6c84cf39 100644 --- a/xla/python/profiler/xplane_visitor.cc +++ b/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 #include // 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 xspace) +ProfileEvent::ProfileEvent(const XEvent* event, int64_t line_timestamp_ns, + const XPlane* plane, + std::shared_ptr 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 XEventVisitor::stats_begin() { +VisitorIterator ProfileEvent::stats_begin() { return VisitorIterator( &event_->stats(), [this](const XStat& stat) { return stats_to_tuple(stat, plane_); }); } -VisitorIterator XEventVisitor::stats_end() { +VisitorIterator ProfileEvent::stats_end() { return VisitorIterator( &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 xspace) +ProfileLine::ProfileLine(const XLine* line, const XPlane* plane, + std::shared_ptr 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 XLineVisitor::events_begin() { - return VisitorIterator( +VisitorIterator ProfileLine::events_begin() { + return VisitorIterator( &line_->events(), [this](const XEvent& event) { - return XEventVisitor(&event, line_->timestamp_ns(), plane_, xspace_); + return ProfileEvent(&event, line_->timestamp_ns(), plane_, xspace_); }); } -VisitorIterator XLineVisitor::events_end() { - return VisitorIterator( +VisitorIterator ProfileLine::events_end() { + return VisitorIterator( &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 xspace) +ProfilePlane::ProfilePlane(const XPlane* plane, + std::shared_ptr 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 XPlaneVisitor::lines_begin() { - return VisitorIterator( +VisitorIterator ProfilePlane::lines_begin() { + return VisitorIterator( &plane_->lines(), [this](const XLine& line) { - return XLineVisitor(&line, plane_, xspace_); + return ProfileLine(&line, plane_, xspace_); }); } -VisitorIterator XPlaneVisitor::lines_end() { - return VisitorIterator( +VisitorIterator ProfilePlane::lines_end() { + return VisitorIterator( &plane_->lines(), - [this](const XLine& line) { - return XLineVisitor(&line, plane_, xspace_); - }, + [this](const XLine& line) { return ProfileLine(&line, plane_, xspace_); }, plane_->lines().size()); } -VisitorIterator XPlaneVisitor::stats_begin() { +VisitorIterator ProfilePlane::stats_begin() { return VisitorIterator( &plane_->stats(), [this](const XStat& stat) { return stats_to_tuple(stat, plane_); }); } -VisitorIterator XPlaneVisitor::stats_end() { +VisitorIterator ProfilePlane::stats_end() { return VisitorIterator( &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(capsule.data()); auto proto_ptr = std::shared_ptr(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,11 +206,11 @@ XSpaceVisitor::XSpaceVisitor(const char* serialized_xspace_ptr, CHECK(xspace_->ParseFromArray(serialized_xspace_ptr, serialized_xspace_size)); } -/*explicit*/ XSpaceVisitor::XSpaceVisitor(std::shared_ptr xspace_ptr) { +/*explicit*/ ProfileData::ProfileData(std::shared_ptr 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(); } @@ -220,24 +218,23 @@ XSpaceVisitor::XSpaceVisitor(const char* serialized_xspace_ptr, serialized_xspace.size())); } -VisitorIterator XSpaceVisitor::planes_begin() { - return VisitorIterator( +VisitorIterator ProfileData::planes_begin() { + return VisitorIterator( &xspace_->planes(), - [this](const XPlane& plane) { return XPlaneVisitor(&plane, xspace_); }); + [this](const XPlane& plane) { return ProfilePlane(&plane, xspace_); }); } -VisitorIterator XSpaceVisitor::planes_end() { - return VisitorIterator( +VisitorIterator ProfileData::planes_end() { + return VisitorIterator( &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; diff --git a/xla/python/profiler/xplane_visitor.h b/xla/python/profiler/profile_data.h similarity index 60% rename from xla/python/profiler/xplane_visitor.h rename to xla/python/profiler/profile_data.h index 527ceabda9e03..980c76375ae80 100644 --- a/xla/python/profiler/xplane_visitor.h +++ b/xla/python/profiler/profile_data.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_PYTHON_PROFILER_XPLANE_VISITOR_H_ -#define XLA_PYTHON_PROFILER_XPLANE_VISITOR_H_ +#ifndef XLA_PYTHON_PROFILER_PROFILE_DATA_H_ +#define XLA_PYTHON_PROFILER_PROFILE_DATA_H_ #include @@ -76,14 +76,14 @@ class VisitorIterator int pos_ = 0; }; -class XEventVisitor { +class ProfileEvent { public: - XEventVisitor() = delete; + ProfileEvent() = delete; - XEventVisitor(const tensorflow::profiler::XEvent* event, - int64_t line_timestamp_ns, - const tensorflow::profiler::XPlane* plane, - std::shared_ptr xspace); + ProfileEvent(const tensorflow::profiler::XEvent* event, + int64_t line_timestamp_ns, + const tensorflow::profiler::XPlane* plane, + std::shared_ptr xspace); double start_ns() const; @@ -101,44 +101,44 @@ class XEventVisitor { const XPlane* plane_; const int64_t line_timestamp_ns_; // The actual XSpace protobuf we are wrapping around. A shared ptr is used so - // the different levels of visitors (XSpaceVisitor, XPlaneVisitor, - // XLineVisitor, etc.) don't depend on the lifetime of others. + // the different levels of visitors (ProfileData, ProfilePlane, + // ProfileLine, etc.) don't depend on the lifetime of others. const std::shared_ptr xspace_; }; -class XLineVisitor { +class ProfileLine { public: - XLineVisitor() = delete; + ProfileLine() = delete; - XLineVisitor(const tensorflow::profiler::XLine* line, - const tensorflow::profiler::XPlane* plane, - std::shared_ptr xspace); + ProfileLine(const tensorflow::profiler::XLine* line, + const tensorflow::profiler::XPlane* plane, + std::shared_ptr xspace); const std::string& name() const; - VisitorIterator events_begin(); - VisitorIterator events_end(); + VisitorIterator events_begin(); + VisitorIterator events_end(); private: const XLine* line_; const XPlane* plane_; // The actual XSpace protobuf we are wrapping around. A shared ptr is used so - // the different levels of visitors (XSpaceVisitor, XPlaneVisitor, - // XLineVisitor, etc.) don't depend on the lifetime of others. + // the different levels of visitors (ProfileData, ProfilePlane, + // ProfileLine, etc.) don't depend on the lifetime of others. const std::shared_ptr xspace_; }; -class XPlaneVisitor { +class ProfilePlane { public: - XPlaneVisitor() = delete; + ProfilePlane() = delete; - XPlaneVisitor(const tensorflow::profiler::XPlane* plane, - std::shared_ptr xspace); + ProfilePlane(const tensorflow::profiler::XPlane* plane, + std::shared_ptr xspace); const std::string& name() const; - VisitorIterator lines_begin(); - VisitorIterator lines_end(); + VisitorIterator lines_begin(); + VisitorIterator lines_end(); VisitorIterator stats_begin(); @@ -147,46 +147,44 @@ class XPlaneVisitor { private: const XPlane* plane_; // The actual XSpace protobuf we are wrapping around. A shared ptr is used so - // the different levels of visitors (XSpaceVisitor, XPlaneVisitor, - // XLineVisitor, etc.) don't depend on the lifetime of others. + // the different levels of visitors (ProfileData, ProfilePlane, + // ProfileLine, etc.) don't depend on the lifetime of others. const std::shared_ptr xspace_; }; -class XSpaceVisitor { +class ProfileData { public: - static XSpaceVisitor from_serialized_xspace( - const nb::bytes& serialized_xspace); + static ProfileData from_serialized_xspace(const nb::bytes& serialized_xspace); - static XSpaceVisitor from_file(const std::string& proto_file_path); + static ProfileData from_file(const std::string& proto_file_path); - static XSpaceVisitor from_raw_cpp_ptr(nb::capsule capsule); + static ProfileData from_raw_cpp_ptr(nb::capsule capsule); - XSpaceVisitor() = delete; + ProfileData() = delete; - XSpaceVisitor(const char* serialized_xspace_ptr, - size_t serialized_xspace_size); + ProfileData(const char* serialized_xspace_ptr, size_t serialized_xspace_size); - explicit XSpaceVisitor(std::shared_ptr xspace_ptr); + explicit ProfileData(std::shared_ptr xspace_ptr); - explicit XSpaceVisitor(const nb::bytes& serialized_xspace); + explicit ProfileData(const nb::bytes& serialized_xspace); - VisitorIterator planes_begin(); + VisitorIterator planes_begin(); - VisitorIterator planes_end(); + VisitorIterator planes_end(); - XPlaneVisitor* find_plane_with_name(const std::string& name) const; + ProfilePlane* find_plane_with_name(const std::string& name) const; private: // The actual XSpace protobuf we are wrapping around. A shared ptr is used so - // the different levels of visitors (XSpaceVisitor, XPlaneVisitor, - // XLineVisitor, etc.) don't depend on the lifetime of others. + // the different levels of visitors (ProfileData, ProfilePlane, + // ProfileLine, etc.) don't depend on the lifetime of others. std::shared_ptr xspace_; }; -XSpaceVisitor from_serialized_xspace(const std::string& serialized_xspace); +ProfileData from_serialized_xspace(const std::string& serialized_xspace); -XSpaceVisitor from_file(const std::string& proto_file_path); +ProfileData from_file(const std::string& proto_file_path); } // namespace tensorflow::profiler::python -#endif // XLA_PYTHON_PROFILER_XPLANE_VISITOR_H_ +#endif // XLA_PYTHON_PROFILER_PROFILE_DATA_H_ diff --git a/xla/python/profiler/xplane_visitor.pyi b/xla/python/profiler/profile_data.pyi similarity index 75% rename from xla/python/profiler/xplane_visitor.pyi rename to xla/python/profiler/profile_data.pyi index 7481b9cd61a12..9a84529639a47 100644 --- a/xla/python/profiler/xplane_visitor.pyi +++ b/xla/python/profiler/profile_data.pyi @@ -12,42 +12,42 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utilities for visiting XPlanes.""" +"""Utilities for visiting program execution data.""" from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple -class XSpaceVisitor: - """Parses a XSpace protobuf and provides accessors to its contents.""" +class ProfileData: + """Program execution data captured by jax.profiler functions.""" def __init__(self, serialized_xspace: bytes): ... @classmethod - def from_file(cls, path: str) -> 'XSpaceVisitor': - """Creates an XSpaceVisitor from a serialized XSpace proto file.""" + def from_file(cls, path: str) -> 'ProfileData': + """Creates a ProfileData from a serialized XSpace proto file.""" ... @classmethod - def from_serialized_xspace(cls, serialized_xspace: bytes) -> 'XSpaceVisitor': - """Creates an XSpaceVisitor from a serialized XSpace proto.""" + def from_serialized_xspace(cls, serialized_xspace: bytes) -> 'ProfileData': + """Creates a ProfileData from a serialized XSpace proto.""" ... @classmethod - def from_raw_cpp_ptr(cls, raw_proto_ptr: object) -> 'XSpaceVisitor': - """Creates an XSpaceVisitor from a raw C++ pointer enclosed in a capsule to a XSpace proto.""" + def from_raw_cpp_ptr(cls, raw_proto_ptr: object) -> 'ProfileData': + """Creates a ProfileData from a raw C++ pointer enclosed in a capsule to a XSpace proto.""" ... @property - def planes(self) -> Iterator['XPlaneVisitor']: + def planes(self) -> Iterator['ProfilePlane']: ... - def find_plane_with_name(self, name: str) -> Optional['XPlaneVisitor']: + def find_plane_with_name(self, name: str) -> Optional['ProfilePlane']: """Finds the plane with the given name.""" ... -class XPlaneVisitor: +class ProfilePlane: """Wraps XPlane protobuf and provides accessors to its contents.""" @property @@ -56,7 +56,7 @@ class XPlaneVisitor: ... @property - def lines(self) -> Iterator['XLineVisitor']: + def lines(self) -> Iterator['ProfileLine']: """Lines in the plane.""" ... @@ -72,7 +72,7 @@ class XPlaneVisitor: ... -class XLineVisitor: +class ProfileLine: """Wraps XLine protobuf and provides accessors to its contents.""" @property @@ -81,12 +81,12 @@ class XLineVisitor: ... @property - def events(self) -> Iterator['XEventVisitor']: + def events(self) -> Iterator['ProfileEvent']: """Events in the line.""" ... -class XEventVisitor: +class ProfileEvent: """Wraps XEvent protobuf and provides accessors to its contents.""" @property diff --git a/xla/python/profiler/py_xplane_visitor.cc b/xla/python/profiler/py_profile_data.cc similarity index 59% rename from xla/python/profiler/py_xplane_visitor.cc rename to xla/python/profiler/py_profile_data.cc index 09bd144a85cb3..b8bacfe9500ee 100644 --- a/xla/python/profiler/py_xplane_visitor.cc +++ b/xla/python/profiler/py_profile_data.cc @@ -16,7 +16,7 @@ limitations under the License. #include // For automatic conversion of std::iterator to Python iterable. #include // For automatic conversion of std::string to Python string. -#include "xla/python/profiler/xplane_visitor.h" +#include "xla/python/profiler/profile_data.h" namespace { @@ -26,59 +26,58 @@ using namespace nb::literals; using namespace tensorflow::profiler::python; // NOLINTEND(build/namespaces) -NB_MODULE(xplane_visitor, m) { - nb::class_(m, "XEventVisitor") - .def_prop_ro("start_ns", &XEventVisitor::start_ns) - .def_prop_ro("duration_ns", &XEventVisitor::duration_ns) - .def_prop_ro("end_ns", &XEventVisitor::end_ns) - .def_prop_ro("name", &XEventVisitor::name) +NB_MODULE(profile_data, m) { + nb::class_(m, "ProfileEvent") + .def_prop_ro("start_ns", &ProfileEvent::start_ns) + .def_prop_ro("duration_ns", &ProfileEvent::duration_ns) + .def_prop_ro("end_ns", &ProfileEvent::end_ns) + .def_prop_ro("name", &ProfileEvent::name) .def_prop_ro( "stats", - [](XEventVisitor&& e) { + [](ProfileEvent&& e) { return nb::make_iterator(nb::type(), "event_stats", e.stats_begin(), e.stats_end()); }, nb::keep_alive<0, 1>()); - nb::class_(m, "XLineVisitor") - .def_prop_ro("name", &XLineVisitor::name) + nb::class_(m, "ProfileLine") + .def_prop_ro("name", &ProfileLine::name) .def_prop_ro( "events", - [](XLineVisitor&& l) { - return nb::make_iterator(nb::type(), "events", + [](ProfileLine&& l) { + return nb::make_iterator(nb::type(), "events", l.events_begin(), l.events_end()); }, nb::keep_alive<0, 1>()); - nb::class_(m, "XPlaneVisitor") - .def_prop_ro("name", &XPlaneVisitor::name) + nb::class_(m, "ProfilePlane") + .def_prop_ro("name", &ProfilePlane::name) .def_prop_ro( "lines", - [](XPlaneVisitor&& p) { - return nb::make_iterator(nb::type(), "lines", + [](ProfilePlane&& p) { + return nb::make_iterator(nb::type(), "lines", p.lines_begin(), p.lines_end()); }, nb::keep_alive<0, 1>()) .def_prop_ro( "stats", - [](XPlaneVisitor&& p) { + [](ProfilePlane&& p) { return nb::make_iterator(nb::type(), "plane_stats", p.stats_begin(), p.stats_end()); }, nb::keep_alive<0, 1>()); - nb::class_(m, "XSpaceVisitor") - .def_static("from_raw_cpp_ptr", &XSpaceVisitor::from_raw_cpp_ptr, + nb::class_(m, "ProfileData") + .def_static("from_raw_cpp_ptr", &ProfileData::from_raw_cpp_ptr, "capsule"_a) - .def_static( - "from_file", &XSpaceVisitor::from_file, "proto_file_path"_a, - "Creates an XSpaceVisitor from a serialized XSpace proto file.") + .def_static("from_file", &ProfileData::from_file, "proto_file_path"_a, + "Creates a ProfileData from a serialized XSpace proto file.") .def_static("from_serialized_xspace", - &XSpaceVisitor::from_serialized_xspace, "serialized_xspace"_a) + &ProfileData::from_serialized_xspace, "serialized_xspace"_a) .def(nb::init()) - .def("find_plane_with_name", &XSpaceVisitor::find_plane_with_name, - "name"_a, nb::keep_alive<0, 1>()) + .def("find_plane_with_name", &ProfileData::find_plane_with_name, "name"_a, + nb::keep_alive<0, 1>()) .def_prop_ro( "planes", - [](XSpaceVisitor&& s) { - return nb::make_iterator(nb::type(), "planes", + [](ProfileData&& s) { + return nb::make_iterator(nb::type(), "planes", s.planes_begin(), s.planes_end()); }, nb::keep_alive<0, 1>());