pkgs/top-level/cuda-packages.nix

   1 { lib
   2 , pkgs
   3 , cudaVersion
   4 }:
   5
   6 with lib;
   7
   8 let
   9
  10   scope = makeScope pkgs.newScope (final: {
  11     # Here we put package set configuration and utility functions.
  12     inherit cudaVersion;
  13     cudaMajorVersion = versions.major final.cudaVersion;
  14     cudaMajorMinorVersion = lib.versions.majorMinor final.cudaVersion;
  15     inherit lib pkgs;
  16
  17     addBuildInputs = drv: buildInputs: drv.overrideAttrs (oldAttrs: {
  18       buildInputs = (oldAttrs.buildInputs or []) ++ buildInputs;
  19     });
  20   });
  21
  22   cutensorExtension = final: prev: let
  23     ### CuTensor
  24
  25     buildCuTensorPackage = final.callPackage ../development/libraries/science/math/cutensor/generic.nix;
  26
  27     cuTensorVersions = {
  28       "1.2.2.5" = {
  29         hash = "sha256-lU7iK4DWuC/U3s1Ct/rq2Gr3w4F2U7RYYgpmF05bibY=";
  30       };
  31       "1.5.0.3" = {
  32         hash = "sha256-T96+lPC6OTOkIs/z3QWg73oYVSyidN0SVkBWmT9VRx0=";
  33       };
  34     };
  35
  36     inherit (final) cudaMajorMinorVersion cudaMajorVersion;
  37
  38     cutensor = buildCuTensorPackage rec {
  39       version = if cudaMajorMinorVersion == "10.1" then "1.2.2.5" else "1.5.0.3";
  40       inherit (cuTensorVersions.${version}) hash;
  41       # This can go into generic.nix
  42       libPath = "lib/${if cudaMajorVersion == "10" then cudaMajorMinorVersion else cudaMajorVersion}";
  43     };
  44   in { inherit cutensor; };
  45
  46   extraPackagesExtension = final: prev: {
  47
  48     nccl = final.callPackage ../development/libraries/science/math/nccl { };
  49
  50     nccl-tests = final.callPackage ../development/libraries/science/math/nccl/tests.nix { };
  51
  52     autoAddOpenGLRunpathHook = final.callPackage ( { makeSetupHook, addOpenGLRunpath }:
  53       makeSetupHook {
  54         name = "auto-add-opengl-runpath-hook";
  55         propagatedBuildInputs = [
  56           addOpenGLRunpath
  57         ];
  58       } ../development/compilers/cudatoolkit/auto-add-opengl-runpath-hook.sh
  59     ) {};
  60
  61   };
  62
  63   composedExtension = composeManyExtensions ([
  64     extraPackagesExtension
  65     (import ../development/compilers/cudatoolkit/extension.nix)
  66     (import ../development/compilers/cudatoolkit/redist/extension.nix)
  67     (import ../development/compilers/cudatoolkit/redist/overrides.nix)
  68     (import ../development/libraries/science/math/cudnn/extension.nix)
  69     (import ../development/libraries/science/math/tensorrt/extension.nix)
  70     (import ../test/cuda/cuda-samples/extension.nix)
  71     (import ../test/cuda/cuda-library-samples/extension.nix)
  72     cutensorExtension
  73   ]);
  74
  75 in (scope.overrideScope composedExtension)