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<(Arc<ComputePipeline>, CompilerInfo), (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, repr: Option<AutoRepresentationRef<'_>>, source: &str, 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>, ) -> Self

Create a new server.

Trait Implementations

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

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

Returns the total time of GPU work this sync completes.

type Kernel = Box<dyn CubeTask<C>>

The kernel type defines the computation algorithms.

type Storage = WgpuStorage

The storage type defines how data is stored and accessed.

type MemoryLayoutPolicy = ContiguousMemoryLayoutPolicy

Manages how allocations are performed for a server.

type Info = Backend

Information that can be retrieved for the runtime.

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

Retrieve the server logger.

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

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 get_resource( &mut self, binding: BufferBinding, stream_id: StreamId, ) -> Result<ManagedResource<WgpuResource>, ServerError>

Given a resource handle, returns the storage resource.

unsafe fn launch( &mut self, kernel: Self::Kernel, 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 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 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)

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

Which pipeline stage this service runs on.

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

const SERVER_COMM_ENABLED: bool = false

Indicates whether server-to-server communication is enabled for this implementation.

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.