Support creating HloModule from proto with compilation environments.

PiperOrigin-RevId: 725560495
openxla · Feb 14, 2025 · 98fba1f · 98fba1f
1 parent 6da089d
commit 98fba1f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/xla/hlo/ir/hlo_module.cc b/xla/hlo/ir/hlo_module.cc
@@ -556,7 +556,8 @@ absl::Status HloModule::CheckUniqueNamesAndIdsForComputationsAndInstructions()
 /* static */
 absl::StatusOr<std::unique_ptr<HloModule>> HloModule::CreateFromProto(
     const HloModuleProto& proto, const HloModuleConfig& module_config,
-    bool prohibit_empty_literal) {
+    bool prohibit_empty_literal,
+    std::unique_ptr<CompilationEnvironments> comp_envs) {
   VLOG(2) << "CreateFromProto()";
   XLA_VLOG_LINES(3, proto.DebugString());
 
@@ -609,7 +610,10 @@ absl::StatusOr<std::unique_ptr<HloModule>> HloModule::CreateFromProto(
   }
   TF_RET_CHECK(entry != nullptr);
 
-  auto module = std::make_unique<HloModule>(proto.name(), module_config);
+  auto module = comp_envs
+                    ? std::make_unique<HloModule>(proto.name(), module_config,
+                                                  std::move(comp_envs))
+                    : std::make_unique<HloModule>(proto.name(), module_config);
 
   // Sort the computations in the proto id's order.
   absl::c_sort(computations, [&](const std::unique_ptr<HloComputation>& a,
@@ -803,12 +807,14 @@ absl::StatusOr<HloModuleConfig> HloModule::CreateModuleConfigFromProto(
 }
 
 absl::StatusOr<std::unique_ptr<HloModule>> HloModule::CreateFromProtoWithConfig(
-    const HloModuleProtoWithConfig& proto, bool prohibit_empty_literal) {
+    const HloModuleProtoWithConfig& proto, bool prohibit_empty_literal,
+    std::unique_ptr<CompilationEnvironments> comp_envs) {
   const auto& hlo_module_proto = proto.hlo_module();
   TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModuleConfig> config_ptr,
                       HloModuleConfig::CreateFromProto(proto.config()));
   return HloModule::CreateFromProto(hlo_module_proto, *config_ptr,
-                                    prohibit_empty_literal);
+                                    prohibit_empty_literal,
+                                    std::move(comp_envs));
 }
 
 namespace {

diff --git a/xla/hlo/ir/hlo_module.h b/xla/hlo/ir/hlo_module.h
@@ -421,13 +421,14 @@ class HloModule {
   HloModuleProto ToProto() const;
   static absl::StatusOr<std::unique_ptr<HloModule>> CreateFromProto(
       const HloModuleProto& proto, const HloModuleConfig& module_config,
-      bool prohibit_empty_literal = true);
+      bool prohibit_empty_literal = true,
+      std::unique_ptr<CompilationEnvironments> comp_envs = nullptr);
 
   // Convert an HloModule to or from a proto that includes module configuration
   HloModuleProtoWithConfig ToProtoWithConfig() const;
   static absl::StatusOr<std::unique_ptr<HloModule>> CreateFromProtoWithConfig(
-      const HloModuleProtoWithConfig& proto,
-      bool prohibit_empty_literal = true);
+      const HloModuleProtoWithConfig& proto, bool prohibit_empty_literal = true,
+      std::unique_ptr<CompilationEnvironments> comp_envs = nullptr);
 
   // Creates and returns an HloModuleConfig with an appropriate program shape
   // for the HLO module in the given proto.