[xla:cpu] InProcessCommunicator: compute collective operations in par…

…allel using all ranks PiperOrigin-RevId: 730672408
openxla · Feb 25, 2025 · 8d3bdff · 8d3bdff
1 parent 2efc815
commit 8d3bdff
Show file tree

Hide file tree

Showing 3 changed files with 185 additions and 98 deletions.
diff --git a/xla/backends/cpu/collectives/BUILD b/xla/backends/cpu/collectives/BUILD
@@ -142,14 +142,18 @@ cc_library(
         "//xla/service:global_device_id",
         "//xla/service:rendezvous",
         "//xla/stream_executor:device_memory",
+        "//xla/tsl/lib/math:math_util",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/cleanup",
         "@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/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
     ],