ollama-vulkan: init package (#456459)

pkgs/by-name/ol/ollama/package.nix

@@ -18,6 +18,11 @@
   cudaArches ? cudaPackages.flags.realArches or [ ],
   autoAddDriverRunpath,
   apple-sdk_15,
+  vulkan-tools,
+  vulkan-headers,
+  vulkan-loader,
+  shaderc,
+  ccache,
 
   versionCheckHook,
   writableTmpDirAsHomeHook,
@@ -29,7 +34,7 @@
   ollama-cuda,
 
   config,
-  # one of `[ null false "rocm" "cuda" ]`
+  # one of `[ null false "rocm" "cuda" "vulkan" ]`
   acceleration ? null,
 }:
 
@@ -38,6 +43,7 @@ assert builtins.elem acceleration [
   false
   "rocm"
   "cuda"
+  "vulkan"
 ];
 
 let
@@ -50,9 +56,11 @@ let
 
   rocmRequested = shouldEnable "rocm" config.rocmSupport;
   cudaRequested = shouldEnable "cuda" config.cudaSupport;
+  vulkanRequested = shouldEnable "vulkan" false;
 
   enableRocm = rocmRequested && stdenv.hostPlatform.isLinux;
   enableCuda = cudaRequested && stdenv.hostPlatform.isLinux;
+  enableVulkan = vulkanRequested && stdenv.hostPlatform.isLinux;
 
   rocmLibs = [
     rocmPackages.clr
@@ -75,6 +83,11 @@ let
     cudaPackages.cuda_cccl
   ];
 
+  vulkanLibs = [
+    vulkan-headers
+    vulkan-loader
+  ];
+
   # Extract the major version of CUDA. e.g. 11 12
   cudaMajorVersion = lib.versions.major cudaPackages.cuda_cudart.version;
 
@@ -103,6 +116,9 @@ let
   ]
   ++ lib.optionals enableCuda [
     "--suffix LD_LIBRARY_PATH : '${lib.makeLibraryPath (map lib.getLib cudaLibs)}'"
+  ]
+  ++ lib.optionals enableVulkan [
+    "--suffix LD_LIBRARY_PATH : '${lib.makeLibraryPath (map lib.getLib vulkanLibs)}'"
   ];
   wrapperArgs = builtins.concatStringsSep " " wrapperOptions;
 
@@ -111,6 +127,8 @@ let
       buildGoModule.override { stdenv = cudaPackages.backendStdenv; }
     else if enableRocm then
       buildGoModule.override { stdenv = rocmPackages.stdenv; }
+    else if enableVulkan then
+      buildGoModule.override { stdenv = vulkan-tools.stdenv; }
     else
       buildGoModule;
   inherit (lib) licenses platforms maintainers;
@@ -137,7 +155,8 @@ goBuild (finalAttrs: {
       CFLAGS = "-Wno-c++17-extensions -I${rocmPath}/include";
       CXXFLAGS = "-Wno-c++17-extensions -I${rocmPath}/include";
     }
-    // lib.optionalAttrs enableCuda { CUDA_PATH = cudaPath; };
+    // lib.optionalAttrs enableCuda { CUDA_PATH = cudaPath; }
+    // lib.optionalAttrs enableVulkan { VULKAN_SDK = shaderc.bin; };
 
   nativeBuildInputs = [
     cmake
@@ -151,12 +170,16 @@ goBuild (finalAttrs: {
   ++ lib.optionals (enableRocm || enableCuda) [
     makeBinaryWrapper
     autoAddDriverRunpath
+  ]
+  ++ lib.optionals enableVulkan [
+    ccache
   ];
 
   buildInputs =
     lib.optionals enableRocm (rocmLibs ++ [ libdrm ])
     ++ lib.optionals enableCuda cudaLibs
-    ++ lib.optionals stdenv.hostPlatform.isDarwin [ apple-sdk_15 ];
+    ++ lib.optionals stdenv.hostPlatform.isDarwin [ apple-sdk_15 ]
+    ++ lib.optionals enableVulkan vulkanLibs;
 
   # replace inaccurate version number with actual release version
   postPatch = ''
@@ -261,11 +284,13 @@ goBuild (finalAttrs: {
     description =
       "Get up and running with large language models locally"
       + lib.optionalString rocmRequested ", using ROCm for AMD GPU acceleration"
-      + lib.optionalString cudaRequested ", using CUDA for NVIDIA GPU acceleration";
+      + lib.optionalString cudaRequested ", using CUDA for NVIDIA GPU acceleration"
+      + lib.optionalString vulkanRequested ", using Vulkan for generic GPU acceleration";
     homepage = "https://github.com/ollama/ollama";
     changelog = "https://github.com/ollama/ollama/releases/tag/v${finalAttrs.version}";
     license = licenses.mit;
-    platforms = if (rocmRequested || cudaRequested) then platforms.linux else platforms.unix;
+    platforms =
+      if (rocmRequested || cudaRequested || vulkanRequested) then platforms.linux else platforms.unix;
     mainProgram = "ollama";
     maintainers = with maintainers; [
       abysssol

pkgs/top-level/all-packages.nix

@@ -392,6 +392,7 @@ with pkgs;
 
   ollama-rocm = callPackage ../by-name/ol/ollama/package.nix { acceleration = "rocm"; };
   ollama-cuda = callPackage ../by-name/ol/ollama/package.nix { acceleration = "cuda"; };
+  ollama-vulkan = callPackage ../by-name/ol/ollama/package.nix { acceleration = "vulkan"; };
 
   device-tree_rpi = callPackage ../os-specific/linux/device-tree/raspberrypi.nix { };