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