From 6fc00d2c10125cccae3dd700a4b8810b3ec8440c Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Tue, 5 Nov 2024 11:34:55 -0800
Subject: [PATCH] Add rudimentary barrier

---
 mlx/distributed/sockets/sockets.cpp | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/mlx/distributed/sockets/sockets.cpp b/mlx/distributed/sockets/sockets.cpp
index 0b9b38c58..1cf10a11b 100644
--- a/mlx/distributed/sockets/sockets.cpp
+++ b/mlx/distributed/sockets/sockets.cpp
@@ -385,6 +385,19 @@ Group Group::split(int color, int key) {
   throw std::runtime_error("Splitting not supported yet");
 }
 
+void Group::barrier() {
+  char buff[128];
+  std::memset(buff, 1, 128);
+
+  auto group = std::static_pointer_cast<GroupImpl>(raw_group());
+  int size = group->size();
+  int rank = group->rank();
+
+  for (int distance = 1; distance <= size / 2; distance *= 2) {
+    group->send_recv_sum<char>(buff, 128, rank ^ distance);
+  }
+}
+
 Group init(bool strict /* = false */) {
   static std::shared_ptr<GroupImpl> global_group = nullptr;