From e8ebdebeeb655feaa85a51f6b24ece5b6d5518d1 Mon Sep 17 00:00:00 2001
From: Irakli Salia <65120973+Isalia20@users.noreply.github.com>
Date: Tue, 28 Apr 2026 20:39:56 +0400
Subject: [PATCH] Add barrier to JACCL (#3459)

---
 mlx/distributed/jaccl/lib/README.md           |  5 +++
 .../jaccl/lib/examples/CMakeLists.txt         |  1 +
 .../jaccl/lib/examples/minimal_barrier.cpp    | 42 +++++++++++++++++++
 mlx/distributed/jaccl/lib/jaccl/group.h       |  1 +
 mlx/distributed/jaccl/lib/jaccl/mesh.cpp      |  5 +++
 mlx/distributed/jaccl/lib/jaccl/mesh.h        |  2 +
 mlx/distributed/jaccl/lib/jaccl/ring.cpp      |  5 +++
 mlx/distributed/jaccl/lib/jaccl/ring.h        |  2 +
 8 files changed, 63 insertions(+)
 create mode 100644 mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp

diff --git a/mlx/distributed/jaccl/lib/README.md b/mlx/distributed/jaccl/lib/README.md
index 32b2defa..d682a06b 100644
--- a/mlx/distributed/jaccl/lib/README.md
+++ b/mlx/distributed/jaccl/lib/README.md
@@ -29,6 +29,8 @@ in macOS 26.2.
 - **Point-to-Point Operations**:
   - `send`: Send data to a specific node
   - `recv`: Receive data from a specific node
+- **Synchronization**:
+  - `barrier`: Block until all nodes in the group reach this point
 - **Type Support**: Bool, Int8-64, UInt8-64, Float16, BFloat16, Float32,
   Float64, Complex64
 
@@ -286,6 +288,9 @@ class Group {
   // Simple send/recv primitives.
   virtual void send(const void* input, size_t n_bytes, int dst) = 0;
   virtual void recv(void* output, size_t n_bytes, int src) = 0;
+
+  // Block until every rank reaches this point.
+  virtual void barrier() = 0;
 };
 ```
 
diff --git a/mlx/distributed/jaccl/lib/examples/CMakeLists.txt b/mlx/distributed/jaccl/lib/examples/CMakeLists.txt
index 48b7522d..869f5dc2 100644
--- a/mlx/distributed/jaccl/lib/examples/CMakeLists.txt
+++ b/mlx/distributed/jaccl/lib/examples/CMakeLists.txt
@@ -35,6 +35,7 @@ endfunction()
 # Examples
 build_example(minimal_env.cpp)
 build_example(minimal_cfg.cpp)
+build_example(minimal_barrier.cpp)
 
 # Benchmarks
 build_example(allreduce_bench.cpp)
diff --git a/mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp b/mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp
new file mode 100644
index 00000000..0d94ffb6
--- /dev/null
+++ b/mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp
@@ -0,0 +1,42 @@
+// Copyright © 2026 Apple Inc.
+//
+// Exercises Group::barrier(). Ranks arrive at the barrier at staggered times;
+// after the barrier returns we do a small all_sum to confirm the group is
+// healthy and that barrier() carried the correct fence semantics.
+
+#include <chrono>
+#include <iostream>
+#include <thread>
+
+#include <jaccl/jaccl.h>
+
+int main() {
+  auto group = jaccl::init();
+  if (!group) {
+    std::cerr << "Failed to initialize JACCL" << std::endl;
+    return 1;
+  }
+
+  int rank = group->rank();
+  int size = group->size();
+
+  std::this_thread::sleep_for(std::chrono::milliseconds(100 * rank));
+  std::cout << "rank " << rank << " entering barrier" << std::endl;
+
+  group->barrier();
+
+  std::cout << "rank " << rank << " exited barrier" << std::endl;
+
+  int in = rank + 1;
+  int out = 0;
+  group->all_sum(&in, &out, sizeof(in), jaccl::Int32);
+  int expected = size * (size + 1) / 2;
+  if (out != expected) {
+    std::cerr << "rank " << rank << ": post-barrier all_sum mismatch (got "
+              << out << ", expected " << expected << ")" << std::endl;
+    return 1;
+  }
+  std::cout << "rank " << rank << ": post-barrier all_sum OK (" << out << ")"
+            << std::endl;
+  return 0;
+}
diff --git a/mlx/distributed/jaccl/lib/jaccl/group.h b/mlx/distributed/jaccl/lib/jaccl/group.h
index 495458c5..758109c2 100644
--- a/mlx/distributed/jaccl/lib/jaccl/group.h
+++ b/mlx/distributed/jaccl/lib/jaccl/group.h
@@ -30,6 +30,7 @@ class Group {
 
   virtual void send(const void* input, size_t n_bytes, int dst) = 0;
   virtual void recv(void* output, size_t n_bytes, int src) = 0;
+  virtual void barrier() = 0;
 };
 
 /**
diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp
index 6ea1da58..d1201143 100644
--- a/mlx/distributed/jaccl/lib/jaccl/mesh.cpp
+++ b/mlx/distributed/jaccl/lib/jaccl/mesh.cpp
@@ -184,6 +184,11 @@ void MeshGroup::recv(void* output, size_t n_bytes, int src) {
   mesh_.recv(static_cast<char*>(output), n_bytes, src);
 }
 
+void MeshGroup::barrier() {
+  uint8_t b = 0;
+  all_sum(&b, &b, sizeof(b), Dtype::UInt8);
+}
+
 template <typename T, typename ReduceOp>
 void MeshGroup::all_reduce(
     const void* input,
diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh.h b/mlx/distributed/jaccl/lib/jaccl/mesh.h
index 3184f1fb..2ced2441 100644
--- a/mlx/distributed/jaccl/lib/jaccl/mesh.h
+++ b/mlx/distributed/jaccl/lib/jaccl/mesh.h
@@ -47,6 +47,8 @@ class MeshGroup : public Group {
   void send(const void* input, size_t n_bytes, int dst) override;
   void recv(void* output, size_t n_bytes, int src) override;
 
+  void barrier() override;
+
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(
diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/mlx/distributed/jaccl/lib/jaccl/ring.cpp
index 96ddb0e0..541e825b 100644
--- a/mlx/distributed/jaccl/lib/jaccl/ring.cpp
+++ b/mlx/distributed/jaccl/lib/jaccl/ring.cpp
@@ -190,6 +190,11 @@ void RingGroup::recv(void* output, size_t n_bytes, int src) {
   ring_.recv(static_cast<char*>(output), n_bytes, src, n_conns_);
 }
 
+void RingGroup::barrier() {
+  uint8_t b = 0;
+  all_sum(&b, &b, sizeof(b), Dtype::UInt8);
+}
+
 template <typename T, typename ReduceOp>
 void RingGroup::all_reduce(
     const void* input,
diff --git a/mlx/distributed/jaccl/lib/jaccl/ring.h b/mlx/distributed/jaccl/lib/jaccl/ring.h
index f0bc1ba9..016c6af6 100644
--- a/mlx/distributed/jaccl/lib/jaccl/ring.h
+++ b/mlx/distributed/jaccl/lib/jaccl/ring.h
@@ -48,6 +48,8 @@ class RingGroup : public Group {
   void send(const void* input, size_t n_bytes, int dst) override;
   void recv(void* output, size_t n_bytes, int src) override;
 
+  void barrier() override;
+
  private:
   template <typename T, typename ReduceOp>
   void all_reduce(