Trait ComputeBackend 

pub trait ComputeBackend:
    Send
    + Sync
    + Debug {
    // Required methods
    fn name(&self) -> &str;
    fn init(&mut self) -> BackendResult<()>;
    fn is_initialized(&self) -> bool;
    fn gemm(
        &self,
        trans_a: BackendTranspose,
        trans_b: BackendTranspose,
        m: usize,
        n: usize,
        k: usize,
        alpha: f64,
        a_ptr: u64,
        lda: usize,
        b_ptr: u64,
        ldb: usize,
        beta: f64,
        c_ptr: u64,
        ldc: usize,
    ) -> BackendResult<()>;
    fn conv2d_forward(
        &self,
        input_ptr: u64,
        input_shape: &[usize],
        filter_ptr: u64,
        filter_shape: &[usize],
        output_ptr: u64,
        output_shape: &[usize],
        stride: &[usize],
        padding: &[usize],
    ) -> BackendResult<()>;
    fn attention(
        &self,
        q_ptr: u64,
        k_ptr: u64,
        v_ptr: u64,
        o_ptr: u64,
        batch: usize,
        heads: usize,
        seq_q: usize,
        seq_kv: usize,
        head_dim: usize,
        scale: f64,
        causal: bool,
    ) -> BackendResult<()>;
    fn reduce(
        &self,
        op: ReduceOp,
        input_ptr: u64,
        output_ptr: u64,
        shape: &[usize],
        axis: usize,
    ) -> BackendResult<()>;
    fn unary(
        &self,
        op: UnaryOp,
        input_ptr: u64,
        output_ptr: u64,
        n: usize,
    ) -> BackendResult<()>;
    fn binary(
        &self,
        op: BinaryOp,
        a_ptr: u64,
        b_ptr: u64,
        output_ptr: u64,
        n: usize,
    ) -> BackendResult<()>;
    fn synchronize(&self) -> BackendResult<()>;
    fn alloc(&self, bytes: usize) -> BackendResult<u64>;
    fn free(&self, ptr: u64) -> BackendResult<()>;
    fn copy_htod(&self, dst: u64, src: &[u8]) -> BackendResult<()>;
    fn copy_dtoh(&self, dst: &mut [u8], src: u64) -> BackendResult<()>;

    // Provided methods
    fn capabilities(&self) -> Capabilities { ... }
    fn available_devices(&self) -> BackendResult<Vec<DeviceInfo>> { ... }
    fn recommended_tile_for(&self, m: usize, n: usize, k: usize) -> TileShape { ... }
    fn gemm_mixed_precision(
        &self,
        prec: MixedPrecision,
        trans_a: BackendTranspose,
        trans_b: BackendTranspose,
        m: usize,
        n: usize,
        k: usize,
        alpha: f32,
        a_ptr: u64,
        lda: usize,
        b_ptr: u64,
        ldb: usize,
        beta: f32,
        c_ptr: u64,
        ldc: usize,
    ) -> BackendResult<()> { ... }
    fn conv2d_backward_data(
        &self,
        grad_output_ptr: u64,
        grad_output_shape: &[usize],
        filter_ptr: u64,
        filter_shape: &[usize],
        grad_input_ptr: u64,
        grad_input_shape: &[usize],
        stride: &[usize],
        padding: &[usize],
    ) -> BackendResult<()> { ... }
    fn conv2d_backward_filter(
        &self,
        input_ptr: u64,
        input_shape: &[usize],
        grad_output_ptr: u64,
        grad_output_shape: &[usize],
        grad_filter_ptr: u64,
        grad_filter_shape: &[usize],
        stride: &[usize],
        padding: &[usize],
    ) -> BackendResult<()> { ... }
    fn softmax(
        &self,
        input_ptr: u64,
        output_ptr: u64,
        shape: &[usize],
        axis: usize,
    ) -> BackendResult<()> { ... }
    fn gather(
        &self,
        input_ptr: u64,
        indices: &[usize],
        output_ptr: u64,
        rows: usize,
        cols: usize,
    ) -> BackendResult<()> { ... }
    fn scatter(
        &self,
        input_ptr: u64,
        indices: &[usize],
        output_ptr: u64,
        rows: usize,
        cols: usize,
    ) -> BackendResult<()> { ... }
    fn batched_gemm(
        &self,
        trans_a: BackendTranspose,
        trans_b: BackendTranspose,
        m: usize,
        n: usize,
        k: usize,
        alpha: f64,
        a_ptr: u64,
        lda: usize,
        stride_a: usize,
        b_ptr: u64,
        ldb: usize,
        stride_b: usize,
        beta: f64,
        c_ptr: u64,
        ldc: usize,
        stride_c: usize,
        batch_count: usize,
    ) -> BackendResult<()> { ... }
}

Abstract compute backend trait.

Implementations provide GPU-accelerated compute operations. All operations work with opaque device memory pointers (u64) and explicit shape/stride information, making the trait independent of any particular memory management scheme.

Object Safety

This trait is object-safe and can be used as Box<dyn ComputeBackend> or &dyn ComputeBackend for dynamic dispatch.

Lifecycle

1. Create the backend (CudaBackend::new()).
2. Call init to select a device and create a context.
3. Allocate memory with alloc.
4. Transfer data with copy_htod.
5. Run compute operations (gemm, conv2d_forward, etc.).
6. Read results with copy_dtoh.
7. Free memory with free.

Required Methods

fn name(&self) -> &str

Backend name (e.g., "cuda", "rocm", "metal").

fn init(&mut self) -> BackendResult<()>

Initialize the backend (select device, create context).

Must be called before any other operation. Calling init on an already-initialized backend is a no-op.

fn is_initialized(&self) -> bool

Returns true if the backend is ready for operations.

fn gemm( &self, trans_a: BackendTranspose, trans_b: BackendTranspose, m: usize, n: usize, k: usize, alpha: f64, a_ptr: u64, lda: usize, b_ptr: u64, ldb: usize, beta: f64, c_ptr: u64, ldc: usize, ) -> BackendResult<()>

General matrix multiply: C = alpha * op(A) * op(B) + beta * C.

Arguments

• trans_a, trans_b — transpose modes for A and B.
• m, n, k — matrix dimensions (C is m×n, A is m×k, B is k×n after transpose).
• alpha, beta — scaling factors.
• a_ptr, b_ptr, c_ptr — device pointers to column-major f64 matrices.
• lda, ldb, ldc — leading dimensions.

fn conv2d_forward( &self, input_ptr: u64, input_shape: &[usize], filter_ptr: u64, filter_shape: &[usize], output_ptr: u64, output_shape: &[usize], stride: &[usize], padding: &[usize], ) -> BackendResult<()>

2D convolution forward pass.

Arguments

• input_ptr — device pointer to input tensor (NCHW layout).
• input_shape — [N, C, H, W].
• filter_ptr — device pointer to filter tensor.
• filter_shape — [K, C, Fh, Fw].
• output_ptr — device pointer to output tensor.
• output_shape — [N, K, Oh, Ow].
• stride — [sh, sw].
• padding — [ph, pw].

fn attention( &self, q_ptr: u64, k_ptr: u64, v_ptr: u64, o_ptr: u64, batch: usize, heads: usize, seq_q: usize, seq_kv: usize, head_dim: usize, scale: f64, causal: bool, ) -> BackendResult<()>

Scaled dot-product attention.

Computes softmax(Q * K^T / scale) * V with optional causal masking.

Arguments

• q_ptr, k_ptr, v_ptr — device pointers to query, key, value tensors.
• o_ptr — device pointer to output tensor.
• batch, heads — batch size and number of attention heads.
• seq_q, seq_kv — query and key/value sequence lengths.
• head_dim — dimension of each attention head.
• scale — attention scale factor (typically 1 / sqrt(head_dim)).
• causal — if true, apply causal (lower-triangular) mask.

fn reduce( &self, op: ReduceOp, input_ptr: u64, output_ptr: u64, shape: &[usize], axis: usize, ) -> BackendResult<()>

Reduction along an axis.

Reduces input along axis using the specified op and writes to output.

fn unary( &self, op: UnaryOp, input_ptr: u64, output_ptr: u64, n: usize, ) -> BackendResult<()>

Element-wise unary operation.

Applies op to each of the n elements at input_ptr and writes to output_ptr.

fn binary( &self, op: BinaryOp, a_ptr: u64, b_ptr: u64, output_ptr: u64, n: usize, ) -> BackendResult<()>

Element-wise binary operation.

Applies op element-wise: output[i] = op(a[i], b[i]) for n elements.

fn synchronize(&self) -> BackendResult<()>

Synchronize all pending operations on this backend.

Blocks the host until all previously submitted GPU work completes.

fn alloc(&self, bytes: usize) -> BackendResult<u64>

Allocate device memory.

Returns an opaque device pointer. The caller is responsible for eventually calling free.

fn free(&self, ptr: u64) -> BackendResult<()>

Free device memory previously allocated with alloc.

fn copy_htod(&self, dst: u64, src: &[u8]) -> BackendResult<()>

Copy data from host memory to device memory.

• dst — device pointer (destination).
• src — host byte slice (source).

fn copy_dtoh(&self, dst: &mut [u8], src: u64) -> BackendResult<()>

Copy data from device memory to host memory.

• dst — host byte slice (destination).
• src — device pointer (source).

Provided Methods

fn capabilities(&self) -> Capabilities

Report this backend’s capabilities (precision support, Tensor Cores, unified memory, thread/shared-memory limits, …).

The default is the conservative CPU profile; GPU backends override it with values read from their driver.

fn available_devices(&self) -> BackendResult<Vec<DeviceInfo>>

Enumerate the devices this backend exposes, in a backend-agnostic DeviceInfo shape.

The default returns an empty list (a backend that cannot enumerate devices, e.g. a pure trait stub). Real backends override it; the CpuBackend reports a single host device.

fn recommended_tile_for(&self, m: usize, n: usize, k: usize) -> TileShape

Suggest a GEMM tile shape (tile_m, tile_n, tile_k) for the given problem dimensions, to seed an autotuner.

The default heuristic scales the tile with the problem size and snaps to a WMMA-aligned tile when capabilities reports Tensor Cores. Backends may override with hardware-specific shapes.

fn gemm_mixed_precision( &self, prec: MixedPrecision, trans_a: BackendTranspose, trans_b: BackendTranspose, m: usize, n: usize, k: usize, alpha: f32, a_ptr: u64, lda: usize, b_ptr: u64, ldb: usize, beta: f32, c_ptr: u64, ldc: usize, ) -> BackendResult<()>

Mixed-precision GEMM: C = alpha * op(A) * op(B) + beta * C where the A/B operands are stored in a reduced 16-bit format (MixedPrecision::F16 or MixedPrecision::Bf16) but the dot products accumulate in f32 — the Tensor-Core / WMMA contract.

Unlike gemm (column-major f64), this operates on column-major f32 buffers, matching the f32 storage every GPU uploads to a half/bfloat16 GEMM. The CPU reference emulates the 16-bit storage by rounding each input element to the target format (round-to-nearest, ties-to-even) before the f32 accumulation, so its output equals what a real reduced-precision kernel would produce.

Arguments

• prec — input storage format (f16 or bf16); the accumulator and output C are always f32.
• trans_a, trans_b — transpose modes for A and B.
• m, n, k — matrix dimensions (C is m×n, A is m×k, B is k×n after transpose).
• alpha, beta — scaling factors (applied in f32).
• a_ptr, b_ptr, c_ptr — device pointers to column-major f32 matrices.
• lda, ldb, ldc — leading dimensions.

The default implementation returns BackendError::Unsupported; the CpuBackend implements the reference math, and GPU backends override it with a Tensor-Core kernel.

fn conv2d_backward_data( &self, grad_output_ptr: u64, grad_output_shape: &[usize], filter_ptr: u64, filter_shape: &[usize], grad_input_ptr: u64, grad_input_shape: &[usize], stride: &[usize], padding: &[usize], ) -> BackendResult<()>

Backward pass of conv2d_forward w.r.t. the input (data gradient): given the upstream gradient grad_output, produce grad_input of the same shape as the forward input.

Mathematically this is the full convolution of grad_output with the spatially-flipped filter (equivalently, the transpose of the forward im2col matrix applied to the upstream gradient). All tensors are row-major f32 in NCHW / KCHW layout, matching conv2d_forward.

Arguments

• grad_output_ptr / grad_output_shape — upstream gradient, [N, K, Oh, Ow].
• filter_ptr / filter_shape — forward filter, [K, C, Fh, Fw].
• grad_input_ptr / grad_input_shape — output data gradient, [N, C, H, W].
• stride — [sh, sw]; padding — [ph, pw] (same as the forward pass).

The default returns BackendError::Unsupported; the CpuBackend implements the reference math.

fn conv2d_backward_filter( &self, input_ptr: u64, input_shape: &[usize], grad_output_ptr: u64, grad_output_shape: &[usize], grad_filter_ptr: u64, grad_filter_shape: &[usize], stride: &[usize], padding: &[usize], ) -> BackendResult<()>

Backward pass of conv2d_forward w.r.t. the filter (weight gradient): given the forward input and the upstream gradient grad_output, produce grad_filter of the same shape as the forward filter.

Mathematically this is the correlation of input with grad_output (the forward im2col matrix multiplied by the upstream gradient). All tensors are row-major f32 in NCHW / KCHW layout, matching conv2d_forward.

Arguments

• input_ptr / input_shape — forward input, [N, C, H, W].
• grad_output_ptr / grad_output_shape — upstream gradient, [N, K, Oh, Ow].
• grad_filter_ptr / grad_filter_shape — output weight gradient, [K, C, Fh, Fw].
• stride — [sh, sw]; padding — [ph, pw] (same as the forward pass).

The default returns BackendError::Unsupported; the CpuBackend implements the reference math.

fn softmax( &self, input_ptr: u64, output_ptr: u64, shape: &[usize], axis: usize, ) -> BackendResult<()>

Numerically-stable softmax along axis of the tensor described by shape (row-major, f32).

The default implementation returns BackendError::Unsupported; the CpuBackend implements it directly, and GPU backends override it with a fused kernel. Consumers that need softmax on a backend that does not provide it can still compose it from reduce(Max) + unary(Exp) + reduce(Sum) + binary(Div).

fn gather( &self, input_ptr: u64, indices: &[usize], output_ptr: u64, rows: usize, cols: usize, ) -> BackendResult<()>

Row-gather: copy the rows named by indices out of a rows × cols (f32, row-major) table into a contiguous output of indices.len() × cols.

Needed by embedding tables and MoE routing. The default returns BackendError::Unsupported.

fn scatter( &self, input_ptr: u64, indices: &[usize], output_ptr: u64, rows: usize, cols: usize, ) -> BackendResult<()>

Row-scatter: write each input row (indices.len() × cols, f32) into output at the destination row given by indices, preserving unreferenced rows of the rows × cols output table.

The inverse routing primitive to gather. The default returns BackendError::Unsupported.

fn batched_gemm( &self, trans_a: BackendTranspose, trans_b: BackendTranspose, m: usize, n: usize, k: usize, alpha: f64, a_ptr: u64, lda: usize, stride_a: usize, b_ptr: u64, ldb: usize, stride_b: usize, beta: f64, c_ptr: u64, ldc: usize, stride_c: usize, batch_count: usize, ) -> BackendResult<()>

Strided batched GEMM: for each batch b in 0..batch_count, compute C_b = alpha * op(A_b) * op(B_b) + beta * C_b where A_b starts at a_ptr + b * stride_a * 4 bytes (f32 elements), etc.

Arguments

• trans_a, trans_b — transpose modes for A and B.
• m, n, k — matrix dimensions (C is m×n).
• alpha, beta — scaling factors.
• a_ptr, b_ptr, c_ptr — device pointers to the first matrix in each batch.
• lda, ldb, ldc — leading dimensions.
• stride_a, stride_b, stride_c — element strides between consecutive matrices.
• batch_count — number of GEMM operations in the batch.

The default implementation dispatches batch_count individual gemm calls with pointer offsets.

Dyn Compatibility

This trait is dyn compatible.

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

Implementations on Foreign Types

impl<T: ComputeBackend + ?Sized> ComputeBackend for &mut T

Forward every ComputeBackend method through a mutable reference, so callers holding &mut dyn ComputeBackend (or &mut T) can pass it where a ComputeBackend is expected without re-boxing.

fn name(&self) -> &str

fn init(&mut self) -> BackendResult<()>

fn is_initialized(&self) -> bool

fn capabilities(&self) -> Capabilities

fn available_devices(&self) -> BackendResult<Vec<DeviceInfo>>

fn recommended_tile_for(&self, m: usize, n: usize, k: usize) -> TileShape

fn gemm( &self, trans_a: BackendTranspose, trans_b: BackendTranspose, m: usize, n: usize, k: usize, alpha: f64, a_ptr: u64, lda: usize, b_ptr: u64, ldb: usize, beta: f64, c_ptr: u64, ldc: usize, ) -> BackendResult<()>

fn conv2d_forward( &self, input_ptr: u64, input_shape: &[usize], filter_ptr: u64, filter_shape: &[usize], output_ptr: u64, output_shape: &[usize], stride: &[usize], padding: &[usize], ) -> BackendResult<()>

fn attention( &self, q_ptr: u64, k_ptr: u64, v_ptr: u64, o_ptr: u64, batch: usize, heads: usize, seq_q: usize, seq_kv: usize, head_dim: usize, scale: f64, causal: bool, ) -> BackendResult<()>

fn reduce( &self, op: ReduceOp, input_ptr: u64, output_ptr: u64, shape: &[usize], axis: usize, ) -> BackendResult<()>

fn unary( &self, op: UnaryOp, input_ptr: u64, output_ptr: u64, n: usize, ) -> BackendResult<()>

fn binary( &self, op: BinaryOp, a_ptr: u64, b_ptr: u64, output_ptr: u64, n: usize, ) -> BackendResult<()>

fn gemm_mixed_precision( &self, prec: MixedPrecision, trans_a: BackendTranspose, trans_b: BackendTranspose, m: usize, n: usize, k: usize, alpha: f32, a_ptr: u64, lda: usize, b_ptr: u64, ldb: usize, beta: f32, c_ptr: u64, ldc: usize, ) -> BackendResult<()>

fn conv2d_backward_data( &self, grad_output_ptr: u64, grad_output_shape: &[usize], filter_ptr: u64, filter_shape: &[usize], grad_input_ptr: u64, grad_input_shape: &[usize], stride: &[usize], padding: &[usize], ) -> BackendResult<()>

fn conv2d_backward_filter( &self, input_ptr: u64, input_shape: &[usize], grad_output_ptr: u64, grad_output_shape: &[usize], grad_filter_ptr: u64, grad_filter_shape: &[usize], stride: &[usize], padding: &[usize], ) -> BackendResult<()>

fn softmax( &self, input_ptr: u64, output_ptr: u64, shape: &[usize], axis: usize, ) -> BackendResult<()>

fn gather( &self, input_ptr: u64, indices: &[usize], output_ptr: u64, rows: usize, cols: usize, ) -> BackendResult<()>

fn scatter( &self, input_ptr: u64, indices: &[usize], output_ptr: u64, rows: usize, cols: usize, ) -> BackendResult<()>

fn synchronize(&self) -> BackendResult<()>

fn alloc(&self, bytes: usize) -> BackendResult<u64>

fn free(&self, ptr: u64) -> BackendResult<()>

fn copy_htod(&self, dst: u64, src: &[u8]) -> BackendResult<()>

fn copy_dtoh(&self, dst: &mut [u8], src: u64) -> BackendResult<()>

Implementors

impl ComputeBackend for CpuBackend

impl ComputeBackend for NullBackend