Remove proto serialization and deserialization from HloModuleGroup.

Inline these into their only user.

PiperOrigin-RevId: 855971432

Author: hawkinsp

third_party/xla/xla/hlo/ir/BUILD

@@ -329,12 +329,9 @@ cc_library(
     hdrs = ["hlo_module_group.h"],
     deps = [
         ":hlo",
-        "//xla:status_macros",
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_proto_cc",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",

third_party/xla/xla/hlo/ir/hlo_module_group.cc

@@ -23,15 +23,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/log/check.h"
-#include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
-#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/service/hlo.pb.h"
-#include "xla/service/hlo_module_config.h"
-#include "xla/status_macros.h"
-#include "xla/tsl/platform/statusor.h"
 
 namespace xla {
 
@@ -55,38 +47,6 @@ std::string HloModuleGroup::ToString() const {
   return s.str();
 }
 
-HloModuleGroupProto HloModuleGroup::ToProto() const {
-  HloModuleGroupProto proto;
-  proto.set_name(name());
-  if (module_) {
-    *proto.add_hlo_modules() = module_->ToProto();
-  }
-  return proto;
-}
-
-/* static */ absl::StatusOr<HloModuleGroup> HloModuleGroup::CreateFromProto(
-    const HloModuleGroupProto& proto,
-    absl::Span<const HloModuleConfig> module_configs) {
-  TF_RET_CHECK(!proto.name().empty()) << "Module group name cannot be empty";
-  TF_RET_CHECK(proto.hlo_modules_size() > 0)
-      << "Module group must have at least one HLO module";
-  TF_RET_CHECK(proto.hlo_modules_size() == module_configs.size());
-
-  std::vector<std::unique_ptr<HloModule>> modules;
-  if (proto.hlo_modules_size() != 1) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("HloModuleGroupProto should have exactly one module, "
-                     "but it has ",
-                     proto.hlo_modules_size()));
-  }
-  const HloModuleProto& module_proto = proto.hlo_modules(0);
-  TF_ASSIGN_OR_RETURN(
-      std::unique_ptr<HloModule> module,
-      HloModule::CreateFromProto(module_proto, module_configs[0]));
-
-  return HloModuleGroup(std::move(module));
-}
-
 void HloModuleGroup::AddModule(std::unique_ptr<HloModule> module) {
   CHECK_EQ(module_, nullptr);
   module_ = std::move(module);

third_party/xla/xla/hlo/ir/hlo_module_group.h

@@ -80,11 +80,6 @@ class HloModuleGroup {
     return H::combine(std::move(h), group.module_);
   }
 
-  // Serialize the module group to/from a proto.
-  HloModuleGroupProto ToProto() const;
-  static absl::StatusOr<HloModuleGroup> CreateFromProto(
-      const HloModuleGroupProto& proto,
-      absl::Span<const HloModuleConfig> module_configs);
 
   // Returns the number of modules in the module group.
   int size() const { return module_ ? 1 : 0; }

third_party/xla/xla/service/hlo_module_group_test.cc

@@ -58,13 +58,6 @@ ENTRY %entry (x: f32[], y: f32[]) -> f32[] {
       group.module(0).entry_computation()->instructions(),
       ::testing::ElementsAre(op::Parameter(), op::Parameter(), op::Add()));
 
-  TF_ASSERT_OK_AND_ASSIGN(HloModuleGroup group_copy,
-                          HloModuleGroup::CreateFromProto(
-                              group.ToProto(), {group.module(0).config()}));
-  EXPECT_EQ(group_copy.modules().size(), 1);
-  EXPECT_THAT(
-      group_copy.module(0).entry_computation()->instructions(),
-      ::testing::ElementsAre(op::Parameter(), op::Parameter(), op::Add()));
 
   std::vector<std::unique_ptr<HloModule>> modules = group.ConsumeModules();
   EXPECT_EQ(modules.size(), 1);

third_party/xla/xla/stream_executor/tpu/BUILD

@@ -565,13 +565,10 @@ cc_library(
         ":status_helper",
         ":tpu_executable",
         ":tpu_executor_api",
-        ":tpu_executor_c_api_hdrs",
         ":tpu_executor_hdrs",
         ":tpu_platform_id",
         "//xla:shape_util",
-        "//xla:util",
         "//xla/hlo/ir:hlo",
-        "//xla/hlo/ir:hlo_module_group",
         "//xla/service:compiler",
         "//xla/service:executable",
         "//xla/service:hlo_cost_analysis",

third_party/xla/xla/stream_executor/tpu/tpu_on_demand_compiler.cc

@@ -21,7 +21,6 @@ limitations under the License.
 #include "absl/cleanup/cleanup.h"
 #include "absl/status/statusor.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/hlo/ir/hlo_module_group.h"
 #include "xla/service/compiler.h"
 #include "xla/service/executable.h"
 #include "xla/service/hlo.pb.h"
@@ -35,9 +34,7 @@ limitations under the License.
 #include "xla/stream_executor/tpu/tpu_executable.h"
 #include "xla/stream_executor/tpu/tpu_executor.h"
 #include "xla/stream_executor/tpu/tpu_executor_api.h"
-#include "xla/stream_executor/tpu/tpu_executor_c_api.h"
 #include "xla/stream_executor/tpu/tpu_platform_id.h"
-#include "xla/util.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {
@@ -116,13 +113,14 @@ class TpuCompiler : public Compiler {
       std::unique_ptr<HloModule> hlo_module,
       std::vector<stream_executor::StreamExecutor*> stream_exec,
       const CompileOptions& options) override {
-    HloModuleGroup module_group(std::move(hlo_module));
+    HloModuleGroupProto proto;
+    proto.set_name(hlo_module->name());
+    *proto.add_hlo_modules() = hlo_module->ToProto();
+
     XLA_HloModuleGroup se_module_group;
-    se_module_group.proto =
-        stream_executor::tpu::SerializeProto(module_group.ToProto());
+    se_module_group.proto = stream_executor::tpu::SerializeProto(proto);
     se_module_group.module_config = new XLA_HloModuleConfig[1];
-    se_module_group.module_config[0] =
-        ApiConverter::ToC(module_group.module(0).config());
+    se_module_group.module_config[0] = ApiConverter::ToC(hlo_module->config());
 
     auto cleanup_config = absl::MakeCleanup([&se_module_group]() {
       ApiConverter::Destroy(&se_module_group.module_config[0]);