Skip to content

Commit

Permalink
add Vulkan support to Nix flake
Browse files Browse the repository at this point in the history
  • Loading branch information
mschwaig authored and philiptaron committed Feb 3, 2024
1 parent e920ed3 commit 60ecf09
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 4 deletions.
21 changes: 17 additions & 4 deletions .devops/nix/package.nix
Original file line number Diff line number Diff line change
Expand Up @@ -13,18 +13,22 @@
cudaPackages,
darwin,
rocmPackages,
vulkan-headers,
vulkan-loader,
clblast,
useBlas ? builtins.all (x: !x) [
useCuda
useMetalKit
useOpenCL
useRocm
useVulkan
],
useCuda ? config.cudaSupport,
useMetalKit ? stdenv.isAarch64 && stdenv.isDarwin && !useOpenCL,
useMpi ? false, # Increases the runtime closure size by ~700M
useOpenCL ? false,
useRocm ? config.rocmSupport,
useVulkan ? false,
llamaVersion ? "0.0.0", # Arbitrary version, substituted by the flake
}@inputs:

Expand All @@ -48,7 +52,8 @@ let
++ lib.optionals useMetalKit [ "MetalKit" ]
++ lib.optionals useMpi [ "MPI" ]
++ lib.optionals useOpenCL [ "OpenCL" ]
++ lib.optionals useRocm [ "ROCm" ];
++ lib.optionals useRocm [ "ROCm" ]
++ lib.optionals useVulkan [ "Vulkan" ];

pnameSuffix =
strings.optionalString (suffices != [ ])
Expand Down Expand Up @@ -108,6 +113,11 @@ let
hipblas
rocblas
];

vulkanBuildInputs = [
vulkan-headers
vulkan-loader
];
in

effectiveStdenv.mkDerivation (
Expand Down Expand Up @@ -164,7 +174,8 @@ effectiveStdenv.mkDerivation (
++ optionals useCuda cudaBuildInputs
++ optionals useMpi [ mpi ]
++ optionals useOpenCL [ clblast ]
++ optionals useRocm rocmBuildInputs;
++ optionals useRocm rocmBuildInputs
++ optionals useVulkan vulkanBuildInputs;

cmakeFlags =
[
Expand All @@ -178,6 +189,7 @@ effectiveStdenv.mkDerivation (
(cmakeBool "LLAMA_HIPBLAS" useRocm)
(cmakeBool "LLAMA_METAL" useMetalKit)
(cmakeBool "LLAMA_MPI" useMpi)
(cmakeBool "LLAMA_VULKAN" useVulkan)
]
++ optionals useCuda [
(
Expand Down Expand Up @@ -218,6 +230,7 @@ effectiveStdenv.mkDerivation (
useMpi
useOpenCL
useRocm
useVulkan
;

shell = mkShell {
Expand All @@ -242,11 +255,11 @@ effectiveStdenv.mkDerivation (
# Configurations we don't want even the CI to evaluate. Results in the
# "unsupported platform" messages. This is mostly a no-op, because
# cudaPackages would've refused to evaluate anyway.
badPlatforms = optionals (useCuda || useOpenCL) lib.platforms.darwin;
badPlatforms = optionals (useCuda || useOpenCL || useVulkan) lib.platforms.darwin;

# Configurations that are known to result in build failures. Can be
# overridden by importing Nixpkgs with `allowBroken = true`.
broken = (useMetalKit && !effectiveStdenv.isDarwin);
broken = (useMetalKit && !effectiveStdenv.isDarwin) || (useVulkan && effectiveStdenv.isDarwin);

description = "Inference of LLaMA model in pure C/C++${descriptionSuffix}";
homepage = "https://github.com/ggerganov/llama.cpp/";
Expand Down
1 change: 1 addition & 0 deletions flake.nix
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,7 @@

mpi-cpu = config.packages.default.override { useMpi = true; };
mpi-cuda = config.packages.default.override { useMpi = true; };
vulkan = config.packages.default.override { useVulkan = true; };
}
// lib.optionalAttrs (system == "x86_64-linux") {
rocm = config.legacyPackages.llamaPackagesRocm.llama-cpp;
Expand Down

0 comments on commit 60ecf09

Please sign in to comment.