[CUDA] Fallback Event impl when there is no hardware cpu/gpu coherency (#3070)

Author: zcbenz

mlx/backend/cuda/device.cpp

@@ -42,6 +42,12 @@ Device::Device(int device) : device_(device) {
       &concurrent_managed_access_,
       cudaDevAttrConcurrentManagedAccess,
       device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &host_native_atomic_, cudaDevAttrHostNativeAtomicSupported, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &managed_memory_, cudaDevAttrManagedMemory, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_));
 }
 
 Device::~Device() {

mlx/backend/cuda/device.h

@@ -164,12 +164,24 @@ class Device {
   bool concurrent_managed_access() const {
     return concurrent_managed_access_ == 1;
   }
+  bool host_native_atomic() const {
+    return host_native_atomic_ == 1;
+  }
+  bool managed_memory() const {
+    return managed_memory_ == 1;
+  }
+  bool memory_pools() const {
+    return memory_pools_ == 1;
+  }
 
  private:
   int device_;
   int compute_capability_major_;
   int compute_capability_minor_;
   int concurrent_managed_access_;
+  int host_native_atomic_;
+  int managed_memory_;
+  int memory_pools_;
   std::string device_name_;
   cublasLtHandle_t cublaslt_handle_{nullptr};
   cudnnHandle_t cudnn_handle_{nullptr};

mlx/backend/cuda/event.cu

@@ -174,62 +174,94 @@ class CopyableCudaEvent {
 // AtomicEvent implementations
 ///////////////////////////////////////////////////////////////////////////////
 
-__host__ __device__ void event_wait(AtomicEvent::Atomic* ac, uint64_t value) {
-  uint64_t current;
-  while ((current = ac->load()) < value) {
-    ac->wait(current);
+__host__ __device__ void event_wait(uint32_t* ptr, uint32_t value) {
+  cuda::atomic_ref<uint32_t> ac(*ptr);
+  uint32_t current;
+  while ((current = ac.load()) < value) {
+    ac.wait(current);
   }
 }
 
-__host__ __device__ void event_signal(AtomicEvent::Atomic* ac, uint64_t value) {
-  ac->store(value);
-  ac->notify_all();
+__host__ __device__ void event_signal(uint32_t* ptr, uint32_t value) {
+  cuda::atomic_ref<uint32_t> ac(*ptr);
+  ac.store(value);
+  ac.notify_all();
 }
 
-__global__ void event_wait_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
-  event_wait(ac, value);
+__global__ void event_wait_kernel(uint32_t* ptr, uint32_t value) {
+  event_wait(ptr, value);
 }
 
-__global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
-  event_signal(ac, value);
+__global__ void event_signal_kernel(uint32_t* ptr, uint32_t value) {
+  __threadfence_system();
+  event_signal(ptr, value);
+  __threadfence_system();
 }
 
-bool supports_concurrent_managed_access() {
-  static bool concurrent_managed_access = []() {
+auto check_gpu_coherency() {
+  static auto coherency = []() {
     int device_count = gpu::device_count();
+    bool concurrent_managed_access = true;
+    bool host_native_atomic = true;
     for (int i = 0; i < device_count; ++i) {
-      if (!cu::device(i).concurrent_managed_access()) {
-        return false;
-      }
+      auto& d = cu::device(i);
+      concurrent_managed_access &= d.concurrent_managed_access();
+      host_native_atomic &= d.host_native_atomic();
     }
-    return true;
+    return std::make_tuple(concurrent_managed_access, host_native_atomic);
   }();
-  return concurrent_managed_access;
+  return coherency;
 }
 
 AtomicEvent::AtomicEvent() {
-  if (!supports_concurrent_managed_access()) {
-    throw std::runtime_error(
-        "Device does not support synchronization in managed memory.");
+  void* buf;
+  cudaError_t (*cuda_free)(void*);
+  // There are 3 kinds of systems we are implementing for:
+  // 1. concurrentManagedAccess == true
+  //    => use cuda::atom_ref on managed memory
+  // 2. hostNativeAtomicSupported == true
+  //    => use cuda::atom_ref on pinned host memory
+  // 2. no hardware cpu/gpu coherency
+  //    => use cuda::atom_ref on device memory
+  auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency();
+  if (concurrent_managed_access) {
+    CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = true;
+  } else if (host_native_atomic) {
+    CHECK_CUDA_ERROR(cudaMallocHost(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFreeHost;
+    coherent_ = true;
+  } else {
+    CHECK_CUDA_ERROR(cudaMalloc(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = false;
+  }
+  buf_ = std::shared_ptr<void>(
+      buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); });
+  if (coherent_) {
+    *ptr() = 0;
+  } else {
+    CHECK_CUDA_ERROR(cudaMemset(buf, 0, sizeof(uint32_t)));
   }
-  buf_ = std::shared_ptr<Buffer>(
-      new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) {
-        allocator().free(*ptr);
-        delete ptr;
-      });
-  *static_cast<uint64_t*>(buf_->raw_ptr()) = 0;
 }
 
-void AtomicEvent::wait(uint64_t value) {
+void AtomicEvent::wait(uint32_t value) {
   nvtx3::scoped_range r("cu::AtomicEvent::wait");
-  event_wait(atomic(), value);
+  if (coherent_) {
+    event_wait(ptr(), value);
+  } else {
+    while (!is_signaled(value)) {
+      std::this_thread::yield();
+    }
+  }
 }
 
-void AtomicEvent::wait(cudaStream_t stream, uint64_t value) {
-  event_wait_kernel<<<1, 1, 0, stream>>>(atomic(), value);
+void AtomicEvent::wait(cudaStream_t stream, uint32_t value) {
+  event_wait_kernel<<<1, 1, 0, stream>>>(ptr(), value);
 }
 
-void AtomicEvent::wait(Stream s, uint64_t value) {
+void AtomicEvent::wait(Stream s, uint32_t value) {
   nvtx3::scoped_range r("cu::AtomicEvent::wait(s)");
   if (s.device == mlx::core::Device::cpu) {
     scheduler::enqueue(s, [*this, value]() mutable { wait(value); });
@@ -241,22 +273,26 @@ void AtomicEvent::wait(Stream s, uint64_t value) {
   }
 }
 
-void AtomicEvent::signal(uint64_t value) {
+void AtomicEvent::signal(uint32_t value) {
   nvtx3::scoped_range r("cu::AtomicEvent::signal");
-  event_signal(atomic(), value);
+  if (coherent_) {
+    event_signal(ptr(), value);
+  } else {
+    signal(signal_stream(), value);
+  }
 }
 
-void AtomicEvent::signal(cudaStream_t stream, uint64_t value) {
-  event_signal_kernel<<<1, 1, 0, stream>>>(atomic(), value);
+void AtomicEvent::signal(cudaStream_t stream, uint32_t value) {
+  event_signal_kernel<<<1, 1, 0, stream>>>(ptr(), value);
 }
 
-void AtomicEvent::signal(Stream s, uint64_t value) {
+void AtomicEvent::signal(Stream s, uint32_t value) {
   nvtx3::scoped_range r("cu::AtomicEvent::signal(s)");
   if (s.device == mlx::core::Device::cpu) {
     // Signal through a GPU stream so the atomic is updated in GPU - updating
     // the atomic in CPU sometimes does not get GPU notified.
-    static CudaStream stream(device(mlx::core::Device::gpu));
-    scheduler::enqueue(s, [*this, value]() mutable { signal(stream, value); });
+    scheduler::enqueue(
+        s, [*this, value]() mutable { signal(signal_stream(), value); });
   } else {
     auto& encoder = get_command_encoder(s);
     encoder.commit();
@@ -265,14 +301,26 @@ void AtomicEvent::signal(Stream s, uint64_t value) {
   }
 }
 
-bool AtomicEvent::is_signaled(uint64_t value) const {
-  nvtx3::scoped_range r("cu::AtomicEvent::is_signaled");
-  return atomic()->load() >= value;
+bool AtomicEvent::is_signaled(uint32_t val) const {
+  return value() >= val;
 }
 
-uint64_t AtomicEvent::value() const {
+uint32_t AtomicEvent::value() const {
   nvtx3::scoped_range r("cu::AtomicEvent::value");
-  return atomic()->load();
+  if (coherent_) {
+    cuda::atomic_ref<uint32_t> ac(*ptr());
+    return ac.load();
+  } else {
+    uint32_t val;
+    CHECK_CUDA_ERROR(
+        cudaMemcpy(&val, ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost));
+    return val;
+  }
+}
+
+const CudaStream& AtomicEvent::signal_stream() {
+  static CudaStream stream(device(0));
+  return stream;
 }
 
 } // namespace cu

mlx/backend/cuda/event.h

@@ -54,25 +54,26 @@ class CudaEvent {
 // CudaEvent so the latter should always be preferred when possible.
 class AtomicEvent {
  public:
-  using Atomic = cuda::atomic<uint64_t>;
-
   AtomicEvent();
 
-  void wait(uint64_t value);
-  void wait(cudaStream_t stream, uint64_t value);
-  void wait(Stream s, uint64_t value);
-  void signal(uint64_t value);
-  void signal(cudaStream_t stream, uint64_t value);
-  void signal(Stream s, uint64_t value);
-  bool is_signaled(uint64_t value) const;
-  uint64_t value() const;
+  void wait(uint32_t value);
+  void wait(cudaStream_t stream, uint32_t value);
+  void wait(Stream s, uint32_t value);
+  void signal(uint32_t value);
+  void signal(cudaStream_t stream, uint32_t value);
+  void signal(Stream s, uint32_t value);
+  bool is_signaled(uint32_t value) const;
+  uint32_t value() const;
 
  private:
-  Atomic* atomic() const {
-    return static_cast<AtomicEvent::Atomic*>(buf_->raw_ptr());
+  const CudaStream& signal_stream();
+
+  uint32_t* ptr() const {
+    return static_cast<uint32_t*>(buf_.get());
   }
 
-  std::shared_ptr<allocator::Buffer> buf_;
+  bool coherent_;
+  std::shared_ptr<void> buf_;
 };
 
 } // namespace mlx::core::cu