knok-build 0.4.0

//! Build-script frontend for `knok` static graph tracing.
//!
//! This crate is used from `build.rs`, not from the final target crate. A build
//! script defines graph functions with [`graph`], executes them with traced
//! tensor values through [`compile_graphs`] or [`grad_graphs`], and writes
//! generated wrappers plus VMFB artifacts into `OUT_DIR`.
//!
//! ```ignore
//! use knok_build::prelude::*;
//!
//! #[knok_build::graph(backend = Backend::LlvmCpu)]
//! fn forward(x: T2<f32, 2, 2>) -> T2<f32, 2, 2> {
//!     relu(matmul(x.clone(), x) + 1.0)
//! }
//!
//! fn main() {
//!     knok_build::compile_graphs!(forward);
//! }
//! ```
//!
//! Experimental value-and-gradient wrappers can be emitted for scalar-loss
//! graphs with [`grad_graphs`]. The generated module is named
//! `{graph_name}_value_and_grad` and returns the loss followed by gradients for
//! the graph's floating-point inputs.
//!
//! Target crates import the generated wrappers with
//! `knok::generated_graphs!(pub mod graphs);`.

mod codegen;
mod trace;

use std::{
    collections::BTreeMap,
    env, fs,
    path::{Path, PathBuf},
};

use anyhow::Context;
pub use knok_compile::CompileOptions;
use knok_compile::{compile_graph_with_options, compile_mlir_source_with_options};
use knok_core::{value_and_grad_graph_with_registry, Input, TensorType, TypedGraph};

pub use knok_build_macros::{
    compile_graphs, compile_graphs_with_options, compile_mlir_models,
    compile_mlir_models_with_options, grad_graphs, grad_graphs_with_options, graph,
};
pub use trace::*;

/// Result type used by build-script tracing and code generation.
pub type Result<T> = anyhow::Result<T>;

/// IREE target backend selected by a traced graph.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Backend {
    /// CPU backend compiled through IREE's LLVM CPU target.
    LlvmCpu,
    /// Apple Metal backend compiled through IREE's Metal/SPIR-V path on macOS.
    #[cfg(any(target_os = "macos", doc))]
    MetalSpirv,
    /// Vulkan backend compiled through IREE's Vulkan/SPIR-V path.
    #[cfg(feature = "vulkan")]
    VulkanSpirv,
    /// NVIDIA CUDA backend compiled through IREE's CUDA target.
    #[cfg(feature = "cuda")]
    Cuda,
}

impl Backend {
    /// Returns the IREE target backend flag value.
    pub const fn name(self) -> &'static str {
        match self {
            Self::LlvmCpu => "llvm-cpu",
            #[cfg(any(target_os = "macos", doc))]
            Self::MetalSpirv => "metal-spirv",
            #[cfg(feature = "vulkan")]
            Self::VulkanSpirv => "vulkan-spirv",
            #[cfg(feature = "cuda")]
            Self::Cuda => "cuda",
        }
    }

    /// Returns the default IREE runtime driver used by generated wrappers.
    pub const fn default_driver(self) -> &'static str {
        match self {
            Self::LlvmCpu => "local-task",
            #[cfg(any(target_os = "macos", doc))]
            Self::MetalSpirv => "metal",
            #[cfg(feature = "vulkan")]
            Self::VulkanSpirv => "vulkan",
            #[cfg(feature = "cuda")]
            Self::Cuda => "cuda",
        }
    }
}

/// Options controlling generated wrapper output from build-script graph macros.
#[derive(Clone, Debug)]
pub struct BuildOptions {
    output_file: String,
    stub_artifacts: bool,
    compiler: CompileOptions,
}

impl BuildOptions {
    /// Sets IREE compiler tuning for every emitted graph, including gradients and imports.
    pub fn compiler_options(mut self, options: CompileOptions) -> Self {
        self.compiler = options;
        self
    }
    /// Creates placeholder VMFB files for check-only builds.
    ///
    /// Stub artifacts let no-std or cross-target fixture crates typecheck
    /// generated wrappers without requiring a runnable IREE compiler artifact.
    /// The generated artifacts are intentionally not executable.
    pub fn stub_artifacts_for_check() -> Self {
        Self {
            stub_artifacts: true,
            ..Self::default()
        }
    }

    /// Writes generated Rust wrappers to `name` inside `OUT_DIR`.
    ///
    /// Target crates must pass the same filename to
    /// `knok::generated_graphs!(pub mod graphs, "...")`.
    pub fn output_file(mut self, name: impl Into<String>) -> Self {
        self.output_file = name.into();
        self
    }
}

impl Default for BuildOptions {
    fn default() -> Self {
        Self {
            output_file: "knok_graphs.rs".into(),
            stub_artifacts: false,
            compiler: CompileOptions::default(),
        }
    }
}

/// Registry populated by generated `#[graph]` registration glue.
#[derive(Default)]
pub struct GraphRegistry {
    graphs: Vec<RegisteredGraph>,
}

/// One traced graph registered for build-time compilation.
pub struct RegisteredGraph {
    graph: TypedGraph,
    backend: Backend,
}

/// One external MLIR module compiled and wrapped by `compile_mlir_models!`.
#[derive(Clone, Debug)]
pub struct MlirModel {
    /// Generated Rust module name.
    pub name: String,
    /// Path to the `.mlir` file, relative to the build-script crate unless absolute.
    pub path: PathBuf,
    /// IREE VM function name, such as `imported.forward`.
    pub function_name: String,
    /// IREE compile backend.
    pub backend: Backend,
    /// Input tensor names and types used by the generated wrapper.
    pub inputs: Vec<Input>,
    /// Output tensor types used by the generated wrapper.
    pub outputs: Vec<TensorType>,
}

impl MlirModel {
    /// Creates an external MLIR model descriptor.
    pub fn new(
        name: impl Into<String>,
        path: impl Into<PathBuf>,
        function_name: impl Into<String>,
        backend: Backend,
        inputs: Vec<Input>,
        outputs: Vec<TensorType>,
    ) -> Self {
        Self {
            name: name.into(),
            path: path.into(),
            function_name: function_name.into(),
            backend,
            inputs,
            outputs,
        }
    }
}

impl GraphRegistry {
    /// Creates an empty graph registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Traces one graph body and stores the resulting typed graph.
    ///
    /// Most users call this indirectly through [`compile_graphs`] or
    /// [`compile_graphs_with_options`].
    pub fn trace<F, O>(&mut self, name: &str, backend: Backend, build: F) -> Result<()>
    where
        F: FnOnce(&mut TraceContext) -> O,
        O: TraceOutput,
    {
        let mut context = TraceContext::default();
        let output = build(&mut context);
        let graph = context.finish(name, backend.name(), output)?;
        self.graphs.push(RegisteredGraph { graph, backend });
        Ok(())
    }
}

/// Emits wrappers and VMFB artifacts for all graphs in `registry`.
pub fn emit_registered_graphs(registry: GraphRegistry) -> Result<()> {
    emit_registered_graphs_with_options(registry, BuildOptions::default())
}

/// Emits wrappers and VMFB artifacts with explicit build options.
pub fn emit_registered_graphs_with_options(
    registry: GraphRegistry,
    options: BuildOptions,
) -> Result<()> {
    let out_dir = PathBuf::from(env::var("OUT_DIR")?);
    fs::create_dir_all(&out_dir)?;
    if env::var_os("KNOK_CACHE_DIR").is_none() {
        env::set_var("KNOK_CACHE_DIR", out_dir.join("knok-cache"));
    }
    println!("cargo:rerun-if-changed=build.rs");
    track_compiler_environment();

    let mut generated = String::new();
    let graphs = registry.graphs;
    let graph_registry = graphs
        .iter()
        .map(|registered| (registered.graph.name.clone(), registered.graph.clone()))
        .collect::<BTreeMap<_, _>>();

    for registered in &graphs {
        let vmfb_name = graph_vmfb_name(&registered.graph.name);
        let vmfb_path = out_dir.join(&vmfb_name);
        if options.stub_artifacts {
            fs::write(&vmfb_path, [0u8; 16])?;
        } else {
            let compiled =
                compile_registered_graph(&registered.graph, &graph_registry, &options.compiler)?;
            fs::write(&vmfb_path, compiled.vmfb)?;
        }
        generated.push_str(&codegen::graph_module(
            &registered.graph,
            registered.backend,
            &vmfb_name,
            &options.compiler.flags(registered.backend.name())?,
        )?);
        generated.push('\n');
    }

    fs::write(out_dir.join(options.output_file), generated)?;
    Ok(())
}

/// Emits value-and-gradient wrappers and VMFB artifacts for all graphs in `registry`.
pub fn emit_registered_grad_graphs(registry: GraphRegistry) -> Result<()> {
    emit_registered_grad_graphs_with_options(registry, BuildOptions::default())
}

/// Emits value-and-gradient wrappers and VMFB artifacts with explicit build options.
pub fn emit_registered_grad_graphs_with_options(
    registry: GraphRegistry,
    options: BuildOptions,
) -> Result<()> {
    let out_dir = PathBuf::from(env::var("OUT_DIR")?);
    fs::create_dir_all(&out_dir)?;
    if env::var_os("KNOK_CACHE_DIR").is_none() {
        env::set_var("KNOK_CACHE_DIR", out_dir.join("knok-cache"));
    }
    println!("cargo:rerun-if-changed=build.rs");
    track_compiler_environment();

    let mut generated = String::new();
    let graphs = registry.graphs;
    let graph_registry = graphs
        .iter()
        .map(|registered| (registered.graph.name.clone(), registered.graph.clone()))
        .collect::<BTreeMap<_, _>>();

    for registered in &graphs {
        let graph = value_and_grad_graph_with_registry(&registered.graph, &graph_registry)?;
        let vmfb_name = graph_vmfb_name(&graph.name);
        let vmfb_path = out_dir.join(&vmfb_name);
        if options.stub_artifacts {
            fs::write(&vmfb_path, [0u8; 16])?;
        } else {
            let compiled = compile_registered_graph(&graph, &graph_registry, &options.compiler)?;
            fs::write(&vmfb_path, compiled.vmfb)?;
        }
        generated.push_str(&codegen::graph_module(
            &graph,
            registered.backend,
            &vmfb_name,
            &options.compiler.flags(registered.backend.name())?,
        )?);
        generated.push('\n');
    }

    fs::write(out_dir.join(options.output_file), generated)?;
    Ok(())
}

/// Emits wrappers and VMFB artifacts for external MLIR models.
pub fn emit_mlir_models(models: Vec<MlirModel>) -> Result<()> {
    emit_mlir_models_with_options(models, BuildOptions::default())
}

/// Emits external MLIR model wrappers and VMFB artifacts with explicit options.
pub fn emit_mlir_models_with_options(models: Vec<MlirModel>, options: BuildOptions) -> Result<()> {
    let out_dir = PathBuf::from(env::var("OUT_DIR")?);
    fs::create_dir_all(&out_dir)?;
    if env::var_os("KNOK_CACHE_DIR").is_none() {
        env::set_var("KNOK_CACHE_DIR", out_dir.join("knok-cache"));
    }
    println!("cargo:rerun-if-changed=build.rs");
    track_compiler_environment();

    let mut generated = String::new();
    for model in &models {
        let source_path = resolve_model_path(&model.path)?;
        println!("cargo:rerun-if-changed={}", source_path.display());
        let vmfb_name = mlir_model_vmfb_name(&model.name);
        let vmfb_path = out_dir.join(&vmfb_name);
        if options.stub_artifacts {
            fs::write(&vmfb_path, [0u8; 16])?;
        } else {
            let source = fs::read_to_string(&source_path)?;
            let vmfb =
                compile_mlir_source_with_options(model.backend.name(), &source, &options.compiler)?;
            fs::write(&vmfb_path, vmfb)?;
        }
        generated.push_str(&codegen::mlir_model_module(
            model,
            &vmfb_name,
            &options.compiler.flags(model.backend.name())?,
        )?);
        generated.push('\n');
    }

    fs::write(out_dir.join(options.output_file), generated)?;
    Ok(())
}

fn track_compiler_environment() {
    for variable in [
        "EERIE_PYTHON",
        "KNOK_IREE_COMPILER",
        "KNOK_CACHE_DIR",
        "PATH",
        "PYTHONPATH",
    ] {
        println!("cargo:rerun-if-env-changed={variable}");
    }
}

fn resolve_model_path(path: &Path) -> Result<PathBuf> {
    if path.is_absolute() {
        return Ok(path.to_path_buf());
    }
    Ok(PathBuf::from(env::var("CARGO_MANIFEST_DIR")?).join(path))
}

fn graph_vmfb_name(name: &str) -> String {
    format!("{name}.vmfb")
}

fn mlir_model_vmfb_name(name: &str) -> String {
    format!("mlir-model-{name}.vmfb")
}

fn compile_registered_graph(
    graph: &TypedGraph,
    graphs: &BTreeMap<String, TypedGraph>,
    options: &CompileOptions,
) -> Result<knok_compile::CompiledGraph> {
    compile_graph_with_options(graph, graphs, options)
        .with_context(|| format!("failed to compile graph `{}`", graph.name))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn external_mlir_artifacts_do_not_collide_with_traced_graph_artifacts() {
        assert_eq!(graph_vmfb_name("forward"), "forward.vmfb");
        assert_eq!(mlir_model_vmfb_name("forward"), "mlir-model-forward.vmfb");
        assert_ne!(graph_vmfb_name("forward"), mlir_model_vmfb_name("forward"));
    }

    #[test]
    fn backend_names_match_iree_driver_pairs() {
        assert_eq!(Backend::LlvmCpu.name(), "llvm-cpu");
        assert_eq!(Backend::LlvmCpu.default_driver(), "local-task");
        #[cfg(any(target_os = "macos", doc))]
        {
            assert_eq!(Backend::MetalSpirv.name(), "metal-spirv");
            assert_eq!(Backend::MetalSpirv.default_driver(), "metal");
        }

        #[cfg(feature = "vulkan")]
        {
            assert_eq!(Backend::VulkanSpirv.name(), "vulkan-spirv");
            assert_eq!(Backend::VulkanSpirv.default_driver(), "vulkan");
        }

        #[cfg(feature = "cuda")]
        {
            assert_eq!(Backend::Cuda.name(), "cuda");
            assert_eq!(Backend::Cuda.default_driver(), "cuda");
        }
    }
}

/// Common build-script imports for graph definitions.
pub mod prelude {
    pub use crate::{
        abs, all, all_axis, amax, amax_axis, amin, amin_axis, any, any_axis, arange, arange_step,
        arange_to, argmax, argmax_axis, argmin, argmin_axis, avg_pool2d, avg_pool2d_options,
        broadcast, ceil, clip, compile_graphs, compile_graphs_with_options, compile_mlir_models,
        compile_mlir_models_with_options, concat, conv2d, conv2d_options, cos, detach, diagonal,
        diagonal_axes, dot, equal, exp, exp2, expm1, eye, flip, flip_axes, floor, full_like,
        gather, grad_graphs, grad_graphs_with_options, graph, greater, greater_equal, identity,
        inner, isnan, less, less_equal, linspace, log, log10, log1p, log2, logical_and,
        logical_not, logical_or, logical_xor, matmul, max, max_axis, max_pool2d,
        max_pool2d_options, maximum, mean, mean_axis, min, min_axis, minimum, moveaxis, not_equal,
        ones_like, outer, pad, permute, permute_dims, pow, prod, prod_axis, ptp, ptp_axis, r#where,
        rand, randint, randint_static, reciprocal, relu, repeat, reshape, rint, roll, round,
        sigmoid, sin, slice, softmax, softmax_axis, split, sqrt, square, squeeze, stack, std,
        std_axis, sum, sum_axis, swapaxes, take, take_along_axis, tan, tanh, tile, trace,
        trace_axes, transpose, transpose_axes, uniform, uniform_static, unsqueeze, var, var_axis,
        vecdot, vecdot_axis, zeros_like, Backend, BuildOptions, CompileOptions, Conv2dOptions,
        MlirModel, Pool2dOptions, Tensor0, Tensor1, Tensor2, Tensor3, Tensor4, Tensor5, Tensor6,
        T0, T1, T2, T3, T4, T5, T6,
    };
    pub use half::{bf16, f16};
}

#[doc(hidden)]
pub mod __private {
    pub use knok_core::{ElementType, Input, TensorType};
}