Struct Client 

pub struct Client { /* private fields */ }

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

Implementations

impl Client

pub fn name(&self) -> &'static str

The runtime name on this device, as logs and cache keys show it.

pub fn init<S: ServerStorage>(device_id: DeviceId, server: S) -> Self

Create a new client with a new server.

pub fn load<S: ServerStorage>(device_id: DeviceId) -> Self

Load the client for the given device, starting a server of type S there if none runs yet.

pub fn service_id(&self) -> ServiceId

The service this client reaches: what its handles are stamped with.

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 Client::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 Client::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 Client::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<S: ServerStorage>( &self, handle: Handle, ) -> Result<ManagedResource<<S::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>

Run task with this device to itself, so nothing else is scheduled against it for the duration.

Errors

The device could not be taken exclusively — another holder has it, or its runner is gone. Nothing ran, so the caller may retry.

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

Run task with every allocation it makes routed to the persistent pool, then restore the previous mode.

Persistent slices are exact-fit and are not reclaimed by the ordinary sweep, which is what weights want: allocated once, alive for the process, and stable enough for a graph capture to record against.

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 Client::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 Client::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 Client::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 Client::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 Client::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 Client::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.

src must be this client’s. The bytes go device to device when both clients are of the same runtime and it has a collective transport; otherwise, and always across runtimes, they go through the host.

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

Perform an all_reduce operation on the given devices.

pub fn has_device_transport(&self) -> bool

Whether this runtime moves data between its devices itself. Without it, to_client copies through the host and the collectives refuse.

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: Box<dyn CubeKernel>, count: CubeCount, bindings: KernelArguments, )

Launches the kernel with the given bindings.

pub fn check<'a>( &self, handles: impl IntoIterator<Item = &'a Handle>, ) -> Result<(), ServerError>

Whether the bytes behind handles can be trusted, right now and with no barrier: the claim check a read makes, without the read. One lookup per handle, so a fusion layer or an autotuner can recover per tensor instead of tearing down a device.

Instant means enqueue-time failures only — a compile or binding failure is visible here immediately, a device fault is not until the queue drains. sync_buffers is the complete answer; read_one is that plus the copy.

Errors

ServerError::Several naming every failure these buffers carry, each once however many carry it. The bytes are gone, so there is nothing to retry: this is the answer, not a hint.

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 Server::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 Server::begin_capture). Pin the client to a dedicated stream with set_stream, then graph_prepare and warm up first.

Between this and stop_capture the window records launches and nothing else: reading, syncing or profiling the stream is refused, and so is writing to a handle — a recorded graph cannot carry a host copy, so feed fresh inputs by writing between replays instead. A refused write is reported late, by failing stop_capture, rather than handing back a graph that silently skips it. Fresh allocation inside the window is fatal on a hardware-graph backend and merely wasteful on a software-graph one, which is what the warmup run exists to avoid.

Returns an error on backends without graph support.

pub fn stop_capture(&self) -> Result<Graph, 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.

The barrier alone, which also reports a device fault — the only failure left that no buffer can report. A launch failure is not this sync’s to report: it lives on the buffers the launch never wrote and surfaces on any read, check or sync_buffers of those.

pub fn sync_buffers<'a>( &self, handles: impl IntoIterator<Item = &'a Handle>, ) -> DynFut<Result<(), ServerError>>

The barrier, and then an answer for handles.

sync first, so a device fault counts, and then the claim check a read would have made — a read without the read, for the caller that needs to know its work produced something trustworthy and does not want to pull it to the host to find out.

Errors

The device fault the barrier found, or ServerError::Several naming every failure these buffers carry.

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_shared(&self) -> Arc<DeviceProperties>

The device properties, shared: what a kernel keeps to expand itself on the device thread without holding the client.

pub fn target_properties(&self) -> &TargetProperties

What the target this client compiles for guarantees about its own instructions, resolved once when the device came up.

pub fn target_properties_shared(&self) -> Arc<TargetProperties>

The target properties, shared: the other half of what a kernel keeps to expand itself on the device thread without naming a runtime.

Cloning this is one atomic increment, which is why the generated launch functions can afford to do it per launch where calling Runtime::target_properties again would not be.

pub fn memory_usage(&self) -> MemoryUsage

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 memory_report(&self) -> MemoryReport

Structured per-pool report of the calling stream’s main GPU memory: each pool’s shape, usage, and high-water marks, in allocation-routing order.

The read side of a measured memory plan — install a layout with install_memory_pools, measure under a DryRun, cap at the observed peaks; the full cycle is on MemoryReport.

Unlike memory_usage, which aggregates across streams, this reads one stream: pools are per stream, and a plan is measured and installed on the stream that runs the workload.

pub fn record_memory(&self, label: &str)

Write a snapshot of the calling stream’s memory report to the environment’s records, under label. Nothing is read when the environment records nothing.

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 install_memory_pools( &self, pools: &MemoryPoolsConfig, ) -> Result<(), InstallMemoryPoolsError>

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

This replaces the pools themselves, not just a setting they read. It lands in two places:

• The calling stream’s pools are rebuilt in place, discarding the old ones — which is why it only happens when nothing is live in them, and why the high-water marks in memory_report start over.
• The layout becomes the one every stream created afterwards is built with. Other streams that already exist keep theirs; memory is per stream, and rebuilding a stream this call is not synchronized with would swap pools under its live slices.

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), so install at a quiescent point such as right after unloading a model. Auxiliary pools (pinned CPU, staging, uniforms) and the persistent pool are never affected.

Errors

PoolsInUse when the current stream kept its old layout because something was still live in its pools — e.g. a garbage-collection task that has not released its cross-stream pins yet, which can lag behind an explicit memory_cleanup. Nothing is disturbed, the layout still applies to streams created afterwards, and retrying after the remaining work drains rebuilds the current stream too.

Unsupported from a runtime with no configurable pools, where retrying will never succeed.

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) — that is a bad layout literal rather than a runtime condition, and an explicit layout that cannot be honored must not be silently replaced.

pub fn profile_start(&self) -> Result<ProfileWindow, ProfileError>

Open a profiling window at the current position of the calling stream.

Prefer the bracketed profile, which also holds the device for the closure. This pair is for a caller that cannot bracket the work in a closure — a lazy queue drained on another thread, say — and only knows when on the stream its window opens and closes.

The window keeps the stream it was opened on, and profile_end closes it there whichever thread calls it. Nothing keeps other streams’ work out of the window.

An open window costs something on every backend and stays open until it is ended or abandoned, so a caller that bails out between the two calls has to abandon it.

pub fn profile_end( &self, window: ProfileWindow, ) -> Result<ProfileDuration, ProfileError>

Close window at the current position of the stream it was opened on.

pub fn profile_abandon(&self, window: ProfileWindow)

Drop window without measuring it, for a caller that will never reach profile_end, such as an error path between the two calls.

Does not wait for the server to drop it, but does flush, because this is usually a caller’s last word: an abandon left sitting in the queue holds the window open for exactly as long as it is the only thing in there, which is the case it exists for.

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, probe: impl FnOnce() -> Result<ThroughputValue, ThroughputError>, ) -> Result<ThroughputValue, ThroughputError>

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

Errors

Whatever probe reports.

Trait Implementations

impl Clone for Client

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.