diff --git a/pkgs/development/libraries/science/math/tiny-cuda-nn/default.nix b/pkgs/development/libraries/science/math/tiny-cuda-nn/default.nix new file mode 100644 index 0000000000000..ef501dfd09a04 --- /dev/null +++ b/pkgs/development/libraries/science/math/tiny-cuda-nn/default.nix @@ -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; + }; +}) diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix index bfecf795b2092..366c7316d4bf6 100644 --- a/pkgs/top-level/all-packages.nix +++ b/pkgs/top-level/all-packages.nix @@ -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; diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index de3f2ba61f0fe..def6dd5472404 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -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 { };