Skip to main content

Module registry

Module registry 

Source
Expand description

Kernel registry for runtime dispatch of compute operations.

A KernelRegistry maps graph-level operations (OpKind) to concrete compute implementations (Kernel). This allows the execution engine to resolve an operation at runtime without hardcoding kernel types.

Kernels are stored as trait objects (Box<dyn Kernel>), enabling different kernel implementations to be registered and overridden dynamically.

§Examples

Register and retrieve a kernel:

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

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

let kernel = reg.get(&OpKind::Add).unwrap();

let shape = vec![1, 4];
let a = Tensor::from_vec(shape.clone(), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
let b = Tensor::from_vec(shape.clone(), vec![10.0, 20.0, 30.0, 40.0]).unwrap();
let mut out = Tensor::zeros(shape).unwrap();

kernel.compute(&[&a, &b], &mut out).unwrap();
assert_eq!(out.data(), &[11.0, 22.0, 33.0, 44.0]);

Overwrite an existing kernel:

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

let mut reg = KernelRegistry::new();

assert!(reg.register(OpKind::Add, Box::new(AddKernel)).is_none());
assert!(reg.register(OpKind::Add, Box::new(MatMulKernel)).is_some());

Structs§

KernelRegistry
Registry mapping OpKind to runtime-executable Kernel implementations.