From 895217f25b64e813c710a4bb3e061d0d7f3ca014 Mon Sep 17 00:00:00 2001
From: David Koski <46639364+davidkoski@users.noreply.github.com>
Date: Mon, 27 Oct 2025 07:52:03 -0700
Subject: [PATCH] optionally load metallib from framework (#2702)

* optionally load metallib from framework

* pre-commit

* adjust logic
---
 mlx/backend/metal/device.cpp | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp
index bd9d16b15..e82d734a2 100644
--- a/mlx/backend/metal/device.cpp
+++ b/mlx/backend/metal/device.cpp
@@ -72,6 +72,19 @@ MTL::Library* try_load_bundle(
   }
   return nullptr;
 }
+
+MTL::Library* try_load_framework(
+    MTL::Device* device,
+    NS::URL* url,
+    const std::string& lib_name) {
+  std::string resource_path = std::string(url->fileSystemRepresentation()) +
+      "/" + lib_name + ".metallib";
+  auto [lib, error] = load_library_from_path(device, resource_path.c_str());
+  if (lib) {
+    return lib;
+  }
+  return nullptr;
+}
 #endif
 
 // Firstly, search for the metallib in the same path as this binary
@@ -103,6 +116,17 @@ std::pair<MTL::Library*, NS::Error*> load_swiftpm_library(
       return {library, nullptr};
     }
   }
+  // if SWIFTPM_BUNDLE is a framework identifier, try loading from that
+  auto frameworks = NS::Bundle::allFrameworks();
+  for (int i = 0, c = (int)frameworks->count(); i < c; i++) {
+    auto bundle = reinterpret_cast<NS::Bundle*>(frameworks->object(i));
+    if (!strcmp(bundle->bundleIdentifier()->utf8String(), SWIFTPM_BUNDLE)) {
+      library = try_load_framework(device, bundle->resourceURL(), lib_name);
+      if (library != nullptr) {
+        return {library, nullptr};
+      }
+    }
+  }
 #endif
   return {nullptr, nullptr};
 }