Do not require ConcurrentManagedAccess when not used (#3062)

mlx/backend/cuda/allocator.cpp

@@ -3,7 +3,9 @@
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/memory.h"
+#include "mlx/scheduler.h"
 #include "mlx/utils.h"
 
 #include <cuda_runtime.h>
@@ -45,12 +47,13 @@ SmallSizePool::SmallSizePool() {
 
   CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size));
 
-  int device_count = 0;
-  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
+  int device_count = gpu::device_count();
   for (int i = 0; i < device_count; ++i) {
-    auto loc = cuda_mem_loc(i);
-    CHECK_CUDA_ERROR(
-        cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
+    if (cu::device(i).concurrent_managed_access()) {
+      auto loc = cuda_mem_loc(i);
+      CHECK_CUDA_ERROR(cudaMemAdvise(
+          data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
+    }
   }
 
   auto curr = next_free_;
@@ -294,10 +297,14 @@ void CudaAllocator::clear_cache() {
 }
 
 CudaAllocator& allocator() {
-  // By creating the |allocator_| on heap, the destructor of CudaAllocator
-  // will not be called on exit and buffers in the cache will be leaked. This
-  // can save some time at program exit.
-  static CudaAllocator* allocator_ = new CudaAllocator;
+  static auto* allocator_ = []() {
+    // Ensure scheduler is created before allocator.
+    scheduler::scheduler();
+    // By creating the |allocator_| on heap, the destructor of CudaAllocator
+    // will not be called on exit and buffers in the cache will be leaked. This
+    // can save some time at program exit.
+    return new CudaAllocator();
+  }();
   return *allocator_;
 }
 

mlx/backend/cuda/conv.cpp

@@ -103,7 +103,7 @@ std::optional<DnnGraph> build_conv_graph(
     const std::vector<int64_t>& dilation) {
   auto compute_dtype =
       (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
-  DnnGraph graph(encoder.device().cudnn_handle(), dtype, compute_dtype);
+  DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype);
   auto x_ = graph.tensor_nchw("X", 'x', x);
   auto w_ = graph.tensor_nchw("W", 'w', w);
 

mlx/backend/cuda/cublas_utils.cpp

@@ -98,7 +98,7 @@ void CublasMatmulBase::init_base(
   M_ = a_rows;
   N_ = b_cols;
   scale_type_ = scale_type;
-  handle_ = device.lt_handle();
+  handle_ = device.get_cublaslt_handle();
   pref_ = cublas_utils::get_preference(device);
   heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
 

mlx/backend/cuda/device.cpp

@@ -3,6 +3,7 @@
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/jit_module.h"
 #include "mlx/backend/cuda/worker.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/utils.h"
 
 #include <fmt/format.h>
@@ -37,31 +38,19 @@ Device::Device(int device) : device_(device) {
       &compute_capability_major_, cudaDevAttrComputeCapabilityMajor, device_));
   CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
       &compute_capability_minor_, cudaDevAttrComputeCapabilityMinor, device_));
-  // Validate the requirements of device.
-  int attr = 0;
   CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
-      &attr, cudaDevAttrConcurrentManagedAccess, device_));
-  if (attr != 1) {
-    throw std::runtime_error(
-        fmt::format(
-            "Device {} does not support synchronization in managed memory.",
-            device_));
-  }
-
-  // The cublasLt handle is used by matmul.
-  make_current();
-  CHECK_CUBLAS_ERROR(cublasLtCreate(&lt_));
-  // The cudnn handle is used by Convolution.
-  CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_));
-
-  // Initialize the jit module cache here ensures it is not
-  // unloaded before any evaluation is done
-  get_jit_module_cache();
+      &concurrent_managed_access_,
+      cudaDevAttrConcurrentManagedAccess,
+      device_));
 }
 
 Device::~Device() {
-  CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_));
-  CHECK_CUBLAS_ERROR(cublasLtDestroy(lt_));
+  if (cudnn_handle_) {
+    CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_handle_));
+  }
+  if (cublaslt_handle_) {
+    CHECK_CUBLAS_ERROR(cublasLtDestroy(cublaslt_handle_));
+  }
 }
 
 void Device::make_current() {
@@ -82,6 +71,22 @@ CommandEncoder& Device::get_command_encoder(Stream s) {
   return it->second;
 }
 
+cublasLtHandle_t Device::get_cublaslt_handle() {
+  if (!cublaslt_handle_) {
+    make_current();
+    CHECK_CUBLAS_ERROR(cublasLtCreate(&cublaslt_handle_));
+  }
+  return cublaslt_handle_;
+}
+
+cudnnHandle_t Device::get_cudnn_handle() {
+  if (!cudnn_handle_) {
+    make_current();
+    CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_handle_));
+  }
+  return cudnn_handle_;
+}
+
 CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {
   enc.device().make_current();
   if (!use_cuda_graphs()) {
@@ -491,13 +496,23 @@ void CommandEncoder::synchronize() {
   f.wait();
 }
 
-Device& device(mlx::core::Device device) {
-  static std::unordered_map<int, Device> devices;
-  auto it = devices.find(device.index);
-  if (it == devices.end()) {
-    it = devices.try_emplace(device.index, device.index).first;
-  }
-  return it->second;
+Device& device(int cuda_device) {
+  static auto devices = []() {
+    std::vector<Device> devices;
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      devices.emplace_back(i);
+    }
+    // Initialize the jit module cache here ensures it is not unloaded before
+    // any evaluation is done.
+    get_jit_module_cache();
+    return devices;
+  }();
+  return devices.at(cuda_device);
+}
+
+Device& device(mlx::core::Device d) {
+  return device(d.index);
 }
 
 CommandEncoder& get_command_encoder(Stream s) {

mlx/backend/cuda/device.h

@@ -142,6 +142,7 @@ class Device {
   explicit Device(int device);
   ~Device();
 
+  Device(Device&&) = default;
   Device(const Device&) = delete;
   Device& operator=(const Device&) = delete;
 
@@ -149,6 +150,8 @@ class Device {
   void make_current();
 
   CommandEncoder& get_command_encoder(Stream s);
+  cublasLtHandle_t get_cublaslt_handle();
+  cudnnHandle_t get_cudnn_handle();
 
   int cuda_device() const {
     return device_;
@@ -159,24 +162,23 @@ class Device {
   int compute_capability_minor() const {
     return compute_capability_minor_;
   }
-  cublasLtHandle_t lt_handle() const {
-    return lt_;
-  }
-  cudnnHandle_t cudnn_handle() const {
-    return cudnn_;
+  bool concurrent_managed_access() const {
+    return concurrent_managed_access_ == 1;
   }
 
  private:
   int device_;
   int compute_capability_major_;
   int compute_capability_minor_;
+  int concurrent_managed_access_;
   std::string device_name_;
-  cublasLtHandle_t lt_;
-  cudnnHandle_t cudnn_;
+  cublasLtHandle_t cublaslt_handle_{nullptr};
+  cudnnHandle_t cudnn_handle_{nullptr};
   std::unordered_map<int, CommandEncoder> encoders_;
 };
 
-Device& device(mlx::core::Device device);
+Device& device(int cuda_device);
+Device& device(mlx::core::Device d);
 CommandEncoder& get_command_encoder(Stream s);
 
 // Return an execution policy that does not sync for result.

mlx/backend/cuda/event.cu

@@ -3,6 +3,7 @@
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/event.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/event.h"
 #include "mlx/scheduler.h"
 
@@ -193,7 +194,24 @@ __global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
   event_signal(ac, value);
 }
 
+bool supports_concurrent_managed_access() {
+  static bool concurrent_managed_access = []() {
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      if (!cu::device(i).concurrent_managed_access()) {
+        return false;
+      }
+    }
+    return true;
+  }();
+  return concurrent_managed_access;
+}
+
 AtomicEvent::AtomicEvent() {
+  if (!supports_concurrent_managed_access()) {
+    throw std::runtime_error(
+        "Device does not support synchronization in managed memory.");
+  }
   buf_ = std::shared_ptr<Buffer>(
       new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) {
         allocator().free(*ptr);

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -269,7 +269,7 @@ void sdpa_cudnn(
     bool output_logsumexp,
     Stream s) {
   auto& encoder = cu::get_command_encoder(s);
-  auto handle = encoder.device().cudnn_handle();
+  auto handle = encoder.device().get_cudnn_handle();
 
   malloc_with_same_layout(encoder, o, q);
 
@@ -327,7 +327,7 @@ void sdpa_backward_cudnn(
     array& d_v,
     Stream s) {
   auto& encoder = cu::get_command_encoder(s);
-  auto handle = encoder.device().cudnn_handle();
+  auto handle = encoder.device().get_cudnn_handle();
 
   malloc_with_same_layout(encoder, d_q, q);
   malloc_with_same_layout(encoder, d_k, k);