makepad/libs/ai/metal/build.rs

use std::env;
use std::fs;
use std::process::Command;

// The metallib build for THE Metal shader store (aiarch.md §1 + §8b ring 1,
// lane T5). Absorbed verbatim from libs/ggml/build.rs::build_metallib —
// same env var names (MAKEPAD_GGML_METAL_PRECOMPILE, still read directly
// off the process environment; build scripts see it regardless of which
// crate's Cargo.toml declares it), same feature name (`metal-precompile`,
// forwarded here from makepad-ggml via
// `metal-precompile = ["makepad-ai-metal/metal-precompile"]` in
// libs/ggml/Cargo.toml so CARGO_FEATURE_METAL_PRECOMPILE lands on THIS
// build script), same three-mode behavior (precompile disabled /
// precompile attempted but failed / precompile succeeded), same output
// filename (`ggml-default.metallib`).
//
// This crate sets `links = "makepad_ai_metal"` in its Cargo.toml, so on
// success (in EVERY mode — the metallib file always exists, empty or not)
// it emits the links-metadata line `cargo:metallib=<path>`, which flows to
// immediate dependents (libs/ggml) as `DEP_MAKEPAD_AI_METAL_METALLIB`.
// libs/ggml/build.rs reads that env var and re-emits
// `cargo:rustc-env=MAKEPAD_GGML_METALLIB=<path>` (same env name as before
// the move), so `libs/ggml/src/backend/metal/{runtime,compat}.rs`'s
// `include_bytes!(env!("MAKEPAD_GGML_METALLIB"))` sites need zero edits.
//
// Two of the include_str! sites in those same files (ggml-metal.metal,
// ggml-common.h, ggml-metal-impl.h) used to reach the shader tree by a
// RELATIVE path (it used to sit right next to them). Now that the shader
// tree lives in this crate, we also emit `cargo:shader_dir=<dir>`, which
// flows to libs/ggml as `DEP_MAKEPAD_AI_METAL_SHADER_DIR`; libs/ggml's
// build.rs re-emits that as `cargo:rustc-env=MAKEPAD_GGML_METAL_SHADER_DIR=<dir>`
// and the include sites switch to
// `include_str!(concat!(env!("MAKEPAD_GGML_METAL_SHADER_DIR"), "/ggml-metal.metal"))`
// — file contents stay byte-identical, only how the path is found changes.
fn main() {
    // Apple BLAS for the RIFE host-path GEMMs (rife.rs): cblas_sgemm off
    // the Accelerate framework — hundreds of GFLOP/s with zero dispatch
    // overhead, exactly what fifty small convolutions per frame pair want.
    #[allow(clippy::needless_return)]
    {
        let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
        if target_os == "macos" || target_os == "ios" {
            println!("cargo:rustc-link-lib=framework=Accelerate");
        }
    }

    println!("cargo:rerun-if-env-changed=MAKEPAD_GGML_METAL_PRECOMPILE");

    // shim.rs's non-macos CUDA bounce only compiles when makepad-ai-cuda
    // actually built kernels (same links-metadata handshake models/common
    // uses); otherwise a None-returning stub takes its place.
    println!("cargo:rustc-check-cfg=cfg(makepad_ai_cuda_kernels)");
    if env::var("DEP_MAKEPAD_AI_CUDA_KERNELS").as_deref() == Ok("1") {
        println!("cargo:rustc-cfg=makepad_ai_cuda_kernels");
    }

    let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
    if target_os == "macos" {
        build_metallib();
    } else {
        // env! sites in macos-only modules still need a value if rustc
        // type-checks them. Point at this crate's shader tree.
        let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
        let metal_dir = format!("{manifest_dir}/shaders/ggml");
        println!("cargo:rustc-env=MAKEPAD_METAL_SHADER_DIR={metal_dir}");
        println!("cargo:rustc-env=MAKEPAD_GGML_METAL_SHADER_DIR={metal_dir}");
        println!("cargo:rustc-env=MAKEPAD_METALLIB=/dev/null");
        println!("cargo:rustc-env=MAKEPAD_GGML_METALLIB=/dev/null");
    }
    // non-macos targets: nothing to do, no metallib, no shader_dir line —
    // dependents see DEP_MAKEPAD_AI_METAL_METALLIB / _SHADER_DIR unset.
}

fn build_metallib() {
    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
    let metal_dir = format!("{}/shaders/ggml", manifest_dir);
    let metal_src = format!("{}/ggml-metal.metal", metal_dir);
    let common_h = format!("{}/ggml-common.h", metal_dir);
    let impl_h = format!("{}/ggml-metal-impl.h", metal_dir);
    println!("cargo:rerun-if-changed={}", metal_src);
    println!("cargo:rerun-if-changed={}", common_h);
    println!("cargo:rerun-if-changed={}", impl_h);

    // links-metadata handshake: flows to immediate dependents (libs/ggml)
    // as DEP_MAKEPAD_AI_METAL_SHADER_DIR=<metal_dir>. Emitted unconditionally
    // (independent of precompile mode) because the include_str! source
    // fallback in libs/ggml needs it regardless of whether a real metallib
    // was baked.
    println!("cargo:shader_dir={}", metal_dir);
    println!("cargo:rustc-env=MAKEPAD_METAL_SHADER_DIR={metal_dir}");
    println!("cargo:rustc-env=MAKEPAD_GGML_METAL_SHADER_DIR={metal_dir}");

    let precompile_default = env::var_os("CARGO_FEATURE_METAL_PRECOMPILE").is_some();
    let precompile_enabled = env::var("MAKEPAD_GGML_METAL_PRECOMPILE")
        .ok()
        .map(|v| {
            let v = v.trim().to_ascii_lowercase();
            !(v.is_empty() || v == "0" || v == "false" || v == "no" || v == "off")
        })
        .unwrap_or(precompile_default);

    let out_dir = env::var("OUT_DIR").unwrap();
    let _ = fs::create_dir_all(&out_dir);
    let air_path = format!("{}/ggml-metal.air", out_dir);
    let metallib_path = format!("{}/ggml-default.metallib", out_dir);

    if !precompile_enabled {
        let _ = fs::write(&metallib_path, []);
        emit_metallib_env(&metallib_path);
        return;
    }

    let metal_status = Command::new("xcrun")
        .args([
            "--sdk",
            "macosx",
            "metal",
            "-O3",
            "-fno-fast-math",
            // The runtime source compile defines this on bfloat-capable
            // devices; the precompiled library must carry the same kernels
            // (the shader drops the define itself below Metal 3.1).
            "-DGGML_METAL_HAS_BF16=1",
            "-c",
            &metal_src,
            "-I",
            &metal_dir,
            "-o",
            &air_path,
        ])
        .status();

    let ok = metal_status.as_ref().is_ok_and(|s| s.success());
    if !ok {
        println!(
            "cargo:warning=failed to compile ggml-metal.metal to AIR; runtime source compile will be used"
        );
        let _ = fs::write(&metallib_path, []);
        emit_metallib_env(&metallib_path);
        return;
    }

    let mut metallib = Command::new("xcrun");
    metallib.args(["--sdk", "macosx", "metallib", &air_path]);
    let metallib_status = metallib.arg("-o").arg(&metallib_path).status();

    let ok = metallib_status.as_ref().is_ok_and(|s| s.success());
    if !ok {
        println!(
            "cargo:warning=failed to build ggml default metallib; runtime source compile will be used"
        );
        let _ = fs::write(&metallib_path, []);
    }

    emit_metallib_env(&metallib_path);
}

fn emit_metallib_env(metallib_path: &str) {
    // Own crate's rustc-env: harmless (this crate has no Rust code that
    // reads it today), kept for parity/future use.
    println!("cargo:rustc-env=MAKEPAD_METALLIB={metallib_path}");
    println!("cargo:rustc-env=MAKEPAD_GGML_METALLIB={metallib_path}");
    // links-metadata handshake: flows to immediate dependents (libs/ggml)
    // as DEP_MAKEPAD_AI_METAL_METALLIB=<metallib_path>.
    println!("cargo:metallib={}", metallib_path);
}