Crate nuvec_compile 

Function compilation: tracing, graph simplification and kernel fusion.

compile works as MLX’s mx.compile. The first call with a given input signature traces the function on placeholders and simplifies the graph:

• identity views and conversions are removed;
• equal scalar constants are merged;
• equal operations on equal inputs are computed once (common-subexpression elimination).

On the GPU, chains of element-wise operations are then fused into single generated kernels. Later calls replay the simplified graph on the new inputs without running the function again.

use nuvec_compile::compile;
use nuvec_core::{Array, Result, add, multiply};

// x * x + 1, as one kernel on the GPU
let f = compile(|xs: &[Array]| Ok(vec![add(&multiply(&xs[0], &xs[0])?, 1.0)?]));
let x = Array::from_slice(&[1.0f32, 2.0], &[2])?;
assert_eq!(f(&[x])?[0].to_vec::<f32>()?, [2.0, 5.0]);
let y = Array::from_slice(&[3.0f32, 4.0], &[2])?;
assert_eq!(f(&[y])?[0].to_vec::<f32>()?, [10.0, 17.0]); // replayed, not traced again

Status: pre-alpha.

Constants

VERSION

The version of this crate.

Functions

compile

Returns f compiled: a function with the same results that traces f once per input signature (dtypes and shapes) and default stream, and replays the simplified, fused graph on later calls.