Unverified Commit ecae25c4 authored by Pavol Rusnak's avatar Pavol Rusnak Committed by GitHub
Browse files

ollama-vulkan: init package (#456459)

parents 9d68a63a 115a7b7e

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

+30 −5
Original line number Diff line number Diff line
@@ -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

+1 −0
Original line number Diff line number Diff line
@@ -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 { };