From 6e1102ef654f5e11ccde7a654ce4f92207310f68 Mon Sep 17 00:00:00 2001
From: Joseph Stahl <1269177+josephst@users.noreply.github.com>
Date: Mon, 25 Mar 2024 20:51:46 -0400
Subject: [PATCH] nix: make `xcrun` visible in Nix sandbox for precompiling
 Metal shaders (#6118)

* Symlink to /usr/bin/xcrun so that `xcrun` binary
is usable during build (used for compiling Metal shaders)

Fixes https://github.com/ggerganov/llama.cpp/issues/6117

* cmake - copy default.metallib to install directory

When metal files are compiled to default.metallib, Cmake needs to add this to the install directory so that it's visible to llama-cpp

Also, update package.nix to use absolute path for default.metallib (it's not finding the bundle)

* add `precompileMetalShaders` flag (defaults to false) to disable precompilation of metal shader

Precompilation requires Xcode to be installed and requires disable sandbox on nix-darwin
---
 .devops/nix/package.nix | 26 ++++++++++++++++++++++++--
 CMakeLists.txt          |  6 ++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/.devops/nix/package.nix b/.devops/nix/package.nix
index 0330b79d9143a2..b651f9e61f7718 100644
--- a/.devops/nix/package.nix
+++ b/.devops/nix/package.nix
@@ -4,6 +4,7 @@
   config,
   stdenv,
   mkShell,
+  runCommand,
   cmake,
   ninja,
   pkg-config,
@@ -35,7 +36,8 @@
   # It's necessary to consistently use backendStdenv when building with CUDA support,
   # otherwise we get libstdc++ errors downstream.
   effectiveStdenv ? if useCuda then cudaPackages.backendStdenv else stdenv,
-  enableStatic ? effectiveStdenv.hostPlatform.isStatic
+  enableStatic ? effectiveStdenv.hostPlatform.isStatic,
+  precompileMetalShaders ? false
 }@inputs:
 
 let
@@ -87,6 +89,11 @@ let
     ]
   );
 
+  xcrunHost = runCommand "xcrunHost" {} ''
+    mkdir -p $out/bin
+    ln -s /usr/bin/xcrun $out/bin
+  '';
+
   # apple_sdk is supposed to choose sane defaults, no need to handle isAarch64
   # separately
   darwinBuildInputs =
@@ -150,6 +157,8 @@ effectiveStdenv.mkDerivation (
     postPatch = ''
       substituteInPlace ./ggml-metal.m \
         --replace '[bundle pathForResource:@"ggml-metal" ofType:@"metal"];' "@\"$out/bin/ggml-metal.metal\";"
+      substituteInPlace ./ggml-metal.m \
+        --replace '[bundle pathForResource:@"default" ofType:@"metallib"];' "@\"$out/bin/default.metallib\";"
 
       # TODO: Package up each Python script or service appropriately.
       # If we were to migrate to buildPythonPackage and prepare the `pyproject.toml`,
@@ -157,6 +166,14 @@ effectiveStdenv.mkDerivation (
       substituteInPlace ./*.py --replace "/usr/bin/env python" "${llama-python}/bin/python"
     '';
 
+    # With PR#6015 https://github.com/ggerganov/llama.cpp/pull/6015,
+    # `default.metallib` may be compiled with Metal compiler from XCode
+    # and we need to escape sandbox on MacOS to access Metal compiler.
+    # `xcrun` is used find the path of the Metal compiler, which is varible
+    # and not on $PATH
+    # see https://github.com/ggerganov/llama.cpp/pull/6118 for discussion
+    __noChroot = effectiveStdenv.isDarwin && useMetalKit && precompileMetalShaders;
+
     nativeBuildInputs =
       [
         cmake
@@ -173,6 +190,8 @@ effectiveStdenv.mkDerivation (
       ]
       ++ optionals (effectiveStdenv.hostPlatform.isGnu && enableStatic) [
         glibc.static
+      ] ++ optionals (effectiveStdenv.isDarwin && useMetalKit && precompileMetalShaders) [
+        xcrunHost
       ];
 
     buildInputs =
@@ -217,7 +236,10 @@ effectiveStdenv.mkDerivation (
         # Should likely use `rocmPackages.clr.gpuTargets`.
         "-DAMDGPU_TARGETS=gfx803;gfx900;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102"
       ]
-      ++ optionals useMetalKit [ (lib.cmakeFeature "CMAKE_C_FLAGS" "-D__ARM_FEATURE_DOTPROD=1") ];
+      ++ optionals useMetalKit [
+        (lib.cmakeFeature "CMAKE_C_FLAGS" "-D__ARM_FEATURE_DOTPROD=1")
+        (cmakeBool "LLAMA_METAL_EMBED_LIBRARY" (!precompileMetalShaders))
+      ];
 
     # TODO(SomeoneSerge): It's better to add proper install targets at the CMake level,
     # if they haven't been added yet.
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3f23ba4d369336..ed1e776455a118 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1265,6 +1265,12 @@ if (LLAMA_METAL)
             GROUP_READ
             WORLD_READ
         DESTINATION ${CMAKE_INSTALL_BINDIR})
+    if (NOT LLAMA_METAL_EMBED_LIBRARY)
+        install(
+            FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+            DESTINATION ${CMAKE_INSTALL_BINDIR}
+        )
+    endif()
 endif()
 
 #