Add the MLX backend to the Apple frameworks and SwiftPM package

.github/workflows/apple.yml

@@ -192,6 +192,7 @@ jobs:
  "executorch"
  "executorch_llm"
  "backend_coreml"
+ "backend_mlx"
  "backend_xnnpack"
  "kernels_llm"
  "kernels_optimized"
@@ -220,6 +221,13 @@ jobs:
  zip -r "${RUNNER_TEMP}/artifacts/${FRAMEWORK}_debug-${VERSION}.zip" "${FRAMEWORK}_debug.xcframework"
  ) done
 
+ # MLX metallibs are data files, so carry them in the artifact for the SwiftPM job to commit.
+ if [ -d .Package.swift/backend_mlx_resources ]; then
+ mkdir -p "${RUNNER_TEMP}/artifacts/backend_mlx_resources"
+ cp .Package.swift/backend_mlx_resources/*.metallib \
+ "${RUNNER_TEMP}/artifacts/backend_mlx_resources/" 2>/dev/null || true
+ fi
+
  upload-frameworks-ios:
  # NB: Don't run this on fork PRs because they won't have access to the secret and would fail anyway
  if: ${{ !github.event.pull_request.head.repo.fork }}
@@ -328,9 +336,38 @@ jobs:
  exit 1
  fi
  done < "${RUNNER_TEMP}/checksums.txt"
+ # If MLX is being published, its Metal kernels have to come with it.
+ # The copy below is conditional and the capture upstream is best
+ # effort, so without this an incomplete set publishes a manifest
+ # declaring three resources that are not on the branch, and the
+ # consumer only finds out at Metal device init.
+ if grep -q "^backend_mlx " "${RUNNER_TEMP}/checksums.txt"; then
+ for SLICE in mlx-ios mlx-ios-simulator mlx-macos; do
+ METALLIB="${RUNNER_TEMP}/frameworks-ios/backend_mlx_resources/${SLICE}.metallib"
+ if [ ! -s "${METALLIB}" ]; then
+ echo "::error::backend_mlx is being published but ${SLICE}.metallib is missing or empty"
+ exit 1
+ fi
+ done
+ fi
+ # Nothing else evaluates this manifest before consumers do: the checks
+ # above are text greps, the swiftpm branch has no CI of its own, and a
+ # .template is not a manifest. Resolve it here so a Swift-level defect
+ # fails this run rather than every consumer of the published branch.
+ if ! swift package dump-package > /dev/null; then
+ echo "::error::the substituted Package.swift does not evaluate"
+ exit 1
+ fi
 
  git config --global user.name "PyTorch Bot"
  git config --global user.email "pytorchbot@users.noreply.github.com"
+ # Commit the MLX metallibs onto the package branch beside the manifest.
+ if [ -d "${RUNNER_TEMP}/frameworks-ios/backend_mlx_resources" ]; then
+ mkdir -p .Package.swift/backend_mlx_resources
+ cp "${RUNNER_TEMP}"/frameworks-ios/backend_mlx_resources/*.metallib \
+ .Package.swift/backend_mlx_resources/
+ git add -f .Package.swift/backend_mlx_resources/*.metallib
+ fi
  git add Package.swift
  git commit -am "${VERSION}"
  git push -f origin "${BRANCH}"

.gitignore

@@ -101,3 +101,4 @@ zephyr_dev_root.backup.*/
 # Agents
 .claude/*.local.*
 extension/pybindings/mlx.metallib
+.Package.swift/backend_mlx_resources/*.metallib

CMakePresets.json

@@ -45,7 +45,7 @@
  "CMAKE_TOOLCHAIN_FILE": "${sourceDir}/third-party/ios-cmake/ios.toolchain.cmake",
  "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/macos.cmake",
  "PLATFORM": "MAC_ARM64",
- "DEPLOYMENT_TARGET": "12.0",
+ "DEPLOYMENT_TARGET": "14.0",
  "CMAKE_MACOSX_BUNDLE": "OFF"
  },
  "condition": {

Package.swift

@@ -53,6 +53,13 @@ let products = deliverables([
  "sqlite3",
  ],
  ],
+ "backend_mlx": [
+ "frameworks": [
+ "Metal",
+ "Foundation",
+ "QuartzCore",
+ ],
+ ],
  "backend_xnnpack": [
  "targets": [
  "threadpool",
@@ -107,17 +114,49 @@ for (key, value) in products {
  name: key,
  path: "cmake-out/\(key).xcframework"
  ))
- let target: Target = .target(
+ // Split into typed locals to stay under the manifest type-check budget.
+ let targetNames: [String] = [key] + (value["targets"] as? [String] ?? []).map {
+ key.hasSuffix(debug_suffix) ? $0 + debug_suffix : $0
+ }
+ let dependencies: [Target.Dependency] = targetNames.map { .target(name: $0) }
+ let frameworks: [LinkerSetting] =
+ (value["frameworks"] as? [String] ?? []).map { .linkedFramework($0) }
+ let libraries: [LinkerSetting] =
+ (value["libraries"] as? [String] ?? []).map { .linkedLibrary($0) }
+ packageTargets.append(.target(
  name: "\(key)\(dependencies_suffix)",
- dependencies: ([key] + (value["targets"] as? [String] ?? []).map {
- key.hasSuffix(debug_suffix) ? $0 + debug_suffix : $0
- }).map { .target(name: $0) },
+ dependencies: dependencies,
  path: ".Package.swift/\(key)",
- linkerSettings:
- (value["frameworks"] as? [String] ?? []).map { .linkedFramework($0) } +
- (value["libraries"] as? [String] ?? []).map { .linkedLibrary($0) }
- )
- packageTargets.append(target)
+ linkerSettings: frameworks + libraries
+ ))
+}
+
+// One resource bundle shared by both MLX products (release and debug) so they
+// resolve the same bundle name, with a per-slice metallib since one slice's does
+// not load on another. Kept out of the loop above to avoid a per-debug copy.
+let mlxMetallibSlices = ["mlx-ios", "mlx-ios-simulator", "mlx-macos"]
+// Anchor to this file's own directory: as a dependency the manifest runs with the
+// consumer's cwd, where a relative path would miss every slice and ship an empty bundle.
+let mlxResourcesRelDir = ".Package.swift/backend_mlx_resources"
+let mlxResourcesDir =
+ URL(fileURLWithPath: #filePath).deletingLastPathComponent()
+ .appendingPathComponent(mlxResourcesRelDir).path
+if products.keys.contains("backend_mlx") {
+ packageTargets.append(.target(
+ name: "backend_mlx_resources",
+ path: mlxResourcesRelDir,
+ resources: mlxMetallibSlices.compactMap { slice in
+ FileManager.default.fileExists(atPath: "\(mlxResourcesDir)/\(slice).metallib")
+ ? .copy("\(slice).metallib") : nil
+ }
+ ))
+ for suffix in ["", debug_suffix] {
+ if let index = packageTargets.firstIndex(where: {
+ $0.name == "backend_mlx\(suffix)\(dependencies_suffix)"
+ }) {
+ packageTargets[index].dependencies.append(.target(name: "backend_mlx_resources"))
+ }
+ }
 }
 
 // Test fixtures. add_coreml.pte and add_mul_coreml.pte are generated at CI
@@ -151,6 +190,12 @@ if FileManager.default.fileExists(atPath: "\(objcTestsDir)/add_mul_coreml.pte")
 }
 
 let testLinkerSettings: [LinkerSetting] = [
+ // The test targets depend on the executorch binary target directly, which
+ // carries no linker settings, rather than the with-dependencies target that
+ // owns the libc++ link, so they must link libc++ themselves. Below a macOS 13
+ // deployment target a Swift back-deployment shim used to supply it implicitly;
+ // at this package's floor that shim is gone, so name it explicitly here.
+ .linkedLibrary("c++"),
  .unsafeFlags([
  "-Xlinker", "-force_load",
  "-Xlinker", "cmake-out/kernels_optimized.xcframework/macos-arm64/libkernels_optimized_macos.a",
@@ -166,7 +211,7 @@ let package = Package(
  name: "executorch",
  platforms: [
  .iOS(.v17),
- .macOS(.v12),
+ .macOS(.v14),
  ],
  products: packageProducts,
  targets: packageTargets + [

backends/mlx/CMakeLists.txt

@@ -188,16 +188,91 @@ set(_mlx_patches
  ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_jit_sdk_gate.patch
  ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch
  ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_gather_mm_rhs_lda.patch
+ ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_sdk_per_platform.patch
+ ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_swiftpm_metallib_name.patch
 )
+# In a framework build the delegate is static, so MLX cannot find a colocated
+# metallib and instead loads one from a SwiftPM resource bundle. SWIFTPM_BUNDLE
+# is the bundle name and MLX_SWIFTPM_METALLIB_NAME the per-slice file inside it
+# (one per slice, since a slice's metallib does not load on another). PLATFORM
+# is set only for the Apple presets; a plain wheel build leaves it empty and
+# keeps the colocated path.
+set(_mlx_extra_cxx_flags "")
+set(_mlx_cxx_flags_arg "")
+if(PLATFORM)
+ if(PLATFORM STREQUAL "OS64")
+ set(_mlx_metallib_slice "ios")
+ elseif(PLATFORM STREQUAL "SIMULATORARM64")
+ set(_mlx_metallib_slice "ios-simulator")
+ else()
+ set(_mlx_metallib_slice "macos")
+ endif()
+ set(_mlx_extra_cxx_flags
+ "-DSWIFTPM_BUNDLE=\\\"executorch_backend_mlx_resources\\\" -DMLX_SWIFTPM_METALLIB_NAME=\\\"mlx-${_mlx_metallib_slice}\\\""
+ )
+ # Only override the sub-build's CXX flags when there is something to add. An
+ # empty -DCMAKE_CXX_FLAGS= on the command line beats the environment, so
+ # passing it unconditionally would silently drop a wheel build's CXXFLAGS for
+ # MLX only.
+ set(_mlx_cxx_flags_arg "-DCMAKE_CXX_FLAGS=${_mlx_extra_cxx_flags}")
+endif()
+
+# Prefer the preset's per-slice DEPLOYMENT_TARGET; fall back to the toolchain
+# value.
+if(DEPLOYMENT_TARGET)
+ set(_mlx_osx_deployment_target ${DEPLOYMENT_TARGET})
+else()
+ set(_mlx_osx_deployment_target ${CMAKE_OSX_DEPLOYMENT_TARGET})
+endif()
+
+# The sub-build is forced to a single-config generator (Unix Makefiles), which
+# reads CMAKE_BUILD_TYPE, not the --config passed to the parent. Follow the
+# active configuration through $<CONFIG>, which resolves to the config being
+# built under a multi-config parent (Xcode) and to CMAKE_BUILD_TYPE under a
+# single-config one. It is empty only for a bare single-config parent that set
+# no type, so default that case to Release rather than compile MLX with no
+# optimization.
+set(_mlx_build_type "$<IF:$<BOOL:$<CONFIG>>,$<CONFIG>,Release>")
+
+# Decide the MLX sub-build's build command, consumed by BUILD_COMMAND below. An
+# empty value leaves ExternalProject's default in place, which is only the right
+# choice for a Makefiles parent: ExternalProject emits `$(MAKE)` (inheriting the
+# parent's jobserver) only when the PARENT generator matches "Make". Under Xcode
+# or Ninja it emits a bare `cmake --build`, with no -j and no jobserver, so MLX
+# would compile on one core. Name a job count there; a bare --parallel with no
+# number would expand to `make -j`, unbounded.
+cmake_host_system_information(
+ RESULT _mlx_host_cores QUERY NUMBER_OF_LOGICAL_CORES
+)
+if(CMAKE_GENERATOR MATCHES "Make")
+ set(_mlx_build_command "")
+else()
+ set(_mlx_build_command ${CMAKE_COMMAND} --build <BINARY_DIR> --parallel
+ ${_mlx_host_cores}
+ )
+endif()
+
 ExternalProject_Add(
  mlx_external
  SOURCE_DIR ${MLX_SOURCE_DIR}
  BINARY_DIR ${_mlx_binary_dir}
- PATCH_COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/patches/apply.sh
- ${MLX_SOURCE_DIR} ${_mlx_patches}
- CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
+ # Force a single-config generator for the sub-build. The parent Apple build
+ # uses the Xcode generator, which is multi-config and would place libmlx.a in
+ # a per-config subdirectory (mlx/Debug/libmlx.a). The paths below that consume
+ # the archive expect it flat at ${_mlx_binary_dir}/libmlx.a, the way a
+ # single-config generator emits it, which is also how the wheel build already
+ # produces it.
+ CMAKE_GENERATOR "Unix Makefiles"
+ CMAKE_ARGS "-DCMAKE_BUILD_TYPE=${_mlx_build_type}"
  -DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD}
- -DCMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}
+ ${_mlx_cxx_flags_arg}
+ # The preset's DEPLOYMENT_TARGET carries the correct per-slice
+ # minimum (iOS 17, macOS 14); the ios.toolchain leaves
+ # CMAKE_OSX_DEPLOYMENT_TARGET at 12.0 on the iOS slices and only
+ # syncs it for MAC. The shader-flag patch reads
+ # CMAKE_OSX_DEPLOYMENT_TARGET, so feed the preset value in as that,
+ # or the iOS metallib is stamped -mios-version-min=12.0.
+ -DCMAKE_OSX_DEPLOYMENT_TARGET=${_mlx_osx_deployment_target}
  -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE}
  -DPLATFORM=${PLATFORM}
  -DDEPLOYMENT_TARGET=${DEPLOYMENT_TARGET}
@@ -217,6 +292,14 @@ ExternalProject_Add(
  # MLX's own install() does not emit libmlx.a where we consume it or the
  # metallib at all, so skip the install step and read both from the build tree.
  INSTALL_COMMAND ""
+ # Build in parallel, but only where it is needed and always with a job count.
+ # Under the multi-config Xcode parent the generated sub-build step is `cmake
+ # --build . --config <cfg>` with no -j and no jobserver to inherit, so MLX
+ # would compile fully serially. Under a single-config parent the default build
+ # command is `$(MAKE)`, which already inherits the parent's jobserver, so
+ # overriding it there would break that and serialize MLX instead. A bare
+ # --parallel with no number expands to `make -j`, unbounded, so name a count.
+ BUILD_COMMAND ${_mlx_build_command}
  # ExternalProject stamps its build, so a bare MLX submodule bump (git
  # submodule update) would not invalidate the stamp and we'd link a stale
  # libmlx.a with no signal. BUILD_ALWAYS reruns the build step every configure;
@@ -226,6 +309,25 @@ ExternalProject_Add(
  BUILD_BYPRODUCTS ${_mlx_static_lib} ${_mlx_metallib}
 )
 
+# ExternalProject stamps the patch step and BUILD_ALWAYS does not re-run it, so
+# a reused build directory whose MLX source was reset (patches reverted) would
+# recompile an unpatched MLX and silently drop the iOS Metal SDK selection and
+# the SwiftPM metallib name. Re-apply the patches on every build; apply.sh is
+# idempotent (it reverse-checks each patch, skipping those already applied).
+# DEPENDERS configure, not build: mlx_metal_sdk_per_platform.patch edits the
+# sub-project's own CMake to pick the Metal SDK from PLATFORM, so it must land
+# before the sub-configure runs, or the shaders are built against the macOS SDK
+# for every slice. That does mean the sub-configure re-runs on each build, which
+# is the price of the ordering.
+ExternalProject_Add_Step(
+ mlx_external reapply_patches
+ COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/patches/apply.sh ${MLX_SOURCE_DIR}
+ ${_mlx_patches}
+ DEPENDEES patch
+ DEPENDERS configure
+ ALWAYS 1
+)
+
 # Imported target for the MLX static library produced by mlx_external. A static
 # libmlx.a carries no transitive link deps, so re-add the frameworks MLX itself
 # links (mirrors third-party/mlx/CMakeLists.txt:209). CPU is OFF, so Accelerate
@@ -283,6 +385,21 @@ add_dependencies(mlxdelegate mlx_schema)
 # Depend on mlx_external directly so libmlx.a exists before mlxdelegate links.
 add_dependencies(mlxdelegate mlx_external)
 
+# The Apple framework build reads each static library straight from the target
+# output directory (it merges them without an install step), but libmlx.a is
+# produced by the MLX sub-build in its own binary dir, not next to the delegate.
+# Copy it beside libmlxdelegate.a after the delegate builds so the framework
+# packaging finds it where it expects. Under the multi-config Xcode generator
+# the Apple presets use, TARGET_FILE_DIR resolves to the per-config directory
+# the packaging reads (<preset>/Release), which is what this relies on.
+add_custom_command(
+ TARGET mlxdelegate
+ POST_BUILD
+ COMMAND ${CMAKE_COMMAND} -E copy_if_different ${_mlx_static_lib}
+ $<TARGET_FILE_DIR:mlxdelegate>/libmlx.a
+ VERBATIM
+)
+
 # Add logging flag if enabled
 if(ET_MLX_ENABLE_OP_LOGGING)
  target_compile_definitions(mlxdelegate PRIVATE ET_MLX_ENABLE_OP_LOGGING=1)
@@ -333,6 +450,9 @@ endif()
 
 executorch_target_link_options_shared_lib(mlxdelegate)
 target_compile_options(mlxdelegate PRIVATE ${_common_compile_options})
+# Core tensor headers carry pre-existing narrowing conversions that trip Xcode's
+# -Wshorten-64-to-32 -Werror; suppress it here as XNNPACK and abseil already do.
+target_compile_options(mlxdelegate PRIVATE -Wno-shorten-64-to-32)
 if(EXECUTORCH_MLX_ENABLE_SANITIZERS)
  target_link_options(mlxdelegate PRIVATE ${_mlx_sanitizer_link_options})
 endif()