nix mlx compilation for better portability

.github/workflows/pipeline.yml

@@ -113,6 +113,28 @@ jobs:
         with:
           lfs: false
 
+      - name: Select Xcode
+        if: startsWith(matrix.runner, 'macos-')
+        run: |
+          XCODE_BASEDIR="$(printf '%s\n' /Applications/Xcode_*.app | sort -V | tail -n 1)"
+          [[ -z "$XCODE_BASEDIR" ]] && exit 1
+          sudo mv "$XCODE_BASEDIR" /Applications/Xcode.app
+
+          ls -ld "/Applications/Xcode.app"
+          sudo /usr/bin/xcode-select -s "/Applications/Xcode.app"
+          /usr/bin/xcode-select -p || true
+          /usr/bin/xcrun --toolchain default --find xcodebuild || true
+
+      - name: Install Metal toolchain component
+        if: startsWith(matrix.runner, 'macos-')
+        run: |
+          set -e
+          if ! xcrun --find metal >/dev/null 2>&1; then
+            sudo xcodebuild -downloadComponent MetalToolchain
+          fi
+          xcrun --find metal
+          xcrun --find metallib
+
       - uses: cachix/install-nix-action@v31
         with:
           nix_path: nixpkgs=channel:nixos-unstable

flake.nix

@@ -70,7 +70,12 @@
             touch $out
           '';
 
-          devShells.default = with pkgs; pkgs.mkShell {
+          packages =
+            if pkgs.stdenv.isDarwin then {
+              metal = pkgs.callPackage ./nix/metalWrapper.nix { metalVersion = "230"; };
+            } else { };
+
+          devShells.default = with pkgs; mkShellNoCC {
             packages =
               [
                 # FORMATTING
@@ -102,7 +107,7 @@
                 just
                 jq
               ]
-              ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [
+              ++ (lib.optionals stdenv.isLinux [
                 # IFCONFIG
                 unixtools.ifconfig
 
@@ -110,11 +115,12 @@
                 pkg-config
                 openssl
               ])
-              ++ (pkgs.lib.optionals pkgs.stdenv.isDarwin [
+              ++ (lib.optionals stdenv.isDarwin [
                 # MACMON
                 macmon
               ]);
 
+
             shellHook = ''
               # PYTHON
               export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${pkgs.python313}/lib"

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
+       "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+       ${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 ${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"
+ 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)

nix/metalWrapper.nix

@@ -0,0 +1,22 @@
+{ stdenv
+, metalVersion
+, xcodeBaseDir ? "/Applications/Xcode.app"
+}:
+assert stdenv.isDarwin;
+stdenv.mkDerivation {
+  pname = "metal-wrapper-impure";
+  version = metalVersion;
+
+  __noChroot = true;
+  buildCommand = ''
+    DEVELOPER_DIR=${xcodeBaseDir}/Contents/Developer
+    [[ -x "$DEVELOPER_DIR/usr/bin/xcodebuild" ]] || (echo "Missing xcodebuild at $DEVELOPER_DIR/usr/bin/xcodebuild" && exit 1)
+    SDKROOT=${xcodeBaseDir}/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk
+    [[ -d "$SDKROOT" ]] || (echo "Missing SDKROOT at $SDKROOT" && exit 1)
+    export DEVELOPER_DIR SDKROOT
+    mkdir -p $out/bin && cd $out/bin
+    ln -s $(/usr/bin/xcrun --sdk macosx -f metal)
+    ln -s $(/usr/bin/xcrun --sdk macosx -f metallib)
+    [[ -f $out/bin/metal ]] && [[ -f $out/bin/metallib ]] || exit 1
+  '';
+}

nix/mlx.nix

@@ -0,0 +1,154 @@
+{ stdenv
+, lib
+, buildPythonPackage
+, fetchFromGitHub
+, replaceVars
+, fetchzip
+, setuptools
+, cmake
+, nanobind
+, pybind11
+, nlohmann_json
+, apple-sdk_26
+, metal
+, numpy
+, pytestCheckHook
+, python
+, runCommand
+, fmt
+}:
+assert stdenv.isDarwin;
+let
+  # 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=";
+  };
+
+  mlx = buildPythonPackage rec {
+    pname = "mlx";
+    version = "0.30.1";
+    pyproject = true;
+
+    src = fetchFromGitHub {
+      owner = "ml-explore";
+      repo = "mlx";
+      tag = "v${version}";
+      hash = "sha256-Vt0RH+70VBwUjXSfPTsNdRS3g0ookJHhzf2kvgEtgH8=";
+    };
+
+    patches = [
+      (replaceVars ./darwin-build-fixes.patch {
+        sdkVersion = apple-sdk_26.version;
+        metalVersion = metal.version;
+      })
+    ];
+
+    postPatch = ''
+      substituteInPlace pyproject.toml \
+        --replace-fail "nanobind==2.10.2" "nanobind"
+
+      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;
+      # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
+      # this is what the xcode wrapper is for - it patches in the system metal cli
+      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.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+        (lib.cmakeBool "MLX_BUILD_METAL" true)
+        (lib.cmakeOptionType "filepath" "METAL_LIB"
+          "${metal}/Metal.framework")
+        (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 = [
+      setuptools
+    ];
+
+    nativeBuildInputs = [
+      cmake
+      metal
+    ];
+
+    buildInputs = [
+      fmt
+      gguf-tools
+      nanobind
+      pybind11
+      apple-sdk_26
+    ];
+
+    pythonImportsCheck = [ "mlx" ];
+
+    # Run the mlx Python test suite.
+    nativeCheckInputs = [
+      numpy
+      pytestCheckHook
+    ];
+
+    enabledTestPaths = [
+      "python/tests/"
+    ];
+
+    # Additional testing by executing the example Python scripts supplied with mlx
+    # using the version of the library we've built.
+    passthru.tests = {
+      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 = [ "x86_64-linux" "aarch64-linux" "aarch64-darwin" ];
+    };
+  };
+in
+mlx