Faster synchronization Fence primitive

benchmarks/python/synchronize_bench.py

@@ -0,0 +1,32 @@
+import time
+
+import mlx.core as mx
+
+world = mx.distributed.init()
+
+a = mx.ones((5, 5), mx.int32)
+its = 10
+its_per_eval = 100
+
+
+def fn(x):
+    for _ in range(its_per_eval):
+        x = mx.distributed.all_sum(x)
+        x = x - 1
+    return x
+
+
+# warmup
+for _ in range(5):
+    x = fn(a)
+    assert mx.array_equal(x, mx.ones_like(x))
+
+tic = time.perf_counter()
+
+for _ in range(its):
+    x = fn(a)
+    mx.eval(x)
+
+toc = time.perf_counter()
+ms = 1000 * (toc - tic) / (its * its_per_eval)
+print(f"Time per iteration {ms:.6f} (ms)")

mlx/backend/metal/CMakeLists.txt

@@ -89,6 +89,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp

mlx/backend/metal/device.cpp

@@ -134,6 +134,13 @@ CommandEncoder::~CommandEncoder() {
   enc_->release();
 }
 
+void CommandEncoder::set_buffer(
+    const MTL::Buffer* buf,
+    int idx,
+    int64_t offset /* = 0 */) {
+  enc_->setBuffer(buf, offset, idx);
+}
+
 void CommandEncoder::set_input_array(
     const array& a,
     int idx,
@@ -155,6 +162,10 @@ void CommandEncoder::set_output_array(
     int64_t offset /* = 0 */) {
   // Add barriers before adding the output to the output set
   set_input_array(a, idx, offset);
+  register_output_array(a);
+}
+
+void CommandEncoder::register_output_array(array& a) {
   all_outputs_.insert(a.buffer().ptr());
   auto buf = static_cast<MTL::Resource*>(a.buffer().ptr());
   if (concurrent_) {
@@ -189,6 +200,10 @@ void CommandEncoder::dispatch_threads(
   enc_->dispatchThreads(grid_dims, group_dims);
 }
 
+void CommandEncoder::barrier() {
+  enc_->memoryBarrier(MTL::BarrierScopeBuffers);
+}
+
 Device::Device() {
   auto pool = new_scoped_memory_pool();
   device_ = load_device();

mlx/backend/metal/device.h

@@ -60,9 +60,11 @@ struct CommandEncoder {
 
   void set_input_array(const array& a, int idx, int64_t offset = 0);
   void set_output_array(array& a, int idx, int64_t offset = 0);
+  void register_output_array(array& a);
   void dispatch_threadgroups(MTL::Size grid_dims, MTL::Size group_dims);
   void dispatch_threads(MTL::Size grid_dims, MTL::Size group_dims);
   void maybeInsertBarrier();
+  void set_buffer(const MTL::Buffer* buf, int idx, int64_t offset = 0);
 
   void set_compute_pipeline_state(MTL::ComputePipelineState* kernel) {
     enc_->setComputePipelineState(kernel);
@@ -110,6 +112,8 @@ struct CommandEncoder {
     return all_outputs_;
   };
 
+  void barrier();
+
  private:
   MTL::ComputeCommandEncoder* enc_;
   bool needs_barrier_{false};

mlx/backend/metal/distributed.cpp

@@ -6,6 +6,7 @@
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/event.h"
+#include "mlx/backend/metal/fence.h"
 #include "mlx/distributed/ops.h"
 #include "mlx/distributed/primitives.h"
 #include "mlx/scheduler.h"
@@ -26,23 +27,28 @@ void AllReduce::eval_gpu(
   assert(outputs.size() == 1);
 
   auto& in = inputs[0];
+
+  Fence f{stream()};
+
+  if (in.event().valid()) {
+    f.update_gpu(in);
+  }
+
   auto& out = outputs[0];
   if (in.is_donatable()) {
     out.move_shared_buffer(in);
   } else {
     out.set_data(allocator::malloc_or_wait(out.nbytes()));
   }
+  f.wait_gpu(out);
 
-  auto e = Event(stream());
-  e.set_value(1);
-  signal_and_wait(in.event(), e);
   auto task = [in = in,
                out = out,
-               e = std::move(e),
+               f = std::move(f),
                reduce_type = reduce_type_,
                group = group()]() mutable {
     if (in.event().valid()) {
-      in.event().wait();
+      f.wait();
     }
     switch (reduce_type) {
       case Sum:
@@ -52,7 +58,7 @@ void AllReduce::eval_gpu(
       default:
         throw std::runtime_error("Only all reduce sum is supported for now");
     }
-    e.signal();
+    f.update();
   };
   scheduler::enqueue(detail::communication_stream(), std::move(task));
 }
@@ -67,17 +73,20 @@ void AllGather::eval_gpu(
 
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
-  auto e = Event(stream());
-  e.set_value(1);
-  signal_and_wait(in.event(), e);
+  Fence f{stream()};
+
+  if (in.event().valid()) {
+    f.update_gpu(in);
+  }
+  f.wait_gpu(out);
 
   auto task =
-      [in = in, out = out, e = std::move(e), group = group()]() mutable {
+      [in = in, out = out, f = std::move(f), group = group()]() mutable {
         if (in.event().valid()) {
-          in.event().wait();
+          f.wait();
         }
         distributed::detail::all_gather(group, in, out);
-        e.signal();
+        f.update();
       };
   scheduler::enqueue(detail::communication_stream(), std::move(task));
 }
@@ -89,22 +98,28 @@ void Send::eval_gpu(
   assert(outputs.size() == 1);
 
   auto& in = inputs[0];
+
+  // Encode a signal event for the input
+  Fence f{stream()};
+  if (in.event().valid()) {
+    f.update_gpu(in);
+  }
+
   auto& out = outputs[0];
   move_or_copy(in, out);
 
   // Schedule an async send on the comm stream
-  auto task = [in = in, out = out, group = group(), dst = dst_]() mutable {
+  auto task = [in = in,
+               out = out,
+               f = std::move(f),
+               group = group(),
+               dst = dst_]() mutable {
     if (in.event().valid()) {
-      in.event().wait();
+      f.wait();
     }
     distributed::detail::send(group, out, dst);
   };
   scheduler::enqueue(detail::communication_stream(), std::move(task));
-
-  // Encode a signal event for the input
-  if (in.event().valid()) {
-    encode_signal(in.event());
-  }
 }
 
 void Recv::eval_gpu(
@@ -117,16 +132,14 @@ void Recv::eval_gpu(
 
   out.set_data(allocator::malloc_or_wait(out.nbytes()));
 
-  auto e = Event(stream());
-  e.set_value(1);
-
-  encode_wait(e);
+  Fence f{stream()};
+  f.wait_gpu(out);
 
   // Schedule an async recv on the comm stream
   auto task =
-      [out = out, e = std::move(e), group = group(), src = src_]() mutable {
+      [out = out, f = std::move(f), group = group(), src = src_]() mutable {
         distributed::detail::recv(group, out, src);
-        e.signal();
+        f.update();
       };
   scheduler::enqueue(detail::communication_stream(), std::move(task));
 }

mlx/backend/metal/fence.cpp

@@ -0,0 +1,145 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/fence.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/event.h"
+#include "mlx/backend/metal/metal_impl.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+Fence::Fence(const Stream& stream) : stream_(stream) {
+  auto d = metal::device(stream.device).mtl_device();
+  if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
+    use_fast_ = false;
+  } else if (__builtin_available(macOS 15, iOS 18, *)) {
+    use_fast_ = env::metal_fast_synch();
+  }
+  if (!use_fast_) {
+    // Wraps Metal SharedEvent
+    auto dtor = [](void* ptr) {
+      auto p = metal::new_scoped_memory_pool();
+      static_cast<MTL::SharedEvent*>(ptr)->release();
+    };
+    auto p = metal::new_scoped_memory_pool();
+    fence_ = std::shared_ptr<void>(
+        metal::device(stream.device).mtl_device()->newSharedEvent(), dtor);
+  } else {
+    auto dtor = [](void* buf) {
+      allocator::free(static_cast<MTL::Buffer*>(buf));
+    };
+    auto buf = allocator::malloc_or_wait(sizeof(uint32_t)).ptr();
+    fence_ = std::shared_ptr<void>(buf, dtor);
+    cpu_value()[0] = 0;
+  }
+}
+
+void Fence::wait_gpu(array& x) {
+  gpu_count_++;
+  auto& d = metal::device(stream_.device);
+  auto idx = stream_.index;
+
+  if (!use_fast_) {
+    d.end_encoding(idx);
+    auto command_buffer = d.get_command_buffer(idx);
+    command_buffer->encodeWait(
+        static_cast<MTL::Event*>(fence_.get()), gpu_count_);
+    command_buffer->addCompletedHandler(
+        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+    return;
+  }
+
+  auto& compute_encoder = d.get_command_encoder(idx);
+
+  // Register the output to ensure that no kernels which depends on the
+  // output starts before this one is done
+  compute_encoder.register_output_array(x);
+
+  auto kernel = d.get_kernel("fence_wait");
+  MTL::Size kernel_dims = MTL::Size(1, 1, 1);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  auto buf = static_cast<MTL::Buffer*>(fence_.get());
+  compute_encoder.set_buffer(buf, 0);
+  compute_encoder.set_bytes(gpu_count_, 1);
+  compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
+
+  d.get_command_buffer(idx)->addCompletedHandler(
+      [fence = fence_](MTL::CommandBuffer* cbuf) {});
+}
+
+void Fence::update_gpu(const array& x) {
+  gpu_count_++;
+  auto& d = metal::device(stream_.device);
+  auto idx = stream_.index;
+
+  if (!use_fast_) {
+    d.end_encoding(idx);
+    auto command_buffer = d.get_command_buffer(idx);
+    command_buffer->encodeSignalEvent(
+        static_cast<MTL::Event*>(fence_.get()), gpu_count_);
+    command_buffer->addCompletedHandler(
+        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+    return;
+  }
+
+  // Launch input visibility kernel
+  auto& compute_encoder = d.get_command_encoder(idx);
+  auto kernel = d.get_kernel("input_coherent");
+  uint32_t nthreads =
+      (x.data_size() * x.itemsize() + sizeof(uint32_t) - 1) / sizeof(uint32_t);
+
+  MTL::Size group_dims = MTL::Size(1024, 1, 1);
+  MTL::Size grid_dims = MTL::Size((nthreads + 1024 - 1) / 1024, 1, 1);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(x, 0);
+  compute_encoder.set_bytes(nthreads, 1);
+  compute_encoder.dispatch_threadgroups(group_dims, grid_dims);
+
+  // Barrier on previous kernel
+  compute_encoder.barrier();
+
+  // Launch value update kernel
+  kernel = d.get_kernel("fence_update");
+  MTL::Size kernel_dims = MTL::Size(1, 1, 1);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  auto buf = static_cast<MTL::Buffer*>(fence_.get());
+  compute_encoder.set_buffer(buf, 0);
+  compute_encoder.set_bytes(gpu_count_, 1);
+  compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
+
+  d.get_command_buffer(idx)->addCompletedHandler(
+      [fence = fence_](MTL::CommandBuffer* cbuf) {});
+}
+
+void Fence::wait() {
+  cpu_count_++;
+  if (!use_fast_) {
+    if (!static_cast<MTL::SharedEvent*>(fence_.get())
+             ->waitUntilSignaledValue(cpu_count_, -1)) {
+      throw std::runtime_error("[Event::wait] Timed out");
+    }
+    return;
+  }
+  while (cpu_value()[0] < cpu_count_) {
+  }
+}
+
+void Fence::update() {
+  cpu_count_++;
+
+  if (!use_fast_) {
+    static_cast<MTL::SharedEvent*>(fence_.get())->setSignaledValue(cpu_count_);
+    return;
+  }
+
+  cpu_value()[0] = cpu_count_;
+}
+
+std::atomic_uint* Fence::cpu_value() {
+  return static_cast<std::atomic_uint*>(
+      static_cast<MTL::Buffer*>(fence_.get())->contents());
+}
+
+} // namespace mlx::core