From ca2924ff26bf4e45f5bbdf7dca72496220cbd0cf Mon Sep 17 00:00:00 2001 From: Antigravity Date: Tue, 19 May 2026 02:03:26 +0200 Subject: [PATCH 1/2] Build SwiftPM default metallib resource --- Package.swift | 22 +++++- Plugins/BuildSwiftPMMetalLibrary/plugin.swift | 55 ++++++++++++++ tools/build-swiftpm-metallib.sh | 76 +++++++++++++++++++ 3 files changed, 150 insertions(+), 3 deletions(-) create mode 100644 Plugins/BuildSwiftPMMetalLibrary/plugin.swift create mode 100755 tools/build-swiftpm-metallib.sh diff --git a/Package.swift b/Package.swift index e36e0d86..7664b171 100644 --- a/Package.swift +++ b/Package.swift @@ -83,6 +83,7 @@ let noCudaCmlxExcludes = [ let cxxSettings: [CXXSetting] let linkerSettings: [LinkerSetting] let mlxSwiftExcludes: [String] + let cmlxPlugins: [Target.PluginUsage] if Context.environment["SPM_CUDA"] != "0" { // Linux with CUDA @@ -140,6 +141,10 @@ let noCudaCmlxExcludes = [ "GPU+Metal.swift", "MLXArray+Metal.swift", ] + + cmlxPlugins = [ + .plugin(name: "CudaBuild") + ] } else { // Linux without CUDA (CPU only) @@ -175,6 +180,10 @@ let noCudaCmlxExcludes = [ "MLXFast.swift", "MLXFastKernel.swift", ] + + cmlxPlugins = [ + .plugin(name: "CudaBuild") + ] } #else // Apple's platforms with Metal @@ -212,6 +221,11 @@ let noCudaCmlxExcludes = [ let mlxSwiftExcludes: [String] = [ "GPU+CUDA.swift" ] + + let cmlxPlugins: [Target.PluginUsage] = [ + .plugin(name: "CudaBuild"), + .plugin(name: "BuildSwiftPMMetalLibrary"), + ] #endif let cmlx = Target.target( @@ -289,9 +303,7 @@ let cmlx = Target.target( .define("MLX_VERSION", to: "\"0.31.1\""), ], linkerSettings: linkerSettings, - plugins: [ - .plugin(name: "CudaBuild") - ], + plugins: cmlxPlugins ) let package = Package( @@ -321,6 +333,10 @@ let package = Package( ], targets: [ cmlx, + .plugin( + name: "BuildSwiftPMMetalLibrary", + capability: .buildTool() + ), .testTarget( name: "CmlxTests", dependencies: ["Cmlx"] diff --git a/Plugins/BuildSwiftPMMetalLibrary/plugin.swift b/Plugins/BuildSwiftPMMetalLibrary/plugin.swift new file mode 100644 index 00000000..b61460e1 --- /dev/null +++ b/Plugins/BuildSwiftPMMetalLibrary/plugin.swift @@ -0,0 +1,55 @@ +import Foundation +import PackagePlugin + +@main +struct BuildSwiftPMMetalLibrary: BuildToolPlugin { + func createBuildCommands(context: PluginContext, target: any Target) async throws -> [Command] { + #if os(Linux) + return [] + #else + let packageRoot = context.package.directory + let script = packageRoot.appending("tools", "build-swiftpm-metallib.sh") + let output = context.pluginWorkDirectory.appending("default.metallib") + + return [ + .buildCommand( + displayName: "Build SwiftPM default.metallib", + executable: Path("/bin/bash"), + arguments: [script, output], + inputFiles: inputFiles(packageRoot: packageRoot, script: script), + outputFiles: [output] + ) + ] + #endif + } + + #if !os(Linux) + private func inputFiles(packageRoot: Path, script: Path) -> [Path] { + let kernelsDirectory = packageRoot.appending( + "Source", + "Cmlx", + "mlx", + "mlx", + "backend", + "metal", + "kernels" + ) + var files = [script] + files.append(contentsOf: recursivelyCollectedMetalInputs(in: kernelsDirectory)) + return files + } + + private func recursivelyCollectedMetalInputs(in directory: Path) -> [Path] { + let fileManager = FileManager.default + guard let enumerator = fileManager.enumerator(atPath: directory.string) else { + return [] + } + + return enumerator.compactMap { entry -> Path? in + guard let entry = entry as? String else { return nil } + guard entry.hasSuffix(".metal") || entry.hasSuffix(".h") else { return nil } + return directory.appending(subpath: entry) + }.sorted { $0.string < $1.string } + } + #endif +} diff --git a/tools/build-swiftpm-metallib.sh b/tools/build-swiftpm-metallib.sh new file mode 100755 index 00000000..08103bdd --- /dev/null +++ b/tools/build-swiftpm-metallib.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# Build the default Metal library resource used by SwiftPM Cmlx builds. + +set -euo pipefail + +if [[ $# -ne 1 ]]; then + echo "usage: $0 OUTPUT_METALLIB" >&2 + exit 64 +fi + +OUTPUT="$1" +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) +ROOT_DIR=$(realpath "${SCRIPT_DIR}/..") +KERNELS_DIR="${ROOT_DIR}/Source/Cmlx/mlx/mlx/backend/metal/kernels" + +METAL=$(xcrun -sdk macosx -find metal) +METALLIB=$(xcrun -sdk macosx -find metallib) +TMP_DIR=$(mktemp -d) +trap 'rm -rf "${TMP_DIR}"' EXIT + +DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-14.0}" + +metal_version=$( + printf '%s\n' '__METAL_VERSION__' | + "${METAL}" "-mmacosx-version-min=${DEPLOYMENT_TARGET}" -E -x metal -P - | + tail -1 | + tr -d '[:space:]' +) +metal_version=${metal_version:-0} + +kernels=( + "arg_reduce" + "conv" + "gemv" + "layer_norm" + "random" + "rms_norm" + "rope" + "scaled_dot_product_attention" +) + +if (( metal_version >= 320 )); then + kernels+=("fence") +fi + +metal_flags=( + -x metal + -Wall + -Wextra + -fno-fast-math + -Wno-c++17-extensions + -Wno-c++20-extensions + -mmacosx-version-min="${DEPLOYMENT_TARGET}" +) + +if (( metal_version >= 400 )); then + metal_flags+=(-std=metal4.0) +elif (( metal_version >= 320 )); then + metal_flags+=(-std=metal3.2) +elif (( metal_version >= 310 )); then + metal_flags+=(-std=metal3.1) +elif (( metal_version >= 300 )); then + metal_flags+=(-std=metal3.0) +fi + +air_files=() +for kernel in "${kernels[@]}"; do + source="${KERNELS_DIR}/${kernel}.metal" + air="${TMP_DIR}/${kernel}.air" + "${METAL}" "${metal_flags[@]}" -c "${source}" -I"${ROOT_DIR}/Source/Cmlx/mlx" -o "${air}" + air_files+=("${air}") +done + +mkdir -p "$(dirname "${OUTPUT}")" +"${METALLIB}" "${air_files[@]}" -o "${TMP_DIR}/default.metallib" +mv "${TMP_DIR}/default.metallib" "${OUTPUT}" From 2b9dda12f9c1a9681e54b38aa718a9437eb1e13e Mon Sep 17 00:00:00 2001 From: RNT56 Date: Thu, 25 Jun 2026 15:52:17 +0200 Subject: [PATCH 2/2] Harden SwiftPM Metal resource build --- Plugins/BuildSwiftPMMetalLibrary/plugin.swift | 42 +++++++------ tools/build-swiftpm-metallib.sh | 59 +++++++++++++++++-- 2 files changed, 73 insertions(+), 28 deletions(-) diff --git a/Plugins/BuildSwiftPMMetalLibrary/plugin.swift b/Plugins/BuildSwiftPMMetalLibrary/plugin.swift index b61460e1..bed4a079 100644 --- a/Plugins/BuildSwiftPMMetalLibrary/plugin.swift +++ b/Plugins/BuildSwiftPMMetalLibrary/plugin.swift @@ -7,15 +7,15 @@ struct BuildSwiftPMMetalLibrary: BuildToolPlugin { #if os(Linux) return [] #else - let packageRoot = context.package.directory - let script = packageRoot.appending("tools", "build-swiftpm-metallib.sh") - let output = context.pluginWorkDirectory.appending("default.metallib") + let packageRoot = context.package.directoryURL + let script = packageRoot.appendingPathComponent("tools/build-swiftpm-metallib.sh") + let output = context.pluginWorkDirectoryURL.appendingPathComponent("default.metallib") return [ .buildCommand( displayName: "Build SwiftPM default.metallib", - executable: Path("/bin/bash"), - arguments: [script, output], + executable: URL(fileURLWithPath: "/bin/bash"), + arguments: [script.path, output.path], inputFiles: inputFiles(packageRoot: packageRoot, script: script), outputFiles: [output] ) @@ -24,32 +24,30 @@ struct BuildSwiftPMMetalLibrary: BuildToolPlugin { } #if !os(Linux) - private func inputFiles(packageRoot: Path, script: Path) -> [Path] { - let kernelsDirectory = packageRoot.appending( - "Source", - "Cmlx", - "mlx", - "mlx", - "backend", - "metal", - "kernels" - ) + private func inputFiles(packageRoot: URL, script: URL) -> [URL] { + let kernelsDirectory = packageRoot.appendingPathComponent( + "Source/Cmlx/mlx/mlx/backend/metal/kernels") var files = [script] files.append(contentsOf: recursivelyCollectedMetalInputs(in: kernelsDirectory)) return files } - private func recursivelyCollectedMetalInputs(in directory: Path) -> [Path] { + private func recursivelyCollectedMetalInputs(in directory: URL) -> [URL] { let fileManager = FileManager.default - guard let enumerator = fileManager.enumerator(atPath: directory.string) else { + guard + let enumerator = fileManager.enumerator( + at: directory, includingPropertiesForKeys: nil) + else { return [] } - return enumerator.compactMap { entry -> Path? in - guard let entry = entry as? String else { return nil } - guard entry.hasSuffix(".metal") || entry.hasSuffix(".h") else { return nil } - return directory.appending(subpath: entry) - }.sorted { $0.string < $1.string } + return enumerator.compactMap { entry -> URL? in + guard let url = entry as? URL else { return nil } + guard url.pathExtension == "metal" || url.pathExtension == "h" else { + return nil + } + return url + }.sorted { $0.path < $1.path } } #endif } diff --git a/tools/build-swiftpm-metallib.sh b/tools/build-swiftpm-metallib.sh index 08103bdd..8b69a5e9 100755 --- a/tools/build-swiftpm-metallib.sh +++ b/tools/build-swiftpm-metallib.sh @@ -13,16 +13,63 @@ SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) ROOT_DIR=$(realpath "${SCRIPT_DIR}/..") KERNELS_DIR="${ROOT_DIR}/Source/Cmlx/mlx/mlx/backend/metal/kernels" -METAL=$(xcrun -sdk macosx -find metal) -METALLIB=$(xcrun -sdk macosx -find metallib) +normalize_sdk_name() { + local raw="$1" + raw=$(basename "${raw}") + raw=$(printf '%s' "${raw}" | tr '[:upper:]' '[:lower:]') + case "${raw}" in + macosx*) echo "macosx" ;; + iphoneos*) echo "iphoneos" ;; + iphonesimulator*) echo "iphonesimulator" ;; + appletvos*) echo "appletvos" ;; + appletvsimulator*) echo "appletvsimulator" ;; + xros* | visionos*) echo "xros" ;; + xrsimulator* | visionsimulator*) echo "xrsimulator" ;; + *) echo "${raw}" ;; + esac +} + +requested_sdk="${SDK_NAME:-${PLATFORM_NAME:-}}" +if [[ -z "${requested_sdk}" && -n "${SDKROOT:-}" ]]; then + requested_sdk=$(basename "${SDKROOT}") +fi +SDK=$(normalize_sdk_name "${requested_sdk:-macosx}") + +case "${SDK}" in + macosx) + DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-14.0}" + deployment_flag=("-mmacosx-version-min=${DEPLOYMENT_TARGET}") + ;; + iphoneos | iphonesimulator) + DEPLOYMENT_TARGET="${IPHONEOS_DEPLOYMENT_TARGET:-${IOS_DEPLOYMENT_TARGET:-17.0}}" + deployment_flag=("-mios-version-min=${DEPLOYMENT_TARGET}") + ;; + appletvos | appletvsimulator) + DEPLOYMENT_TARGET="${TVOS_DEPLOYMENT_TARGET:-17.0}" + deployment_flag=("-mtvos-version-min=${DEPLOYMENT_TARGET}") + ;; + xros) + DEPLOYMENT_TARGET="${XROS_DEPLOYMENT_TARGET:-${VISIONOS_DEPLOYMENT_TARGET:-1.0}}" + deployment_flag=("-mtargetos=xros${DEPLOYMENT_TARGET}") + ;; + xrsimulator) + DEPLOYMENT_TARGET="${XROS_DEPLOYMENT_TARGET:-${VISIONOS_DEPLOYMENT_TARGET:-1.0}}" + deployment_flag=("-mtargetos=xros${DEPLOYMENT_TARGET}-simulator") + ;; + *) + echo "unsupported SDK '${SDK}'" >&2 + exit 65 + ;; +esac + +METAL=$(xcrun -sdk "${SDK}" -find metal) +METALLIB=$(xcrun -sdk "${SDK}" -find metallib) TMP_DIR=$(mktemp -d) trap 'rm -rf "${TMP_DIR}"' EXIT -DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-14.0}" - metal_version=$( printf '%s\n' '__METAL_VERSION__' | - "${METAL}" "-mmacosx-version-min=${DEPLOYMENT_TARGET}" -E -x metal -P - | + "${METAL}" "${deployment_flag[@]}" -E -x metal -P - | tail -1 | tr -d '[:space:]' ) @@ -50,7 +97,7 @@ metal_flags=( -fno-fast-math -Wno-c++17-extensions -Wno-c++20-extensions - -mmacosx-version-min="${DEPLOYMENT_TARGET}" + "${deployment_flag[@]}" ) if (( metal_version >= 400 )); then