Struct WgpuServer 

pub struct WgpuServer<C: WgpuCompiler> {
    pub compilation_options: WgpuCompilationOptions,
    /* private fields */
}

Wgpu compute server.

Fields

compilation_options: WgpuCompilationOptions

Implementations

impl<C: WgpuCompiler> WgpuServer<C>

pub fn load_cached_pipeline( &mut self, kernel_id: &KernelId, bindings: &KernelArguments, mode: ExecutionMode, ) -> Result<Option<Result<PipelineEntry, (u64, KernelCacheKey)>>, CompilationError>

Loads a cached kernel if present and creates the pipeline for it. Returns None if the cache isn’t enabled, Some(Ok(pipeline)) if a cache entry was found, and Some(Err(cache_key)) if the cache is enabled but doesn’t contain this kernel.

pub fn create_module( &self, entrypoint_name: &str, cube_dim: CubeDim, source: ModuleSource<'_>, mode: ExecutionMode, ) -> Result<ShaderModule, CompilationError>

pub fn create_pipeline( &self, entrypoint_name: &str, repr: Option<AutoRepresentationRef<'_>>, module: ShaderModule, bindings: &KernelArguments, ) -> Arc<ComputePipeline>

impl<C: WgpuCompiler> WgpuServer<C>

pub fn new( memory_properties: MemoryDeviceProperties, memory_config: MemoryConfiguration, compilation_options: WgpuCompilationOptions, device: Device, queue: Queue, tasks_max: usize, backend: Backend, timing_method: TimingMethod, utilities: ServerUtilities, ) -> Self

Create a new server.

Trait Implementations

impl<C: Debug + WgpuCompiler> Debug for WgpuServer<C>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<C: WgpuCompiler> DeviceService for WgpuServer<C>

fn init(device_id: DeviceId) -> Self

Initializes the service. It is only called once per device.

fn utilities(&self) -> ServerUtilitiesHandle

Get the service utilities.

fn stage() -> DeviceServiceStage

where Self: Sized,

Which pipeline stage this service runs on.

impl<C: WgpuCompiler> Server for WgpuServer<C>

fn sync( &mut self, handles: Vec<BufferBinding>, stream_id: StreamId, ) -> DynFut<Result<(), ServerError>>

Returns the total time of GPU work this sync completes.

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

Retrieve the server logger.

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

Retrieve the server utilities.

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 initialize_memory( &mut self, memory: ManagedMemoryHandle, size: u64, stream_id: StreamId, )

Initializes memory on the given stream with the given size.

fn read( &mut self, descriptors: Vec<CopyDescriptor>, stream_id: StreamId, ) -> DynFut<Result<Vec<Bytes>, ServerError>>

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 check( &mut self, handles: Vec<BufferBinding>, _stream_id: StreamId, ) -> Result<(), ServerError>

Whether the bytes the handles name can be trusted, right now and with no barrier: the claim check a read makes, without the read. Instant — enqueue-time failures only. A device fault needs sync, which drains first.

unsafe fn launch( &mut self, kernel: Box<dyn CubeKernel>, count: CubeCount, args: KernelArguments, stream_id: StreamId, launch_mode: LaunchMode, )

Executes the kernel over the given memory handles.

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

Flush all outstanding tasks in the server.

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 abandon_profile(&mut self, stream_id: StreamId, token: ProfilingToken)

Drop the window token opened without measuring it, for a caller that will never close it with end_profile.

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

Memory usage of the given stream.

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

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 in cubecl-server.

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

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

fn memory_cleanup(&mut self, stream_id: StreamId)

Ask the server to release memory that it can release.

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

Update the memory mode of allocation in the server.

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.

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.

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.

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, ) -> Result<(), ServerError>

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.

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.

impl<C: WgpuCompiler> ServerCommunication for WgpuServer<C>

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

Ensure that all queued collective operations have been executed.

fn comm_init(&mut self, device_ids: Vec<DeviceId>) -> Result<(), ServerError>

Initialize the communication between the devices in device_ids.

fn all_reduce( &mut self, src: BufferBinding, dst: BufferBinding, dtype: ElemType, stream_id: StreamId, op: ReduceOperation, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Performs an all_reduce operation on the input data and writes it to the output buffer. see https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce

fn send( &mut self, desc: CopyDescriptor, dtype: ElemType, stream_id: StreamId, device_id_dst: DeviceId, ) -> Result<(), ServerError>

Sends data from this server to a destination server.

fn recv( &mut self, handle: Handle, dtype: ElemType, stream_id: StreamId, device_id_src: DeviceId, ) -> Result<(), ServerError>

Receive data from another server.

impl<C: WgpuCompiler> ServerStorage for WgpuServer<C>

type Storage = WgpuStorage

The storage type defines how data is stored and accessed.

fn get_resource( &mut self, binding: BufferBinding, stream_id: StreamId, ) -> Result<ManagedResource<WgpuResource>, ServerError>

Given a resource handle, returns the storage resource.

impl<C: WgpuCompiler> WriteScoped for WgpuServer<C>

type Streams = SchedulerMultiStream<ScheduledWgpuBackend>

The multi-stream driver the server keeps.

fn write_streams(&mut self) -> &mut Self::Streams

The driver, split-borrowed from the server so the scope can claim on the way in and settle on the way out while body holds the rest.

fn on_failure(&mut self, stream: StreamId, error: &ServerError)

Told about every failure a scope settles, so a measurement in flight on stream is invalidated wherever the failure happened.

fn capturing(&mut self, stream: StreamId) -> Option<&mut StreamCapture>

The graph capture of the pooled stream stream folds onto, for a server that records captures. Defaults to None, for one that does not.

fn write_set(&mut self) -> Vec<BufferBinding>

An empty write set for the caller to fill with what the work it is about to run writes.