cubecl_runtime/channel/
base.rs

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

cubecl_runtime/channel/
base.rs
