pkgs/by-name/ti/tiny-cuda-nn/package.nix

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