metal: dispatch/commit/sync/wait counters (gemma4-qat#181)

mlx/backend/cuda/eval.cpp

@@ -81,7 +81,7 @@ void finalize(Stream s) {
   cu::get_command_encoder(s).commit();
 }
 
-void synchronize(Stream s) {
+void synchronize(Stream s, bool) {
   nvtx3::scoped_range r("gpu::synchronize");
   cu::get_command_encoder(s).synchronize();
 }

mlx/backend/gpu/eval.h

@@ -15,7 +15,8 @@ void new_stream(Stream s);
 void new_thread_unsafe_stream(Stream s);
 void eval(array& arr);
 void finalize(Stream s);
-void synchronize(Stream s);
+// explicit_sync is false for internal flushes (e.g. eval error recovery).
+void synchronize(Stream s, bool explicit_sync = true);
 void clear_streams();
 
 } // namespace mlx::core::gpu

mlx/backend/metal/device.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 
+#include <atomic>
 #include <cstdlib>
 #include <sstream>
 
@@ -14,6 +15,7 @@
 #include "mlx/backend/metal/event.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/backend/metal/utils.h"
+#include "mlx/scheduler.h"
 #include "mlx/utils.h"
 
 namespace std {
@@ -32,6 +34,12 @@ namespace mlx::core::metal {
 
 namespace {
 
+// Statistics for counters() (metal.h), not used for synchronization.
+std::atomic<uint64_t> dispatch_count{0};
+std::atomic<uint64_t> commit_count{0};
+std::atomic<uint64_t> sync_count{0};
+std::atomic<uint64_t> wait_count{0};
+
 constexpr const char* default_mtllib_path = METAL_PATH;
 
 void set_compile_options(
@@ -418,6 +426,7 @@ void CommandEncoder::dispatch_threadgroups(
     MTL::Size group_dims) {
   maybeInsertBarrier();
   buffer_ops_++;
+  dispatch_count.fetch_add(1, std::memory_order_relaxed);
   get_command_encoder()->dispatchThreadgroups(grid_dims, group_dims);
 }
 
@@ -426,6 +435,7 @@ void CommandEncoder::dispatch_threads(
     MTL::Size group_dims) {
   maybeInsertBarrier();
   buffer_ops_++;
+  dispatch_count.fetch_add(1, std::memory_order_relaxed);
   get_command_encoder()->dispatchThreads(grid_dims, group_dims);
 }
 
@@ -564,16 +574,24 @@ void CommandEncoder::commit(std::function<void()> completion) {
         }
       });
   buffer_->commit();
+  commit_count.fetch_add(1, std::memory_order_relaxed);
   buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences());
   buffer_ops_ = 0;
   buffer_sizes_ = 0;
 }
 
-void CommandEncoder::synchronize() {
+void CommandEncoder::synchronize(bool explicit_sync) {
   auto pool = new_scoped_memory_pool();
   auto cbuf = buffer_; // retained
   end_encoding();
   commit();
+  if (explicit_sync) {
+    sync_count.fetch_add(1, std::memory_order_relaxed);
+  }
+  // Encoder teardown (clear_streams(), thread exit) is not a host wait.
+  if (!exiting_) {
+    count_host_wait();
+  }
   cbuf->waitUntilCompleted();
 
   if (!exiting_) {
@@ -980,4 +998,25 @@ bool is_nax_available() {
 #endif
 }
 
+void count_host_wait() {
+  wait_count.fetch_add(1, std::memory_order_relaxed);
+}
+
+Counters counters() {
+  return {
+      dispatch_count.load(std::memory_order_relaxed),
+      commit_count.load(std::memory_order_relaxed),
+      sync_count.load(std::memory_order_relaxed),
+      wait_count.load(std::memory_order_relaxed) +
+          scheduler::scheduler().gpu_waits()};
+}
+
+void reset() {
+  dispatch_count.store(0, std::memory_order_relaxed);
+  commit_count.store(0, std::memory_order_relaxed);
+  sync_count.store(0, std::memory_order_relaxed);
+  wait_count.store(0, std::memory_order_relaxed);
+  scheduler::scheduler().reset_gpu_waits();
+}
+
 } // namespace mlx::core::metal

mlx/backend/metal/device.h

@@ -94,7 +94,7 @@ class MLX_API CommandEncoder {
   void signal_event(Event event, uint64_t value);
   bool needs_commit() const;
   void commit(std::function<void()> completion = nullptr);
-  void synchronize();
+  void synchronize(bool explicit_sync = false);
 
   MTL::CommandBuffer* get_command_buffer() const {
     return buffer_.get();
@@ -257,4 +257,7 @@ NS::SharedPtr<NS::AutoreleasePool> new_scoped_memory_pool();
 
 bool is_nax_available();
 
+// Call before each host wait on GPU work.
+void count_host_wait();
+
 } // namespace mlx::core::metal

mlx/backend/metal/eval.cpp

@@ -76,8 +76,8 @@ void finalize(Stream s) {
   encoder.commit();
 }
 
-void synchronize(Stream s) {
-  metal::get_command_encoder(s).synchronize();
+void synchronize(Stream s, bool explicit_sync) {
+  metal::get_command_encoder(s).synchronize(explicit_sync);
 }
 
 void clear_streams() {

mlx/backend/metal/event.cpp

@@ -45,15 +45,24 @@ Event::Event(Stream stream) : stream_(stream) {
 
 void Event::wait() {
   check_error();
+  if (stream_.device == Device::gpu) {
+    metal::count_host_wait();
+  }
   cast<metal::EventImpl>().wait(value());
   check_error();
 }
 
 void Event::wait(Stream stream) {
   if (stream.device == Device::cpu) {
-    scheduler::wait_event(stream, *this, [value = value()](Event& self) {
-      self.cast<metal::EventImpl>().wait(value);
-    });
+    scheduler::wait_event(
+        stream,
+        *this,
+        [value = value(), gpu = stream_.device == Device::gpu](Event& self) {
+          if (gpu) {
+            metal::count_host_wait();
+          }
+          self.cast<metal::EventImpl>().wait(value);
+        });
   } else {
     auto& encoder = metal::get_command_encoder(stream);
     encoder.wait_event(*this, value());

mlx/backend/metal/fence.cpp

@@ -7,7 +7,7 @@
 namespace mlx::core {
 
 struct FenceImpl {
-  FenceImpl(Stream stream) {
+  FenceImpl(Stream stream) : device(stream.device) {
     auto d = metal::device(stream.device).mtl_device();
     if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
       use_fast = false;
@@ -29,6 +29,8 @@ struct FenceImpl {
       allocator::free(allocator::Buffer{static_cast<MTL::Buffer*>(fence)});
     }
   }
+  // Device of the stream that updates (signals) the fence.
+  Device device;
   bool use_fast{false};
   uint32_t count{0};
   void* fence;
@@ -55,6 +57,9 @@ void Fence::wait(Stream stream, const array& x) {
   if (stream.device == Device::cpu) {
     scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
       auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (f.device == Device::gpu) {
+        metal::count_host_wait();
+      }
       while (f.cpu_value()[0] < count) {
       }
     });

mlx/backend/metal/metal.h

@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include <cstdint>
 #include <string>
 #include <unordered_map>
 #include <variant>
@@ -27,4 +28,15 @@ MLX_API const
 MLX_API void set_metallib_path(const std::string& path);
 MLX_API const std::string& get_metallib_path();
 
+// Process-wide Metal work counters. syncs counts mx::synchronize() calls,
+// waits counts host waits on GPU work. Read or reset only while no eval runs.
+struct Counters {
+  uint64_t dispatches;
+  uint64_t commits;
+  uint64_t syncs;
+  uint64_t waits;
+};
+MLX_API Counters counters();
+MLX_API void reset();
+
 } // namespace mlx::core::metal

mlx/backend/metal/no_metal.cpp

@@ -28,6 +28,11 @@ const std::string& get_metallib_path() {
       "[metal::get_metallib_path] Cannot get metallib path without metal backend");
 }
 
+Counters counters() {
+  return {0, 0, 0, 0};
+}
+void reset() {}
+
 } // namespace metal
 
 } // namespace mlx::core

mlx/backend/no_gpu/eval.cpp

@@ -27,7 +27,7 @@ void finalize(Stream) {
   throw std::runtime_error("[gpu::finalize] GPU backend is not available");
 }
 
-void synchronize(Stream) {
+void synchronize(Stream, bool) {
   throw std::runtime_error("[gpu::synchronize] GPU backend is not available");
 }
 

mlx/scheduler.h

@@ -37,6 +37,9 @@ class MLX_API Scheduler {
     {
       std::lock_guard<std::mutex> lk(mtx);
       n_active_tasks_++;
+      if (stream.device == Device::gpu) {
+        n_active_gpu_tasks_++;
+      }
     }
     completion_cv.notify_all();
   }
@@ -45,6 +48,9 @@ class MLX_API Scheduler {
     {
       std::lock_guard<std::mutex> lk(mtx);
       n_active_tasks_--;
+      if (stream.device == Device::gpu) {
+        n_active_gpu_tasks_--;
+      }
     }
     completion_cv.notify_all();
   }
@@ -57,18 +63,34 @@ class MLX_API Scheduler {
     std::unique_lock<std::mutex> lk(mtx);
     int n_tasks_old = n_active_tasks();
     if (n_tasks_old > 1) {
+      int n_gpu_tasks_old = n_active_gpu_tasks_;
       completion_cv.wait(lk, [this, n_tasks_old] {
         return this->n_active_tasks() < n_tasks_old;
       });
+      // Count only when a GPU task completed during the wait.
+      if (n_active_gpu_tasks_ < n_gpu_tasks_old) {
+        gpu_waits_.fetch_add(1, std::memory_order_relaxed);
+      }
     }
   }
 
+  // Blocking wait_for_one() calls that ended on a GPU task completion.
+  uint64_t gpu_waits() const {
+    return gpu_waits_.load(std::memory_order_relaxed);
+  }
+
+  void reset_gpu_waits() {
+    gpu_waits_.store(0, std::memory_order_relaxed);
+  }
+
  private:
   friend Stream mlx::core::new_stream(Device d);
 
   StreamThread& get_thread(Stream s);
 
   int n_active_tasks_{0};
+  int n_active_gpu_tasks_{0};
+  std::atomic<uint64_t> gpu_waits_{0};
   std::unordered_map<int, std::unique_ptr<StreamThread>> threads_;
   std::shared_mutex threads_mtx_;
   std::condition_variable completion_cv;

mlx/transforms.cpp

@@ -326,7 +326,11 @@ array eval_impl(std::vector<array> outputs, bool async) {
     }
     for (auto& s : open_streams) {
       try {
-        synchronize(s);
+        if (s.device == Device::gpu) {
+          gpu::synchronize(s, /* explicit_sync = */ false);
+        } else {
+          synchronize(s);
+        }
       } catch (...) {
         // Preserve the original exception.
       }

tests/CMakeLists.txt

@@ -11,7 +11,7 @@ if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
 endif()
 
 if(MLX_BUILD_METAL)
-  list(APPEND METAL_TEST_SOURCES residency_tests.cpp)
+  list(APPEND METAL_TEST_SOURCES residency_tests.cpp metal_counters_tests.cpp)
 endif()
 
 include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake)