Struct KernelRegistry 

pub struct KernelRegistry { /* private fields */ }

Registry mapping OpKind to runtime-executable Kernel implementations.

This type is used by the execution engine to dispatch operations to the correct kernel at runtime.

Implementations

impl KernelRegistry

pub fn new() -> Self

Creates an empty kernel registry.

Examples

use tensor_forge::registry::KernelRegistry;
use tensor_forge::op::OpKind;

let reg = KernelRegistry::new();

assert!(reg.get(&OpKind::Add).is_none());
assert!(reg.get(&OpKind::MatMul).is_none());
assert!(reg.get(&OpKind::ReLU).is_none());

Examples found in repository

examples/custom_kernel.rs (line 100)

fn main() {
    // Build a tiny graph:
    //
    //   out = add(a, b)
    //
    // The graph describes *what* should be computed, not how.
    let mut graph = Graph::new();

    let a = graph.input_node(vec![2, 2]);
    let b = graph.input_node(vec![2, 2]);
    let out = graph
        .add(a, b)
        .expect("Adding valid input nodes should succeed");

    graph
        .set_output_node(out)
        .expect("Setting output node should succeed");

    // Create a custom registry.
    //
    // Start from an empty registry and explicitly register only the kernel(s)
    // needed by this graph.
    let mut registry = KernelRegistry::new();

    // Register our custom Add kernel.
    //
    // `register(...)` returns the previous mapping if one existed.
    let old = registry.register(OpKind::Add, Box::new(CustomAddKernel));
    assert!(
        old.is_none(),
        "First Add registration should not replace an existing kernel"
    );

    // Construct the executor with the custom registry.
    let exec = Executor::new(registry);

    // Bind runtime inputs.
    //
    // These are ordinary tensors supplied for the graph input nodes.
    let a_tensor = Tensor::from_vec(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0])
        .expect("Tensor construction should succeed");
    let b_tensor = Tensor::from_vec(vec![2, 2], vec![10.0, 20.0, 30.0, 40.0])
        .expect("Tensor construction should succeed");

    // Execute the graph.
    //
    // During execution:
    // - the executor validates input bindings,
    // - walks the graph in topological order,
    // - sees an `OpKind::Add` node,
    // - looks up `OpKind::Add` in the registry,
    // - dispatches to `CustomAddKernel::compute(...)`.
    let outputs = exec
        .execute(&graph, vec![(a, a_tensor), (b, b_tensor)])
        .expect("Execution should succeed");

    let result = outputs
        .get(&out)
        .expect("Declared output should be present in executor results");

    println!("Computed output for node {:?}: {:?}", out, result);
}

pub fn register( &mut self, op: OpKind, kernel: Box<dyn Kernel>, ) -> Option<Box<dyn Kernel>>

Registers kernel as the implementation for op.

Returns the previously registered kernel if one existed.

Examples

use tensor_forge::kernel::{AddKernel, MatMulKernel};
use tensor_forge::op::OpKind;
use tensor_forge::registry::KernelRegistry;

let mut reg = KernelRegistry::new();

let old_kernel = reg.register(OpKind::Add, Box::new(AddKernel));
assert!(old_kernel.is_none());
assert!(reg.get(&OpKind::Add).is_some());

// Add conflicting mapping
let old_kernel = reg.register(OpKind::Add, Box::new(MatMulKernel));
assert!(old_kernel.is_some());             // returns old AddKernel Mapping.
assert!(reg.get(&OpKind::Add).is_some());  // `OpKind::Add` now maps to MatMulKernel.

Examples found in repository

examples/custom_kernel.rs (line 105)

fn main() {
    // Build a tiny graph:
    //
    //   out = add(a, b)
    //
    // The graph describes *what* should be computed, not how.
    let mut graph = Graph::new();

    let a = graph.input_node(vec![2, 2]);
    let b = graph.input_node(vec![2, 2]);
    let out = graph
        .add(a, b)
        .expect("Adding valid input nodes should succeed");

    graph
        .set_output_node(out)
        .expect("Setting output node should succeed");

    // Create a custom registry.
    //
    // Start from an empty registry and explicitly register only the kernel(s)
    // needed by this graph.
    let mut registry = KernelRegistry::new();

    // Register our custom Add kernel.
    //
    // `register(...)` returns the previous mapping if one existed.
    let old = registry.register(OpKind::Add, Box::new(CustomAddKernel));
    assert!(
        old.is_none(),
        "First Add registration should not replace an existing kernel"
    );

    // Construct the executor with the custom registry.
    let exec = Executor::new(registry);

    // Bind runtime inputs.
    //
    // These are ordinary tensors supplied for the graph input nodes.
    let a_tensor = Tensor::from_vec(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0])
        .expect("Tensor construction should succeed");
    let b_tensor = Tensor::from_vec(vec![2, 2], vec![10.0, 20.0, 30.0, 40.0])
        .expect("Tensor construction should succeed");

    // Execute the graph.
    //
    // During execution:
    // - the executor validates input bindings,
    // - walks the graph in topological order,
    // - sees an `OpKind::Add` node,
    // - looks up `OpKind::Add` in the registry,
    // - dispatches to `CustomAddKernel::compute(...)`.
    let outputs = exec
        .execute(&graph, vec![(a, a_tensor), (b, b_tensor)])
        .expect("Execution should succeed");

    let result = outputs
        .get(&out)
        .expect("Declared output should be present in executor results");

    println!("Computed output for node {:?}: {:?}", out, result);
}

pub fn get(&self, op: &OpKind) -> Option<&dyn Kernel>

Returns the kernel registered for op, if present.

Examples

use tensor_forge::kernel::AddKernel;
use tensor_forge::op::OpKind;
use tensor_forge::registry::KernelRegistry;

let mut reg = KernelRegistry::new();
reg.register(OpKind::Add, Box::new(AddKernel));

assert!(reg.get(&OpKind::Add).is_some());
assert!(reg.get(&OpKind::MatMul).is_none());

Trait Implementations

impl Default for KernelRegistry

fn default() -> Self

Creates a registry populated with the built-in kernels.

Registers:

• OpKind::Add → crate::kernel::AddKernel
• OpKind::MatMul → crate::kernel::MatMulKernel
• OpKind::ReLU → crate::kernel::ReluKernel

Examples

use tensor_forge::op::OpKind;
use tensor_forge::registry::KernelRegistry;

let reg = KernelRegistry::default();
assert!(reg.get(&OpKind::Add).is_some());
assert!(reg.get(&OpKind::MatMul).is_some());
assert!(reg.get(&OpKind::ReLU).is_some());