[XLA:GPU] Use HloHardwareIndependentTestBase::CreateDefaultDeviceAssignment.

PiperOrigin-RevId: 839349267

Author: olegshyshkov

xla/tests/BUILD

@@ -3007,14 +3007,12 @@ xla_test(
         ":literal_test_util",
         ":test_utils",
         ":xla_internal_test_main",
-        "//xla:array",
         "//xla:error_spec",
         "//xla:literal",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
         "//xla/hlo/testlib:verified_hlo_module",
-        "//xla/service:computation_placer_hdr",
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_runner",
         "//xla/service/gpu:backend_configs_cc",

xla/tests/collective_ops_e2e_test_base.cc

@@ -22,13 +22,11 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
-#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
@@ -147,7 +145,7 @@ CollectiveOpsE2ETestBase::ExecuteReplicated(
   CHECK(num_partitions > 0 && "expect at least one partition");
 
   DeviceAssignment device_assignment =
-      MakeDeviceAssignment(num_replicas, num_partitions);
+      GetDefaultDeviceAssignment(num_replicas, num_partitions);
   int64_t num_devices = num_replicas * num_partitions;
 
   CHECK(num_devices == arguments.size() &&

xla/tests/collective_ops_e2e_test_base.h

@@ -20,19 +20,15 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
-#include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/types/span.h"
-#include "xla/array.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/literal.h"
 #include "xla/service/backend.h"
-#include "xla/service/computation_placer.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/hlo_runner.h"
 #include "xla/service/hlo_runner_interface.h"
@@ -78,19 +74,6 @@ class CollectiveOpsE2ETestBase : public HloHardwareIndependentTestBase {
            Capability().cuda_compute_capability()->IsAtLeastHopper();
   }
 
-  // Makes an iota device assignment.
-  DeviceAssignment MakeDeviceAssignment(int64_t num_replicas,
-                                        int64_t num_partitions = 1) {
-    DeviceAssignment assn(/*replica_count=*/num_replicas,
-                          /*computation_count=*/num_partitions);
-    for (int64_t i = 0; i < num_replicas; ++i) {
-      for (int64_t j = 0; j < num_partitions; ++j) {
-        assn(i, j) = i * num_partitions + j;
-      }
-    }
-    return assn;
-  }
-
  protected:
   std::unique_ptr<HloRunner> hlo_runner_;
   std::unique_ptr<HloRunner> reference_hlo_runner_;

xla/tests/collective_ops_sharded_unsharded_e2e_test.cc

@@ -26,14 +26,12 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "xla/array.h"
 #include "xla/error_spec.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_sharding.h"
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/literal.h"
-#include "xla/service/computation_placer.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_runner.h"