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.

Author: olilarkin

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

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();