Add barrier to JACCL (ml-explore#3459)

Author: Isalia20

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;
 };
 ```
 

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)

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;
+}

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;
 };
 
 /**

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,

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(

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,

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(