[XLA:GPU] Remove unused ExecuteReplicated that require OpaqueExecutable.

olegshyshkov · Google-ML-Automation · commit 2e19703ab56f · 2025-12-02T09:21:09.000-08:00
PiperOrigin-RevId: 839297753
diff --git a/xla/hlo/testlib/hlo_hardware_independent_test_base.cc b/xla/hlo/testlib/hlo_hardware_independent_test_base.cc
@@ -37,8 +37,8 @@ limitations under the License.
 #include "xla/debug_options_flags.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/hlo/ir/hlo_module_group.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/ir/hlo_print_options.h"
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/hlo/testlib/filecheck.h"
diff --git a/xla/tests/collective_ops_e2e_test.cc b/xla/tests/collective_ops_e2e_test.cc
@@ -2792,27 +2792,23 @@ TEST_F(CollectiveOpsTestE2E, OptimizedSubByteAllGatherOnDim0OutputIsCorrect) {
     e {
       a = s4[2,4]{1,0:E(4)} constant({{0,1,2,3},{4,5,5,4}})
       b = s4[4,4]{1,0:E(4)} all-gather(a), dimensions={0}
-    })"));
+    })",
+                                                       kNumReplicas));
 
-  TF_ASSERT_OK_AND_ASSIGN(auto executable, hlo_runner_->CreateExecutable(
-                                               std::move(unoptimized_module),
-                                               /*run_hlo_passes=*/true));
-
-  TF_ASSERT_OK_AND_ASSIGN(const HloModule* const module,
-                          hlo_runner_->HloModuleFromWrapped(executable.get()));
+  TF_ASSERT_OK_AND_ASSIGN(ExecutionResult execution_result,
+                          ExecuteReplicated(std::move(unoptimized_module)));
 
+  const HloModule* module = execution_result.optimized_module;
   EXPECT_THAT(module->entry_computation()->root_instruction(),
               GmockMatch(m::Bitcast(m::AllGatherDone().WithShape(S8, {4, 2}))));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::vector<Literal> result,
-                          ExecuteReplicated(executable.get(), kNumReplicas));
-
   const Literal expected_result =
       LiteralUtil::CreateR2<s4>({{s4(0), s4(1), s4(2), s4(3)},
                                  {s4(4), s4(5), s4(5), s4(4)},
                                  {s4(0), s4(1), s4(2), s4(3)},
                                  {s4(4), s4(5), s4(5), s4(4)}});
 
+  const std::vector<Literal>& result = execution_result.results;
   ASSERT_EQ(result.size(), kNumReplicas);
   for (int i = 0; i < kNumReplicas; ++i) {
     EXPECT_TRUE(LiteralTestUtil::Equal(expected_result, result[i]))
@@ -2833,30 +2829,26 @@ TEST_F(CollectiveOpsTestE2E, OptimizedSubByteAllGatherOnDim1OutputIsCorrect) {
     e {
       a = s4[4,2]{1,0:E(4)} constant({{0,1},{2,3},{4,5},{5,4}})
       b = s4[4,4]{1,0:E(4)} all-gather(a), dimensions={1}
-    })"));
+    })",
+                                                       kNumReplicas));
 
-  TF_ASSERT_OK_AND_ASSIGN(auto executable, hlo_runner_->CreateExecutable(
-                                               std::move(unoptimized_module),
-                                               /*run_hlo_passes=*/true));
-
-  TF_ASSERT_OK_AND_ASSIGN(const HloModule* const module,
-                          hlo_runner_->HloModuleFromWrapped(executable.get()));
+  TF_ASSERT_OK_AND_ASSIGN(ExecutionResult execution_result,
+                          ExecuteReplicated(std::move(unoptimized_module)));
 
+  const HloModule* module = execution_result.optimized_module;
   const HloInstruction* root = module->entry_computation()->root_instruction();
   EXPECT_THAT(root, GmockMatch(m::Fusion(
                         m::Bitcast(m::AllGatherDone().WithShape(S8, {2, 4})))));
   EXPECT_THAT(root->fused_expression_root(),
               GmockMatch(m::Transpose(m::Parameter())));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::vector<Literal> result,
-                          ExecuteReplicated(executable.get(), kNumReplicas));
-
   const Literal expected_result =
       LiteralUtil::CreateR2<s4>({{s4(0), s4(1), s4(0), s4(1)},
                                  {s4(2), s4(3), s4(2), s4(3)},
                                  {s4(4), s4(5), s4(4), s4(5)},
                                  {s4(5), s4(4), s4(5), s4(4)}});
 
+  const std::vector<Literal>& result = execution_result.results;
   ASSERT_EQ(result.size(), kNumReplicas);
   for (int i = 0; i < kNumReplicas; ++i) {
     EXPECT_TRUE(LiteralTestUtil::Equal(expected_result, result[i]))
diff --git a/xla/tests/collective_ops_e2e_test_base.cc b/xla/tests/collective_ops_e2e_test_base.cc
@@ -116,67 +116,6 @@ CollectiveOpsE2ETestBase::CollectiveOpsE2ETestBase() {
       reference_platform, /*intra_op_parallelism_threads=*/0);
 }
 
-absl::StatusOr<std::vector<Literal>>
-CollectiveOpsE2ETestBase::ExecuteReplicated(
-    absl::AnyInvocable<OpaqueExecutable*(int64_t)> executable_provider,
-    absl::AnyInvocable<int64_t(int64_t)> argument_count_provider,
-    absl::AnyInvocable<const Literal*(int64_t, int64_t)> argument_provider,
-    const int64_t num_replicas, const bool run_hlo_passes,
-    DeviceAssignment* const device_assignment) {
-  // TODO(b/441865120): Use designated initializers this once XLA moves to
-  // C++20.
-  HloRunnerInterface::ReplicatedExecuteOptions options;
-  options.num_replicas = num_replicas;
-  options.run_hlo_passes = run_hlo_passes;
-  options.use_threads = true;
-
-  return hlo_runner_->ExecuteReplicated(
-      std::move(executable_provider), std::move(argument_count_provider),
-      std::move(argument_provider), std::move(options), device_assignment);
-}
-
-absl::StatusOr<std::vector<Literal>>
-CollectiveOpsE2ETestBase::ExecuteReplicated(
-    std::unique_ptr<HloModule> module,
-    const absl::Span<const Literal* const> arguments,
-    const int64_t num_replicas, DeviceAssignment* const device_assignment,
-    const bool run_hlo_passes, const bool use_threads) {
-  // TODO(b/441865120): Use designated initializers this once XLA moves to
-  // C++20.
-  HloRunnerInterface::ReplicatedExecuteOptions options;
-  options.num_replicas = num_replicas;
-  options.arguments = {arguments.begin(), arguments.end()};
-  options.run_hlo_passes = run_hlo_passes;
-  options.use_threads = use_threads;
-
-  return hlo_runner_->ExecuteReplicated(std::move(module), std::move(options),
-                                        device_assignment);
-}
-
-absl::StatusOr<std::vector<Literal>>
-CollectiveOpsE2ETestBase::ExecuteReplicated(
-    std::unique_ptr<HloModule> module,
-    const std::vector<std::vector<Literal*>> arguments,
-    DeviceAssignment* const device_assignment, const int64_t num_replicas,
-    const bool run_hlo_passes) {
-  CHECK(num_replicas > 0 && "expect at least one replica");
-  CHECK(num_replicas == arguments.size() &&
-        "expect arguments for each replica");
-  int64_t argument_count = arguments.front().size();
-  TF_ASSIGN_OR_RETURN(
-      const std::unique_ptr<OpaqueExecutable> executable,
-      hlo_runner_->CreateExecutable(std::move(module), run_hlo_passes));
-  return ExecuteReplicated(
-      /*executable_provider=*/[&](int64_t) { return executable.get(); },
-      /*argument_count_provider=*/[&](int64_t) { return argument_count; },
-      /*argument_provider=*/
-      [&](int64_t replica_idx, int64_t argument_idx) -> const Literal* {
-        return arguments[replica_idx][argument_idx];
-      },
-      num_replicas, /*run_hlo_passes=*/run_hlo_passes,
-      /*device_assignment=*/device_assignment);
-}
-
 absl::StatusOr<CollectiveOpsE2ETestBase::ExecutionResult>
 CollectiveOpsE2ETestBase::ExecuteReplicated(std::unique_ptr<HloModule> module) {
   return ExecuteReplicated(std::move(module),
@@ -224,9 +163,16 @@ CollectiveOpsE2ETestBase::ExecuteReplicated(
       execution_result.optimized_module,
       hlo_runner_->HloModuleFromWrapped(execution_result.executable.get()));
 
+  // TODO(b/441865120): Use designated initializers this once XLA moves to
+  // C++20.
+  HloRunnerInterface::ReplicatedExecuteOptions options;
+  options.num_replicas = num_devices;
+  options.run_hlo_passes = run_hlo_passes;
+  options.use_threads = true;
+
   TF_ASSIGN_OR_RETURN(
       execution_result.results,
-      ExecuteReplicated(
+      hlo_runner_->ExecuteReplicated(
           /*executable_provider=*/
           [&](int64_t) { return execution_result.executable.get(); },
           /*argument_count_provider=*/
@@ -235,21 +181,10 @@ CollectiveOpsE2ETestBase::ExecuteReplicated(
           [&](int64_t replica_idx, int64_t argument_idx) -> const Literal* {
             return arguments[replica_idx][argument_idx];
           },
-          /*num_replicas=*/num_devices,
-          /*run_hlo_passes=*/run_hlo_passes,
+          std::move(options),
           /*device_assignment=*/&device_assignment));
-  return execution_result;
-}
 
-absl::StatusOr<std::vector<Literal>>
-CollectiveOpsE2ETestBase::ExecuteReplicated(OpaqueExecutable* executable,
-                                            int64_t num_replicas) {
-  DeviceAssignment device_assignment = MakeDeviceAssignment(num_replicas);
-  return ExecuteReplicated(
-      /*executable_provider*/ [&](int64_t) { return executable; },
-      /*argument_count_provider*/ [](int64_t) { return 0; },
-      /*argument_provider*/ [](int64_t, int64_t) { return nullptr; },
-      num_replicas, /*run_hlo_passes=*/false, &device_assignment);
+  return execution_result;
 }
 
 DebugOptions CollectiveOpsWithFlagsBase::GetDebugOptionsForTest() const {
diff --git a/xla/tests/collective_ops_e2e_test_base.h b/xla/tests/collective_ops_e2e_test_base.h
@@ -54,27 +54,6 @@ class CollectiveOpsE2ETestBase : public HloHardwareIndependentTestBase {
     const HloModule* optimized_module;
   };
 
-  absl::StatusOr<std::vector<Literal>> ExecuteReplicated(
-      absl::AnyInvocable<OpaqueExecutable*(int64_t)> executable_provider,
-      absl::AnyInvocable<int64_t(int64_t)> argument_count_provider,
-      absl::AnyInvocable<const Literal*(int64_t, int64_t)> argument_provider,
-      int64_t num_replicas, bool run_hlo_passes,
-      DeviceAssignment* device_assignment);
-
-  absl::StatusOr<std::vector<Literal>> ExecuteReplicated(
-      std::unique_ptr<HloModule> module,
-      absl::Span<const Literal* const> arguments, int64_t num_replicas,
-      DeviceAssignment* vice_assignment, bool run_hlo_passes, bool use_threads);
-
-  absl::StatusOr<std::vector<Literal>> ExecuteReplicated(
-      std::unique_ptr<HloModule> module,
-      std::vector<std::vector<Literal*>> arguments,
-      DeviceAssignment* device_assignment, int64_t num_replicas,
-      bool run_hlo_passes);
-
-  absl::StatusOr<std::vector<Literal>> ExecuteReplicated(
-      OpaqueExecutable* executable, int64_t num_replicas);
-
   absl::StatusOr<ExecutionResult> ExecuteReplicated(
       std::unique_ptr<HloModule> module);
 
diff --git a/xla/tests/collective_ops_sharded_unsharded_e2e_test.cc b/xla/tests/collective_ops_sharded_unsharded_e2e_test.cc
@@ -61,13 +61,15 @@ class CollectiveOpsTestE2EShardedUnsharded : public CollectiveOpsE2ETestBase {
                    << " available)";
     }
 
-    TF_ASSERT_OK_AND_ASSIGN(std::vector<Literal> ref_results,
+    TF_ASSERT_OK_AND_ASSIGN(ExecutionResult ref_execution_result,
                             ExecuteUnsharded(hlo_text));
+    const std::vector<Literal>& ref_results = ref_execution_result.results;
     ASSERT_EQ(ref_results.size(), 1);
 
     TF_ASSERT_OK_AND_ASSIGN(
-        std::vector<Literal> results,
+        ExecutionResult execution_result,
         ExecuteSharded(hlo_text, num_partitions, enable_enzyme_comms_opt));
+    const std::vector<Literal>& results = execution_result.results;
     ASSERT_EQ(results.size(), num_partitions);
 
     ErrorSpec error_spec{1e-4, 1e-4};
@@ -77,7 +79,7 @@ class CollectiveOpsTestE2EShardedUnsharded : public CollectiveOpsE2ETestBase {
 
  private:
   // Execute the unsharded case.
-  absl::StatusOr<std::vector<Literal>> ExecuteUnsharded(
+  absl::StatusOr<ExecutionResult> ExecuteUnsharded(
       const std::string& hlo_text) {
     // Create the unsharded reference case by removing the sharding metadata
     // from the HLO string.
@@ -90,9 +92,12 @@ class CollectiveOpsTestE2EShardedUnsharded : public CollectiveOpsE2ETestBase {
     DebugOptions ref_opts = GetDebugOptionsForTest();
     ref_opts.set_xla_gpu_enable_triton_gemm(false);
     ref_config.set_debug_options(ref_opts);
-    ref_config.set_num_partitions(1);
     TF_ASSIGN_OR_RETURN(std::unique_ptr<VerifiedHloModule> ref_module,
                         ParseAndReturnVerifiedModule(hlo_text_ref, ref_config));
+
+    ref_module->mutable_config().set_replica_count(1);
+    ref_module->mutable_config().set_num_partitions(1);
+
     const int64_t num_params =
         ref_module->entry_computation()->num_parameters();
 
@@ -102,17 +107,11 @@ class CollectiveOpsTestE2EShardedUnsharded : public CollectiveOpsE2ETestBase {
       ref_fake_ptrs[i] = &fake_args[i];
     }
 
-    DeviceAssignment ref_assn(/*replica_count=*/1,
-                              /*computation_count=*/1);
-    ref_assn(0, 0) = 0;
-    return ExecuteReplicated(std::move(ref_module), ref_fake_ptrs,
-                             /*num_replicas=*/1, &ref_assn,
-                             /*run_hlo_passes=*/true,
-                             /*use_threads=*/true);
+    return ExecuteReplicated(std::move(ref_module), ref_fake_ptrs);
   }
 
   // Execute the sharded case.
-  absl::StatusOr<std::vector<Literal>> ExecuteSharded(
+  absl::StatusOr<ExecutionResult> ExecuteSharded(
       const std::string& hlo_text, int64_t num_partitions,
       bool enable_enzyme_comms_opt = false) {
     HloModuleConfig config = GetModuleConfigForTest();
@@ -182,21 +181,14 @@ class CollectiveOpsTestE2EShardedUnsharded : public CollectiveOpsE2ETestBase {
       }
     }
 
-    DeviceAssignment assn(/*replica_count=*/1,
-                          /*computation_count=*/num_partitions);
-    for (int64_t i = 0; i < num_partitions; ++i) {
-      assn(0, i) = i;
-    }
-    return ExecuteReplicated(std::move(module), fake_ptrs, &assn,
-                             num_partitions,
-                             /*run_hlo_passes=*/true);
+    return ExecuteReplicated(std::move(module), fake_ptrs);
   }
 
   // Slice the unsharded reference results and compare to the sharded case.
   void CompareShardedUnsharded(const std::string& hlo_text,
                                int64_t num_partitions,
-                               std::vector<Literal>& ref_results,
-                               std::vector<Literal>& results,
+                               const std::vector<Literal>& ref_results,
+                               const std::vector<Literal>& results,
                                ErrorSpec& error_spec) {
     HloModuleConfig config = GetModuleConfigForTest();
     DebugOptions opts = GetDebugOptionsForTest();
