[xla:collectives] NFC: Remove NcclCliqueKey alias

PiperOrigin-RevId: 702392147

xla/backends/gpu/collectives/BUILD

@@ -1,12 +1,13 @@
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured")
 load("//xla:xla.bzl", "xla_cc_test")
 load("//xla/stream_executor:build_defs.bzl", "if_gpu_is_configured")
+load("//xla/tsl:tsl.bzl", "internal_visibility")
 load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
 load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
-    default_visibility = [":friends"],
+    default_visibility = internal_visibility([":friends"]),
     licenses = ["notice"],
 )
 

xla/pjrt/gpu/BUILD

@@ -219,12 +219,12 @@ cc_library(
     deps = [
         "//xla:status_macros",
         "//xla:util",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:clique_id",
         "//xla/core/collectives:clique_key",
         "//xla/pjrt/distributed:key_value_store_interface",
         "//xla/service:global_device_id",
         "//xla/service/gpu/runtime:nccl_api",
-        "//xla/service/gpu/runtime:nccl_clique_key",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status:statusor",

xla/pjrt/gpu/nccl_id_store.cc

@@ -21,10 +21,10 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
+#include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/core/collectives/clique_id.h"
 #include "xla/core/collectives/clique_key.h"
 #include "xla/service/gpu/runtime/nccl_api.h"
-#include "xla/service/gpu/runtime/nccl_clique_key.h"
 #include "xla/status_macros.h"
 #include "xla/util.h"
 #include "tsl/platform/casts.h"
@@ -34,7 +34,7 @@ limitations under the License.
 namespace xla {
 
 absl::StatusOr<CliqueId> NcclIdStore::GetNcclUniqueId(const CliqueKey& key) {
-  auto* gpu_key = tsl::down_cast<const gpu::NcclCliqueKey*>(&key);
+  auto* gpu_key = tsl::down_cast<const gpu::GpuCliqueKey*>(&key);
   if (gpu_key == nullptr) {
     return InvalidArgument("Expected GPU clique key");
   }

xla/pjrt/gpu/nccl_id_store.h

@@ -23,15 +23,15 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
+#include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/core/collectives/clique_id.h"
 #include "xla/core/collectives/clique_key.h"
 #include "xla/pjrt/distributed/key_value_store_interface.h"
 #include "xla/service/global_device_id.h"
-#include "xla/service/gpu/runtime/nccl_clique_key.h"
 
 namespace xla {
 
-// A table mapping NcclCliqueKeys to NcclCliqueIds. In a distributed setup the
+// A table mapping GpuCliqueKeys to CliqueIds. In a distributed setup the
 // table of NCCL IDs is kept on the master node (node 0). The node of the first
 // participating device will create the unique id.
 class NcclIdStore {
@@ -51,7 +51,7 @@ class NcclIdStore {
   const std::shared_ptr<KeyValueStoreInterface> kv_store_;
 
   absl::Mutex mu_;
-  absl::flat_hash_map<gpu::NcclCliqueKey, CliqueId> cache_ ABSL_GUARDED_BY(mu_);
+  absl::flat_hash_map<gpu::GpuCliqueKey, CliqueId> cache_ ABSL_GUARDED_BY(mu_);
 };
 
 }  // namespace xla

xla/service/gpu/BUILD

@@ -97,10 +97,10 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//xla:executable_run_options",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:clique_id",
         "//xla/core/collectives:clique_key",
         "//xla/service:global_device_id",
-        "//xla/service/gpu/runtime:nccl_clique_key",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
     ],
@@ -339,6 +339,7 @@ cc_library(
         "//xla:status_macros",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/ffi:attribute_map",
         "//xla/ffi:ffi_api",
         "//xla/ffi/api:c_api",
@@ -375,7 +376,6 @@ cc_library(
         "//xla/service/gpu/runtime:nccl_all_to_all_thunk",
         "//xla/service/gpu/runtime:nccl_api",
         "//xla/service/gpu/runtime:nccl_clique",
-        "//xla/service/gpu/runtime:nccl_clique_key",
         "//xla/service/gpu/runtime:nccl_collective_broadcast_thunk",
         "//xla/service/gpu/runtime:nccl_collective_permute_thunk",
         "//xla/service/gpu/runtime:nccl_collective_thunk",
@@ -575,6 +575,7 @@ cc_library(
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla:util",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
         "//xla/service:executable",
@@ -587,7 +588,6 @@ cc_library(
         "//xla/service/gpu/runtime:annotation",
         "//xla/service/gpu/runtime:for_all_thunks",
         "//xla/service/gpu/runtime:nccl_clique",
-        "//xla/service/gpu/runtime:nccl_clique_key",
         "//xla/service/gpu/runtime:sequential_thunk",
         "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor:device_description",

xla/service/gpu/fusions/BUILD

@@ -69,6 +69,7 @@ cc_library(
         "//xla:status_macros",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/ffi:attribute_map",
         "//xla/ffi:ffi_api",
         "//xla/hlo/ir:hlo",
@@ -95,7 +96,6 @@ cc_library(
         "//xla/service/gpu/runtime:kernel_thunk",
         "//xla/service/gpu/runtime:nccl_all_reduce_thunk",
         "//xla/service/gpu/runtime:nccl_api",
-        "//xla/service/gpu/runtime:nccl_clique_key",
         "//xla/service/gpu/runtime:nccl_collective_thunk",
         "//xla/service/gpu/runtime:thunk",
         "@com_google_absl//absl/algorithm:container",

xla/service/gpu/fusions/custom.cc

@@ -35,6 +35,7 @@ limitations under the License.
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Support/LLVM.h"
+#include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/ffi/attribute_map.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
@@ -64,7 +65,6 @@ limitations under the License.
 #include "xla/service/gpu/runtime/kernel_thunk.h"
 #include "xla/service/gpu/runtime/nccl_all_reduce_thunk.h"
 #include "xla/service/gpu/runtime/nccl_api.h"
-#include "xla/service/gpu/runtime/nccl_clique_key.h"
 #include "xla/service/gpu/runtime/nccl_collective_thunk.h"
 #include "xla/service/gpu/runtime/thunk.h"
 #include "xla/service/gpu/stream_executor_util.h"

xla/service/gpu/gpu_executable.cc

@@ -37,6 +37,7 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/executable_run_options.h"
 #include "xla/hlo/ir/hlo_input_output_alias_config.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -51,7 +52,6 @@ limitations under the License.
 #include "xla/service/gpu/runtime/annotation.h"
 #include "xla/service/gpu/runtime/for_all_thunks.h"
 #include "xla/service/gpu/runtime/nccl_clique.h"
-#include "xla/service/gpu/runtime/nccl_clique_key.h"
 #include "xla/service/gpu/runtime/sequential_thunk.h"
 #include "xla/service/gpu/runtime/thunk.h"
 #include "xla/service/gpu/stream_executor_util.h"
@@ -187,7 +187,7 @@ namespace {
 // Shared resources required for thunk initialization and execution.
 class ResourceRequests : public Thunk::ResourceRequests {
  public:
-  absl::Status AddClique(const NcclCliqueKey& clique_key,
+  absl::Status AddClique(const GpuCliqueKey& clique_key,
                          int32_t num_local_participants) final {
     VLOG(5) << "Add collective clique request: " << clique_key.ToString()
             << "; num_local_participants: " << num_local_participants;
@@ -286,7 +286,7 @@ class ResourceRequests : public Thunk::ResourceRequests {
 
  private:
   struct CliqueRequest {
-    NcclCliqueKey key;
+    GpuCliqueKey key;
     int64_t num_local_participants;
     int64_t id;
   };
@@ -326,7 +326,7 @@ class ResourceRequests : public Thunk::ResourceRequests {
     return cliques;
   }
 
-  absl::flat_hash_map<NcclCliqueKey, CliqueRequest> cliques_;
+  absl::flat_hash_map<GpuCliqueKey, CliqueRequest> cliques_;
 };
 
 absl::Status MaybeSyncAndProfile(const ServiceExecutableRunOptions* run_options,

xla/service/gpu/ir_emitter_unnested.cc

@@ -68,6 +68,7 @@ limitations under the License.
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
+#include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/ffi/api/c_api.h"
 #include "xla/ffi/attribute_map.h"
 #include "xla/ffi/ffi_api.h"
@@ -126,7 +127,6 @@ limitations under the License.
 #include "xla/service/gpu/runtime/nccl_all_to_all_thunk.h"
 #include "xla/service/gpu/runtime/nccl_api.h"
 #include "xla/service/gpu/runtime/nccl_clique.h"
-#include "xla/service/gpu/runtime/nccl_clique_key.h"
 #include "xla/service/gpu/runtime/nccl_collective_broadcast_thunk.h"
 #include "xla/service/gpu/runtime/nccl_collective_permute_thunk.h"
 #include "xla/service/gpu/runtime/nccl_collective_thunk.h"

xla/service/gpu/runtime/BUILD

@@ -61,7 +61,6 @@ cc_library(
         ":nccl_all_reduce_thunk",
         ":nccl_all_to_all_thunk",
         ":nccl_api",
-        ":nccl_clique_key",
         ":nccl_collective_broadcast_thunk",
         ":nccl_collective_thunk",
         ":thunk",
@@ -72,6 +71,7 @@ cc_library(
         "//xla:status_macros",
         "//xla:types",
         "//xla:util",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/backends/gpu/collectives:gpu_collectives",
         "//xla/core/collectives:communicator",
         "//xla/ffi:call_frame",
@@ -212,10 +212,10 @@ cc_library(
     hdrs = ["nccl_api.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
-        ":nccl_clique_key",
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/backends/gpu/collectives:gpu_collectives",
         "//xla/backends/gpu/collectives:nccl_communicator",
         "//xla/core/collectives:clique_id",
@@ -257,9 +257,9 @@ cc_library(
     hdrs = ["nccl_api.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
-        ":nccl_clique_key",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/backends/gpu/collectives:gpu_collectives",
         "//xla/core/collectives:clique_id",
         "//xla/core/collectives:communicator",
@@ -282,12 +282,12 @@ cc_library(
     hdrs = ["nccl_clique.h"],
     deps = [
         ":nccl_api",
-        ":nccl_clique_key",
         "//xla:debug_options_flags",
         "//xla:executable_run_options",
         "//xla:status_macros",
         "//xla:types",
         "//xla:util",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:clique_id",
         "//xla/core/collectives:clique_key",
         "//xla/core/collectives:communicator",
@@ -317,32 +317,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "nccl_clique_key",
-    hdrs = ["nccl_clique_key.h"],
-    compatible_with = get_compatible_with_portable(),
-    deps = [
-        "//xla/backends/gpu/collectives:gpu_clique_key",
-        "//xla/backends/gpu/collectives:gpu_collectives",
-        "//xla/core/collectives:clique_id",
-        "//xla/core/collectives:clique_key",
-        "//xla/core/collectives:rank_id",
-        "//xla/service:global_device_id",
-        "//xla/tsl/lib/gtl:int_type",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/crc:crc32c",
-        "@com_google_absl//absl/hash",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/types:span",
-        "@tsl//tsl/platform:casts",
-        "@tsl//tsl/platform:logging",
-    ],
-)
-
 #===-------------------------------------------------------------------------------------------===//
 # XLA Thunks Runtime
 #===-------------------------------------------------------------------------------------------===//
@@ -868,11 +842,11 @@ cc_library(
     hdrs = ["nccl_all_to_all_thunk.h"],
     deps = [
         ":nccl_api",
-        ":nccl_clique_key",
         ":nccl_collective_thunk",
         ":thunk",
         "//xla:shape_util",
         "//xla:status_macros",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/backends/gpu/collectives:gpu_collectives",
         "//xla/core/collectives:communicator",
         "//xla/hlo/ir:hlo",
@@ -946,12 +920,12 @@ cc_library(
     deps = [
         ":nccl_api",
         ":nccl_clique",
-        ":nccl_clique_key",
         ":thunk",
         "//xla:debug_options_flags",
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/backends/gpu/collectives:gpu_collectives",
         "//xla/core/collectives:communicator",
         "//xla/core/collectives:rank_id",
@@ -995,12 +969,12 @@ cc_library(
     srcs = ["nccl_p2p_thunk_common.cc"],
     hdrs = ["nccl_p2p_thunk_common.h"],
     deps = [
-        ":nccl_clique_key",
         ":nccl_collective_thunk",
         "//xla:executable_run_options",
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:communicator",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
@@ -1023,11 +997,11 @@ cc_library(
     hdrs = ["nccl_recv_thunk.h"],
     deps = [
         ":nccl_api",
-        ":nccl_clique_key",
         ":nccl_collective_thunk",
         ":nccl_p2p_thunk_common",
         ":thunk",
         "//xla:status_macros",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:communicator",
         "//xla/hlo/ir:hlo",
         "//xla/service:collective_ops_utils",
@@ -1049,11 +1023,11 @@ cc_library(
     hdrs = ["nccl_send_thunk.h"],
     deps = [
         ":nccl_api",
-        ":nccl_clique_key",
         ":nccl_collective_thunk",
         ":nccl_p2p_thunk_common",
         ":thunk",
         "//xla:status_macros",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:communicator",
         "//xla/hlo/ir:hlo",
         "//xla/service:collective_ops_utils",
@@ -1075,9 +1049,9 @@ cc_library(
     hdrs = ["nccl_group_thunk.h"],
     deps = [
         ":nccl_api",
-        ":nccl_clique_key",
         ":nccl_collective_thunk",
         ":thunk",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/hlo/ir:hlo",
         "//xla/stream_executor:stream",
         "@com_google_absl//absl/status",
@@ -1189,8 +1163,8 @@ cc_library(
     hdrs = ["thunk.h"],
     deps = [
         ":nccl_clique",
-        ":nccl_clique_key",
         "//xla:executable_run_options",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:communicator",
         "//xla/core/collectives:rank_id",
         "//xla/ffi:execution_context",