Add a new memory_profile tool to xprof.

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,
 )

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));

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 <string>
+
+#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

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 <string>
+#include <vector>
+
+#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<std::string> 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<std::string>& map_output_files) override {
+    return absl::UnimplementedError(
+        "Reduce not implemented for UnifiedMemoryProfileProcessor");
+  }
+};
+
+}  // namespace xprof
+
+#endif  // THIRD_PARTY_XPROF_CONVERT_UNIFIED_MEMORY_PROFILE_PROCESSOR_H_

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