onnxruntime-ep-mlx 0.3.0

MLX-native ONNX Runtime execution provider (plugin EP) for Apple Silicon — binds mlx-c directly, no mlx-rs.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use std::env;
use std::path::PathBuf;
use std::process::Command;

fn brew_prefix(pkg: &str) -> String {
    let out = Command::new("brew")
        .arg("--prefix")
        .arg(pkg)
        .output()
        .unwrap_or_else(|_| panic!("failed to run `brew --prefix {pkg}`"));
    String::from_utf8(out.stdout).unwrap().trim().to_string()
}

/// Discover the ONNX Runtime C-API include directory for a standalone checkout.
///
/// Resolution order:
///   1. `ORT_INCLUDE_DIR` — an explicit include dir (highest precedence).
///   2. `$ORT_HOME/include` — the standard layout of an ORT release tarball
///      (`onnxruntime-osx-arm64-<ver>/include`), which is what CI provisions.
///
/// Whichever we pick, we verify `onnxruntime_c_api.h` is actually present so the
/// error surfaces here (with a clear message) rather than deep inside bindgen.
fn resolve_ort_include() -> String {
    let candidate = env::var("ORT_INCLUDE_DIR").ok().or_else(|| {
        env::var("ORT_HOME")
            .ok()
            .map(|home| format!("{home}/include"))
    });

    match candidate {
        Some(dir) if PathBuf::from(&dir).join("onnxruntime_c_api.h").is_file() => dir,
        Some(dir) => panic!(
            "ORT include dir '{dir}' does not contain onnxruntime_c_api.h. \
             Set ORT_INCLUDE_DIR to the ONNX Runtime C-API include directory, \
             or set ORT_HOME to an ONNX Runtime release root (expects \
             $ORT_HOME/include/onnxruntime_c_api.h)."
        ),
        None => panic!(
            "Could not locate the ONNX Runtime headers. Set ORT_INCLUDE_DIR to \
             the ORT C-API include directory, or set ORT_HOME to an ONNX Runtime \
             release root (expects $ORT_HOME/include/onnxruntime_c_api.h)."
        ),
    }
}

fn main() {
    let ort_inc = resolve_ort_include();
    let mlxc = brew_prefix("mlx-c");
    let mlx = brew_prefix("mlx");
    let out = PathBuf::from(env::var("OUT_DIR").unwrap());

    println!("cargo:rerun-if-changed=wrapper_ort.h");
    println!("cargo:rerun-if-changed=wrapper_mlx.h");
    println!("cargo:rerun-if-env-changed=ORT_INCLUDE_DIR");
    println!("cargo:rerun-if-env-changed=ORT_HOME");

    // --- ORT plugin-EP C ABI bindings (pure C header; pulls in onnxruntime_ep_c_api.h) ---
    bindgen::Builder::default()
        .header("wrapper_ort.h")
        .clang_arg(format!("-I{ort_inc}"))
        .allowlist_type("Ort.*")
        .allowlist_type("ONNX.*")
        .allowlist_function("Ort.*")
        .wrap_unsafe_ops(true)
        .generate()
        .expect("ORT bindgen failed")
        .write_to_file(out.join("ort.rs"))
        .unwrap();

    // --- mlx-c bindings (bind the C API DIRECTLY, no mlx-rs crate) ---
    bindgen::Builder::default()
        .header("wrapper_mlx.h")
        .clang_arg(format!("-I{mlxc}/include"))
        .allowlist_function("mlx_.*")
        .allowlist_type("mlx_.*")
        .allowlist_var("MLX_.*")
        .wrap_unsafe_ops(true)
        .generate()
        .expect("mlx-c bindgen failed")
        .write_to_file(out.join("mlx.rs"))
        .unwrap();

    // We call ORT purely through the OrtApi function-pointer table handed to
    // CreateEpFactories, so we do NOT link libonnxruntime. Only mlx-c + mlx + frameworks.
    println!("cargo:rustc-link-search=native={mlxc}/lib");
    println!("cargo:rustc-link-search=native={mlx}/lib");
    println!("cargo:rustc-link-lib=dylib=mlxc");
    println!("cargo:rustc-link-lib=dylib=mlx");
    println!("cargo:rustc-link-arg=-Wl,-rpath,{mlxc}/lib");
    println!("cargo:rustc-link-arg=-Wl,-rpath,{mlx}/lib");
    for fw in ["Metal", "Foundation", "QuartzCore", "Accelerate"] {
        println!("cargo:rustc-link-lib=framework={fw}");
    }
}