Cuda device change optimization (#864)

Author: nathanielsimard

crates/cubecl-common/src/device.rs

@@ -0,0 +1,45 @@
+use core::cmp::Ordering;
+
+/// The device id.
+#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy, new)]
+pub struct DeviceId {
+    /// The type id identifies the type of the device.
+    pub type_id: u16,
+    /// The index id identifies the device number.
+    pub index_id: u32,
+}
+
+/// Device trait for all cubecl devices.
+pub trait Device: Default + Clone + core::fmt::Debug + Send + Sync {
+    /// Create a device from its [id](DeviceId).
+    fn from_id(device_id: DeviceId) -> Self;
+    /// Retrieve the [device id](DeviceId) from the device.
+    fn to_id(&self) -> DeviceId;
+    /// Returns the number of devices available under the provided type id.
+    fn device_count(type_id: u16) -> usize;
+    /// Returns the total number of devices that can be handled by the runtime.
+    fn device_count_total() -> usize {
+        Self::device_count(0)
+    }
+}
+
+impl core::fmt::Display for DeviceId {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_fmt(format_args!("{self:?}"))
+    }
+}
+
+impl Ord for DeviceId {
+    fn cmp(&self, other: &Self) -> Ordering {
+        match self.type_id.cmp(&other.type_id) {
+            Ordering::Equal => self.index_id.cmp(&other.index_id),
+            other => other,
+        }
+    }
+}
+
+impl PartialOrd for DeviceId {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}

crates/cubecl-common/src/lib.rs

@@ -12,6 +12,9 @@ extern crate derive_new;
 /// std environments.
 pub mod rand;
 
+/// Device module.
+pub mod device;
+
 /// Stub module contains types for stubs for non-std environments and for std environments.
 pub mod stub;
 

crates/cubecl-core/src/id.rs

@@ -1,4 +1,5 @@
-use cubecl_runtime::{client::ComputeClient, id::DeviceId};
+use cubecl_common::device::{Device, DeviceId};
+use cubecl_runtime::client::ComputeClient;
 
 /// ID used to identify a Just-in-Time environment.
 #[derive(Hash, PartialEq, Eq, Debug, Clone)]
@@ -14,7 +15,7 @@ impl CubeTuneId {
         device: &R::Device,
     ) -> Self {
         Self {
-            device: R::device_id(device),
+            device: device.to_id(),
             name: R::name(client),
         }
     }

crates/cubecl-core/src/runtime.rs

@@ -1,7 +1,7 @@
 use crate::codegen::Compiler;
 use crate::compute::CubeTask;
+use cubecl_common::device::Device;
 use cubecl_ir::{StorageType, TargetProperties};
-use cubecl_runtime::id::DeviceId;
 use cubecl_runtime::{channel::ComputeChannel, client::ComputeClient, server::ComputeServer};
 
 pub use cubecl_runtime::channel;
@@ -18,10 +18,7 @@ pub trait Runtime: Send + Sync + 'static + core::fmt::Debug {
     /// The channel used to communicate with the compute server.
     type Channel: ComputeChannel<Self::Server>;
     /// The device used to retrieve the compute client.
-    type Device: Default + Clone + core::fmt::Debug + Send + Sync;
-
-    /// Fetch the id for the given device.
-    fn device_id(device: &Self::Device) -> DeviceId;
+    type Device: Device;
 
     /// Retrieve the compute client from the runtime device.
     fn client(device: &Self::Device) -> ComputeClient<Self::Server, Self::Channel>;
@@ -50,8 +47,7 @@ pub trait Runtime: Send + Sync + 'static + core::fmt::Debug {
 
     fn can_read_tensor(shape: &[usize], strides: &[usize]) -> bool;
 
-    fn device_count() -> usize;
-
+    /// Returns the properties of the target hardware architecture.
     fn target_properties() -> TargetProperties;
 }
 

crates/cubecl-core/src/runtime_tests/mod.rs

@@ -22,6 +22,7 @@ pub mod slice;
 pub mod synchronization;
 pub mod tensor;
 pub mod tensormap;
+pub mod to_client;
 pub mod topology;
 pub mod traits;
 pub mod unary;
@@ -135,6 +136,8 @@ macro_rules! testgen_untyped {
 
         cubecl_core::testgen_enums!();
         cubecl_core::testgen_comparison!();
+
+        cubecl_core::testgen_to_client!();
     };
 }
 

crates/cubecl-core/src/runtime_tests/to_client.rs

@@ -0,0 +1,58 @@
+use cubecl_common::device::{Device, DeviceId};
+
+use crate::Runtime;
+use crate::prelude::*;
+
+pub fn test_to_client<R: Runtime>() {
+    let type_id = 0;
+    let device_count = R::Device::device_count(type_id);
+
+    if device_count < 2 {
+        return;
+    }
+
+    for (device_0, device_1) in num_combination(type_id, device_count as u32) {
+        let device_0 = R::Device::from_id(device_0);
+        let device_1 = R::Device::from_id(device_1);
+
+        println!("Moving data from {device_0:?} to {device_1:?} ...");
+
+        let client_0 = R::client(&device_0);
+        let client_1 = R::client(&device_1);
+
+        let expected = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0];
+        let input = client_0.create(f32::as_bytes(&expected));
+
+        let output = client_0.to_client(input, &client_1).handle;
+
+        let actual = client_1.read_one(output);
+        let actual = f32::from_bytes(&actual);
+
+        assert_eq!(actual, expected);
+    }
+}
+
+fn num_combination(type_id: u16, n: u32) -> Vec<(DeviceId, DeviceId)> {
+    let mut results = Vec::new();
+
+    for i in 0..n {
+        for j in i + 1..n {
+            results.push((DeviceId::new(type_id, i), DeviceId::new(type_id, j)));
+        }
+    }
+
+    results
+}
+
+#[allow(missing_docs)]
+#[macro_export]
+macro_rules! testgen_to_client {
+    () => {
+        use super::*;
+
+        #[test]
+        fn test_to_client() {
+            cubecl_core::runtime_tests::to_client::test_to_client::<TestRuntime>();
+        }
+    };
+}

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, Handle,
-        IoError, ProfileError, ProfilingToken,
+        Allocation, AllocationDescriptor, Binding, Bindings, ComputeServer, CopyDescriptor,
+        DataTransferService, Handle, IoError, ProfileError, ProfilingToken,
     },
 };
 use cubecl_runtime::{
@@ -28,6 +28,8 @@ pub struct CpuServer {
     logger: ServerLogger,
 }
 
+impl DataTransferService for CpuServer {}
+
 impl CpuServer {
     pub fn new(ctx: CpuContext) -> Self {
         Self {

crates/cubecl-cpu/src/device.rs

@@ -1,2 +1,21 @@
+use cubecl_common::device::{Device, DeviceId};
+
 #[derive(new, Clone, PartialEq, Eq, Default, Hash, Debug)]
 pub struct CpuDevice;
+
+impl Device for CpuDevice {
+    fn from_id(_device_id: DeviceId) -> Self {
+        Self
+    }
+
+    fn to_id(&self) -> DeviceId {
+        DeviceId {
+            type_id: 0,
+            index_id: 0,
+        }
+    }
+
+    fn device_count(_type_id: u16) -> usize {
+        1
+    }
+}

crates/cubecl-cpu/src/runtime.rs

@@ -7,7 +7,6 @@ use cubecl_core::{
 };
 use cubecl_runtime::{
     ComputeRuntime, DeviceProperties,
-    id::DeviceId,
     memory_management::{HardwareProperties, MemoryDeviceProperties, MemoryManagement},
     storage::BytesStorage,
 };
@@ -63,6 +62,7 @@ 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 =
@@ -106,18 +106,10 @@ impl Runtime for CpuRuntime {
         (u32::MAX, u32::MAX, u32::MAX)
     }
 
-    fn device_id(_device: &Self::Device) -> DeviceId {
-        DeviceId::new(0, 0)
-    }
-
     fn can_read_tensor(shape: &[usize], strides: &[usize]) -> bool {
         is_contiguous(shape, strides)
     }
 
-    fn device_count() -> usize {
-        1
-    }
-
     fn target_properties() -> TargetProperties {
         TargetProperties {
             // Values are irrelevant, since no wgsl backends currently support manual mma

crates/cubecl-cuda/Cargo.toml

@@ -18,12 +18,13 @@ default = [
     "cubecl-common/default",
     "cubecl-core/default",
     "cudarc/dynamic-loading",
-    "cudarc/cuda-12050",
+    "cuda-12050",
 ]
 ptx-wmma = []
 std = ["cubecl-runtime/std", "cubecl-common/std", "cubecl-core/std"]
 
-cuda-12080 = []
+cuda-12080 = ["cudarc/cuda-12080"]
+cuda-12050 = ["cudarc/cuda-12050"]
 
 conv_tests = ["cubecl-convolution/conv_tests"]
 matmul_tests_all = [
@@ -96,9 +97,6 @@ half = { workspace = true }
 log = { workspace = true }
 serde = { workspace = true }
 
-[build-dependencies]
-cudarc = { workspace = true }
-
 [dev-dependencies]
 cubecl-convolution = { path = "../cubecl-convolution", version = "0.7.0", features = [
     "export_tests",