From d70b4df686a714f4a7f97bdf67eda1473f87707a Mon Sep 17 00:00:00 2001 From: Aidan Gauland Date: Tue, 21 Jun 2022 15:36:41 +1200 Subject: [PATCH] tensorrt: init at 8.4.0.6 Add derivation for TensorRT 8, a high-performance deep learning interface SDK from NVIDIA, which is at this point non-redistributable. The current version aldo requires CUDA 11, so this is left out of the cudaPackages_10* scopes. --- .../libraries/science/math/tensorrt/8.nix | 79 +++++++++++++++++++ .../python-modules/tensorrt/default.nix | 52 ++++++++++++ pkgs/top-level/cuda-packages.nix | 15 +++- pkgs/top-level/python-packages.nix | 2 + 4 files changed, 146 insertions(+), 2 deletions(-) create mode 100644 pkgs/development/libraries/science/math/tensorrt/8.nix create mode 100644 pkgs/development/python-modules/tensorrt/default.nix diff --git a/pkgs/development/libraries/science/math/tensorrt/8.nix b/pkgs/development/libraries/science/math/tensorrt/8.nix new file mode 100644 index 000000000000..e80d5f3a6fe0 --- /dev/null +++ b/pkgs/development/libraries/science/math/tensorrt/8.nix @@ -0,0 +1,79 @@ +{ lib +, stdenv +, requireFile +, autoPatchelfHook +, autoAddOpenGLRunpathHook +, cudaVersion +, cudatoolkit +, cudnn +}: + +assert lib.assertMsg (lib.strings.versionAtLeast cudaVersion "11.0") + "This version of TensorRT requires at least CUDA 11.0 (current version is ${cudaVersion})"; +assert lib.assertMsg (lib.strings.versionAtLeast cudnn.version "8.3") + "This version of TensorRT requires at least cuDNN 8.3 (current version is ${cudnn.version})"; + +stdenv.mkDerivation rec { + pname = "cudatoolkit-${cudatoolkit.majorVersion}-tensorrt"; + version = "8.4.0.6"; + src = requireFile rec { + name = "TensorRT-${version}.Linux.x86_64-gnu.cuda-11.6.cudnn8.3.tar.gz"; + sha256 = "sha256-DNgHHXF/G4cK2nnOWImrPXAkOcNW6Wy+8j0LRpAH/LQ="; + message = '' + To use the TensorRT derivation, you must join the NVIDIA Developer Program + and download the ${version} Linux x86_64 TAR package from + ${meta.homepage}. + + Once you have downloaded the file, add it to the store with the following + command, and try building this derivation again. + + $ nix-store --add-fixed sha256 ${name} + ''; + }; + + outputs = [ "out" "dev" ]; + + nativeBuildInputs = [ + autoPatchelfHook + autoAddOpenGLRunpathHook + ]; + + # Used by autoPatchelfHook + buildInputs = [ + cudatoolkit.cc.cc.lib # libstdc++ + cudatoolkit + cudnn + ]; + + sourceRoot = "TensorRT-${version}"; + + installPhase = '' + install --directory "$dev" "$out" + mv include "$dev" + mv targets/x86_64-linux-gnu/lib "$out" + install -D --target-directory="$out/bin" targets/x86_64-linux-gnu/bin/trtexec + ''; + + # Tell autoPatchelf about runtime dependencies. + # (postFixup phase is run before autoPatchelfHook.) + postFixup = + let + mostOfVersion = builtins.concatStringsSep "." + (lib.take 3 (lib.versions.splitVersion version)); + in + '' + echo 'Patching RPATH of libnvinfer libs' + patchelf --debug --add-needed libnvinfer.so \ + "$out/lib/libnvinfer.so.${mostOfVersion}" \ + "$out/lib/libnvinfer_plugin.so.${mostOfVersion}" \ + "$out/lib/libnvinfer_builder_resource.so.${mostOfVersion}" + ''; + + meta = with lib; { + description = "TensorRT: a high-performance deep learning interface"; + homepage = "https://developer.nvidia.com/tensorrt"; + license = licenses.unfree; + platforms = [ "x86_64-linux" ]; + maintainers = with maintainers; [ aidalgol ]; + }; +} diff --git a/pkgs/development/python-modules/tensorrt/default.nix b/pkgs/development/python-modules/tensorrt/default.nix new file mode 100644 index 000000000000..30da346c810e --- /dev/null +++ b/pkgs/development/python-modules/tensorrt/default.nix @@ -0,0 +1,52 @@ +{ lib +, python +, buildPythonPackage +, autoPatchelfHook +, unzip +, cudaPackages +}: + +let + pyVersion = "${lib.versions.major python.version}${lib.versions.minor python.version}"; +in +buildPythonPackage rec { + pname = "tensorrt"; + version = cudaPackages.tensorrt.version; + + src = cudaPackages.tensorrt.src; + + format = "wheel"; + # We unpack the wheel ourselves because of the odd packaging. + dontUseWheelUnpack = true; + + nativeBuildInputs = [ + unzip + autoPatchelfHook + cudaPackages.autoAddOpenGLRunpathHook + ]; + + preUnpack = '' + mkdir -p dist + tar --strip-components=2 -xf "$src" --directory=dist \ + "TensorRT-${version}/python/tensorrt-${version}-cp${pyVersion}-none-linux_x86_64.whl" + ''; + + sourceRoot = "."; + + buildInputs = [ + cudaPackages.cudnn + cudaPackages.tensorrt + ]; + + pythonCheckImports = [ + "tensorrt" + ]; + + meta = with lib; { + description = "Python bindings for TensorRT, a high-performance deep learning interface"; + homepage = "https://developer.nvidia.com/tensorrt"; + license = licenses.unfree; + platforms = [ "x86_64-linux" ]; + maintainers = with maintainers; [ aidalgol ]; + }; +} diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix index 211540260d10..af8beb9b58c3 100644 --- a/pkgs/top-level/cuda-packages.nix +++ b/pkgs/top-level/cuda-packages.nix @@ -43,6 +43,16 @@ let }; in { inherit cutensor; }; + tensorrtExtension = final: prev: let + ### Tensorrt + + inherit (final) cudaMajorMinorVersion cudaMajorVersion; + + # TODO: Add derivations for TensorRT versions that support older CUDA versions. + + tensorrt = final.callPackage ../development/libraries/science/math/tensorrt/8.nix { }; + in { inherit tensorrt; }; + extraPackagesExtension = final: prev: { nccl = final.callPackage ../development/libraries/science/math/nccl { }; @@ -58,7 +68,7 @@ let }; - composedExtension = composeManyExtensions [ + composedExtension = composeManyExtensions ([ extraPackagesExtension (import ../development/compilers/cudatoolkit/extension.nix) (import ../development/compilers/cudatoolkit/redist/extension.nix) @@ -67,6 +77,7 @@ let (import ../test/cuda/cuda-samples/extension.nix) (import ../test/cuda/cuda-library-samples/extension.nix) cutensorExtension - ]; + ] ++ (lib.optional (lib.strings.versionAtLeast cudaVersion "11.0") tensorrtExtension)); + # We only package the current version of TensorRT, which requires CUDA 11. in (scope.overrideScope' composedExtension) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index d958958d0618..9f525e31bd40 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -10384,6 +10384,8 @@ in { tensorly = callPackage ../development/python-modules/tensorly { }; + tensorrt = callPackage ../development/python-modules/tensorrt { }; + tellduslive = callPackage ../development/python-modules/tellduslive { }; termcolor = callPackage ../development/python-modules/termcolor { };