Skip to main content

Kernel

Trait Kernel 

Source
pub trait Kernel: Send {
    // Required method
    fn execute(
        &self,
        inputs: &[TensorView<'_>],
        outputs: &mut [TensorMut<'_>],
    ) -> Result<()>;

    // Provided methods
    fn estimated_flops(&self) -> Option<u64> { ... }
    fn supports_strided_input(&self, input_idx: usize) -> bool { ... }
    fn preferred_output_layout(&self) -> Option<TensorLayout> { ... }
    fn cuda_graph_compatible(&self) -> bool { ... }
}
Expand description

A kernel ready to execute a specific op with specific shapes (§4.2).

Required Methods§

Source

fn execute( &self, inputs: &[TensorView<'_>], outputs: &mut [TensorMut<'_>], ) -> Result<()>

Execute over device-resident inputs/outputs.

Provided Methods§

Source

fn estimated_flops(&self) -> Option<u64>

Estimated FLOPs, if known (for the cost model).

Source

fn supports_strided_input(&self, input_idx: usize) -> bool

Whether the kernel accepts a non-contiguous (strided) input at idx.

Source

fn preferred_output_layout(&self) -> Option<TensorLayout>

The layout the kernel writes most efficiently, if it has a preference.

Source

fn cuda_graph_compatible(&self) -> bool

Whether this kernel can be captured inside a CUDA graph.

Dyn Compatibility§

This trait is dyn compatible.

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

Implementors§