cubecl_runtime/channel/
base.rs

1use core::future::Future;
2use cubecl_common::benchmark::TimestampsResult;
3
4use crate::{
5    server::{Binding, ComputeServer, CubeCount, Handle},
6    storage::BindingResource,
7    ExecutionMode,
8};
9use alloc::vec::Vec;
10
11/// The ComputeChannel trait links the ComputeClient to the ComputeServer
12/// while ensuring thread-safety
13pub trait ComputeChannel<Server: ComputeServer>: Clone + core::fmt::Debug + Send + Sync {
14    /// Given bindings, returns owned resources as bytes
15    fn read(&self, bindings: Vec<Binding>) -> impl Future<Output = Vec<Vec<u8>>> + Send;
16
17    /// Given a resource handle, return the storage resource.
18    fn get_resource(&self, binding: Binding) -> BindingResource<Server>;
19
20    /// Given a resource as bytes, stores it and returns the resource handle
21    fn create(&self, data: &[u8]) -> Handle;
22
23    /// Reserves `size` bytes in the storage, and returns a handle over them
24    fn empty(&self, size: usize) -> Handle;
25
26    /// Executes the `kernel` over the given `bindings`.
27    ///
28    /// # Safety
29    ///
30    /// When executing with mode [ExecutionMode::Unchecked], out-of-bound reads and writes can happen.
31    unsafe fn execute(
32        &self,
33        kernel: Server::Kernel,
34        count: CubeCount,
35        bindings: Vec<Binding>,
36        mode: ExecutionMode,
37    );
38
39    /// Flush outstanding work of the server.
40    fn flush(&self);
41
42    /// Wait for the completion of every task in the server.
43    fn sync(&self) -> impl Future<Output = ()> + Send;
44
45    /// Wait for the completion of every task in the server.
46    ///
47    /// Returns the (approximate) total amount of GPU work done since the last sync.
48    fn sync_elapsed(&self) -> impl Future<Output = TimestampsResult> + Send;
49
50    /// Get the current memory usage of the server.
51    fn memory_usage(&self) -> crate::memory_management::MemoryUsage;
52
53    /// Enable collecting timestamps.
54    fn enable_timestamps(&self);
55
56    /// Disable collecting timestamps.
57    fn disable_timestamps(&self);
58}
cubecl_runtime/channel/base.rs

cubecl_runtime/channel/
base.rs
