[CUDA] Set current device before allocating memory (#3110)

mlx/backend/cuda/allocator.cpp

@@ -168,9 +168,6 @@ CudaAllocator::CudaAllocator()
   free_limit_ = total_memory_ - memory_limit_;
   max_pool_size_ = memory_limit_;
 
-  int curr;
-  CHECK_CUDA_ERROR(cudaGetDevice(&curr));
-
   int device_count = gpu::device_count();
   free_streams_.resize(device_count);
   mem_pools_.resize(device_count);
@@ -181,7 +178,6 @@ CudaAllocator::CudaAllocator()
       CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i));
     }
   }
-  CHECK_CUDA_ERROR(cudaSetDevice(curr));
 }
 
 Buffer
@@ -223,6 +219,7 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
       if (device == -1) {
         data = unified_malloc(size);
       } else {
+        cu::device(device).make_current();
         if (mem_pools_[device]) { // supports memory pools
           CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
         } else {

mlx/backend/cuda/event.cu

@@ -213,7 +213,7 @@ auto check_gpu_coherency() {
   return coherency;
 }
 
-AtomicEvent::AtomicEvent() {
+AtomicEvent::AtomicEvent(Device& d) {
   void* buf;
   cudaError_t (*cuda_free)(void*);
   // There are 3 kinds of systems we are implementing for:
@@ -223,6 +223,7 @@ AtomicEvent::AtomicEvent() {
   //    => use cuda::atom_ref on pinned host memory
   // 2. no hardware cpu/gpu coherency
   //    => use cuda::atom_ref on device memory
+  d.make_current();
   auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency();
   if (concurrent_managed_access) {
     CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t)));
@@ -347,11 +348,12 @@ struct EventImpl {
     if (is_created()) {
       return;
     }
+    auto& d = cu::device(s.device);
     if (s.device == mlx::core::Device::cpu || signal_value > 1) {
       nvtx3::mark("Using slow AtomicEvent");
-      atomic = std::make_unique<cu::AtomicEvent>();
+      atomic = std::make_unique<cu::AtomicEvent>(d);
     } else {
-      cuda = std::make_unique<cu::CopyableCudaEvent>(cu::device(s.device));
+      cuda = std::make_unique<cu::CopyableCudaEvent>(d);
     }
   }
 };

mlx/backend/cuda/event.h

@@ -54,7 +54,7 @@ class CudaEvent {
 // CudaEvent so the latter should always be preferred when possible.
 class AtomicEvent {
  public:
-  AtomicEvent();
+  AtomicEvent(Device& d);
 
   void wait(uint32_t value);
   void wait(cudaStream_t stream, uint32_t value);

mlx/backend/cuda/fence.cpp

@@ -14,7 +14,8 @@ struct FenceImpl {
 
 Fence::Fence(Stream s) {
   fence_ = std::shared_ptr<void>(
-      new FenceImpl{0}, [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
+      new FenceImpl{0, cu::device(s.device)},
+      [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
 }
 
 void Fence::wait(Stream s, const array&) {