Trait ComputeServer 

pub trait ComputeServer:
    Send
    + Debug
    + ServerCommunication
    + DeviceService
    + 'static
    + Sized {
    type Kernel: KernelMetadata;
    type Info: Debug + Send + Sync;
    type MemoryLayoutPolicy: MemoryLayoutPolicy;
    type Storage: ComputeStorage;

    // Required methods
    fn initialize_memory(
        &mut self,
        memory: ManagedMemoryHandle,
        size: u64,
        stream_id: StreamId,
    );
    fn logger(&self) -> Arc<ServerLogger> ;
    fn utilities(&self) -> Arc<ServerUtilities<Self>> ;
    fn read(
        &mut self,
        descriptors: Vec<CopyDescriptor>,
        stream_id: StreamId,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<Bytes>, ServerError>> + Send>>;
    fn write(
        &mut self,
        descriptors: Vec<(CopyDescriptor, Bytes)>,
        stream_id: StreamId,
    );
    fn sync(
        &mut self,
        stream_id: StreamId,
    ) -> Pin<Box<dyn Future<Output = Result<(), ServerError>> + Send>>;
    fn get_resource(
        &mut self,
        binding: BufferBinding,
        stream_id: StreamId,
    ) -> Result<ManagedResource<<Self::Storage as ComputeStorage>::Resource>, ServerError>;
    unsafe fn launch(
        &mut self,
        kernel: Self::Kernel,
        count: CubeCount,
        bindings: KernelArguments,
        stream_id: StreamId,
        launch_mode: LaunchMode,
    );
    fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>;
    fn memory_usage(
        &mut self,
        stream_id: StreamId,
    ) -> Result<MemoryUsage, ServerError>;
    fn memory_report(
        &mut self,
        stream_id: StreamId,
    ) -> Result<MemoryReport, ServerError>;
    fn memory_cleanup(&mut self, stream_id: StreamId);
    fn start_profile(
        &mut self,
        stream_id: StreamId,
    ) -> Result<ProfilingToken, ServerError>;
    fn end_profile(
        &mut self,
        stream_id: StreamId,
        token: ProfilingToken,
    ) -> Result<ProfileDuration, ProfileError>;
    fn allocation_mode(
        &mut self,
        mode: MemoryAllocationMode,
        stream_id: StreamId,
    );

    // Provided methods
    fn staging(
        &mut self,
        _sizes: &[usize],
        _stream_id: StreamId,
    ) -> Result<Vec<Bytes>, ServerError> { ... }
    fn graph_prepare(&mut self, stream_id: StreamId) -> Result<(), ServerError> { ... }
    fn begin_capture(&mut self, stream_id: StreamId) -> Result<(), ServerError> { ... }
    fn end_capture(
        &mut self,
        stream_id: StreamId,
    ) -> Result<GraphId, ServerError> { ... }
    fn replay(&mut self, graph: GraphId, stream_id: StreamId) { ... }
    fn graph_destroy(&mut self, graph: GraphId, stream_id: StreamId) { ... }
    fn stream_ids(&self) -> Vec<StreamId> { ... }
    fn install_memory_pools(
        &mut self,
        config: MemoryConfiguration,
        stream_id: StreamId,
    ) -> Result<(), InstallMemoryPoolsError> { ... }
}

The compute server is responsible for handling resources and computations over resources.

Everything in the server is mutable, therefore it should be solely accessed through the ComputeClient for thread safety.

Required Associated Types

type Kernel: KernelMetadata

The kernel type defines the computation algorithms.

type Info: Debug + Send + Sync

Information that can be retrieved for the runtime.

type MemoryLayoutPolicy: MemoryLayoutPolicy

Manages how allocations are performed for a server.

type Storage: ComputeStorage

The storage type defines how data is stored and accessed.

Required Methods

fn initialize_memory( &mut self, memory: ManagedMemoryHandle, size: u64, stream_id: StreamId, )

Initializes memory on the given stream with the given size.

fn logger(&self) -> Arc<ServerLogger>

Retrieve the server logger.

fn utilities(&self) -> Arc<ServerUtilities<Self>>

Retrieve the server utilities.

fn read( &mut self, descriptors: Vec<CopyDescriptor>, stream_id: StreamId, ) -> Pin<Box<dyn Future<Output = Result<Vec<Bytes>, ServerError>> + Send>>

Given bindings, returns the owned resources as bytes.

fn write( &mut self, descriptors: Vec<(CopyDescriptor, Bytes)>, stream_id: StreamId, )

Writes the specified bytes into the buffers given

fn sync( &mut self, stream_id: StreamId, ) -> Pin<Box<dyn Future<Output = Result<(), ServerError>> + Send>>

Wait for the completion of every task in the server.

fn get_resource( &mut self, binding: BufferBinding, stream_id: StreamId, ) -> Result<ManagedResource<<Self::Storage as ComputeStorage>::Resource>, ServerError>

Given a resource handle, returns the storage resource.

unsafe fn launch( &mut self, kernel: Self::Kernel, count: CubeCount, bindings: KernelArguments, stream_id: StreamId, launch_mode: LaunchMode, )

Executes the kernel over the given memory handles.

Kernels have mutable access to every resource they are given and are responsible of determining which should be read or written.

launch_mode says whether the kernel actually runs. On LaunchMode::Skip the server must still do everything a first launch does short of dispatching — expand, compile, validate, fill its caches — and then drop the launch; skipping the compilation instead would defeat the whole point of a dry run.

Safety

When executing with mode [ExecutionMode::Unchecked], out-of-bound reads and writes can happen.

fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Flush all outstanding tasks in the server.

fn memory_usage( &mut self, stream_id: StreamId, ) -> Result<MemoryUsage, ServerError>

Memory usage of the given stream.

fn memory_report( &mut self, stream_id: StreamId, ) -> Result<MemoryReport, ServerError>

Structured per-pool report of the given 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 — see MemoryManagement::memory_report.

fn memory_cleanup(&mut self, stream_id: StreamId)

Ask the server to release memory that it can release.

fn start_profile( &mut self, stream_id: StreamId, ) -> Result<ProfilingToken, ServerError>

Enable collecting timestamps.

fn end_profile( &mut self, stream_id: StreamId, token: ProfilingToken, ) -> Result<ProfileDuration, ProfileError>

Disable collecting timestamps.

fn allocation_mode(&mut self, mode: MemoryAllocationMode, stream_id: StreamId)

Update the memory mode of allocation in the server.

Provided Methods

fn staging( &mut self, _sizes: &[usize], _stream_id: StreamId, ) -> Result<Vec<Bytes>, ServerError>

Reserves N Bytes of the provided sizes to be used as staging to load data.

fn graph_prepare(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Prepare stream_id for an upcoming graph capture: route allocations into a stable pool and snapshot it, so every buffer allocated between here and end_capture can be pinned for the graph’s lifetime. Call this before the warmup run so the capture window reuses the slices warmup left in the pool rather than allocating its own — which a hardware-graph backend cannot do at all (a device malloc inside the capture is illegal there), and which on any backend would grow the memory a graph pins beyond what it replays against.

Prefer having kernels already autotuned before this call: any transient benchmark buffers autotune allocates while the window is armed are forced into the persistent pool and pinned to the graph, so a graph captured over a cold autotune cache retains more device memory than it replays against. Warm the autotune cache first, then graph_prepare and warm up only to populate the pool.

A no-op by default (harmless on backends without graph support); a backend with graph support enables its persistent pool + capture recording.

fn begin_capture(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Begin recording the launches issued on stream_id into a graph instead of executing them, so the sequence can later be replayed without paying the launch path again. Call graph_prepare and warm up first.

Between this call and end_capture the stream must not synchronize — a read, a sync or a profile either aborts the capture or is refused — and should not allocate fresh device memory, which graph_prepare plus a warmup run is what avoids. Whether an operation the window cannot record fails the call or fails end_capture, and whether a mid-window allocation is fatal, is the backend’s to say; see StreamCaptureState::Capture.

The default is unsupported. Two shapes of backend override it: a hardware graph (CUDA, HIP), where the driver records a replayable graph object, and a software graph (wgpu), where the runtime records fully-resolved dispatches and re-encodes them on replay.

fn end_capture(&mut self, stream_id: StreamId) -> Result<GraphId, ServerError>

Stop recording (see begin_capture), store the captured graph in the backend’s registry, and return its GraphId, ready to replay.

fn replay(&mut self, graph: GraphId, stream_id: StreamId)

Replay the graph identified by graph on stream_id, re-running the whole recorded launch sequence against its original buffers. A hardware graph replays as a single dispatch; a software graph re-encodes the recorded dispatches, which is still far cheaper than the launch path but stays O(n) in recorded launches.

Fire-and-forget, like launch: the call enqueues the dispatch and returns without waiting, so a failure is not returned here — it is pushed onto the stream’s error queue and surfaces on the next flush/sync, which leaves the server unhealthy until drained. A no-op by default: a GraphId can only come from end_capture, unsupported here.

fn graph_destroy(&mut self, graph: GraphId, stream_id: StreamId)

Release the graph identified by graph, destroying whatever it recorded and unpinning the buffers it retained. Replay returns at enqueue time, so the backend must guarantee no in-flight replay can still read those buffers once they return to the pool — by syncing stream_id where nothing weaker will do (CUDA, HIP), or by relying on the queue ordering that already places a submitted replay ahead of any later write (wgpu). A no-op by default and for an unknown id.

fn stream_ids(&self) -> Vec<StreamId>

Stream ids the client should iterate to aggregate across the device.

Default is just the calling stream, which is correct for non-multi-stream backends; multi-stream backends override to return one id per initialized stream pool slot.

fn install_memory_pools( &mut self, config: MemoryConfiguration, stream_id: StreamId, ) -> Result<(), InstallMemoryPoolsError>

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

The calling stream’s pools are rebuilt in place (see MemoryManagement::install_pools — a rebuild only happens when nothing is live in them), and the layout becomes the one every stream created afterwards is built with. Pool layouts are a purely programmatic, runtime setting — there is no config-file pathway — so callers size them per workload (e.g. per model, just before loading it).

Errors

PoolsInUse when the calling 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. The layout still applies to streams created afterwards; retry to rebuild the calling stream too.

StreamUnavailable when the calling stream is already in an error state, so its pools could not be reached. Future streams still get the layout.

Unsupported from servers without configurable pools, which is the default implementation.

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors