Backend Refactor

Remove MetaOperation trait: In it's current form, the MetaOperation trait was tightly coupled with the underlying kernel that would run on the GPU. This isn't great when you want to have multiple kernels backing a single Operation. For example MatMul.

Introduce GPUOperation trait:

pub trait GPUOperation: Operation {
/// # Kernel Selection
/// Enumeration of all possible kernels that can be used for this operation.
type KernelEnum: Kernel;
fn select_kernel(self) -> Self::KernelEnum;
...
}

This trait allows for multiple kernels to back a single operation, and removes most of the entanglement between Kernel and Operation.

Introduce the Kernel trait:

pub trait Kernel {
/// There are 2 flavours of metadata:
/// 1. Static Metadata - the structure is known at compile time, so we get `render` for free
///    with the derive macro: `WgslMetadata`. The author still needs to implement
///    `write_metadata`.
/// 2. Dynamic Metadata - the structure is not known at compile time, so the author must
///   implement both `render` and `write_metadata`.
type Metadata: KernelMetadata + 'static;

fn metadata(
    &self,
    dst: &Tensor,
    kernel_element: &KernelElement,
) -> Result<Self::Metadata, OperationError>;

/// # Calculate Dispatch
///
/// Determine required amount of workgroups to execute the operation.
fn calculate_dispatch(&self, dst: &Tensor) -> Result<Workload, OperationError>;

/// # Kernel Key
///
/// Construct a unique cache key for a kernel.
/// If the key is registered in the compute module pool, the module is reused.
///
/// Default implementation is provided, but care must be taken to ensure that the key is
/// unique via the `additional` parameter.
fn kernel_key(
    &self,
    workgroup_size: &WorkgroupSize,
    inplace: bool,
    dst: &Tensor,
    kernel_element: &KernelElement,
) -> KernelKey {
    KernelKey::new(
        &self.kernel_name(),
        &self.srcs(),
        dst,
        workgroup_size,
        inplace,
        kernel_element,
        None,
    )
}

/// # Kernel Element
///
/// Determine the largest possible unit data type that can be used (e.g f32, vec2<f32>, vec4<f32>)
fn kernel_element(&self, dst: &Tensor) -> KernelElement;

fn build_kernel(
    &self,
    inplace: bool,
    dst: &Tensor,
    workgroup_size: &WorkgroupSize,
) -> Result<KernelSource, OperationError>;
}

Kernel is a specific kernel that will be dispatched to the GPU. It also formalizes the Metadata back to being an assoc type. We previously removed this because the Concat operation introduced metadata where the memory layout wasn't known at compile time. This is fixed with DynMetadata.

Introduce KernelRenderable, a trait that ensures that all kernels implement the standard behaviour for Kernels.

/// This trait is focused on the generation of the kernel source code.
pub trait KernelRenderable {
///Every WGSL must have 1 or more bindings registered
///
///This method writes the GPU source of the bindings to the kernel builder.
///This will be called within `build()`
fn register_bindings<P: WgslPrimitive>(
    &self,
    builder: &mut WgslKernelBuilder,
    inplace: bool,
) -> Result<(), OperationError>;

///Generate the full kernel source
fn render<P: WgslPrimitive>(
    &self,
    inplace: bool,
    dst: &Tensor,
    workgroup_size: &WorkgroupSize,
) -> Result<KernelSource, OperationError>;
}

huggingface / ratchet

Refactor backend #226

Backend Refactor