mlx 0.30.7

data/mlx/mlx/backend/metal/device.h

@@ -0,0 +1,289 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <Metal/Metal.hpp>
+#include <functional>
+#include <mutex>
+#include <shared_mutex>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/array.h"
+#include "mlx/device.h"
+
+namespace mlx::core::metal {
+
+using MTLFCList =
+    std::vector<std::tuple<const void*, MTL::DataType, NS::UInteger>>;
+
+struct DeviceStream;
+
+struct MLX_API CommandEncoder {
+  explicit CommandEncoder(DeviceStream& stream);
+  CommandEncoder(const CommandEncoder&) = delete;
+  CommandEncoder& operator=(const CommandEncoder&) = delete;
+
+  struct ConcurrentContext {
+    ConcurrentContext(CommandEncoder& enc) : enc(enc) {
+      enc.concurrent_ = true;
+    }
+    ~ConcurrentContext() {
+      enc.concurrent_ = false;
+      enc.prev_outputs_.insert(
+          enc.concurrent_outputs_.begin(), enc.concurrent_outputs_.end());
+      enc.concurrent_outputs_.clear();
+    }
+
+   private:
+    CommandEncoder& enc;
+  };
+
+  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(const 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);
+  }
+
+  void wait_for_fence(MTL::Fence* fence) {
+    enc_->waitForFence(fence);
+  }
+
+  void update_fence(MTL::Fence* fence) {
+    enc_->updateFence(fence);
+  }
+
+  template <typename Vec, typename = std::enable_if_t<is_vector_v<Vec>>>
+  void set_vector_bytes(const Vec& vec, size_t nelems, int idx) {
+    enc_->setBytes(vec.data(), nelems * sizeof(typename Vec::value_type), idx);
+  }
+  template <typename Vec, typename = std::enable_if_t<is_vector_v<Vec>>>
+  void set_vector_bytes(const Vec& vec, int idx) {
+    return set_vector_bytes(vec, vec.size(), idx);
+  }
+
+  template <typename T>
+  void set_bytes(const T* v, int n, int idx) {
+    return enc_->setBytes(v, n * sizeof(T), idx);
+  }
+
+  template <typename T>
+  void set_bytes(const T& v, int idx) {
+    return enc_->setBytes(&v, sizeof(T), idx);
+  }
+
+  void set_threadgroup_memory_length(size_t length, int idx) {
+    enc_->setThreadgroupMemoryLength(length, idx);
+  }
+
+  ConcurrentContext start_concurrent() {
+    return ConcurrentContext(*this);
+  }
+  ~CommandEncoder();
+
+  // Inputs to all kernels in the encoder including temporaries
+  std::unordered_set<const void*>& inputs() {
+    return all_inputs_;
+  };
+
+  // Outputs of all kernels in the encoder including temporaries
+  std::unordered_set<const void*>& outputs() {
+    return all_outputs_;
+  };
+
+  void barrier();
+
+ private:
+  DeviceStream& stream_;
+  MTL::ComputeCommandEncoder* enc_;
+  bool needs_barrier_{false};
+  bool concurrent_{false};
+  std::unordered_set<MTL::Resource*> prev_outputs_;
+  std::unordered_set<MTL::Resource*> next_outputs_;
+  std::unordered_set<MTL::Resource*> concurrent_outputs_;
+  std::unordered_set<const void*> all_inputs_;
+  std::unordered_set<const void*> all_outputs_;
+};
+
+struct Fence {
+  Fence(MTL::Fence* fence) : fence(fence) {}
+  ~Fence() {
+    fence->release();
+  }
+  MTL::Fence* fence;
+};
+
+struct DeviceStream {
+  DeviceStream(MTL::CommandQueue* queue) : queue(queue) {};
+  ~DeviceStream() {
+    queue->release();
+    if (buffer != nullptr) {
+      buffer->release();
+    }
+  };
+  MTL::CommandQueue* queue;
+  // A map of prior command encoder outputs to their corresponding fence
+  std::unordered_map<const void*, std::shared_ptr<Fence>> outputs;
+  // Used to allow thread-safe access to the outputs map
+  std::mutex fence_mtx;
+
+  // Data updated between command buffers
+  MTL::CommandBuffer* buffer{nullptr};
+  int buffer_ops{0};
+  size_t buffer_sizes{0};
+
+  // The command encoder, fence, and temporaries are updated between command
+  // encoders
+  std::unique_ptr<CommandEncoder> encoder{nullptr};
+  std::shared_ptr<Fence> fence;
+  std::vector<array> temporaries;
+};
+
+class MLX_API Device {
+ public:
+  Device();
+  Device(const Device&) = delete;
+  Device& operator=(const Device&) = delete;
+  ~Device();
+
+  MTL::Device* mtl_device() {
+    return device_;
+  };
+
+  const std::string& get_architecture() {
+    return arch_;
+  }
+
+  int get_architecture_gen() const {
+    return arch_gen_;
+  }
+
+  void new_queue(int index);
+
+  MTL::CommandQueue* get_queue(Stream stream);
+
+  MTL::CommandBuffer* get_command_buffer(int index);
+  bool command_buffer_needs_commit(int index);
+  void commit_command_buffer(int index);
+  CommandEncoder& get_command_encoder(int index);
+  void end_encoding(int index);
+
+  MTL::Library* get_library(
+      const std::string& name,
+      const std::string& path = "");
+
+  MTL::Library* get_library(
+      const std::string& name,
+      const std::function<std::string(void)>& builder);
+
+  void clear_library(const std::string& name);
+
+  MTL::ComputePipelineState* get_kernel(
+      const std::string& base_name,
+      MTL::Library* mtl_lib,
+      const std::string& hash_name = "",
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::ComputePipelineState* get_kernel(
+      const std::string& base_name,
+      const std::string& hash_name = "",
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::ArgumentEncoder* argument_encoder(
+      const std::vector<MTL::ArgumentDescriptor*>& arg_descs) const;
+
+  // Record temporary arrays for the given stream index
+  void add_temporary(array arr, int index);
+  void add_temporaries(std::vector<array> arrays, int index);
+
+  void set_residency_set(const MTL::ResidencySet* residency_set);
+
+ private:
+  DeviceStream& get_stream_(int index) {
+    return stream_map_.find(index)->second;
+  }
+  MTL::Library* get_library_cache_(const std::string& name);
+
+  MTL::Library* get_library_(const std::string& name);
+  MTL::Library* build_library_(const std::string& source_string);
+
+  MTL::Function* get_function_(const std::string& name, MTL::Library* mtl_lib);
+
+  MTL::Function* get_function_(
+      const std::string& name,
+      const std::string& specialized_name,
+      const MTLFCList& func_consts,
+      MTL::Library* mtl_lib);
+
+  MTL::LinkedFunctions* get_linked_functions_(
+      const std::vector<MTL::Function*>& funcs);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& name,
+      const MTL::Function* mtl_function);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& name,
+      const MTL::Function* mtl_function,
+      const MTL::LinkedFunctions* linked_functions);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& base_name,
+      MTL::Library* mtl_lib,
+      const std::string& hash_name,
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::Device* device_;
+  std::unordered_map<int32_t, DeviceStream> stream_map_;
+
+  std::shared_mutex kernel_mtx_;
+  std::shared_mutex library_mtx_;
+  std::unordered_map<std::string, MTL::Library*> library_map_;
+  MTL::Library* default_library_;
+  std::unordered_map<
+      MTL::Library*,
+      std::unordered_map<std::string, MTL::ComputePipelineState*>>
+      library_kernels_;
+  const MTL::ResidencySet* residency_set_{nullptr};
+  std::string arch_;
+  int arch_gen_;
+  int max_ops_per_buffer_;
+  int max_mb_per_buffer_;
+};
+
+MLX_API Device& device(mlx::core::Device);
+
+std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
+
+inline bool is_nax_available() {
+#ifdef MLX_METAL_NO_NAX
+  return false;
+#else
+  auto _check_nax = []() {
+    bool can_use_nax = false;
+    if (__builtin_available(
+            macOS 26.2, iOS 26.2, tvOS 26.2, visionOS 26.2, *)) {
+      can_use_nax = true;
+    }
+    auto& d = metal::device(mlx::core::Device::gpu);
+    auto arch = d.get_architecture().back();
+    auto gen = d.get_architecture_gen();
+    can_use_nax &= gen >= (arch == 'p' ? 18 : 17);
+    return can_use_nax;
+  };
+  static bool is_nax_available_ = _check_nax();
+  return is_nax_available_;
+#endif
+}
+
+} // namespace mlx::core::metal

data/mlx/mlx/backend/metal/device_info.cpp

@@ -0,0 +1,58 @@
+// Copyright © 2026 Apple Inc.
+
+#include <sys/sysctl.h>
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/metal.h"
+
+namespace mlx::core::gpu {
+
+bool is_available() {
+  return metal::is_available();
+}
+
+int device_count() {
+  return 1;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index) {
+  auto init_device_info = []()
+      -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
+    auto pool = metal::new_scoped_memory_pool();
+    auto raw_device = metal::device(mlx::core::Device::gpu).mtl_device();
+    auto name = std::string(raw_device->name()->utf8String());
+    auto arch = std::string(raw_device->architecture()->name()->utf8String());
+
+    size_t memsize = 0;
+    size_t length = sizeof(memsize);
+    sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
+
+    size_t rsrc_limit = 0;
+    sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0);
+    if (rsrc_limit == 0) {
+      rsrc_limit = 499000;
+    }
+
+    return {
+        {"device_name", name},
+        {"architecture", arch},
+        {"max_buffer_length", raw_device->maxBufferLength()},
+        {"max_recommended_working_set_size",
+         raw_device->recommendedMaxWorkingSetSize()},
+        {"memory_size", memsize},
+        {"resource_limit", rsrc_limit}};
+  };
+  static auto device_info_ = init_device_info();
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+
+  if (device_index == 0) {
+    return device_info_;
+  } else {
+    return empty;
+  }
+}
+
+} // namespace mlx::core::gpu

data/mlx/mlx/backend/metal/distributed.cpp

@@ -0,0 +1,38 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/distributed/primitives.h"
+#include "mlx/fence.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core::distributed {
+
+void AllReduce::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[AllReduce::eval_gpu] has no GPU implementation.");
+}
+
+void AllGather::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[AllGather::eval_gpu] has no GPU implementation.");
+}
+
+void Send::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[Send::eval_gpu] has no GPU implementation.");
+}
+
+void Recv::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[Recv::eval_gpu] has no GPU implementation.");
+}
+
+void ReduceScatter::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error(
+      "[ReduceScatter::eval_gpu] has no GPU implementation.");
+}
+
+} // namespace mlx::core::distributed

data/mlx/mlx/backend/metal/eval.cpp

@@ -0,0 +1,97 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <memory>
+
+#include "mlx/backend/gpu/eval.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core::gpu {
+
+void new_stream(Stream stream) {
+  if (stream.device == mlx::core::Device::gpu) {
+    metal::device(stream.device).new_queue(stream.index);
+  }
+}
+
+inline void check_error(MTL::CommandBuffer* cbuf) {
+  if (cbuf->status() == MTL::CommandBufferStatusError) {
+    std::ostringstream msg;
+    msg << "[METAL] Command buffer execution failed: "
+        << cbuf->error()->localizedDescription()->utf8String();
+    throw std::runtime_error(msg.str());
+  }
+}
+
+void eval(array& arr) {
+  auto pool = metal::new_scoped_memory_pool();
+  auto s = arr.primitive().stream();
+  auto& d = metal::device(s.device);
+  auto command_buffer = d.get_command_buffer(s.index);
+
+  auto outputs = arr.outputs();
+  {
+    // If the array is a tracer hold a reference
+    // to its inputs so they don't get donated
+    std::vector<array> inputs;
+    if (arr.is_tracer()) {
+      inputs = arr.inputs();
+    }
+
+    debug_set_primitive_buffer_label(command_buffer, arr.primitive());
+    arr.primitive().eval_gpu(arr.inputs(), outputs);
+  }
+  std::unordered_set<std::shared_ptr<array::Data>> buffers;
+  for (auto& in : arr.inputs()) {
+    buffers.insert(in.data_shared_ptr());
+  }
+  for (auto& s : arr.siblings()) {
+    buffers.insert(s.data_shared_ptr());
+  }
+  // Remove the output if it was donated to by an input
+  if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
+    buffers.erase(it);
+  }
+
+  if (d.command_buffer_needs_commit(s.index)) {
+    d.end_encoding(s.index);
+    scheduler::notify_new_task(s);
+    command_buffer->addCompletedHandler(
+        [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          scheduler::notify_task_completion(s);
+          check_error(cbuf);
+        });
+    d.commit_command_buffer(s.index);
+    d.get_command_buffer(s.index);
+  } else {
+    command_buffer->addCompletedHandler(
+        [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          check_error(cbuf);
+        });
+  }
+}
+
+void finalize(Stream s) {
+  auto pool = metal::new_scoped_memory_pool();
+  auto& d = metal::device(s.device);
+  auto cb = d.get_command_buffer(s.index);
+  d.end_encoding(s.index);
+  cb->addCompletedHandler([](MTL::CommandBuffer* cbuf) { check_error(cbuf); });
+  d.commit_command_buffer(s.index);
+  d.get_command_buffer(s.index);
+}
+
+void synchronize(Stream s) {
+  auto pool = metal::new_scoped_memory_pool();
+  auto& d = metal::device(s.device);
+  auto cb = d.get_command_buffer(s.index);
+  cb->retain();
+  d.end_encoding(s.index);
+  d.commit_command_buffer(s.index);
+  cb->waitUntilCompleted();
+  check_error(cb);
+  cb->release();
+}
+
+} // namespace mlx::core::gpu

data/mlx/mlx/backend/metal/event.cpp

@@ -0,0 +1,62 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/event.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core {
+
+Event::Event(Stream stream) : stream_(stream) {
+  auto dtor = [](void* ptr) {
+    auto p = metal::new_scoped_memory_pool();
+    static_cast<MTL::SharedEvent*>(ptr)->release();
+  };
+  auto p = metal::new_scoped_memory_pool();
+  event_ = std::shared_ptr<void>(
+      metal::device(Device::gpu).mtl_device()->newSharedEvent(), dtor);
+  if (event_ == nullptr) {
+    throw std::runtime_error(
+        "[Event::Event] Failed to create Metal shared event.");
+  }
+}
+
+void Event::wait() {
+  if (!static_cast<MTL::SharedEvent*>(event_.get())
+           ->waitUntilSignaledValue(value(), -1)) {
+    throw std::runtime_error("[Event::wait] Timed out");
+  }
+}
+
+void Event::wait(Stream stream) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [*this]() mutable { wait(); });
+  } else {
+    auto& d = metal::device(stream.device);
+    d.end_encoding(stream.index);
+    auto command_buffer = d.get_command_buffer(stream.index);
+    command_buffer->encodeWait(static_cast<MTL::Event*>(event_.get()), value());
+    command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {});
+  }
+}
+
+void Event::signal(Stream stream) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [*this]() mutable {
+      static_cast<MTL::SharedEvent*>(event_.get())->setSignaledValue(value());
+    });
+  } else {
+    auto& d = metal::device(stream.device);
+    d.end_encoding(stream.index);
+    auto command_buffer = d.get_command_buffer(stream.index);
+    command_buffer->encodeSignalEvent(
+        static_cast<MTL::Event*>(event_.get()), value());
+    command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {});
+  }
+}
+
+bool Event::is_signaled() const {
+  return static_cast<MTL::SharedEvent*>(event_.get())->signaledValue() >=
+      value();
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/fence.cpp

@@ -0,0 +1,162 @@
+// Copyright © 2024 Apple Inc.
+#include "mlx/fence.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/scheduler.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+struct FenceImpl {
+  FenceImpl() {
+    auto d = metal::device(Device::gpu).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) {
+      auto p = metal::new_scoped_memory_pool();
+      fence = static_cast<void*>(d->newSharedEvent());
+    } else {
+      auto buf = allocator::malloc(sizeof(uint32_t)).ptr();
+      fence = static_cast<void*>(buf);
+      cpu_value()[0] = 0;
+    }
+  }
+
+  ~FenceImpl() {
+    if (!use_fast) {
+      // Wraps Metal SharedEvent
+      auto p = metal::new_scoped_memory_pool();
+      static_cast<MTL::SharedEvent*>(fence)->release();
+    } else {
+      allocator::free(allocator::Buffer{static_cast<MTL::Buffer*>(fence)});
+    }
+  }
+  bool use_fast{false};
+  uint32_t count{0};
+  void* fence;
+
+  std::atomic_uint* cpu_value() {
+    return static_cast<std::atomic_uint*>(
+        static_cast<MTL::Buffer*>(fence)->contents());
+  }
+};
+
+Fence::Fence(Stream) {
+  auto dtor = [](void* ptr) { delete static_cast<FenceImpl*>(ptr); };
+  fence_ = std::shared_ptr<void>(new FenceImpl{}, dtor);
+}
+
+void Fence::wait(Stream stream, const array& x) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (!f.use_fast) {
+        if (!static_cast<MTL::SharedEvent*>(f.fence)->waitUntilSignaledValue(
+                count, -1)) {
+          throw std::runtime_error("[Fence::wait] Timed out");
+        }
+        return;
+      }
+      while (f.cpu_value()[0] < count) {
+      }
+    });
+    return;
+  }
+
+  auto& d = metal::device(stream.device);
+  auto idx = stream.index;
+
+  if (!f.use_fast) {
+    d.end_encoding(idx);
+    auto command_buffer = d.get_command_buffer(idx);
+    command_buffer->encodeWait(static_cast<MTL::Event*>(f.fence), f.count);
+    command_buffer->addCompletedHandler(
+        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+    return;
+  }
+
+  auto& compute_encoder = d.get_command_encoder(idx);
+
+  // Register outputs 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*>(f.fence);
+  compute_encoder.set_buffer(buf, 0);
+  compute_encoder.set_bytes(f.count, 1);
+  compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
+
+  d.get_command_buffer(idx)->addCompletedHandler(
+      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+}
+
+void Fence::update(Stream stream, const array& x, bool cross_device) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+  f.count++;
+
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (!f.use_fast) {
+        static_cast<MTL::SharedEvent*>(f.fence)->setSignaledValue(count);
+        return;
+      }
+
+      f.cpu_value()[0] = count;
+    });
+    return;
+  }
+
+  auto& d = metal::device(stream.device);
+  auto idx = stream.index;
+  if (!f.use_fast) {
+    d.end_encoding(idx);
+    auto command_buffer = d.get_command_buffer(idx);
+    command_buffer->encodeSignalEvent(
+        static_cast<MTL::Event*>(f.fence), f.count);
+    command_buffer->addCompletedHandler(
+        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+    return;
+  }
+
+  // Launch input visibility kernels
+  auto& compute_encoder = d.get_command_encoder(idx);
+  if (cross_device) {
+    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(grid_dims, group_dims);
+  }
+
+  // Barrier on previous kernels
+  compute_encoder.barrier();
+
+  // Launch value update kernel
+  auto 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*>(f.fence);
+  compute_encoder.set_buffer(buf, 0);
+  compute_encoder.set_bytes(f.count, 1);
+  compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
+
+  d.get_command_buffer(idx)->addCompletedHandler(
+      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+}
+
+} // namespace mlx::core