cubecl_runtime/channel/
base.rs

1use core::future::Future;
2use cubecl_common::{ExecutionMode, benchmark::ProfileDuration};
3
4use crate::{
5    server::{Binding, BindingWithMeta, Bindings, ComputeServer, CubeCount, Handle},
6    storage::{BindingResource, ComputeStorage},
7};
8use alloc::vec::Vec;
9
10/// The ComputeChannel trait links the ComputeClient to the ComputeServer
11/// while ensuring thread-safety
12pub trait ComputeChannel<Server: ComputeServer>: Clone + core::fmt::Debug + Send + Sync {
13    /// Given bindings, returns owned resources as bytes
14    fn read(&self, bindings: Vec<Binding>) -> impl Future<Output = Vec<Vec<u8>>> + Send;
15
16    /// Given bindings, returns owned resources as bytes
17    fn read_tensor(
18        &self,
19        bindings: Vec<BindingWithMeta>,
20    ) -> impl Future<Output = Vec<Vec<u8>>> + Send;
21
22    /// Given a resource handle, return the storage resource.
23    fn get_resource(
24        &self,
25        binding: Binding,
26    ) -> BindingResource<<Server::Storage as ComputeStorage>::Resource>;
27
28    /// Given a resource as bytes, stores it and returns the resource handle
29    fn create(&self, data: &[u8]) -> Handle;
30
31    /// Given a resource as bytes and a shape, stores it and returns the tensor handle
32    fn create_tensor(&self, data: &[u8], shape: &[usize], elem_size: usize)
33    -> (Handle, Vec<usize>);
34
35    /// Reserves `size` bytes in the storage, and returns a handle over them
36    fn empty(&self, size: usize) -> Handle;
37
38    /// Reserves a tensor with `shape` in the storage, and returns a handle to it
39    fn empty_tensor(&self, shape: &[usize], elem_size: usize) -> (Handle, Vec<usize>);
40
41    /// Executes the `kernel` over the given `bindings`.
42    ///
43    /// # Safety
44    ///
45    /// When executing with mode [ExecutionMode::Unchecked], out-of-bound reads and writes can happen.
46    unsafe fn execute(
47        &self,
48        kernel: Server::Kernel,
49        count: CubeCount,
50        bindings: Bindings,
51        mode: ExecutionMode,
52    );
53
54    /// Flush outstanding work of the server.
55    fn flush(&self);
56
57    /// Wait for the completion of every task in the server.
58    fn sync(&self) -> impl Future<Output = ()> + Send;
59
60    /// Get the current memory usage of the server.
61    fn memory_usage(&self) -> crate::memory_management::MemoryUsage;
62
63    /// Ask the server to release memory that it can release.
64    fn memory_cleanup(&self);
65
66    /// Start a profile on the server. This allows you to profile kernels.
67    ///
68    /// This will measure execution time either by measuring the 'full' execution time by synchronizing
69    /// the execution at the start and the end of the profile, or 'device' time by using device timestamps.
70    /// This function will handle any required synchronization.
71    ///
72    /// Recursive profiling is not allowed and will panic.
73    fn start_profile(&self);
74
75    /// End the profile and return a [`ProfileDuration`].
76    ///
77    /// You can retrieve the Duration of the client profile asynchronously. This function will handle any required synchronization.
78    fn end_profile(&self) -> ProfileDuration;
79}
cubecl_runtime/channel/base.rs

cubecl_runtime/channel/
base.rs
