Struct ComputeClient 

pub struct ComputeClient<R: Runtime> { /* private fields */ }

The ComputeClient is the entry point to require tasks from the ComputeServer. It should be obtained for a specific device via the Compute struct.

Implementations

impl<R: Runtime> ComputeClient<R>

pub fn info(&self) -> &<R::Server as ComputeServer>::Info

Get the info of the current backend.

pub fn init<D: Device>(device: &D, server: R::Server) -> Self

Create a new client with a new server.

pub fn load<D: Device>(device: &D) -> Self

Load the client for the given device.

pub unsafe fn set_stream(&mut self, stream_id: StreamId)

Set the stream in which the current client is operating on.

Safety

This is highly unsafe and should probably only be used by the CubeCL/Burn projects for now.

pub fn read_async( &self, handles: Vec<Handle>, ) -> impl Future<Output = Result<Vec<Bytes>, ServerError>> + Send

Given bindings, returns owned resources as bytes.

pub fn read(&self, handles: Vec<Handle>) -> Vec<Bytes>

Given bindings, returns owned resources as bytes.

Remarks

Panics if the read operation fails.

pub fn read_one(&self, handle: Handle) -> Result<Bytes, ServerError>

Given a binding, returns owned resource as bytes.

pub fn read_one_unchecked(&self, handle: Handle) -> Bytes

Given a binding, returns owned resource as bytes.

Remarks

Panics if the read operation fails. Useful for tests.

pub fn read_tensor_async( &self, descriptors: Vec<CopyDescriptor>, ) -> impl Future<Output = Result<Vec<Bytes>, ServerError>> + Send

Given bindings, returns owned resources as bytes.

pub fn read_tensor(&self, descriptors: Vec<CopyDescriptor>) -> Vec<Bytes>

Given bindings, returns owned resources as bytes.

Remarks

Panics if the read operation fails.

The tensor must be in the same layout as created by the runtime, or more strict. Contiguous tensors are always fine, strided tensors are only ok if the stride is similar to the one created by the runtime (i.e. padded on only the last dimension). A way to check stride compatibility on the runtime will be added in the future.

Also see ComputeClient::create_tensor.

pub fn read_one_tensor_async( &self, descriptor: CopyDescriptor, ) -> impl Future<Output = Result<Bytes, ServerError>> + Send

Given a binding, returns owned resource as bytes. See ComputeClient::read_tensor

pub fn read_one_unchecked_tensor(&self, descriptor: CopyDescriptor) -> Bytes

Given a binding, returns owned resource as bytes.

Remarks

Panics if the read operation fails. See ComputeClient::read_tensor

pub fn read_lazy(&self, descriptor: CopyDescriptor) -> Bytes

Reads the device resource described by descriptor lazily.

The returned Bytes only performs the device-to-host copy on first access (e.g. during serialization), keeping the source allocation alive until then. This lets a large number of device tensors be serialized without materializing them all in host memory at once: drain the Bytes sequentially rather than holding them all alive.

The data reflects the device state at first access, so the buffer must not be mutated between this call and the first read.

pub fn read_lazy_async( &self, descriptor: CopyDescriptor, ) -> impl Future<Output = Result<Bytes, ServerError>> + Send

Reads the device resource described by descriptor lazily, async variant.

On native targets the returned future is immediately ready and yields a lazy Bytes whose device-to-host copy is deferred to first access (see read_lazy).

pub fn get_resource( &self, handle: Handle, ) -> Result<ManagedResource<<<R::Server as ComputeServer>::Storage as ComputeStorage>::Resource>, ServerError>

Given a resource handle, returns the storage resource.

pub fn create_from_slice(&self, slice: &[u8]) -> Handle

Returns a resource handle containing the given data.

Notes

Prefer using the more efficient Self::create function.

pub fn exclusive<'a, Re: Send + 'static, F: FnOnce() -> Re + Send + 'a>( &'a self, task: F, ) -> Result<Re, ServerError>

todo: docs

pub fn memory_persistent_allocation<'a, Re: Send, Input: Send, F: FnOnce(Input) -> Re + Send + 'a>( &'a self, input: Input, task: F, ) -> Result<Re, ServerError>

dodo: Docs

pub fn write(&self, handle: &Handle, data: Bytes)

Write data into an existing allocation, in place (same device pointer).

This is how a captured Graph’s inputs are refreshed between replays: the graph records raw device pointers, so new input bytes must land in the very buffer the capture read from. Issue it from the capture stream (see the stream-ordering notes on Graph) so the write orders against the replays instead of racing them.

Non-blocking: the write is enqueued on this client’s current stream.

pub fn create(&self, data: Bytes) -> Handle

Returns a resource handle containing the given Bytes.

pub fn create_tensor_from_slice( &self, slice: &[u8], shape: Shape, elem_size: usize, ) -> MemoryLayout

Given a resource and shape, stores it and returns the tensor handle and strides. This may or may not return contiguous strides. The layout is up to the runtime, and care should be taken when indexing.

Currently the tensor may either be contiguous (most runtimes), or “pitched”, to use the CUDA terminology. This means the last (contiguous) dimension is padded to fit a certain alignment, and the strides are adjusted accordingly. This can make memory accesses significantly faster since all rows are aligned to at least 16 bytes (the maximum load width), meaning the GPU can load as much data as possible in a single instruction. It may be aligned even more to also take cache lines into account.

However, the stride must be taken into account when indexing and reading the tensor (also see ComputeClient::read_tensor).

Notes

Prefer using Self::create_tensor for better performance.

pub fn create_tensor( &self, bytes: Bytes, shape: Shape, elem_size: usize, ) -> MemoryLayout

Given a resource and shape, stores it and returns the tensor handle and strides. This may or may not return contiguous strides. The layout is up to the runtime, and care should be taken when indexing.

Currently the tensor may either be contiguous (most runtimes), or “pitched”, to use the CUDA terminology. This means the last (contiguous) dimension is padded to fit a certain alignment, and the strides are adjusted accordingly. This can make memory accesses significantly faster since all rows are aligned to at least 16 bytes (the maximum load width), meaning the GPU can load as much data as possible in a single instruction. It may be aligned even more to also take cache lines into account.

However, the stride must be taken into account when indexing and reading the tensor (also see ComputeClient::read_tensor).

pub fn create_tensors_from_slices( &self, descriptors: Vec<(MemoryLayoutDescriptor, &[u8])>, ) -> Vec<MemoryLayout>

Reserves all shapes in a single storage buffer, copies the corresponding data into each handle, and returns the handles for them. See ComputeClient::create_tensor

Notes

Prefer using Self::create_tensors for better performance.

pub fn create_tensors( &self, descriptors: Vec<(MemoryLayoutDescriptor, Bytes)>, ) -> Vec<MemoryLayout>

Reserves all shapes in a single storage buffer, copies the corresponding data into each handle, and returns the handles for them. See ComputeClient::create_tensor

pub fn empty(&self, size: usize) -> Handle

Reserves size bytes in the storage, and returns a handle over them.

pub fn empty_tensor(&self, shape: Shape, elem_size: usize) -> MemoryLayout

Reserves shape in the storage, and returns a tensor handle for it. See ComputeClient::create_tensor

pub fn empty_tensors( &self, descriptors: Vec<MemoryLayoutDescriptor>, ) -> Vec<MemoryLayout>

Reserves all shapes in a single storage buffer, and returns the handles for them. See ComputeClient::create_tensor

pub fn staging<'a, I>(&self, bytes: I, file_only: bool)

where I: Iterator<Item = &'a mut Bytes>,

Marks the given Bytes as being a staging buffer, maybe transferring it to pinned memory for faster data transfer with compute device.

TODO: This blocks the compute queue, so it will drop the compute utilization.

pub fn to_client( &mut self, src: Handle, dst_server: &Self, dtype: ElemType, ) -> Handle

Transfer data from one client to another

pub fn ensure_init_collective(&mut self, device_ids: Vec<DeviceId>)

Perform an all_reduce operation on the given devices.

pub fn sync_collective(&self)

Wait on the communication stream.

pub fn all_reduce( &mut self, src: Handle, dst: Handle, dtype: ElemType, device_ids: Vec<DeviceId>, op: ReduceOperation, )

Perform an all_reduce operation on the given devices.

pub fn to_client_tensor( &mut self, src_descriptor: CopyDescriptor, dst_server: &Self, dtype: ElemType, ) -> Handle

Transfer data from one client to another

Make sure the source description can be read in a contiguous manner.

pub fn launch( &self, kernel: <R::Server as ComputeServer>::Kernel, count: CubeCount, bindings: KernelArguments, )

Launches the kernel with the given bindings.

pub unsafe fn launch_unchecked( &self, kernel: <R::Server as ComputeServer>::Kernel, count: CubeCount, bindings: KernelArguments, )

Launches the kernel with the given bindings without performing any bound checks.

Safety

To ensure this is safe, you must verify your kernel:

• Has no out-of-bound reads and writes that can happen.
• Has no infinite loops that might never terminate.

pub fn flush(&self) -> Result<(), ServerError>

Flush all outstanding commands.

pub fn graph_prepare(&self) -> Result<(), ServerError>

Prepare this client’s stream for a graph capture (see ComputeServer::graph_prepare) — enable the persistent pool + capture recording. Call this before the warmup run, then start_capture around the run to record.

pub fn start_capture(&self) -> Result<(), ServerError>

Begin recording launches on this client’s stream into a graph rather than executing them (see ComputeServer::begin_capture). Pin the client to a dedicated stream with set_stream, then graph_prepare and warm up first; between this and stop_capture no sync or fresh allocation may happen. Returns an error on backends without graph support.

pub fn stop_capture(&self) -> Result<Graph<R>, ServerError>

Stop recording and return the captured graph, ready to replay.

pub fn sync(&self) -> DynFut<Result<(), ServerError>>

Wait for the completion of every task in the server.

pub fn properties(&self) -> &DeviceProperties

Get the features supported by the compute server.

pub fn features(&self) -> &Features

Get the features supported by the compute server.

pub fn properties_mut(&mut self) -> Option<&mut DeviceProperties>

Warning

For private use only.

pub fn memory_usage(&self) -> Result<MemoryUsage, ServerError>

Total memory usage across all streams on this client’s device.

The closure iterates the server’s stream_ids() and folds each per-stream memory_usage(id) with MemoryUsage::combine, so the result is correct regardless of which thread queries it.

pub fn enumerate_devices(&self, type_id: u16) -> Vec<DeviceId>

Get all devices of a specific type available to this runtime

pub fn enumerate_all_devices(&self) -> Vec<DeviceId>

Get all devices available to this runtime

pub fn device_count(&self, type_id: u16) -> usize

Get the number of devices of a specific type available to this runtime

pub fn device_count_total(&self) -> usize

Get the number of devices of a specific type available to this runtime

pub unsafe fn allocation_mode(&self, mode: MemoryAllocationMode)

Change the memory allocation mode.

Safety

This function isn’t thread safe and might create memory leaks.

pub fn memory_cleanup(&self)

Ask the client to release memory that it can release.

Nb: Results will vary on what the memory allocator deems beneficial, so it’s not guaranteed any memory is freed.

pub fn configure_memory_pools(&self, pools: &MemoryPoolsConfig) -> bool

Install a new dynamic-pool layout for the device’s main GPU memory.

Pool layouts are a purely programmatic, runtime setting — there is no config-file pathway — sized per workload (e.g. per model, just before loading it). The current stream’s pools are rebuilt in place when nothing is live in them (reconfigure at a quiescent point, e.g. right after unloading a model), and the layout applies to every stream created afterwards. Auxiliary pools (pinned CPU, staging, uniforms) and the persistent pool are never affected.

Returns true when the current stream’s pools were rebuilt now. Returns false when they kept the old layout because something was still live in them — e.g. a garbage-collection task that has not released its cross-stream pins yet, which can lag behind an explicit memory_cleanup. The layout still applies to streams created afterwards; retry after the remaining work drains to rebuild the current stream too.

Panics

Panics if the layout is invalid (empty list, too many pools, zero page size, slice larger than page, cap smaller than page, unavailable preset) — an explicit layout that cannot be honored must not be silently replaced.

pub fn profile<O: Send + 'static>( &self, func: impl FnOnce() -> O + Send, func_name: &str, ) -> Result<(O, ProfileDuration), ProfileError>

Measure the execution time of some inner operations.

pub fn io_optimized_vector_sizes( &self, size: usize, ) -> impl Iterator<Item = VectorSize> + Clone

Returns all vector sizes that are useful to perform optimal IO operation on the given element.

pub fn measure_throughput( &self, key: ThroughputKey, kernel_config: KernelConfig, ) -> ThroughputValue

Calculates the maximum throughput of the device given the given config (like tensor core with certain sizes and dtypes, or just arithmetic by dtype)

pub fn measure_launch_overhead(&self, sample: impl Fn() -> Duration) -> Duration

Calculates the launch overhead of the device by sampling.

Trait Implementations

impl<R: Runtime> Clone for ComputeClient<R>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.