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: Binding,
        stream_id: StreamId,
    ) -> Result<ManagedResource<<Self::Storage as ComputeStorage>::Resource>, ServerError>;
    unsafe fn launch(
        &mut self,
        kernel: Self::Kernel,
        count: CubeCount,
        bindings: KernelArguments,
        kind: ExecutionMode,
        stream_id: StreamId,
    );
    fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>;
    fn memory_usage(
        &mut self,
        stream_id: StreamId,
    ) -> Result<MemoryUsage, 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 configure_memory_pools(
        &mut self,
        config: MemoryConfiguration,
        stream_id: StreamId,
    ) -> bool { ... }
}

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: Binding, 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, kind: ExecutionMode, stream_id: StreamId, )

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.

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_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 itself needs no fresh device allocation — a device malloc inside the capture is illegal.

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 hardware-graph backend 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 as a single dispatch. Between this call and end_capture the stream must not synchronize or allocate fresh device memory — call graph_prepare and warm up first.

The default is unsupported; a backend with hardware graph support (CUDA, HIP) overrides these methods.

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 — one dispatch that re-runs the whole recorded launch sequence against its original buffers.

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 its executable and unpinning the buffers it retained. The backend must ensure any in-flight replay on stream_id has completed first (replay returns at enqueue time). 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 configure_memory_pools( &mut self, config: MemoryConfiguration, stream_id: StreamId, ) -> bool

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

The calling stream’s pools are rebuilt in place (see MemoryManagement::configure — 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).

Returns true when the calling stream’s pools were rebuilt now, and false when they kept the old layout (something was still live in them); the layout still applies to streams created afterwards.

The default is a no-op returning false for servers without configurable pools.

Dyn Compatibility

This trait is not dyn compatible.

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

Implementors