Merge pull request #218166 from ConnorBaker/feat/tiny-cuda-nn
tiny-cuda-nn,python3Packages.tiny-cuda-nn: init 1.6
This commit is contained in:
commit
49fd723a69
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;
|
||||||
|
};
|
||||||
|
})
|
|
@ -3822,6 +3822,8 @@ with pkgs;
|
||||||
|
|
||||||
tensorflow-lite = callPackage ../development/libraries/science/math/tensorflow-lite { };
|
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 { };
|
tezos-rust-libs = callPackage ../development/libraries/tezos-rust-libs { };
|
||||||
|
|
||||||
behave = with python3Packages; toPythonApplication behave;
|
behave = with python3Packages; toPythonApplication behave;
|
||||||
|
|
|
@ -11702,6 +11702,12 @@ self: super: with self; {
|
||||||
|
|
||||||
timm = callPackage ../development/python-modules/timm { };
|
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 { };
|
tinycss2 = callPackage ../development/python-modules/tinycss2 { };
|
||||||
|
|
||||||
tinycss = callPackage ../development/python-modules/tinycss { };
|
tinycss = callPackage ../development/python-modules/tinycss { };
|
||||||
|
|
Loading…
Reference in a new issue