Remove dependency on -hlo-import-all-computations flag when exporti…

…ng `entry_computation_layout`. PiperOrigin-RevId: 671527962
openxla · Sep 17, 2024 · 2cac71b · 2cac71b
1 parent 161e7db
commit 2cac71b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/xla/translate/hlo_to_mhlo/hlo_module_importer.cc b/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
@@ -55,6 +55,8 @@ absl::Status HloModuleImporter::Import(const HloModule& hlo_module) {
 
   ImportCrossProgramPrefetches(hlo_module, module,
                                flatten_computation_args_result_, builder_);
+  if (!flatten_computation_args_result_)
+    ImportEntryComputationLayoutAndTiles(hlo_module, module, builder_);
   ImportFrontendAttributes(hlo_module, module, builder_);
   ImportInputOutputAlias(hlo_module, module, builder_);
   ImportIsDynamic(hlo_module, module, builder_);
@@ -71,19 +73,18 @@ absl::Status HloModuleImporter::Import(const HloModule& hlo_module) {
     return HloFunctionImporter::ImportAsFunc(
                *hlo_module.entry_computation(), symbol_table_, &function_map_,
                &builder_,
-               /*is_main*/ true, flatten_computation_args_result_)
+               /*is_main=*/true, flatten_computation_args_result_)
         .status();
 
   auto* module_entry_computation = hlo_module.entry_computation();
   for (const auto* computation : hlo_module.computations())
     TF_RETURN_IF_ERROR(HloFunctionImporter::ImportAsFunc(
                            *computation, symbol_table_, &function_map_,
                            &builder_,
-                           /*is_main*/ computation == module_entry_computation,
+                           /*is_main=*/computation == module_entry_computation,
                            flatten_computation_args_result_)
                            .status());
 
-  ImportEntryComputationLayoutAndTiles(hlo_module, module, builder_);
   return absl::OkStatus();
 }
 

diff --git a/xla/translate/hlo_to_mhlo/tests/import_entry_computation_layout.hlo b/xla/translate/hlo_to_mhlo/tests/import_entry_computation_layout.hlo
@@ -1,4 +1,4 @@
-// RUN: xla-translate -hlo-text-to-mlir-hlo -with-layouts -print-layouts -hlo-import-all-computations -split-input-file -verify-diagnostics %s -o - | FileCheck %s
+// RUN: xla-translate -hlo-text-to-mlir-hlo -with-layouts -print-layouts -split-input-file -verify-diagnostics %s -o - | FileCheck %s
 
 HloModule entry, entry_computation_layout={(
   f32[2,3,4]{0,1,2:T(2,128)},