Struct WgpuServer 

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

Wgpu compute server.

Implementations

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

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 compile_queued(&mut self)

Compiles every kernel a LaunchMode::Queue launch queued, now rather than inside the next launch, so a measurement that follows times its own kernels only.

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_report(&mut self, stream_id: StreamId) -> Option<StreamMemoryReport>

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

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 graph_prepare(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Prepare stream_id for an upcoming graph capture. Call this before the warmup run: the capture window allocates nothing, so it reuses what the warmup run left in the pools, and every page the recording touches is guarded for the graph’s lifetime.

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.

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

Enable collecting timestamps from the next kernel launched on the stream, for a window that holds one launch.

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, reduction: Reduction, stream_id: StreamId, device_ids: Vec<DeviceId>, wait: ComputeWait, ) -> 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 all_gather( &mut self, src: BufferBinding, dst: BufferBinding, stream_id: StreamId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Writes every rank’s src, concatenated in rank order, to dst, which holds one src per rank of device_ids.

fn reduce_scatter( &mut self, src: BufferBinding, dst: BufferBinding, reduction: Reduction, stream_id: StreamId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Reduces every rank’s src and writes this rank’s chunk of the result to dst, src holding one dst per rank of device_ids.

fn broadcast( &mut self, src: Option<BufferBinding>, dst: BufferBinding, stream_id: StreamId, root: DeviceId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Writes root’s src to dst on every rank of device_ids. Only root has a src.

fn all_to_all( &mut self, src: BufferBinding, dst: BufferBinding, stream_id: StreamId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Sends chunk r of src to rank r and writes the chunk received from rank r to chunk r of dst, both holding one chunk per rank of device_ids.

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

Sends data from this server to a destination server.

fn recv( &mut self, handle: Handle, 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.