-
-
Notifications
You must be signed in to change notification settings - Fork 13.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
1eeea1f
commit 45726d9
Showing
3 changed files
with
162 additions
and
0 deletions.
There are no files selected for viewing
155 changes: 155 additions & 0 deletions
155
pkgs/development/libraries/science/math/tiny-cuda-nn/default.nix
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,155 @@ | ||
{ cmake | ||
, fetchFromGitHub | ||
, lib | ||
, ninja | ||
, pythonPackages | ||
, pythonSupport ? false | ||
, stdenv | ||
, symlinkJoin | ||
, which | ||
}: | ||
assert lib.asserts.assertMsg | ||
(pythonPackages.torch.cudaSupport) | ||
"tiny-cuda-nn requires torch to be built with cudaSupport"; let | ||
inherit (lib) lists strings; | ||
inherit (pythonPackages.torch) cudaPackages cudaSupport; | ||
inherit (cudaPackages) cudaFlags cudatoolkit; | ||
|
||
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 pythonPackages; [ | ||
pip | ||
setuptools | ||
wheel | ||
]); | ||
|
||
buildInputs = [ | ||
cuda-redist | ||
] ++ lib.optionals pythonSupport ( | ||
with pythonPackages; [ | ||
pybind11 | ||
python | ||
] | ||
); | ||
|
||
propagatedBuildInputs = lib.optionals pythonSupport ( | ||
with pythonPackages; [ | ||
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=${cudatoolkit.cc}/bin/cc | ||
export CXX=${cudatoolkit.cc}/bin/c++ | ||
export CUDAHOSTCXX=${cudatoolkit.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; | ||
broken = !cudaSupport; | ||
}; | ||
}) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters