diff --git a/xprof/convert/BUILD b/xprof/convert/BUILD index 7fd7b589f..0d057cc2f 100644 --- a/xprof/convert/BUILD +++ b/xprof/convert/BUILD @@ -271,6 +271,45 @@ cc_library( ], ) +cc_library( + name = "unified_utilization_viewer_processor", + srcs = ["unified_utilization_viewer_processor.cc"], + hdrs = ["unified_utilization_viewer_processor.h"], + deps = [ + ":tool_options", + ":unified_profile_processor", + ":unified_session_snapshot", + ":xplane_to_utilization_viewer", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@com_google_protobuf//:protobuf", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + ], + alwayslink = 1, +) + +cc_test( + name = "unified_utilization_viewer_processor_test", + srcs = ["unified_utilization_viewer_processor_test.cc"], + deps = [ + ":file_utils", + ":repository", + ":tool_options", + ":unified_profile_processor", + ":unified_profile_processor_factory", + ":unified_tools_registration", + ":unified_utilization_viewer_processor", + "//file/base", + "//file/base:path", + "//third_party/jsoncpp:json_compat", + "@com_google_absl//absl/status", + "@com_google_googletest//:gtest_main", + "@jsoncpp_git//:jsoncpp", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + ], +) + cc_library( name = "overview_page_processor", srcs = ["overview_page_processor.cc"], @@ -789,6 +828,25 @@ cc_test( ], ) +cc_test( + name = "unified_memory_profile_processor_test", + srcs = ["unified_memory_profile_processor_test.cc"], + deps = [ + ":tool_options", + ":unified_memory_profile_processor", + ":unified_memory_profile_processor_registration", + ":unified_profile_processor", + ":unified_profile_processor_factory", + ":unified_session_snapshot", + "//net/proto2/contrib/parse_proto:parse_text_proto", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + "@com_google_protobuf//:protobuf", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + ], +) + cc_test( name = "unified_memory_viewer_processor_test", srcs = ["unified_memory_viewer_processor_test.cc"], @@ -1824,12 +1882,9 @@ cc_library( ":profile_processor_factory", ":repository", ":tool_options", - ":unified_hlo_stats_processor", - ":unified_memory_viewer_processor", - ":unified_op_profile_processor", - ":unified_overview_page_processor", ":unified_profile_processor", ":unified_profile_processor_factory", + ":unified_tools_registration", "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/log", @@ -2994,6 +3049,37 @@ cc_test( ], ) +cc_library( + name = "unified_memory_profile_processor", + srcs = ["unified_memory_profile_processor.cc"], + hdrs = ["unified_memory_profile_processor.h"], + deps = [ + ":preprocess_single_host_xplane", + ":tool_options", + ":unified_profile_processor", + ":unified_session_snapshot", + ":xplane_to_memory_profile", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_protobuf//:protobuf", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + "@xla//xla/tsl/platform:errors", + "@xla//xla/tsl/platform:statusor", + ], +) + +cc_library( + name = "unified_memory_profile_processor_registration", + srcs = ["unified_memory_profile_processor_registration.cc"], + deps = [ + ":unified_memory_profile_processor", + ":unified_profile_processor_factory", + ], + alwayslink = 1, +) + cc_library( name = "unified_memory_viewer_processor", srcs = ["unified_memory_viewer_processor.cc"], @@ -3015,12 +3101,14 @@ cc_library( cc_library( name = "unified_tools_registration", srcs = ["unified_tools_registration.cc"], + hdrs = ["unified_tools_registration.h"], deps = [ ":unified_hlo_stats_processor", ":unified_memory_viewer_processor", ":unified_op_profile_processor", ":unified_overview_page_processor", ":unified_profile_processor_factory", + ":unified_utilization_viewer_processor", ], alwayslink = 1, ) diff --git a/xprof/convert/unified_hlo_stats_processor_test.cc b/xprof/convert/unified_hlo_stats_processor_test.cc index b5676e884..e62955f69 100644 --- a/xprof/convert/unified_hlo_stats_processor_test.cc +++ b/xprof/convert/unified_hlo_stats_processor_test.cc @@ -32,6 +32,7 @@ #include "xprof/convert/repository.h" #include "xprof/convert/tool_options.h" #include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_tools_registration.h" #include "plugin/xprof/protobuf/op_stats.pb.h" namespace xprof { @@ -48,6 +49,7 @@ class UnifiedHloStatsProcessorTest : public ::testing::Test { file::JoinPath(testing::TempDir(), "unified_hlo_stats_processor_test"); file::RecursivelyDelete(session_dir_, file::Defaults()).IgnoreError(); ASSERT_OK(file::CreateDir(session_dir_, file::Defaults())); + RegisterUnifiedToolRegistrations(); } void TearDown() override { @@ -65,7 +67,7 @@ TEST_F(UnifiedHloStatsProcessorTest, MinimalTest) { std::string xspace_path = file::JoinPath(session_dir_, "test_host.xplane.pb"); XSpace dummy_space; - ASSERT_OK(xprof::WriteBinaryProto(xspace_path, dummy_space)); + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); ASSERT_OK_AND_ASSIGN(SessionSnapshot session_snapshot, SessionSnapshot::Create({xspace_path}, std::nullopt)); @@ -83,7 +85,7 @@ TEST_F(UnifiedHloStatsProcessorTest, ProcessCombinedOpStatsTest) { std::string xspace_path = file::JoinPath(session_dir_, "test_host.xplane.pb"); XSpace dummy_space; - ASSERT_OK(xprof::WriteBinaryProto(xspace_path, dummy_space)); + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); ASSERT_OK_AND_ASSIGN(SessionSnapshot session_snapshot, SessionSnapshot::Create({xspace_path}, std::nullopt)); diff --git a/xprof/convert/unified_memory_profile_processor.cc b/xprof/convert/unified_memory_profile_processor.cc new file mode 100644 index 000000000..e9da9c7d6 --- /dev/null +++ b/xprof/convert/unified_memory_profile_processor.cc @@ -0,0 +1,59 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xprof/convert/unified_memory_profile_processor.h" + +#include + +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "google/protobuf/arena.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/preprocess_single_host_xplane.h" +#include "xprof/convert/tool_options.h" +#include "xprof/convert/unified_session_snapshot.h" +#include "xprof/convert/xplane_to_memory_profile.h" + +namespace xprof { + +absl::Status UnifiedMemoryProfileProcessor::ProcessSession( + const XprofSessionSnapshot& session_snapshot, + const tensorflow::profiler::ToolOptions& options) { + if (session_snapshot.XSpaceSize() != 1) { + return absl::InvalidArgumentError( + absl::StrCat("Memory profile tool expects only 1 XSpace path but gets ", + session_snapshot.XSpaceSize())); + } + + LOG(INFO) << "Processing memory profile for host: " + << session_snapshot.GetHostname(0); + + std::string memory_profile_json; + google::protobuf::Arena arena; + TF_ASSIGN_OR_RETURN(tensorflow::profiler::XSpace * xspace, + session_snapshot.GetXSpace(0, &arena)); + PreprocessSingleHostXSpace(xspace, /*step_grouping=*/true, + /*derived_timeline=*/false); + TF_RETURN_IF_ERROR( + ConvertXSpaceToMemoryProfileJson(*xspace, &memory_profile_json)); + + SetOutput(memory_profile_json, "application/json"); + return absl::OkStatus(); +} + +} // namespace xprof diff --git a/xprof/convert/unified_memory_profile_processor.h b/xprof/convert/unified_memory_profile_processor.h new file mode 100644 index 000000000..4044784a0 --- /dev/null +++ b/xprof/convert/unified_memory_profile_processor.h @@ -0,0 +1,59 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef THIRD_PARTY_XPROF_CONVERT_UNIFIED_MEMORY_PROFILE_PROCESSOR_H_ +#define THIRD_PARTY_XPROF_CONVERT_UNIFIED_MEMORY_PROFILE_PROCESSOR_H_ + +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/tool_options.h" +#include "xprof/convert/unified_profile_processor.h" +#include "xprof/convert/unified_session_snapshot.h" + +namespace xprof { + +class UnifiedMemoryProfileProcessor : public virtual UnifiedProfileProcessor { + public: + explicit UnifiedMemoryProfileProcessor( + const tensorflow::profiler::ToolOptions& options) {} + ~UnifiedMemoryProfileProcessor() override = default; + + absl::Status ProcessSession( + const XprofSessionSnapshot& session_snapshot, + const tensorflow::profiler::ToolOptions& options) override; + + absl::StatusOr Map( + const XprofSessionSnapshot& session_snapshot, absl::string_view hostname, + const tensorflow::profiler::XSpace& xspace) override { + return absl::UnimplementedError( + "Map not implemented for UnifiedMemoryProfileProcessor"); + } + + absl::Status Reduce( + const XprofSessionSnapshot& session_snapshot, + const std::vector& map_output_files) override { + return absl::UnimplementedError( + "Reduce not implemented for UnifiedMemoryProfileProcessor"); + } +}; + +} // namespace xprof + +#endif // THIRD_PARTY_XPROF_CONVERT_UNIFIED_MEMORY_PROFILE_PROCESSOR_H_ diff --git a/xprof/convert/unified_memory_profile_processor_registration.cc b/xprof/convert/unified_memory_profile_processor_registration.cc new file mode 100644 index 000000000..fc2c92e61 --- /dev/null +++ b/xprof/convert/unified_memory_profile_processor_registration.cc @@ -0,0 +1,24 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xprof/convert/unified_memory_profile_processor.h" +#include "xprof/convert/unified_profile_processor_factory.h" + +namespace xprof { + +REGISTER_UNIFIED_PROFILE_PROCESSOR("memory_profile", + UnifiedMemoryProfileProcessor); + +} // namespace xprof diff --git a/xprof/convert/unified_memory_profile_processor_test.cc b/xprof/convert/unified_memory_profile_processor_test.cc new file mode 100644 index 000000000..b59bc8579 --- /dev/null +++ b/xprof/convert/unified_memory_profile_processor_test.cc @@ -0,0 +1,84 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xprof/convert/unified_memory_profile_processor.h" + +#include +#include +#include + +#include "testing/base/public/gmock.h" +#include "" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "google/protobuf/arena.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/tool_options.h" +#include "xprof/convert/unified_profile_processor.h" +#include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_session_snapshot.h" + +namespace xprof { +namespace { + +using ::testing::IsEmpty; +using ::testing::Not; + +class MockXprofSessionSnapshot : public XprofSessionSnapshot { + public: + MOCK_METHOD(size_t, XSpaceSize, (), (const, override)); + MOCK_METHOD((absl::StatusOr), GetXSpace, + (size_t index, google::protobuf::Arena* arena), (const, override)); + MOCK_METHOD(std::string, GetHostname, (size_t index), (const, override)); + MOCK_METHOD(absl::string_view, GetSessionRunDir, (), (const, override)); + MOCK_METHOD(absl::StatusOr, GetHostDataFileName, + (tensorflow::profiler::StoredDataType data_type, + absl::string_view host), + (const, override)); +}; + +TEST(UnifiedMemoryProfileProcessorTest, ProcessSessionTest) { + tensorflow::profiler::ToolOptions options; + std::unique_ptr processor = + UnifiedProfileProcessorFactory::GetInstance().Create("memory_profile", + options); + ASSERT_NE(processor, nullptr); + + auto* memory_profile_processor = + dynamic_cast(processor.get()); + ASSERT_NE(memory_profile_processor, nullptr); + + tensorflow::profiler::XSpace xspace_obj; + auto* plane = xspace_obj.add_planes(); + plane->set_name("/host:CPU"); + auto* xspace = &xspace_obj; + + testing::NiceMock session_snapshot; + EXPECT_CALL(session_snapshot, XSpaceSize()) + .WillRepeatedly(testing::Return(1)); + EXPECT_CALL(session_snapshot, GetHostname(0)) + .WillRepeatedly(testing::Return("test_host")); + EXPECT_CALL(session_snapshot, GetXSpace(0, testing::_)) + .WillRepeatedly(testing::Return(xspace)); + + ASSERT_OK( + memory_profile_processor->ProcessSession(session_snapshot, options)); + + EXPECT_EQ(memory_profile_processor->GetContentType(), "application/json"); + EXPECT_THAT(memory_profile_processor->GetData(), Not(IsEmpty())); +} + +} // namespace +} // namespace xprof diff --git a/xprof/convert/unified_memory_viewer_processor_test.cc b/xprof/convert/unified_memory_viewer_processor_test.cc index 2ce1ea226..c7e2a7f4d 100644 --- a/xprof/convert/unified_memory_viewer_processor_test.cc +++ b/xprof/convert/unified_memory_viewer_processor_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "xprof/convert/tool_options.h" #include "xprof/convert/unified_profile_processor.h" #include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_tools_registration.h" #include "xprof/convert/unified_session_snapshot.h" namespace xprof { @@ -51,6 +52,7 @@ class MockXprofSessionSnapshot : public XprofSessionSnapshot { }; TEST(UnifiedMemoryViewerProcessorTest, ProcessHloJsonTest) { + RegisterUnifiedToolRegistrations(); tensorflow::profiler::ToolOptions options; std::unique_ptr processor = UnifiedProfileProcessorFactory::GetInstance().Create("memory_viewer", @@ -76,6 +78,7 @@ TEST(UnifiedMemoryViewerProcessorTest, ProcessHloJsonTest) { } TEST(UnifiedMemoryViewerProcessorTest, ProcessHloHtmlTest) { + RegisterUnifiedToolRegistrations(); tensorflow::profiler::ToolOptions options; options["view_memory_allocation_timeline"] = true; std::unique_ptr processor = diff --git a/xprof/convert/unified_op_profile_processor_test.cc b/xprof/convert/unified_op_profile_processor_test.cc index 083bc650d..3a4a352b9 100644 --- a/xprof/convert/unified_op_profile_processor_test.cc +++ b/xprof/convert/unified_op_profile_processor_test.cc @@ -27,6 +27,7 @@ limitations under the License. #include "xprof/convert/repository.h" #include "xprof/convert/tool_options.h" #include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_tools_registration.h" namespace xprof { namespace { @@ -36,6 +37,7 @@ using ::tensorflow::profiler::ToolOptions; using ::tensorflow::profiler::XSpace; TEST(UnifiedOpProfileProcessorTest, MinimalTest) { + RegisterUnifiedToolRegistrations(); ToolOptions options; auto processor = UnifiedProfileProcessorFactory::GetInstance().Create( "op_profile", options); @@ -47,7 +49,7 @@ TEST(UnifiedOpProfileProcessorTest, MinimalTest) { std::string xspace_path = tsl::io::JoinPath(session_dir, "test_host.xplane.pb"); XSpace dummy_space; - ASSERT_OK(xprof::WriteBinaryProto(xspace_path, dummy_space)); + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); std::vector xspace_paths = {xspace_path}; ASSERT_OK_AND_ASSIGN( diff --git a/xprof/convert/unified_overview_page_processor_test.cc b/xprof/convert/unified_overview_page_processor_test.cc index 611dc570e..6788b1516 100644 --- a/xprof/convert/unified_overview_page_processor_test.cc +++ b/xprof/convert/unified_overview_page_processor_test.cc @@ -33,6 +33,7 @@ #include "xprof/convert/repository.h" #include "xprof/convert/tool_options.h" #include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_tools_registration.h" #include "plugin/xprof/protobuf/op_stats.pb.h" namespace xprof { @@ -53,6 +54,7 @@ class UnifiedOverviewPageProcessorTest : public ::testing::Test { "unified_overview_page_processor_test"); file::RecursivelyDelete(session_dir_, file::Defaults()).IgnoreError(); ASSERT_OK(file::CreateDir(session_dir_, file::Defaults())); + RegisterUnifiedToolRegistrations(); } void TearDown() override { @@ -70,7 +72,7 @@ TEST_F(UnifiedOverviewPageProcessorTest, MinimalTest) { std::string xspace_path = file::JoinPath(session_dir_, "test_host.xplane.pb"); XSpace dummy_space; - ASSERT_OK(xprof::WriteBinaryProto(xspace_path, dummy_space)); + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); ASSERT_OK_AND_ASSIGN(auto session_snapshot, SessionSnapshot::Create({xspace_path}, std::nullopt)); @@ -90,7 +92,7 @@ TEST_F(UnifiedOverviewPageProcessorTest, ProcessCombinedOpStatsTrainingTest) { std::string xspace_path = file::JoinPath(session_dir_, "test_host.xplane.pb"); XSpace dummy_space; - ASSERT_OK(xprof::WriteBinaryProto(xspace_path, dummy_space)); + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); ASSERT_OK_AND_ASSIGN(auto session_snapshot, SessionSnapshot::Create({xspace_path}, std::nullopt)); @@ -128,7 +130,7 @@ TEST_F(UnifiedOverviewPageProcessorTest, ProcessCombinedOpStatsInferenceTest) { std::string xspace_path = file::JoinPath(session_dir_, "test_host.xplane.pb"); XSpace dummy_space; - ASSERT_OK(xprof::WriteBinaryProto(xspace_path, dummy_space)); + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); ASSERT_OK_AND_ASSIGN(auto session_snapshot, SessionSnapshot::Create({xspace_path}, std::nullopt)); diff --git a/xprof/convert/unified_tools_registration.cc b/xprof/convert/unified_tools_registration.cc index 4ac496103..7f5236a8d 100644 --- a/xprof/convert/unified_tools_registration.cc +++ b/xprof/convert/unified_tools_registration.cc @@ -13,19 +13,26 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xprof/convert/unified_tools_registration.h" + #include "xprof/convert/unified_hlo_stats_processor.h" #include "xprof/convert/unified_memory_viewer_processor.h" #include "xprof/convert/unified_op_profile_processor.h" #include "xprof/convert/unified_overview_page_processor.h" #include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_utilization_viewer_processor.h" namespace xprof { -REGISTER_UNIFIED_PROFILE_PROCESSOR("hlo_stats", UnifiedHloStatsProcessor); -REGISTER_UNIFIED_PROFILE_PROCESSOR("memory_viewer", - UnifiedMemoryViewerProcessor); -REGISTER_UNIFIED_PROFILE_PROCESSOR("op_profile", UnifiedOpProfileProcessor); -REGISTER_UNIFIED_PROFILE_PROCESSOR("overview_page", - UnifiedOverviewPageProcessor); +void RegisterUnifiedToolRegistrations() { + REGISTER_UNIFIED_PROFILE_PROCESSOR("hlo_stats", UnifiedHloStatsProcessor); + REGISTER_UNIFIED_PROFILE_PROCESSOR("memory_viewer", + UnifiedMemoryViewerProcessor); + REGISTER_UNIFIED_PROFILE_PROCESSOR("op_profile", UnifiedOpProfileProcessor); + REGISTER_UNIFIED_PROFILE_PROCESSOR("overview_page", + UnifiedOverviewPageProcessor); + REGISTER_UNIFIED_PROFILE_PROCESSOR("utilization_viewer", + UnifiedUtilizationViewerProcessor); +} } // namespace xprof diff --git a/xprof/convert/unified_tools_registration.h b/xprof/convert/unified_tools_registration.h new file mode 100644 index 000000000..084e73b77 --- /dev/null +++ b/xprof/convert/unified_tools_registration.h @@ -0,0 +1,25 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef THIRD_PARTY_XPROF_CONVERT_UNIFIED_TOOLS_REGISTRATION_H_ +#define THIRD_PARTY_XPROF_CONVERT_UNIFIED_TOOLS_REGISTRATION_H_ + +namespace xprof { +// Registers unified profile processors for various profiling tools (e.g., +// hlo_stats, memory_viewer) with the processor factory. +void RegisterUnifiedToolRegistrations(); +} + +#endif // THIRD_PARTY_XPROF_CONVERT_UNIFIED_TOOLS_REGISTRATION_H_ diff --git a/xprof/convert/unified_utilization_viewer_processor.cc b/xprof/convert/unified_utilization_viewer_processor.cc new file mode 100644 index 000000000..367e5f96c --- /dev/null +++ b/xprof/convert/unified_utilization_viewer_processor.cc @@ -0,0 +1,57 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xprof/convert/unified_utilization_viewer_processor.h" + +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "google/protobuf/arena.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/tool_options.h" +#include "xprof/convert/unified_session_snapshot.h" +#include "xprof/convert/xplane_to_utilization_viewer.h" + +namespace xprof { + +absl::Status UnifiedUtilizationViewerProcessor::ProcessSession( + const XprofSessionSnapshot& session_snapshot, + const tensorflow::profiler::ToolOptions& options) { + // TODO(b/something): Support multiple hosts properly if needed, + // or unify the data aggregation logic between 1P and 3P here. + // For now, process the first host's XSpace to establish the structure. + if (session_snapshot.XSpaceSize() == 0) { + return absl::NotFoundError("No XSpace found in the session."); + } + + google::protobuf::Arena arena; + absl::StatusOr xspace = + session_snapshot.GetXSpace(0, &arena); + if (!xspace.ok()) { + return xspace.status(); + } + + absl::StatusOr json_output = + ConvertXSpaceToUtilizationViewer(**xspace); + if (!json_output.ok()) { + return json_output.status(); + } + + SetOutput(*json_output, "application/json"); + return absl::OkStatus(); +} + +} // namespace xprof diff --git a/xprof/convert/unified_utilization_viewer_processor.h b/xprof/convert/unified_utilization_viewer_processor.h new file mode 100644 index 000000000..5e6dc5c8e --- /dev/null +++ b/xprof/convert/unified_utilization_viewer_processor.h @@ -0,0 +1,56 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef THIRD_PARTY_XPROF_CONVERT_UNIFIED_UTILIZATION_VIEWER_PROCESSOR_H_ +#define THIRD_PARTY_XPROF_CONVERT_UNIFIED_UTILIZATION_VIEWER_PROCESSOR_H_ + +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/tool_options.h" +#include "xprof/convert/unified_profile_processor.h" +#include "xprof/convert/unified_session_snapshot.h" + +namespace xprof { + +class UnifiedUtilizationViewerProcessor : public UnifiedProfileProcessor { + public: + explicit UnifiedUtilizationViewerProcessor( + const tensorflow::profiler::ToolOptions& options) {} + + absl::Status ProcessSession( + const XprofSessionSnapshot& session_snapshot, + const tensorflow::profiler::ToolOptions& options) override; + + absl::StatusOr Map( + const XprofSessionSnapshot& session_snapshot, absl::string_view hostname, + const tensorflow::profiler::XSpace& xspace) override { + return absl::UnimplementedError("Map not implemented"); + } + + absl::Status Reduce( + const XprofSessionSnapshot& session_snapshot, + const std::vector& map_output_files) override { + return absl::UnimplementedError("Reduce not implemented"); + } +}; + +} // namespace xprof + +#endif // THIRD_PARTY_XPROF_CONVERT_UNIFIED_UTILIZATION_VIEWER_PROCESSOR_H_ diff --git a/xprof/convert/unified_utilization_viewer_processor_test.cc b/xprof/convert/unified_utilization_viewer_processor_test.cc new file mode 100644 index 000000000..7bc075234 --- /dev/null +++ b/xprof/convert/unified_utilization_viewer_processor_test.cc @@ -0,0 +1,88 @@ +/* Copyright 2026 The OpenXLA Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include "file/base/filesystem.h" +#include "file/base/options.h" +#include "file/base/path.h" +#include "testing/base/public/gmock.h" +#include "" +#include "absl/status/status.h" +#include "third_party/jsoncpp/include/json/reader.h" +#include "third_party/jsoncpp/include/json/value.h" +#include "tsl/profiler/protobuf/xplane.pb.h" +#include "xprof/convert/file_utils.h" +#include "xprof/convert/repository.h" +#include "xprof/convert/tool_options.h" +#include "xprof/convert/unified_profile_processor.h" +#include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_tools_registration.h" + +namespace xprof { +namespace { + +using ::tensorflow::profiler::SessionSnapshot; +using ::tensorflow::profiler::ToolOptions; +using ::tensorflow::profiler::XSpace; + +class UnifiedUtilizationViewerProcessorTest : public ::testing::Test { + protected: + void SetUp() override { + session_dir_ = + file::JoinPath(testing::TempDir(), "unified_utilization_viewer_test"); + file::RecursivelyDelete(session_dir_, file::Defaults()).IgnoreError(); + ASSERT_OK(file::CreateDir(session_dir_, file::Defaults())); + RegisterUnifiedToolRegistrations(); + } + + void TearDown() override { + file::RecursivelyDelete(session_dir_, file::Defaults()).IgnoreError(); + } + + std::string session_dir_; + ToolOptions options_; +}; + +TEST_F(UnifiedUtilizationViewerProcessorTest, EmptyXSpaceTest) { + auto processor = UnifiedProfileProcessorFactory::GetInstance().Create( + "utilization_viewer", options_); + ASSERT_NE(processor, nullptr); + + std::string xspace_path = file::JoinPath(session_dir_, "test_host.xplane.pb"); + XSpace dummy_space; + ASSERT_OK(WriteBinaryProto(xspace_path, dummy_space)); + + ASSERT_OK_AND_ASSIGN(SessionSnapshot session_snapshot, + SessionSnapshot::Create({xspace_path}, std::nullopt)); + + ASSERT_OK(processor->ProcessSession(session_snapshot, options_)); + + std::string output_str = processor->GetData(); + EXPECT_FALSE(output_str.empty()); + EXPECT_EQ(processor->GetContentType(), "application/json"); + + Json::Value json; + Json::Reader reader; + ASSERT_TRUE(reader.parse(output_str, json)); + ASSERT_TRUE(json.isMember("rows")); + EXPECT_EQ(json["rows"].size(), 0); +} + +} // namespace +} // namespace xprof diff --git a/xprof/convert/xplane_to_tools_data_with_profile_processor.cc b/xprof/convert/xplane_to_tools_data_with_profile_processor.cc index 794d9074b..137d592d9 100644 --- a/xprof/convert/xplane_to_tools_data_with_profile_processor.cc +++ b/xprof/convert/xplane_to_tools_data_with_profile_processor.cc @@ -26,6 +26,7 @@ #include "xprof/convert/tool_options.h" #include "xprof/convert/unified_profile_processor.h" #include "xprof/convert/unified_profile_processor_factory.h" +#include "xprof/convert/unified_tools_registration.h" #include "plugin/xprof/protobuf/worker_service.pb.h" #include "plugin/xprof/worker/grpc_utils.h" #include "plugin/xprof/worker/stub_factory.h" @@ -193,6 +194,7 @@ absl::StatusOr ConvertMultiXSpacesToToolDataWithProfileProcessor( absl::Time start_time = absl::Now(); if (absl::GetFlag(FLAGS_enable_unified_xprof)) { + xprof::RegisterUnifiedToolRegistrations(); auto unified_processor = xprof::UnifiedProfileProcessorFactory::GetInstance().Create( tool_name, options);