[xla:cpu] Add CpuClique to XLA:CPU collectives and use generic collec…

…tives APIs to acquire communicator in CollectiveThunk

Implement Cliques support for XLA:CPU collectives for consistency with XLA:GPU. Further unification will be in followup CLs.

PiperOrigin-RevId: 713052194

xla/backends/cpu/collectives/BUILD

@@ -14,6 +14,59 @@ package_group(
     ],
 )
 
+cc_library(
+    name = "cpu_clique_key",
+    srcs = ["cpu_clique_key.cc"],
+    hdrs = ["cpu_clique_key.h"],
+    deps = [
+        "//xla/core/collectives:clique_key",
+        "//xla/service:global_device_id",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@tsl//tsl/platform:casts",
+    ],
+)
+
+cc_library(
+    name = "cpu_clique",
+    srcs = ["cpu_clique.cc"],
+    hdrs = ["cpu_clique.h"],
+    deps = [
+        ":cpu_clique_key",
+        "//xla/core/collectives:clique",
+        "//xla/core/collectives:communicator",
+        "//xla/core/collectives:rank_id",
+        "//xla/tsl/platform:logging",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
+cc_library(
+    name = "cpu_cliques",
+    srcs = ["cpu_cliques.cc"],
+    hdrs = ["cpu_cliques.h"],
+    deps = [
+        ":cpu_clique",
+        ":cpu_clique_key",
+        ":cpu_collectives",
+        "//xla:util",
+        "//xla/core/collectives:clique",
+        "//xla/core/collectives:communicator",
+        "//xla/core/collectives:rank_id",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:logging",
+        "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:node_hash_map",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
 cc_library(
     name = "cpu_collectives",
     srcs = ["cpu_collectives.cc"],
@@ -23,14 +76,17 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/core/collectives",
+        "//xla/core/collectives:clique_id",
         "//xla/core/collectives:collectives_registry",
         "//xla/core/collectives:communicator",
+        "//xla/core/collectives:rank_id",
         "//xla/service:collective_ops_utils",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:span",
         "@tsl//tsl/platform:casts",
     ],
 )

xla/backends/cpu/collectives/cpu_clique.cc

@@ -0,0 +1,59 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/backends/cpu/collectives/cpu_clique.h"
+
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "xla/backends/cpu/collectives/cpu_clique_key.h"
+#include "xla/core/collectives/clique.h"
+#include "xla/core/collectives/communicator.h"
+#include "xla/core/collectives/rank_id.h"
+#include "xla/tsl/platform/logging.h"
+
+namespace xla::cpu {
+
+CpuClique::CpuClique(CpuCliqueKey key) : Clique({}), key_(std::move(key)) {}
+
+std::string CpuClique::DebugString() const {
+  std::string out =
+      absl::StrFormat("key: %s; size: %d; communicators: ", key_.ToString(),
+                      num_communicators());
+  int32_t cnt = 0;
+  ForEachComm([&](RankId rank, Communicator* comm) {
+    if (cnt++) absl::StrAppend(&out, ", ");
+    absl::StrAppendFormat(&out, "[rank=%d, comm=%s]", rank.value(),
+                          comm->ToString());
+  });
+  return out;
+}
+
+absl::Status CpuClique::HealthCheck() const {
+  absl::Status health_check = absl::OkStatus();
+  ForEachComm([&health_check](RankId rank, Communicator* comm) {
+    if (auto s = comm->HealthCheck(); !s.ok()) {
+      LOG(ERROR) << "CPU communicator error (rank " << rank << "): " << s;
+      if (health_check.ok()) health_check = std::move(s);  // return first error
+    }
+  });
+  return health_check;
+}
+
+}  // namespace xla::cpu

xla/backends/cpu/collectives/cpu_clique.h

@@ -0,0 +1,42 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_CPU_COLLECTIVES_CPU_CLIQUE_H_
+#define XLA_BACKENDS_CPU_COLLECTIVES_CPU_CLIQUE_H_
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "xla/backends/cpu/collectives/cpu_clique_key.h"
+#include "xla/core/collectives/clique.h"
+
+namespace xla::cpu {
+
+// A group of CPU communicators making up a clique.
+class CpuClique final : public Clique {
+ public:
+  explicit CpuClique(CpuCliqueKey key);
+
+  absl::Status HealthCheck() const final;
+
+  std::string DebugString() const final;
+
+ private:
+  CpuCliqueKey key_;
+};
+
+}  // namespace xla::cpu
+
+#endif  // XLA_BACKENDS_CPU_COLLECTIVES_CPU_CLIQUE_H_

xla/backends/cpu/collectives/cpu_clique_key.cc

@@ -0,0 +1,59 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/backends/cpu/collectives/cpu_clique_key.h"
+
+#include <string>
+#include <utility>
+
+#include "absl/algorithm/container.h"
+#include "absl/hash/hash.h"
+#include "absl/strings/str_format.h"
+#include "xla/core/collectives/clique_key.h"
+#include "xla/service/global_device_id.h"
+#include "tsl/platform/casts.h"
+
+namespace xla::cpu {
+
+bool CpuCliqueKey::IsSubsetOf(const CliqueKey& other) const {
+  auto* other_cpu = tsl::down_cast<const CpuCliqueKey*>(&other);
+  if (other_cpu == nullptr) return false;
+
+  return absl::c_all_of(devices(), [&](GlobalDeviceId id) {
+    return absl::c_linear_search(other_cpu->devices(), id);
+  });
+}
+
+std::string CpuCliqueKey::ToString() const {
+  return absl::StrFormat("devices=[%s]", GlobalDeviceIdsToString(devices()));
+}
+
+void CpuCliqueKey::HashValue(absl::HashState state) const {
+  absl::HashState::combine(std::move(state), devices());
+}
+
+bool operator==(const CpuCliqueKey& a, const CpuCliqueKey& b) {
+  return a.devices() == b.devices();
+}
+
+bool operator<(const CpuCliqueKey& a, const CpuCliqueKey& b) {
+  return a.devices() < b.devices();
+}
+
+bool operator>(const CpuCliqueKey& a, const CpuCliqueKey& b) {
+  return a.devices() > b.devices();
+}
+
+}  // namespace xla::cpu

xla/backends/cpu/collectives/cpu_clique_key.h

@@ -0,0 +1,44 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_CPU_COLLECTIVES_CPU_CLIQUE_KEY_H_
+#define XLA_BACKENDS_CPU_COLLECTIVES_CPU_CLIQUE_KEY_H_
+
+#include <string>
+
+#include "absl/hash/hash.h"
+#include "xla/core/collectives/clique_key.h"
+
+namespace xla::cpu {
+
+// Clique key for identifying a particular CPU collectives clique.
+class CpuCliqueKey final : public CliqueKey {
+ public:
+  using CliqueKey::CliqueKey;
+
+  bool IsSubsetOf(const CliqueKey& other) const final;
+  std::string ToString() const final;
+
+  friend bool operator==(const CpuCliqueKey& a, const CpuCliqueKey& b);
+  friend bool operator<(const CpuCliqueKey& a, const CpuCliqueKey& b);
+  friend bool operator>(const CpuCliqueKey& a, const CpuCliqueKey& b);
+
+ private:
+  void HashValue(absl::HashState state) const final;
+};
+
+}  // namespace xla::cpu
+
+#endif  // XLA_BACKENDS_CPU_COLLECTIVES_CPU_CLIQUE_KEY_H_