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 view_outputs(
        &self,
        inputs: &[TensorView<'_>],
        num_outputs: usize,
    ) -> Option<Vec<ViewOutput>> { ... }
    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 view_outputs( &self, inputs: &[TensorView<'_>], num_outputs: usize, ) -> Option<Vec<ViewOutput>>

Attempt to express this node’s outputs as zero-copy ViewOutputs over its inputs instead of computing bytes (the layout/movement-op fast path).

inputs carries the real (possibly already-strided) input views and num_outputs is the node’s output arity. Returning:

None — the default — means “compute normally”: the executor allocates output buffers and calls Kernel::execute.
Some(specs) means every output is a view; specs.len() MUST equal num_outputs. A kernel that can view some but not all outputs must return None (all-or-nothing) so correctness never regresses.

When Some is returned, Kernel::execute is not invoked.

Source