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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
|
{ cmake
, cudaPackages
, fetchFromGitHub
, lib
, ninja
, pkgs
, python3Packages ? { }
, pythonSupport ? false
, stdenv
, symlinkJoin
, which
}:
let
inherit (lib) lists strings;
inherit (cudaPackages) backendStdenv cudaFlags;
cuda-common-redist = with cudaPackages; [
libcublas # cublas_v2.h
libcusolver # cusolverDn.h
libcusparse # cusparse.h
];
cuda-native-redist = symlinkJoin {
name = "cuda-redist";
paths = with cudaPackages; [
cuda_cudart # cuda_runtime.h
cuda_nvcc
] ++ cuda-common-redist;
};
cuda-redist = symlinkJoin {
name = "cuda-redist";
paths = cuda-common-redist;
};
in
stdenv.mkDerivation (finalAttrs: {
name = "tiny-cuda-nn";
version = "1.6";
format = strings.optionalString pythonSupport "setuptools";
src = fetchFromGitHub {
owner = "NVlabs";
repo = finalAttrs.name;
rev = "v${finalAttrs.version}";
fetchSubmodules = true;
hash = "sha256-qW6Fk2GB71fvZSsfu+mykabSxEKvaikZ/pQQZUycOy0=";
};
nativeBuildInputs = [
cmake
cuda-native-redist
ninja
which
] ++ lists.optionals pythonSupport (with python3Packages; [
pip
setuptools
wheel
]);
buildInputs = [
cuda-redist
] ++ lib.optionals pythonSupport (
with python3Packages; [
pybind11
python
]
);
propagatedBuildInputs = lib.optionals pythonSupport (
with python3Packages; [
torch
]
);
# NOTE: We cannot use pythonImportsCheck for this module because it uses torch to immediately
# initailize CUDA and GPU access is not allowed in the nix build environment.
# NOTE: There are no tests for the C++ library or the python bindings, so we just skip the check
# phase -- we're not missing anything.
doCheck = false;
preConfigure = ''
export TCNN_CUDA_ARCHITECTURES=${
strings.concatStringsSep ";" (lists.map cudaFlags.dropDot cudaFlags.cudaCapabilities)
}
export CUDA_HOME=${cuda-native-redist}
export LIBRARY_PATH=${cuda-native-redist}/lib/stubs:$LIBRARY_PATH
export CC=${backendStdenv.cc}/bin/cc
export CXX=${backendStdenv.cc}/bin/c++
'';
# When building the python bindings, we cannot re-use the artifacts from the C++ build so we
# skip the CMake confurePhase and the buildPhase.
dontUseCmakeConfigure = pythonSupport;
# The configurePhase usually puts you in the build directory, so for the python bindings we
# need to change directories to the source directory.
configurePhase = strings.optionalString pythonSupport ''
runHook preConfigure
mkdir -p $NIX_BUILD_TOP/build
cd $NIX_BUILD_TOP/build
runHook postConfigure
'';
buildPhase = strings.optionalString pythonSupport ''
runHook preBuild
python -m pip wheel \
--no-build-isolation \
--no-clean \
--no-deps \
--no-index \
--verbose \
--wheel-dir $NIX_BUILD_TOP/build \
$NIX_BUILD_TOP/source/bindings/torch
runHook postBuild
'';
installPhase = ''
runHook preInstall
mkdir -p $out/lib
''
# Installing the C++ library just requires copying the static library to the output directory
+ strings.optionalString (!pythonSupport) ''
cp libtiny-cuda-nn.a $out/lib/
''
# Installing the python bindings requires building the wheel and installing it
+ strings.optionalString pythonSupport ''
python -m pip install \
--no-build-isolation \
--no-cache-dir \
--no-deps \
--no-index \
--no-warn-script-location \
--prefix="$out" \
--verbose \
./*.whl
'' + ''
runHook postInstall
'';
passthru = {
inherit cudaPackages;
};
meta = with lib; {
description = "Lightning fast C++/CUDA neural network framework";
homepage = "https://github.com/NVlabs/tiny-cuda-nn";
license = licenses.bsd3;
maintainers = with maintainers; [ connorbaker ];
platforms = platforms.linux;
};
})
|