Struct KernelRegistry 

pub struct KernelRegistry { /* private fields */ }

Registry that lazily compiles and caches Metal compute pipelines from embedded MSL source.

Usage

let mut registry = KernelRegistry::new();
let pipeline = registry.get_pipeline("elementwise_add", device.metal_device())?;
encoder.encode(&pipeline, &buffers, grid, tg);

Thread Safety

KernelRegistry is not Sync by default (it uses &mut self for get_pipeline to allow mutable cache insertion). If you need concurrent access, wrap it in a Mutex or use one registry per thread.

Implementations

impl KernelRegistry

pub fn new() -> Self

Create a new registry with all embedded shader sources pre-registered.

No compilation happens here — shaders are compiled lazily on first use.

pub fn register_source(&mut self, name: impl Into<String>, source: &'static str)

Register a shader source at runtime (useful for testing and dynamic kernel generation).

pub fn get_pipeline( &mut self, name: &str, device: &DeviceRef, ) -> Result<&ComputePipelineState>

Get a compiled compute pipeline for the named kernel function.

On first call for a given name, this compiles the MSL source into a Metal library, extracts the named function, and creates a ComputePipelineState. Subsequent calls return the cached pipeline.

Errors

• MlxError::KernelNotFound — no source registered for this name.
• MlxError::ShaderCompilationError — MSL compilation or pipeline creation failed.

pub fn get_pipeline_with_constants( &mut self, name: &str, device: &DeviceRef, bool_constants: &[(usize, bool)], int_constants: &[(usize, i32)], ) -> Result<&ComputePipelineState>

Get a compiled compute pipeline for the named kernel, specialized with Metal function constants (both bool and i32 in one call).

bool_constants contains (index, value) pairs mapping to [[function_constant(index)]] bool declarations in the MSL shader. int_constants contains (index, value) pairs mapping to [[function_constant(index)]] int (int32_t) declarations in the MSL shader.

Pipelines are cached by a composite key: "<name>|<index>:b<0|1>|...|<index>:i<value>|...". The ‘b’ prefix marks bool entries and the ‘i’ prefix marks i32 entries, making the format unambiguous regardless of constant ordering. Distinct (name, constants) combinations each compile to a separate pipeline; the slow compilation path runs at most once per unique combination.

Errors

• MlxError::KernelNotFound — no source registered for this name.
• MlxError::ShaderCompilationError — MSL compilation, function specialisation, or pipeline creation failed.

pub fn get_pipeline_with_bool_constants( &mut self, name: &str, device: &DeviceRef, bool_constants: &[(usize, bool)], ) -> Result<&ComputePipelineState>

Get a compiled compute pipeline for the named kernel, specialized with Metal bool function constants.

The bool_constants slice contains (index, value) pairs. Each pair maps to a [[function_constant(index)]] declaration in the MSL shader.

This is a thin wrapper around [get_pipeline_with_constants] that passes an empty int_constants slice. Existing callers continue to work without modification; the cache-key format for pure-bool pipelines is compatible (bool entries carry the ‘b’ type marker, which is the only format ever written by this wrapper).

Errors

• MlxError::KernelNotFound — no source registered for this name.
• MlxError::ShaderCompilationError — MSL compilation, function specialisation, or pipeline creation failed.

pub fn is_cached(&self, name: &str) -> bool

Check if a pipeline for the given name is already compiled and cached.

pub fn cached_count(&self) -> usize

Number of compiled pipelines currently in the cache.

pub fn source_count(&self) -> usize

Number of registered shader sources.

Trait Implementations

impl Default for KernelRegistry

fn default() -> Self

Returns the “default value” for a type.