Trait RuntimeOp 

pub trait RuntimeOp: Send + Sync {
    // Required methods
    fn n_activation_inputs(&self) -> usize;
    fn param_shapes(
        &self,
        input_shapes: &[&[usize]],
        output_shape: &[usize],
    ) -> Vec<Vec<usize>>;
    fn pack_args(
        &self,
        inputs: &[(RawPtr, &[usize])],
        params: &[RawPtr],
        output: RawPtr,
        output_shape: &[usize],
        output_row_stride: i32,
        visitor: &mut dyn ArgVisitor,
    );
    fn block(&self) -> [u32; 3];
    fn grid(&self, output_shape: &[usize]) -> [u32; 3];

    // Provided methods
    fn param_names(&self) -> &'static [&'static str] { ... }
    fn forward_output_row_stride(&self, output_shape: &[usize]) -> usize { ... }
    fn param_init_data(&self, _param_idx: usize) -> Option<Vec<u8>> { ... }
    fn compute_concrete_output_shape(
        &self,
        _input_shapes: &[&[usize]],
        resolved: &[usize],
    ) -> Vec<usize> { ... }
    fn n_launches(&self) -> usize { ... }
    fn pack_args_for_launch(
        &self,
        launch_idx: usize,
        inputs: &[(RawPtr, &[usize])],
        params: &[RawPtr],
        output: RawPtr,
        output_shape: &[usize],
        output_row_stride: i32,
        visitor: &mut dyn ArgVisitor,
    ) { ... }
    fn grid_for_launch(
        &self,
        launch_idx: usize,
        input_shapes: &[&[usize]],
        output_shape: &[usize],
    ) -> [u32; 3] { ... }
    fn has_backward(&self) -> bool { ... }
    fn backward_grad_output_row_stride(&self, output_shape: &[usize]) -> usize { ... }
    fn pack_backward_args(
        &self,
        inputs: &[(RawPtr, &[usize])],
        params: &[RawPtr],
        output: RawPtr,
        output_shape: &[usize],
        grad_output: RawPtr,
        grad_output_row_stride: i32,
        grad_inputs: &[RawPtr],
        grad_params: &[RawPtr],
        visitor: &mut dyn ArgVisitor,
    ) { ... }
    fn backward_block(&self) -> [u32; 3] { ... }
    fn backward_grid(
        &self,
        input_shapes: &[&[usize]],
        output_shape: &[usize],
    ) -> [u32; 3] { ... }
    fn n_backward_launches(&self) -> usize { ... }
    fn pack_backward_args_for_launch(
        &self,
        launch_idx: usize,
        inputs: &[(RawPtr, &[usize])],
        params: &[RawPtr],
        output: RawPtr,
        output_shape: &[usize],
        grad_output: RawPtr,
        grad_output_row_stride: i32,
        grad_inputs: &[RawPtr],
        grad_params: &[RawPtr],
        visitor: &mut dyn ArgVisitor,
    ) { ... }
    fn backward_grid_for_launch(
        &self,
        launch_idx: usize,
        input_shapes: &[&[usize]],
        output_shape: &[usize],
    ) -> [u32; 3] { ... }
}

Encapsulates the runtime-dispatch behaviour for a compiled op node: how many activation inputs it takes, what parameter buffers it needs, how to pack kernel arguments, and how to compute the launch grid.

Implementations live in teeny-kernels alongside the kernel structs so that each kernel owns its arg layout. The trait is defined here in teeny-core so that both teeny-kernels (impl) and teeny-cuda (consumer) can share it without a circular dependency.

Required Methods

fn n_activation_inputs(&self) -> usize

Number of activation tensors taken from predecessor DAG nodes.

fn param_shapes( &self, input_shapes: &[&[usize]], output_shape: &[usize], ) -> Vec<Vec<usize>>

Shapes of additional parameter buffers (weights, biases) needed by this op. Called at LoadedModel::load() time to pre-allocate device buffers. input_shapes / output_shape are concrete (batch dim resolved).

fn pack_args( &self, inputs: &[(RawPtr, &[usize])], params: &[RawPtr], output: RawPtr, output_shape: &[usize], output_row_stride: i32, visitor: &mut dyn ArgVisitor, )

Pack all kernel arguments into visitor in the correct order.

• inputs — (ptr, concrete_shape) per activation input
• params — raw pointers to pre-allocated param buffers
• output — raw pointer to the output buffer for this node
• output_shape — concrete output shape (batch dim resolved)
• output_row_stride — actual memory row stride (elements) of the output buffer (may be padded for TMA alignment)

fn block(&self) -> [u32; 3]

Threads-per-CTA for this kernel (x, y, z).

fn grid(&self, output_shape: &[usize]) -> [u32; 3]

Number of CTAs to launch (x, y, z), given the concrete output shape.

Provided Methods

fn param_names(&self) -> &'static [&'static str]

Names of parameter slots returned by [param_shapes], in the same order. Used as the suffix in the dotted key {node_name}.{slot_name}. Return an empty slice for ops that have no named parameters.

fn forward_output_row_stride(&self, output_shape: &[usize]) -> usize

Returns the required row stride (in elements) for the output buffer of this op’s forward kernel. The default is the natural row-major stride (output_shape[-1]). Kernels using TMA must round up to satisfy the 16-byte alignment constraint (e.g. 4 elements for f32).

fn param_init_data(&self, _param_idx: usize) -> Option<Vec<u8>>

Returns raw (little-endian) bytes to pre-populate parameter slot param_idx immediately after device buffer allocation. Return None to leave the slot zero-initialised (the default for trained parameters). Byte count must equal param_shapes()[param_idx].iter().product() * dtype_bytes.

fn compute_concrete_output_shape( &self, _input_shapes: &[&[usize]], resolved: &[usize], ) -> Vec<usize>

Override to compute the true concrete output shape from concrete input shapes.

Called during both load() (for param allocation) and forward() (for buffer allocation). The default just returns the resolved shape (from resolve_shape).

Override this when the output’s first dimension is a multiple of the batch size (e.g. B * H for attention pack/unpack ops) so that resolve_shape’s simple None → batch_size substitution would under-allocate the buffer.

fn n_launches(&self) -> usize

Number of sequential kernel launches this op requires.

Ops like channel-cat scatter N input chunks into one output buffer and need one kernel call per chunk. The executor loops n_launches() times, calling pack_args_for_launch and grid_for_launch on each iteration. The default is 1, which delegates to pack_args / grid.

fn pack_args_for_launch( &self, launch_idx: usize, inputs: &[(RawPtr, &[usize])], params: &[RawPtr], output: RawPtr, output_shape: &[usize], output_row_stride: i32, visitor: &mut dyn ArgVisitor, )

Pack kernel arguments for launch i (0-indexed).

Only called by the executor when n_launches() > 1. The default delegates to pack_args, ignoring launch_idx.

fn grid_for_launch( &self, launch_idx: usize, input_shapes: &[&[usize]], output_shape: &[usize], ) -> [u32; 3]

Grid for launch i. Receives concrete input shapes so that per-chunk grids can be computed without storing them in the op.

Only called when n_launches() > 1. The default delegates to grid.

fn has_backward(&self) -> bool

Returns true if this op has a backward (gradient) kernel.

fn backward_grad_output_row_stride(&self, output_shape: &[usize]) -> usize

Returns the required row stride (in elements) for the grad_output buffer passed to pack_backward_args. The default is the natural row-major stride (output_shape[-1]). Kernels using TMA must round up to satisfy the 16-byte alignment constraint (e.g. 4 elements for f32).

fn pack_backward_args( &self, inputs: &[(RawPtr, &[usize])], params: &[RawPtr], output: RawPtr, output_shape: &[usize], grad_output: RawPtr, grad_output_row_stride: i32, grad_inputs: &[RawPtr], grad_params: &[RawPtr], visitor: &mut dyn ArgVisitor, )

Pack backward kernel arguments.

• inputs — (ptr, shape) per forward activation input (from cache)
• params — raw ptrs to forward param buffers (weights, biases)
• output — forward output buffer (activation cache)
• output_shape — concrete forward output shape
• grad_output — incoming gradient dL/dy from the consumer node
• grad_output_row_stride — actual memory row stride (elements) of the grad_output buffer (may be padded for TMA alignment)
• grad_inputs — output gradient buffers: dL/dx per activation parent
• grad_params — output gradient buffers: dL/dw, dL/db, etc.

fn backward_block(&self) -> [u32; 3]

Threads-per-CTA for the backward kernel.

fn backward_grid( &self, input_shapes: &[&[usize]], output_shape: &[usize], ) -> [u32; 3]

Number of CTAs for the backward kernel.

input_shapes[i] is the concrete shape of the i-th activation input.

fn n_backward_launches(&self) -> usize

Number of sequential kernel launches for the backward pass.

For ops like channel-cat where the backward must write separate gradient buffers per input chunk, this should return n_inputs. Default is 1.

fn pack_backward_args_for_launch( &self, launch_idx: usize, inputs: &[(RawPtr, &[usize])], params: &[RawPtr], output: RawPtr, output_shape: &[usize], grad_output: RawPtr, grad_output_row_stride: i32, grad_inputs: &[RawPtr], grad_params: &[RawPtr], visitor: &mut dyn ArgVisitor, )

Pack backward kernel arguments for launch i (0-indexed).

Only called when n_backward_launches() > 1. The default delegates to pack_backward_args, ignoring launch_idx.

fn backward_grid_for_launch( &self, launch_idx: usize, input_shapes: &[&[usize]], output_shape: &[usize], ) -> [u32; 3]

Grid for backward launch i.

Only called when n_backward_launches() > 1. The default delegates to backward_grid.

Implementors