tiny-cuda-nn: init 1.6
This commit is contained in:
parent
7913ebc8c3
commit
aacbb6d021
152
pkgs/development/libraries/science/math/tiny-cuda-nn/default.nix
Normal file
152
pkgs/development/libraries/science/math/tiny-cuda-nn/default.nix
Normal file
|
@ -0,0 +1,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;
|
||||
};
|
||||
})
|
|
@ -3812,6 +3812,8 @@ with pkgs;
|
|||
|
||||
tensorflow-lite = callPackage ../development/libraries/science/math/tensorflow-lite { };
|
||||
|
||||
tiny-cuda-nn = callPackage ../development/libraries/science/math/tiny-cuda-nn { };
|
||||
|
||||
tezos-rust-libs = callPackage ../development/libraries/tezos-rust-libs { };
|
||||
|
||||
behave = with python3Packages; toPythonApplication behave;
|
||||
|
|
|
@ -11694,6 +11694,12 @@ self: super: with self; {
|
|||
|
||||
timm = callPackage ../development/python-modules/timm { };
|
||||
|
||||
tiny-cuda-nn = toPythonModule (pkgs.tiny-cuda-nn.override {
|
||||
cudaPackages = self.torch.cudaPackages;
|
||||
python3Packages = self;
|
||||
pythonSupport = true;
|
||||
});
|
||||
|
||||
tinycss2 = callPackage ../development/python-modules/tinycss2 { };
|
||||
|
||||
tinycss = callPackage ../development/python-modules/tinycss { };
|
||||
|
|
Loading…
Reference in a new issue