diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml index 7c5768ac..7f6c5f04 100644 --- a/.github/workflows/pipeline.yml +++ b/.github/workflows/pipeline.yml @@ -123,6 +123,63 @@ jobs: name: exo authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" + - name: Build Metal packages (macOS only) + if: runner.os == 'macOS' + run: | + # Try to build metal-toolchain first (may succeed via cachix cache hit) + if nix build .#metal-toolchain 2>/dev/null; then + echo "metal-toolchain built successfully (likely cache hit)" + else + echo "metal-toolchain build failed, extracting from Xcode..." + + NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw=" + NAR_NAME="metal-toolchain-17C48.nar" + + # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS + WORK_DIR="${RUNNER_TEMP}/metal-work" + mkdir -p "$WORK_DIR" + + # Download the Metal toolchain component + xcodebuild -downloadComponent MetalToolchain + + # Find and mount the DMG + DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1) + if [ -z "$DMG_PATH" ]; then + echo "Error: Could not find Metal toolchain DMG" + exit 1 + fi + + echo "Found DMG at: $DMG_PATH" + hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg" + + # Copy the toolchain + cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export" + hdiutil detach "${WORK_DIR}/metal-dmg" + + # Create NAR and add to store + nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}" + STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}") + echo "Added NAR to store: $STORE_PATH" + + # Verify the hash matches + ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}") + if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then + echo "Warning: NAR hash mismatch!" + echo "Expected: $NAR_HASH" + echo "Actual: $ACTUAL_HASH" + echo "The metal-toolchain.nix may need updating" + fi + + # Clean up + rm -rf "$WORK_DIR" + + # Retry the build now that NAR is in store + nix build .#metal-toolchain + fi + + # Build mlx (depends on metal-toolchain) + nix build .#mlx + - name: Build all Nix outputs run: | nix flake show --json | jq -r ' diff --git a/flake.nix b/flake.nix index c4c018c6..0d6d87b9 100644 --- a/flake.nix +++ b/flake.nix @@ -58,6 +58,11 @@ pkgsSwift = import inputs.nixpkgs-swift { inherit system; }; in { + # Allow unfree for metal-toolchain (needed for Darwin Metal packages) + _module.args.pkgs = import inputs.nixpkgs { + inherit system; + config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain"; + }; treefmt = { projectRootFile = "flake.nix"; programs = { @@ -88,6 +93,21 @@ touch $out ''; + packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin ( + let + uvLock = builtins.fromTOML (builtins.readFile ./uv.lock); + mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx") uvLock.package); + uvLockMlxVersion = mlxPackage.version; + in + { + metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { }; + mlx = pkgs.callPackage ./nix/mlx.nix { + metal-toolchain = self'.packages.metal-toolchain; + inherit uvLockMlxVersion; + }; + } + ); + devShells.default = with pkgs; pkgs.mkShell { inputsFrom = [ self'.checks.cargo-build ]; diff --git a/nix/darwin-build-fixes.patch b/nix/darwin-build-fixes.patch new file mode 100644 index 00000000..edce60d0 --- /dev/null +++ b/nix/darwin-build-fixes.patch @@ -0,0 +1,79 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index 0ed30932..d8528132 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -177,11 +177,7 @@ if(MLX_BUILD_METAL) + add_compile_definitions(MLX_METAL_DEBUG) + endif() + +- # Throw an error if xcrun not found +- execute_process( +- COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version" +- OUTPUT_VARIABLE MACOS_SDK_VERSION +- OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) ++ set(MACOS_SDK_VERSION @sdkVersion@) + + if(${MACOS_SDK_VERSION} LESS 14.0) + message( +@@ -199,11 +195,8 @@ if(MLX_BUILD_METAL) + endif() + set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") + endif() +- execute_process( +- COMMAND +- zsh "-c" +- "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'" +- OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY) ++ set( ++ MLX_METAL_VERSION @metalVersion@) + FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL}) + FetchContent_MakeAvailable(metal_cpp) + target_include_directories( +diff --git a/cmake/extension.cmake b/cmake/extension.cmake +index 13db804a..5b385132 100644 +--- a/cmake/extension.cmake ++++ b/cmake/extension.cmake +@@ -36,7 +36,7 @@ macro(mlx_build_metallib) + add_custom_command( + OUTPUT ${MTLLIB_BUILD_TARGET} + COMMAND +- xcrun -sdk macosx metal ++ metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache + "$" + ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET} + DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES} +diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt +index 262b0495..5c7446ad 100644 +--- a/mlx/backend/metal/kernels/CMakeLists.txt ++++ b/mlx/backend/metal/kernels/CMakeLists.txt +@@ -29,7 +29,7 @@ function(build_kernel_base TARGET SRCFILE DEPS) + "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") + endif() + add_custom_command( +- COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE} ++ COMMAND metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache ${METAL_FLAGS} -c ${SRCFILE} + -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air + DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS} + OUTPUT ${TARGET}.air +@@ -170,7 +170,7 @@ endif() + + add_custom_command( + OUTPUT ${MLX_METAL_PATH}/mlx.metallib +- COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o ++ COMMAND metallib ${KERNEL_AIR} -o + ${MLX_METAL_PATH}/mlx.metallib + DEPENDS ${KERNEL_AIR} + COMMENT "Building mlx.metallib" +diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh +index bb55ed3a..94ea7dd7 100644 +--- a/mlx/backend/metal/make_compiled_preamble.sh ++++ b/mlx/backend/metal/make_compiled_preamble.sh +@@ -31,7 +31,7 @@ OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp + mkdir -p "$OUTPUT_DIR" + + # Use the metal compiler to get a list of headers (with depth) +-CCC="xcrun -sdk macosx metal -x metal" ++CCC="metal -x metal -fmodules-cache-path=${OUTPUT_DIR}/metal-cache" + HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null ) + + # Remove any included system frameworks (for MetalPerformancePrimitive headers) diff --git a/nix/metal-toolchain.nix b/nix/metal-toolchain.nix new file mode 100644 index 00000000..8620e43f --- /dev/null +++ b/nix/metal-toolchain.nix @@ -0,0 +1,56 @@ +{ lib, stdenvNoCC, requireFile, nix }: + +let + narFile = requireFile { + name = "metal-toolchain-17C48.nar"; + message = '' + The Metal Toolchain NAR must be available. + + If you have cachix configured for exo.cachix.org, this should be automatic. + + Otherwise: + 1. Install Xcode 26+ from the App Store + 2. Run: xcodebuild -downloadComponent MetalToolchain + 3. Export the toolchain: + hdiutil attach "$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' | head -1)" -mountpoint /tmp/metal-dmg + cp -R /tmp/metal-dmg/Metal.xctoolchain /tmp/metal-export + hdiutil detach /tmp/metal-dmg + 4. Create NAR and add to store: + nix nar pack /tmp/metal-export > /tmp/metal-toolchain-17C48.nar + nix store add --mode flat /tmp/metal-toolchain-17C48.nar + ''; + hash = "sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw="; + }; +in +stdenvNoCC.mkDerivation { + pname = "metal-toolchain"; + version = "17C48"; + + dontUnpack = true; + dontBuild = true; + dontFixup = true; + + nativeBuildInputs = [ nix ]; + + installPhase = '' + runHook preInstall + + nix-store --restore $out < ${narFile} + + # Create bin directory with symlinks for PATH + mkdir -p $out/bin + ln -s $out/usr/bin/metal $out/bin/metal + ln -s $out/usr/bin/metallib $out/bin/metallib + + runHook postInstall + ''; + + # Metal language version for CMake (from: echo __METAL_VERSION__ | metal -E -x metal -P -) + passthru.metalVersion = "400"; + + meta = { + description = "Apple Metal compiler toolchain"; + platforms = [ "aarch64-darwin" ]; + license = lib.licenses.unfree; + }; +} diff --git a/nix/mlx.nix b/nix/mlx.nix new file mode 100644 index 00000000..baa31da8 --- /dev/null +++ b/nix/mlx.nix @@ -0,0 +1,158 @@ +{ stdenv +, lib +, fetchFromGitHub +, replaceVars +, fetchzip +, cmake +, nlohmann_json +, apple-sdk_26 +, metal-toolchain +, runCommand +, fmt +, python313Packages +, uvLockMlxVersion +}: + +assert stdenv.isDarwin; + +let + python = python313Packages.python; + + # Static dependencies included directly during compilation + gguf-tools = fetchFromGitHub { + owner = "antirez"; + repo = "gguf-tools"; + rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848"; + hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA="; + }; + + metal_cpp = fetchzip { + url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip"; + hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8="; + }; + + nanobind = fetchFromGitHub { + owner = "wjakob"; + repo = "nanobind"; + rev = "v2.10.2"; + hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU="; + fetchSubmodules = true; + }; + + mlx = stdenv.mkDerivation rec { + pname = "mlx"; + version = let v = "0.30.4"; in + assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix."; + v; + pyproject = true; + + src = fetchFromGitHub { + owner = "ml-explore"; + repo = "mlx"; + tag = "v${version}"; + hash = "sha256-OJk6jPlbaSlsUdk3ADz3tWcRzTWXRof3/q8Soe1AO6w="; + }; + + patches = [ + (replaceVars ./darwin-build-fixes.patch { + sdkVersion = apple-sdk_26.version; + metalVersion = metal-toolchain.metalVersion; + }) + ]; + + postPatch = '' + substituteInPlace mlx/backend/cpu/jit_compiler.cpp \ + --replace-fail "g++" "$CXX" + ''; + + dontUseCmakeConfigure = true; + + enableParallelBuilding = true; + + # Allows multiple cores to be used in Python builds. + postUnpack = '' + export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}" + ''; + + # Updates the wrong fetcher rev attribute + passthru.skipBulkUpdate = true; + + env = { + DEV_RELEASE = 1; + CMAKE_ARGS = toString [ + (lib.cmakeBool "USE_SYSTEM_FMT" true) + (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}") + (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}") + (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}") + (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true) + (lib.cmakeBool "MLX_BUILD_METAL" true) + (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}") + (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}") + (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}") + ]; + SDKROOT = apple-sdk_26.passthru.sdkroot; + MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version; + }; + + build-system = [ + python313Packages.setuptools + ]; + + nativeBuildInputs = [ + cmake + metal-toolchain + python313Packages.pypaBuildHook + python313Packages.pypaInstallHook + python313Packages.setuptools + python313Packages.typing-extensions + python313Packages.wheel + python313Packages.cmake + python313Packages.ninja + ]; + + buildInputs = [ + fmt + gguf-tools + python313Packages.nanobind + python313Packages.pybind11 + apple-sdk_26 + ]; + + # Tests require Metal GPU access which isn't available in the Nix sandbox. + # To run tests, build with: nix build --option sandbox false .#mlx.passthru.tests.mlxTest + doCheck = false; + + pythonImportsCheck = [ "mlx" ]; + + passthru.tests = { + # Runs example scripts to verify MLX works. Requires --option sandbox false + # since Metal GPU access is needed. + mlxTest = + runCommand "run-mlx-examples" + { + buildInputs = [ mlx ]; + nativeBuildInputs = [ python ]; + } + '' + cp ${src}/examples/python/logistic_regression.py . + ${python.interpreter} logistic_regression.py + rm logistic_regression.py + + cp ${src}/examples/python/linear_regression.py . + ${python.interpreter} linear_regression.py + rm linear_regression.py + + touch $out + ''; + }; + + meta = { + homepage = "https://github.com/ml-explore/mlx"; + description = "Array framework for Apple silicon"; + changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}"; + license = lib.licenses.mit; + platforms = [ "aarch64-darwin" ]; + }; + }; +in +mlx