tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 config,
3 cmake,
4 cudaPackages,
5 fetchFromGitHub,
6 lib,
7 ninja,
8 python3Packages ? { },
9 pythonSupport ? false,
10 stdenv,
11 symlinkJoin,
12 which,
13}:
14let
15 inherit (lib) lists strings;
16 inherit (cudaPackages) backendStdenv cudaAtLeast flags;
17
18 cuda-common-redist = with cudaPackages; [
19 (lib.getDev cuda_cudart) # cuda_runtime.h
20 (lib.getLib cuda_cudart)
21 (lib.getDev cuda_cccl) # <nv/target>
22 (lib.getDev libcublas) # cublas_v2.h
23 (lib.getLib libcublas)
24 (lib.getDev libcusolver) # cusolverDn.h
25 (lib.getLib libcusolver)
26 (lib.getDev libcusparse) # cusparse.h
27 (lib.getLib libcusparse)
28 ];
29
30 cuda-native-redist = symlinkJoin {
31 name = "cuda-redist";
32 paths = with cudaPackages; [ cuda_nvcc ] ++ cuda-common-redist;
33 };
34
35 cuda-redist = symlinkJoin {
36 name = "cuda-redist";
37 paths = cuda-common-redist;
38 };
39
40 unsupportedCudaCapabilities = [
41 "9.0a"
42 ];
43
44 cudaCapabilities = lists.subtractLists unsupportedCudaCapabilities flags.cudaCapabilities;
45
46 cudaArchitecturesString = strings.concatMapStringsSep ";" flags.dropDots cudaCapabilities;
47in
48stdenv.mkDerivation (finalAttrs: {
49 pname = "tiny-cuda-nn";
50 version = "1.6";
51 strictDeps = true;
52
53 format = strings.optionalString pythonSupport "setuptools";
54
55 src = fetchFromGitHub {
56 owner = "NVlabs";
57 repo = "tiny-cuda-nn";
58 rev = "v${finalAttrs.version}";
59 fetchSubmodules = true;
60 hash = "sha256-qW6Fk2GB71fvZSsfu+mykabSxEKvaikZ/pQQZUycOy0=";
61 };
62
63 # Remove this once a release is made with
64 # https://github.com/NVlabs/tiny-cuda-nn/commit/78a14fe8c292a69f54e6d0d47a09f52b777127e1
65 postPatch = lib.optionals (cudaAtLeast "11.0") ''
66 substituteInPlace bindings/torch/setup.py --replace-fail \
67 "-std=c++14" "-std=c++17"
68 '';
69
70 nativeBuildInputs = [
71 cmake
72 cuda-native-redist
73 ninja
74 which
75 ]
76 ++ lists.optionals pythonSupport (
77 with python3Packages;
78 [
79 pip
80 setuptools
81 wheel
82 ]
83 );
84
85 buildInputs = [
86 cuda-redist
87 ]
88 ++ lib.optionals pythonSupport (
89 with python3Packages;
90 [
91 pybind11
92 python
93 ]
94 );
95
96 propagatedBuildInputs = lib.optionals pythonSupport (
97 with python3Packages;
98 [
99 torch
100 ]
101 );
102
103 # NOTE: We cannot use pythonImportsCheck for this module because it uses torch to immediately
104 # initialize CUDA and GPU access is not allowed in the nix build environment.
105 # NOTE: There are no tests for the C++ library or the python bindings, so we just skip the check
106 # phase -- we're not missing anything.
107 doCheck = false;
108
109 preConfigure = ''
110 export TCNN_CUDA_ARCHITECTURES="${cudaArchitecturesString}"
111 export CUDA_HOME="${cuda-native-redist}"
112 export LIBRARY_PATH="${cuda-native-redist}/lib/stubs:$LIBRARY_PATH"
113 export CC="${backendStdenv.cc}/bin/cc"
114 export CXX="${backendStdenv.cc}/bin/c++"
115 '';
116
117 # When building the python bindings, we cannot re-use the artifacts from the C++ build so we
118 # skip the CMake configurePhase and the buildPhase.
119 dontUseCmakeConfigure = pythonSupport;
120
121 # The configurePhase usually puts you in the build directory, so for the python bindings we
122 # need to change directories to the source directory.
123 configurePhase = strings.optionalString pythonSupport ''
124 runHook preConfigure
125 mkdir -p "$NIX_BUILD_TOP/build"
126 cd "$NIX_BUILD_TOP/build"
127 runHook postConfigure
128 '';
129
130 buildPhase = strings.optionalString pythonSupport ''
131 runHook preBuild
132 python -m pip wheel \
133 --no-build-isolation \
134 --no-clean \
135 --no-deps \
136 --no-index \
137 --verbose \
138 --wheel-dir "$NIX_BUILD_TOP/build" \
139 "$NIX_BUILD_TOP/source/bindings/torch"
140 runHook postBuild
141 '';
142
143 installPhase = ''
144 runHook preInstall
145 mkdir -p "$out/lib"
146 ''
147 # Installing the C++ library just requires copying the static library to the output directory
148 + strings.optionalString (!pythonSupport) ''
149 cp libtiny-cuda-nn.a "$out/lib/"
150 ''
151 # Installing the python bindings requires building the wheel and installing it
152 + strings.optionalString pythonSupport ''
153 python -m pip install \
154 --no-build-isolation \
155 --no-cache-dir \
156 --no-deps \
157 --no-index \
158 --no-warn-script-location \
159 --prefix="$out" \
160 --verbose \
161 ./*.whl
162 ''
163 + ''
164 runHook postInstall
165 '';
166
167 passthru = {
168 inherit cudaPackages;
169 };
170
171 meta = with lib; {
172 description = "Lightning fast C++/CUDA neural network framework";
173 homepage = "https://github.com/NVlabs/tiny-cuda-nn";
174 license = licenses.bsd3;
175 maintainers = with maintainers; [ connorbaker ];
176 platforms = platforms.linux;
177 badPlatforms = [
178 # g++: error: unrecognized command-line option '-mf16c'
179 lib.systems.inspect.patterns.isAarch64
180 ];
181 # Requires torch.cuda._is_compiled() == True to build
182 broken = !config.cudaSupport;
183 };
184})