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, definition: &KernelDefinition, 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: Binding, 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, mode: ExecutionMode, 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 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 configure_memory_pools( &mut self, config: MemoryConfiguration, stream_id: StreamId, ) -> bool

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

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.

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.

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.

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: Binding, dst: Binding, 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.