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 prewarm_pipelines(&mut self, device: &DeviceRef, names: &[&str]) -> usize

ADR-033 §Pi Task #20 iter 11 (2026-05-23) — eagerly compile a list of kernel pipelines to move first-call JIT/PSO-creation cost out of the prefill hot path and into the model-load window.

The profiler showed that on Qwen3.6 35B-A3B MoE prefill at seq=553, the FIRST FA layer + FIRST FFN layer take ~40ms each (vs ~14µs warm) — that’s 80ms of the 221ms prefill, dominated by Metal pipeline state creation. Pre-creating these pipelines at load time (when 3.3s is already being spent on model parse + upload) is a strict perf win for measured prefill throughput.

Best-effort: silently skips kernels that aren’t registered (e.g., list contains a kernel name for an arch this build doesn’t use). Logs at debug level on failure to keep load-path quiet.

Returns the count of pipelines successfully prewarmed.

pub fn prewarm_pipelines_with_bool_constants( &mut self, device: &DeviceRef, entries: &[(&str, &[(usize, bool)])], ) -> usize

ADR-033 §Pi Task #20 iter 12 (2026-05-23) — prewarm pipelines that require [[function_constant]] specialization. Each entry is (name, &[(constant_index, bool_value)]). Mirrors prewarm_pipelines but routes through get_pipeline_with_bool_constants so kernels declaring function_constant decls without defaults can be safely prewarmed.

Use case: hot-path kernels like flash_attn_prefill_bf16_d256 (uses bool constants 200/201/300/301/303 for align/mask/causal/blk flags) cannot be safely prewarmed without specialization — Metal validateWithDevice: asserts and aborts the process. Provide the constants production uses and prewarming becomes safe.

Returns count warmed.

pub fn prewarm_all(&mut self, device: &DeviceRef) -> (usize, usize)

ADR-033 §Pi Task #20 iter 11 (2026-05-23) — prewarm every registered kernel source. Useful when the exact set of needed kernels is hard to enumerate (e.g., serving paths that span multiple arches). Total cost is bounded by the number of registered kernels times the per-pipeline PSO creation cost (~5-15ms typical on M-series).

Returns (warmed, skipped) counts.

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.