From 5c716bdc26976df970abe657f39a47f73485a8f1 Mon Sep 17 00:00:00 2001
From: Oli Larkin <olilarkin@googlemail.com>
Date: Mon, 9 Feb 2026 23:40:53 +0100
Subject: [PATCH 1/2] Add set_metallib_path() for custom metallib location

Allows overriding the default metallib search paths before Device
initialization. Needed when MLX is embedded in plugin bundles (AU,
VST3) where the standard search locations don't apply.
---
 mlx/backend/metal/device.cpp | 15 +++++++++++++++
 mlx/backend/metal/metal.h    |  4 ++++
 2 files changed, 19 insertions(+)

diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp
index 7658ce5f5c..f103a6dbaf 100644
--- a/mlx/backend/metal/device.cpp
+++ b/mlx/backend/metal/device.cpp
@@ -30,6 +30,12 @@ struct hash<NS::SharedPtr<T>> {
 
 namespace mlx::core::metal {
 
+static std::string g_metallib_override_path;
+
+void set_metallib_path(const std::string& path) {
+  g_metallib_override_path = path;
+}
+
 namespace {
 
 constexpr const char* default_mtllib_path = METAL_PATH;
@@ -162,6 +168,15 @@ std::pair<MTL::Library*, NS::Error*> load_swiftpm_library(
 }
 
 MTL::Library* load_default_library(MTL::Device* device) {
+  // Check override path first
+  if (!g_metallib_override_path.empty()) {
+    auto [lib, error] =
+        load_library_from_path(device, g_metallib_override_path.c_str());
+    if (lib) {
+      return lib;
+    }
+  }
+
   NS::Error* error[5];
   MTL::Library* lib;
   // First try the colocated mlx.metallib
diff --git a/mlx/backend/metal/metal.h b/mlx/backend/metal/metal.h
index 6662e21ebd..eaea1814c9 100644
--- a/mlx/backend/metal/metal.h
+++ b/mlx/backend/metal/metal.h
@@ -10,6 +10,10 @@
 
 namespace mlx::core::metal {
 
+/* Set a custom path to mlx.metallib. Must be called before any MLX operation.
+ */
+void set_metallib_path(const std::string& path);
+
 /* Check if the Metal backend is available. */
 MLX_API bool is_available();
 

From 6979af077df0cd602681275c54b15b2ba93efa5d Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Sun, 21 Jun 2026 20:30:02 +0900
Subject: [PATCH 2/2] Clean up code

---
 mlx/backend/metal/device.cpp   | 21 ++++++++++-----------
 mlx/backend/metal/metal.cpp    | 14 ++++++++++++++
 mlx/backend/metal/metal.h      |  9 +++++----
 mlx/backend/metal/no_metal.cpp |  7 +++++++
 4 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp
index f103a6dbaf..2a8e15afd7 100644
--- a/mlx/backend/metal/device.cpp
+++ b/mlx/backend/metal/device.cpp
@@ -30,12 +30,6 @@ struct hash<NS::SharedPtr<T>> {
 
 namespace mlx::core::metal {
 
-static std::string g_metallib_override_path;
-
-void set_metallib_path(const std::string& path) {
-  g_metallib_override_path = path;
-}
-
 namespace {
 
 constexpr const char* default_mtllib_path = METAL_PATH;
@@ -168,13 +162,18 @@ std::pair<MTL::Library*, NS::Error*> load_swiftpm_library(
 }
 
 MTL::Library* load_default_library(MTL::Device* device) {
-  // Check override path first
-  if (!g_metallib_override_path.empty()) {
+  // Check override path before automatic lookup
+  if (!get_metallib_path().empty()) {
     auto [lib, error] =
-        load_library_from_path(device, g_metallib_override_path.c_str());
-    if (lib) {
-      return lib;
+        load_library_from_path(device, get_metallib_path().c_str());
+    if (!lib) {
+      throw std::runtime_error(
+          fmt::format(
+              "Can not load metallib from specified location \"{}\": {}.",
+              get_metallib_path(),
+              error->localizedDescription()->utf8String()));
     }
+    return lib;
   }
 
   NS::Error* error[5];
diff --git a/mlx/backend/metal/metal.cpp b/mlx/backend/metal/metal.cpp
index 6bf3b895b9..e95b6f0b9e 100644
--- a/mlx/backend/metal/metal.cpp
+++ b/mlx/backend/metal/metal.cpp
@@ -7,6 +7,12 @@
 
 namespace mlx::core::metal {
 
+namespace {
+
+std::string g_metallib_path;
+
+} // namespace
+
 bool is_available() {
   return true;
 }
@@ -46,4 +52,12 @@ void stop_capture() {
   manager->stopCapture();
 }
 
+void set_metallib_path(const std::string& path) {
+  g_metallib_path = path;
+}
+
+const std::string& get_metallib_path() {
+  return g_metallib_path;
+}
+
 } // namespace mlx::core::metal
diff --git a/mlx/backend/metal/metal.h b/mlx/backend/metal/metal.h
index eaea1814c9..d686c2971a 100644
--- a/mlx/backend/metal/metal.h
+++ b/mlx/backend/metal/metal.h
@@ -10,10 +10,6 @@
 
 namespace mlx::core::metal {
 
-/* Set a custom path to mlx.metallib. Must be called before any MLX operation.
- */
-void set_metallib_path(const std::string& path);
-
 /* Check if the Metal backend is available. */
 MLX_API bool is_available();
 
@@ -26,4 +22,9 @@ MLX_API const
     std::unordered_map<std::string, std::variant<std::string, size_t>>&
     device_info();
 
+/* Set a custom path to mlx.metallib. Must be called before any MLX operation.
+ */
+MLX_API void set_metallib_path(const std::string& path);
+MLX_API const std::string& get_metallib_path();
+
 } // namespace mlx::core::metal
diff --git a/mlx/backend/metal/no_metal.cpp b/mlx/backend/metal/no_metal.cpp
index 2d3414cc50..3b72ad005a 100644
--- a/mlx/backend/metal/no_metal.cpp
+++ b/mlx/backend/metal/no_metal.cpp
@@ -21,6 +21,13 @@ device_info() {
       "[metal::device_info] Cannot get device info without metal backend");
 };
 
+void set_metallib_path(const std::string& path) {}
+
+const std::string& get_metallib_path() {
+  throw std::runtime_error(
+      "[metal::get_metallib_path] Cannot get metallib path without metal backend");
+}
+
 } // namespace metal
 
 } // namespace mlx::core