From 6ed51729d327fa83e8a7ddcef1c1c175a7a5830a Mon Sep 17 00:00:00 2001
From: Kyle Kelley <rgbkrk@gmail.com>
Date: Wed, 18 Mar 2026 18:11:53 -0700
Subject: [PATCH] Fix metallib path for cargo install

Copy mlx.metallib to a stable versioned directory (~/.mlx/lib/v{version}/)
during the CMake build by overriding MLX_METAL_PATH. The compiled-in
METAL_PATH fallback now survives cargo install temp dir cleanup.

Previously, METAL_PATH pointed to the CMake build directory inside
Cargo's temp workspace. When cargo install copies only the binary and
deletes the temp dir, the metallib becomes unreachable at runtime:

  MLX error: Failed to load the default metallib.

Fixes #327
---
 mlx-sys/build.rs | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/mlx-sys/build.rs b/mlx-sys/build.rs
index 11f6300f0..f291c1f60 100644
--- a/mlx-sys/build.rs
+++ b/mlx-sys/build.rs
@@ -1,7 +1,7 @@
 extern crate cmake;
 
 use cmake::Config;
-use std::{env, path::PathBuf, process::Command};
+use std::{env, fs, path::PathBuf, process::Command};
 
 /// Find the clang runtime library path dynamically using xcrun
 fn find_clang_rt_path() -> Option<String> {
@@ -44,6 +44,27 @@ fn find_clang_rt_path() -> Option<String> {
     None
 }
 
+/// Determine a stable directory for the MLX metallib that survives `cargo install`
+/// temp dir cleanup. Uses `~/.mlx/lib/v{version}/`.
+///
+/// When `cargo install` builds a crate, it uses a temporary directory that is
+/// deleted after the binary is copied. The CMake build bakes the metallib path
+/// into the binary via `-DMETAL_PATH=...`. If that path points to the temp dir,
+/// the binary fails at runtime with "Failed to load the default metallib".
+///
+/// By setting MLX_METAL_PATH to a stable home-directory location (and creating
+/// it before CMake runs so CMake can output the metallib there directly), the
+/// compiled-in METAL_PATH remains valid after the temp dir is cleaned up.
+#[cfg(feature = "metal")]
+fn stable_metallib_dir() -> PathBuf {
+    let version = env!("CARGO_PKG_VERSION");
+    let home = env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
+    PathBuf::from(home)
+        .join(".mlx")
+        .join("lib")
+        .join(format!("v{}", version))
+}
+
 fn build_and_link_mlx_c() {
     let mut config = Config::new("src/mlx-c");
     config.very_verbose(true);
@@ -69,6 +90,14 @@ fn build_and_link_mlx_c() {
     #[cfg(feature = "metal")]
     {
         config.define("MLX_BUILD_METAL", "ON");
+
+        // Point MLX_METAL_PATH to a stable location so the compiled-in
+        // METAL_PATH survives `cargo install` temp dir cleanup.
+        // Must create the directory BEFORE CMake runs because CMake
+        // outputs the compiled metallib directly to MLX_METAL_PATH.
+        let metallib_dir = stable_metallib_dir();
+        fs::create_dir_all(&metallib_dir).expect("Failed to create stable metallib directory");
+        config.define("MLX_METAL_PATH", metallib_dir.to_str().unwrap());
     }
 
     #[cfg(feature = "accelerate")]
@@ -90,6 +119,17 @@ fn build_and_link_mlx_c() {
     #[cfg(feature = "metal")]
     {
         println!("cargo:rustc-link-lib=framework=Metal");
+
+        // Verify CMake output the metallib to the stable location.
+        let metallib_dir = stable_metallib_dir();
+        let target_metallib = metallib_dir.join("mlx.metallib");
+        if !target_metallib.exists() {
+            eprintln!(
+                "cargo:warning=mlx.metallib not found at {}. \
+                 Runtime Metal operations may fail.",
+                target_metallib.display()
+            );
+        }
     }
 
     #[cfg(feature = "accelerate")]