Fix/device transfer (#941)

Author: nathanielsimard

Cargo.toml

@@ -86,10 +86,11 @@ tracy-client = { version = "0.18.0" }
 strum = { version = "0.27.1", features = ["derive"] }
 tracel-xtask = { version = "=2.1.8" }
 
-portable-atomic = { version = "1.10", default-features = false, features = [
+
+portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
+portable-atomic = { version = "1.11", default-features = false, features = [
     "serde",
 ] }
-portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 pretty_assertions = "1.4"
 
 # Async

crates/cubecl-common/Cargo.toml

@@ -21,6 +21,7 @@ fp8 = ["float8"]
 serde = ["serde_bytes"]
 std = ["rand/std", "futures-lite", "rand/thread_rng", "serde_json?/std"]
 
+
 [dependencies]
 # ** Please make sure all dependencies support no_std when std is disabled **
 bytemuck = { workspace = true, features = ["derive"] }
@@ -61,8 +62,8 @@ futures-lite = { workspace = true, features = [
 spin = { workspace = true, features = ["mutex", "spin_mutex"] }
 
 [target.'cfg(not(target_has_atomic = "ptr"))'.dependencies]
-portable-atomic = { workspace = true }
 portable-atomic-util = { workspace = true }
+portable-atomic = { workspace = true }
 spin = { workspace = true, features = [
     "mutex",
     "spin_mutex",

crates/cubecl-common/src/map.rs

@@ -1,18 +1,13 @@
-use crate::stub::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
-
-#[cfg(target_has_atomic = "ptr")]
-use alloc::sync::Arc;
-
-#[cfg(not(target_has_atomic = "ptr"))]
-use portable_atomic_util::Arc;
-
+use crate::stub::{Arc, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
 use hashbrown::HashMap;
 
 /// A thread-safe map that allows concurrent access to values using read-write locks.
 pub struct SharedStateMap<K, V> {
-    state: Mutex<Option<HashMap<K, Arc<RwLock<V>>>>>,
+    state: Mutex<Option<State<K, V>>>,
 }
 
+type State<K, V> = HashMap<K, Arc<RwLock<V>>>;
+
 /// A value in the [SharedStateMap] that provides read and write access.
 pub struct SharedState<V> {
     val: Arc<RwLock<V>>,
@@ -53,7 +48,7 @@ where
     /// Retrieves a value associated with the given key, if it exists.
     pub fn get(&self, k: &K) -> Option<SharedState<V>> {
         let mut state = self.state.lock().unwrap();
-        let map = get_or_init(&mut state);
+        let map = get_or_init::<K, V>(&mut state);
 
         match map.get(k) {
             Some(val) => Some(SharedState { val: val.clone() }),
@@ -68,7 +63,7 @@ where
         K: Clone,
     {
         let mut state = self.state.lock().unwrap();
-        let map = get_or_init(&mut state);
+        let map = get_or_init::<K, V>(&mut state);
 
         match map.get(k) {
             Some(val) => SharedState { val: val.clone() },
@@ -84,24 +79,24 @@ where
     /// Inserts a key-value pair into the map.
     pub fn insert(&self, k: K, v: V) {
         let mut state = self.state.lock().unwrap();
-        let map = get_or_init(&mut state);
+        let map = get_or_init::<K, V>(&mut state);
 
         map.insert(k, Arc::new(RwLock::new(v)));
     }
 
     /// Clears the map, removing all key-value pairs.
     pub fn clear(&self) {
         let mut state = self.state.lock().unwrap();
-        let map = get_or_init(&mut state);
+        let map = get_or_init::<K, V>(&mut state);
         map.clear();
     }
 }
 
-fn get_or_init<T: Default>(state: &mut Option<T>) -> &mut T {
+fn get_or_init<K, V>(state: &mut Option<State<K, V>>) -> &mut State<K, V> {
     match state {
         Some(state) => state,
         None => {
-            *state = Some(T::default());
+            *state = Some(State::<K, V>::default());
             state.as_mut().unwrap()
         }
     }

crates/cubecl-common/src/stub.rs

@@ -10,6 +10,12 @@ pub use spin::{RwLockReadGuard, RwLockWriteGuard};
 #[cfg(feature = "std")]
 pub use std::sync::{RwLockReadGuard, RwLockWriteGuard};
 
+#[cfg(target_has_atomic = "ptr")]
+pub use alloc::sync::Arc;
+
+#[cfg(not(target_has_atomic = "ptr"))]
+pub use portable_atomic_util::Arc;
+
 /// A mutual exclusion primitive useful for protecting shared data
 ///
 /// This mutex will block threads waiting for the lock to become available. The

crates/cubecl-cpu/src/compute/server.rs

@@ -6,8 +6,8 @@ use cubecl_core::{
     compute::CubeTask,
     future::DynFut,
     server::{
-        Allocation, AllocationDescriptor, Binding, Bindings, ComputeServer, CopyDescriptor,
-        DataTransferService, Handle, IoError, ProfileError, ProfilingToken,
+        Allocation, AllocationDescriptor, Binding, Bindings, ComputeServer, CopyDescriptor, Handle,
+        IoError, ProfileError, ProfilingToken, ServerCommunication,
     },
 };
 use cubecl_runtime::{
@@ -28,8 +28,6 @@ pub struct CpuServer {
     logger: Arc<ServerLogger>,
 }
 
-impl DataTransferService for CpuServer {}
-
 impl CpuServer {
     pub fn new(ctx: CpuContext) -> Self {
         Self {
@@ -227,6 +225,10 @@ impl ComputeServer for CpuServer {
     }
 }
 
+impl ServerCommunication for CpuServer {
+    const SERVER_COMM_ENABLED: bool = false;
+}
+
 impl CpuServer {
     fn copy_to_binding(&mut self, binding: Binding, data: &[u8]) {
         let mut resource = self

crates/cubecl-cpu/src/runtime.rs

@@ -64,7 +64,6 @@ fn create_client(options: RuntimeOptions) -> ComputeClient<Server, Channel> {
     let mem_properties = MemoryDeviceProperties {
         max_page_size: max_shared_memory_size as u64,
         alignment: ALIGNMENT,
-        data_transfer_async: false,
     };
 
     let memory_management =

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

@@ -1,9 +1,8 @@
 use crate::{
     CudaCompiler,
     compute::{
-        DataTransferItem, DataTransferRuntime, MB, context::CudaContext,
-        io::controller::PinnedMemoryManagedAllocController, storage::gpu::GpuResource,
-        stream::CudaStreamBackend, sync::Fence, valid_strides,
+        MB, context::CudaContext, io::controller::PinnedMemoryManagedAllocController,
+        storage::gpu::GpuResource, stream::CudaStreamBackend, sync::Fence, valid_strides,
     },
 };
 use cubecl_common::{bytes::Bytes, stream_id::StreamId};
@@ -14,7 +13,6 @@ use cubecl_core::{
     server::{Binding, CopyDescriptor, Handle, IoError, ProfileError},
 };
 use cubecl_runtime::{
-    data_service::DataTransferId,
     id::KernelId,
     logging::ServerLogger,
     memory_management::{MemoryAllocationMode, MemoryHandle},
@@ -25,21 +23,13 @@ use cudarc::driver::sys::{
 };
 use std::{ffi::c_void, ops::DerefMut, sync::Arc};
 
-const DEVICE_TO_DEVICE: DeviceTransferStrategy = DeviceTransferStrategy::Serialized;
-
-#[allow(unused)]
-enum DeviceTransferStrategy {
-    Peer,
-    Serialized,
-}
-
 #[derive(new)]
 /// The `Command` struct encapsulates a CUDA context and a set of resolved CUDA streams, providing an
 /// interface for executing GPU-related operations such as memory allocation, data transfers, kernel
 /// registration, and task execution.
 pub struct Command<'a> {
     ctx: &'a mut CudaContext,
-    streams: ResolvedStreams<'a, CudaStreamBackend>,
+    pub(crate) streams: ResolvedStreams<'a, CudaStreamBackend>,
 }
 
 impl<'a> Command<'a> {
@@ -241,7 +231,7 @@ impl<'a> Command<'a> {
         Ok((data, fences))
     }
 
-    fn copy_to_bytes(
+    pub fn copy_to_bytes(
         &mut self,
         descriptor: CopyDescriptor<'_>,
         pinned: bool,
@@ -344,72 +334,6 @@ impl<'a> Command<'a> {
         Ok(handle)
     }
 
-    /// Registers a source for an asynchronous data transfer operation.
-    ///
-    /// # Parameters
-    ///
-    /// * `id` - The unique identifier for the data transfer.
-    /// * `src` - The descriptor for the source GPU memory.
-    pub fn data_transfer_src(&mut self, id: DataTransferId, src: CopyDescriptor<'_>) {
-        let src_resource = self.resource(src.binding.clone()).unwrap();
-        let client = DataTransferRuntime::client();
-        let current = self.streams.current();
-
-        let handle = DataTransferItem {
-            stream: current.sys,
-            context: self.ctx.context,
-            resource: src_resource,
-        };
-
-        match DEVICE_TO_DEVICE {
-            DeviceTransferStrategy::Peer => {
-                let fence = Fence::new(current.sys);
-
-                client.register_src_peer(id, handle, fence);
-            }
-            DeviceTransferStrategy::Serialized => {
-                client.register_src_serialized(id, handle, src.binding);
-            }
-        }
-    }
-
-    /// Registers a destination for an asynchronous data transfer operation.
-    ///
-    /// # Parameters
-    ///
-    /// * `id` - The unique identifier for the data transfer.
-    /// * `dest` - The descriptor for the destination GPU memory.
-    pub fn data_transfer_dest(&mut self, id: DataTransferId, dest: CopyDescriptor<'_>) {
-        let dst_resource = self.resource(dest.binding).unwrap();
-        let current = self.streams.current();
-        let client = DataTransferRuntime::client();
-
-        let item = DataTransferItem {
-            context: self.ctx.context,
-            stream: current.sys,
-            resource: dst_resource,
-        };
-
-        match DEVICE_TO_DEVICE {
-            DeviceTransferStrategy::Peer => {
-                client.register_dest_peer(id, item);
-            }
-            DeviceTransferStrategy::Serialized => {
-                let num_bytes = dest.shape.iter().product::<usize>() * dest.elem_size;
-                let bytes = self.reserve_cpu(num_bytes, true, None);
-
-                client.register_dest_serialized(
-                    id,
-                    item,
-                    bytes,
-                    dest.shape.to_vec(),
-                    dest.strides.to_vec(),
-                    dest.elem_size,
-                );
-            }
-        }
-    }
-
     /// Synchronizes the current stream, ensuring all pending operations are complete.
     ///
     /// # Returns