Trait GpuBackend 

pub trait GpuBackend {
    type TapeBuffers;

    // Required methods
    fn upload_tape(&self, data: &GpuTapeData) -> Self::TapeBuffers;
    fn num_outputs(&self, tape: &Self::TapeBuffers) -> u32;
    fn forward_batch(
        &self,
        tape: &Self::TapeBuffers,
        inputs: &[f32],
        batch_size: u32,
    ) -> Result<Vec<f32>, GpuError>;
    fn gradient_batch(
        &self,
        tape: &Self::TapeBuffers,
        inputs: &[f32],
        batch_size: u32,
    ) -> Result<(Vec<f32>, Vec<f32>), GpuError>;
    fn sparse_jacobian(
        &self,
        tape: &Self::TapeBuffers,
        tape_cpu: &mut BytecodeTape<f32>,
        x: &[f32],
    ) -> Result<(Vec<f32>, JacobianSparsityPattern, Vec<f32>), GpuError>;
    fn hvp_batch(
        &self,
        tape: &Self::TapeBuffers,
        x: &[f32],
        tangent_dirs: &[f32],
        batch_size: u32,
    ) -> Result<(Vec<f32>, Vec<f32>), GpuError>;
    fn sparse_hessian(
        &self,
        tape: &Self::TapeBuffers,
        tape_cpu: &mut BytecodeTape<f32>,
        x: &[f32],
    ) -> Result<(f32, Vec<f32>, SparsityPattern, Vec<f32>), GpuError>;
    fn taylor_forward_kth_batch(
        &self,
        tape: &Self::TapeBuffers,
        primal_inputs: &[f32],
        direction_seeds: &[f32],
        batch_size: u32,
        order: usize,
    ) -> Result<TaylorKthBatchResult<f32>, GpuError>;

    // Provided method
    fn taylor_forward_2nd_batch(
        &self,
        tape: &Self::TapeBuffers,
        primal_inputs: &[f32],
        direction_seeds: &[f32],
        batch_size: u32,
    ) -> Result<TaylorBatchResult<f32>, GpuError> { ... }
}

Available on crate features gpu-cuda or gpu-wgpu only.

Common interface for GPU backends (f32 operations).

Both WgpuContext and CudaContext implement this trait for the f32 operation set. CUDA additionally provides f64 methods as inherent methods on CudaContext directly.

Associated Type

TapeBuffers is the backend-specific opaque handle returned by upload_tape and passed to all dispatch methods. It holds GPU-resident buffers and is not cloneable.

Implementing a New Backend

A backend must implement all six methods. Construction (new()) is not part of the trait — backends may have different initialization requirements.

Required Associated Types

type TapeBuffers

Backend-specific uploaded tape handle.

Required Methods

fn upload_tape(&self, data: &GpuTapeData) -> Self::TapeBuffers

Upload a tape to the GPU.

The returned handle is used for all subsequent operations and holds GPU-resident buffers for the tape’s opcodes, arguments, and constants.

Panics

Panics if data fails GpuTapeData::validate. The kernels index device memory with these fields, so uploading invalid data would be an out-of-bounds device access; data built by GpuTapeData::from_tape from a valid tape always passes.

fn num_outputs(&self, tape: &Self::TapeBuffers) -> u32

Number of declared outputs on the uploaded tape.

Used by estimators like stde_gpu::laplacian_gpu to enforce single-output assumptions whose coefficient layout depends on the tape’s output count.

fn forward_batch( &self, tape: &Self::TapeBuffers, inputs: &[f32], batch_size: u32, ) -> Result<Vec<f32>, GpuError>

Batched forward evaluation.

inputs is [f32; batch_size * num_inputs] (row-major, one point per row). Returns output values [f32; batch_size * num_outputs].

fn gradient_batch( &self, tape: &Self::TapeBuffers, inputs: &[f32], batch_size: u32, ) -> Result<(Vec<f32>, Vec<f32>), GpuError>

Batched gradient (forward + reverse sweep).

Returns (outputs, gradients) where outputs is [f32; batch_size * num_outputs] and gradients is [f32; batch_size * num_inputs].

fn sparse_jacobian( &self, tape: &Self::TapeBuffers, tape_cpu: &mut BytecodeTape<f32>, x: &[f32], ) -> Result<(Vec<f32>, JacobianSparsityPattern, Vec<f32>), GpuError>

GPU-accelerated sparse Jacobian.

CPU detects sparsity and computes coloring; GPU dispatches colored tangent sweeps. Returns (output_values, pattern, jacobian_values).

fn hvp_batch( &self, tape: &Self::TapeBuffers, x: &[f32], tangent_dirs: &[f32], batch_size: u32, ) -> Result<(Vec<f32>, Vec<f32>), GpuError>

Batched Hessian-vector product (forward-over-reverse).

tangent_dirs is [f32; batch_size * num_inputs] — one direction per batch element. Returns (gradients, hvps) each [f32; batch_size * num_inputs].

fn sparse_hessian( &self, tape: &Self::TapeBuffers, tape_cpu: &mut BytecodeTape<f32>, x: &[f32], ) -> Result<(f32, Vec<f32>, SparsityPattern, Vec<f32>), GpuError>

GPU-accelerated sparse Hessian.

CPU detects Hessian sparsity and computes distance-2 coloring; GPU dispatches HVP sweeps. Returns (value, gradient, pattern, hessian_values).

fn taylor_forward_kth_batch( &self, tape: &Self::TapeBuffers, primal_inputs: &[f32], direction_seeds: &[f32], batch_size: u32, order: usize, ) -> Result<TaylorKthBatchResult<f32>, GpuError>

Available on crate feature stde only.

Batched K-th order Taylor forward propagation.

Supports order in 1..=5. Each batch element pushes one direction through the tape, producing K Taylor coefficients (c0, c1, …, c_{K-1}).

primal_inputs is [f32; batch_size * num_inputs] — primals for each element. direction_seeds is [f32; batch_size * num_inputs] — c1 seeds for each element.

Returns TaylorKthBatchResult with coefficients[k] of size [f32; batch_size * num_outputs] for each k in 0..order.

Provided Methods

fn taylor_forward_2nd_batch( &self, tape: &Self::TapeBuffers, primal_inputs: &[f32], direction_seeds: &[f32], batch_size: u32, ) -> Result<TaylorBatchResult<f32>, GpuError>

Available on crate feature stde only.

Batched second-order Taylor forward propagation.

Each batch element pushes one direction through the tape, producing a Taylor jet with 3 coefficients (c0=value, c1=first derivative, c2=second derivative / 2).

primal_inputs is [f32; batch_size * num_inputs] — primals for each element. direction_seeds is [f32; batch_size * num_inputs] — c1 seeds for each element.

Returns TaylorBatchResult with values, c1s, c2s each of size [f32; batch_size * num_outputs]. Batched second-order Taylor forward propagation.

Default implementation delegates to taylor_forward_kth_batch(order=3).

Dyn Compatibility

This trait is dyn compatible.

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

Implementors

impl GpuBackend for WgpuContext

type TapeBuffers = WgpuTapeBuffers