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