Fix deadlock when copy (#1041)

Author: nathanielsimard

crates/cubecl-common/src/device.rs

@@ -184,6 +184,7 @@ mod context {
     /// Handle for accessing a [DeviceState] associated with a specific device.
     pub struct DeviceContext<S: DeviceState> {
         lock: DeviceStateLock,
+        lock_kind: Arc<ReentrantMutex<()>>,
         device_id: DeviceId,
         _phantom: PhantomData<S>,
     }
@@ -195,6 +196,7 @@ mod context {
         fn clone(&self) -> Self {
             Self {
                 lock: self.lock.clone(),
+                lock_kind: self.lock_kind.clone(),
                 _phantom: self._phantom,
                 device_id: self.device_id,
             }
@@ -295,6 +297,14 @@ mod context {
             Ok(lock)
         }
 
+        /// Locks all devices under the same kind.
+        ///
+        /// This is useful when you need mutable access to multiple devices at once, which can lead
+        /// to deadlocks.
+        pub fn lock_device_kind(&self) -> ReentrantMutexGuard<'_, ()> {
+            self.lock_kind.lock()
+        }
+
         /// Locks the current device making sure this device can be used.
         pub fn lock_device(&self) -> DeviceGuard<'_> {
             let state = self.lock.lock.lock();
@@ -367,8 +377,14 @@ mod context {
 
     static GLOBAL: spin::Mutex<DeviceLocator> = spin::Mutex::new(DeviceLocator { state: None });
 
+    #[derive(Default)]
+    struct DeviceLocatorState {
+        device: HashMap<Key, DeviceStateLock>,
+        device_kind: HashMap<TypeId, Arc<ReentrantMutex<()>>>,
+    }
+
     struct DeviceLocator {
-        state: Option<HashMap<Key, DeviceStateLock>>,
+        state: Option<DeviceLocatorState>,
     }
 
     #[derive(Clone)]
@@ -383,18 +399,19 @@ mod context {
     impl DeviceStateLock {
         fn locate<D: Device + 'static, S: DeviceState>(device: &D) -> DeviceContext<S> {
             let id = device.to_id();
+            let kind = TypeId::of::<D>();
             let key = (id, TypeId::of::<D>());
             let mut global = GLOBAL.lock();
 
-            let map = match &mut global.state {
+            let locator_state = match &mut global.state {
                 Some(state) => state,
                 None => {
-                    global.state = Some(HashMap::default());
+                    global.state = Some(Default::default());
                     global.state.as_mut().expect("Just created Option::Some")
                 }
             };
 
-            let lock = match map.get(&key) {
+            let lock = match locator_state.device.get(&key) {
                 Some(value) => value.clone(),
                 None => {
                     let state = DeviceStateMap::new();
@@ -403,13 +420,31 @@ mod context {
                         lock: Arc::new(ReentrantMutex::new(state)),
                     };
 
-                    map.insert(key, value);
-                    map.get(&key).expect("Just inserted the key/value").clone()
+                    locator_state.device.insert(key, value);
+                    locator_state
+                        .device
+                        .get(&key)
+                        .expect("Just inserted the key/value")
+                        .clone()
+                }
+            };
+            let lock_kind = match locator_state.device_kind.get(&kind) {
+                Some(value) => value.clone(),
+                None => {
+                    locator_state
+                        .device_kind
+                        .insert(kind, Arc::new(ReentrantMutex::new(())));
+                    locator_state
+                        .device_kind
+                        .get(&kind)
+                        .expect("Just inserted the key/value")
+                        .clone()
                 }
             };
 
             DeviceContext {
                 lock,
+                lock_kind,
                 device_id: id,
                 _phantom: PhantomData,
             }

crates/cubecl-cuda/src/compute/context.rs

@@ -162,18 +162,19 @@ impl CudaContext {
         let repr = kernel_compiled.repr.unwrap();
 
         if let Some(cache) = &mut self.ptx_cache {
-            cache
-                .insert(
-                    name.unwrap(),
-                    PtxCacheEntry {
-                        entrypoint_name: kernel_compiled.entrypoint_name.clone(),
-                        cube_dim: (cube_dim.x, cube_dim.y, cube_dim.z),
-                        shared_mem_bytes: repr.shared_memory_size(),
-                        cluster_dim: cluster_dim.map(|cluster| (cluster.x, cluster.y, cluster.z)),
-                        ptx: ptx.clone(),
-                    },
-                )
-                .unwrap();
+            let result = cache.insert(
+                name.unwrap(),
+                PtxCacheEntry {
+                    entrypoint_name: kernel_compiled.entrypoint_name.clone(),
+                    cube_dim: (cube_dim.x, cube_dim.y, cube_dim.z),
+                    shared_mem_bytes: repr.shared_memory_size(),
+                    cluster_dim: cluster_dim.map(|cluster| (cluster.x, cluster.y, cluster.z)),
+                    ptx: ptx.clone(),
+                },
+            );
+            if let Err(err) = result {
+                log::warn!("Unable to save the ptx {err:?}");
+            }
         }
 
         self.load_ptx(

crates/cubecl-runtime/src/client.rs

@@ -485,17 +485,22 @@ where
         dst_server: &Self,
     ) -> Allocation {
         if Server::SERVER_COMM_ENABLED {
+            let guard = self.context.lock_device_kind();
             let mut server_src = self.context.lock();
             let mut server_dst = dst_server.context.lock();
 
-            Server::copy(
+            let copied = Server::copy(
                 server_src.deref_mut(),
                 server_dst.deref_mut(),
                 src_descriptor,
                 self.stream_id(),
                 dst_server.stream_id(),
             )
-            .unwrap()
+            .unwrap();
+            core::mem::drop(server_src);
+            core::mem::drop(server_dst);
+            core::mem::drop(guard);
+            copied
         } else {
             let alloc_desc = AllocationDescriptor::new(
                 AllocationKind::Optimized,

crates/cubecl-std/src/tests/trigonometry.rs

@@ -15,7 +15,7 @@ pub fn test_to_degrees<R: Runtime>(client: ComputeClient<R::Server>) {
     let input_data = vec![0.0, PI / 6.0, PI / 4.0, PI / 2.0, PI, TAU];
     let expected = vec![0.0, 30.0, 45.0, 90.0, 180.0, 360.0];
 
-    let input = client.create(f32::as_bytes(&input_data));
+    let input = client.create_from_slice(f32::as_bytes(&input_data));
     let output = client.empty(input_data.len() * core::mem::size_of::<f32>());
 
     unsafe {
@@ -53,7 +53,7 @@ pub fn test_to_radians<R: Runtime>(client: ComputeClient<R::Server>) {
     let input_data = vec![0.0, 30.0, 45.0, 90.0, 180.0, 360.0];
     let expected = vec![0.0, PI / 6.0, PI / 4.0, PI / 2.0, PI, TAU];
 
-    let input = client.create(f32::as_bytes(&input_data));
+    let input = client.create_from_slice(f32::as_bytes(&input_data));
     let output = client.empty(input_data.len() * core::mem::size_of::<f32>());
 
     unsafe {
@@ -92,8 +92,8 @@ pub fn test_hypot<R: Runtime>(client: ComputeClient<R::Server>) {
     let y_data = vec![4.0, 1.0, 1.0, 12.0, 0.0];
     let expected = vec![5.0, 1.0, 1.4142135623730951, 13.0, 0.0];
 
-    let x = client.create(f32::as_bytes(&x_data));
-    let y = client.create(f32::as_bytes(&y_data));
+    let x = client.create_from_slice(f32::as_bytes(&x_data));
+    let y = client.create_from_slice(f32::as_bytes(&y_data));
     let output = client.empty(x_data.len() * core::mem::size_of::<f32>());
 
     unsafe {