From d70b4df686a714f4a7f97bdf67eda1473f87707a Mon Sep 17 00:00:00 2001 From: Aidan Gauland Date: Tue, 21 Jun 2022 15:36:41 +1200 Subject: [PATCH 1/3] 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 { }; From c8fba8254a5f984a215164f84a4fd43c28c6ed82 Mon Sep 17 00:00:00 2001 From: Aidan Gauland Date: Sat, 2 Jul 2022 09:53:27 +1200 Subject: [PATCH 2/3] tensorrt: support multiple CUDA versions Refactor derivation to pick the version that supports the current CUDA version. Based on the implementation of the same concept in the cudnn derivation. --- .../science/math/tensorrt/extension.nix | 63 +++++++++++++++++++ .../math/tensorrt/{8.nix => generic.nix} | 29 ++++++--- pkgs/top-level/cuda-packages.nix | 14 +---- 3 files changed, 85 insertions(+), 21 deletions(-) create mode 100644 pkgs/development/libraries/science/math/tensorrt/extension.nix rename pkgs/development/libraries/science/math/tensorrt/{8.nix => generic.nix} (67%) diff --git a/pkgs/development/libraries/science/math/tensorrt/extension.nix b/pkgs/development/libraries/science/math/tensorrt/extension.nix new file mode 100644 index 000000000000..8e9f64885bad --- /dev/null +++ b/pkgs/development/libraries/science/math/tensorrt/extension.nix @@ -0,0 +1,63 @@ +final: prev: let + + inherit (final) callPackage; + inherit (prev) cudatoolkit cudaVersion lib pkgs; + + ### TensorRT + + buildTensorRTPackage = args: + callPackage ./generic.nix { } args; + + toUnderscore = str: lib.replaceStrings ["."] ["_"] str; + + majorMinorPatch = str: lib.concatStringsSep "." (lib.take 3 (lib.splitVersion str)); + + tensorRTPackages = with lib; let + # Check whether a file is supported for our cuda version + isSupported = fileData: elem cudaVersion fileData.supportedCudaVersions; + # Return the first file that is supported. In practice there should only ever be one anyway. + supportedFile = files: findFirst isSupported null files; + # Supported versions with versions as keys and file as value + supportedVersions = filterAttrs (version: file: file !=null ) (mapAttrs (version: files: supportedFile files) tensorRTVersions); + # Compute versioned attribute name to be used in this package set + computeName = version: "tensorrt_${toUnderscore version}"; + # Add all supported builds as attributes + allBuilds = mapAttrs' (version: file: nameValuePair (computeName version) (buildTensorRTPackage (removeAttrs file ["fileVersionCuda"]))) supportedVersions; + # Set the default attributes, e.g. tensorrt = tensorrt_8_4; + defaultBuild = { "tensorrt" = allBuilds.${computeName tensorRTDefaultVersion}; }; + in allBuilds // defaultBuild; + + tensorRTVersions = { + "8.4.0" = [ + rec { + fileVersionCuda = "11.6"; + fileVersionCudnn = "8.3"; + fullVersion = "8.4.0.6"; + sha256 = "sha256-DNgHHXF/G4cK2nnOWImrPXAkOcNW6Wy+8j0LRpAH/LQ="; + tarball = "TensorRT-${fullVersion}.Linux.x86_64-gnu.cuda-${fileVersionCuda}.cudnn${fileVersionCudnn}.tar.gz"; + supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" ]; + } + rec { + fileVersionCuda = "10.2"; + fileVersionCudnn = "8.3"; + fullVersion = "8.4.0.6"; + sha256 = "sha256-aCzH0ZI6BrJ0v+e5Bnm7b8mNltA7NNuIa8qRKzAQv+I="; + tarball = "TensorRT-${fullVersion}.Linux.x86_64-gnu.cuda-${fileVersionCuda}.cudnn${fileVersionCudnn}.tar.gz"; + supportedCudaVersions = [ "10.2" ]; + } + ]; + }; + + # Default attributes + tensorRTDefaultVersion = { + "10.2" = "8.4.0"; + "11.0" = "8.4.0"; + "11.1" = "8.4.0"; + "11.2" = "8.4.0"; + "11.3" = "8.4.0"; + "11.4" = "8.4.0"; + "11.5" = "8.4.0"; + "11.6" = "8.4.0"; + }.${cudaVersion}; + +in tensorRTPackages diff --git a/pkgs/development/libraries/science/math/tensorrt/8.nix b/pkgs/development/libraries/science/math/tensorrt/generic.nix similarity index 67% rename from pkgs/development/libraries/science/math/tensorrt/8.nix rename to pkgs/development/libraries/science/math/tensorrt/generic.nix index e80d5f3a6fe0..3447087051f1 100644 --- a/pkgs/development/libraries/science/math/tensorrt/8.nix +++ b/pkgs/development/libraries/science/math/tensorrt/generic.nix @@ -8,20 +8,25 @@ , 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})"; +{ fullVersion +, fileVersionCudnn +, tarball +, sha256 +, supportedCudaVersions ? [ ] +}: + +assert lib.assertMsg (lib.strings.versionAtLeast cudnn.version fileVersionCudnn) + "This version of TensorRT requires at least cuDNN ${fileVersionCudnn} (current version is ${cudnn.version})"; stdenv.mkDerivation rec { pname = "cudatoolkit-${cudatoolkit.majorVersion}-tensorrt"; - version = "8.4.0.6"; + version = fullVersion; src = requireFile rec { - name = "TensorRT-${version}.Linux.x86_64-gnu.cuda-11.6.cudnn8.3.tar.gz"; - sha256 = "sha256-DNgHHXF/G4cK2nnOWImrPXAkOcNW6Wy+8j0LRpAH/LQ="; + name = tarball; + inherit sha256; message = '' - To use the TensorRT derivation, you must join the NVIDIA Developer Program - and download the ${version} Linux x86_64 TAR package from + To use the TensorRT derivation, you must join the NVIDIA Developer Program and + download the ${version} Linux x86_64 TAR package for CUDA ${cudaVersion} from ${meta.homepage}. Once you have downloaded the file, add it to the store with the following @@ -70,6 +75,12 @@ stdenv.mkDerivation rec { ''; meta = with lib; { + # Check that the cudatoolkit version satisfies our min/max constraints (both + # inclusive). We mark the package as broken if it fails to satisfies the + # official version constraints (as recorded in default.nix). In some cases + # you _may_ be able to smudge version constraints, just know that you're + # embarking into unknown and unsupported territory when doing so. + broken = !(elem cudaVersion supportedCudaVersions); description = "TensorRT: a high-performance deep learning interface"; homepage = "https://developer.nvidia.com/tensorrt"; license = licenses.unfree; diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix index af8beb9b58c3..70d57672a0f5 100644 --- a/pkgs/top-level/cuda-packages.nix +++ b/pkgs/top-level/cuda-packages.nix @@ -43,16 +43,6 @@ 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 { }; @@ -74,10 +64,10 @@ let (import ../development/compilers/cudatoolkit/redist/extension.nix) (import ../development/compilers/cudatoolkit/redist/overrides.nix) (import ../development/libraries/science/math/cudnn/extension.nix) + (import ../development/libraries/science/math/tensorrt/extension.nix) (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) From e28a8e07a6a518645625d0041d72493e3a0e2764 Mon Sep 17 00:00:00 2001 From: Aidan Gauland Date: Sat, 2 Jul 2022 11:09:45 +1200 Subject: [PATCH 3/3] cudnn: fix incorrect hash Hash mismatch for cudnn_8_3_2 discovered when building cudaPackages_10_2.tensorrt, which depends on this version of cudnn. --- pkgs/development/libraries/science/math/cudnn/extension.nix | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkgs/development/libraries/science/math/cudnn/extension.nix b/pkgs/development/libraries/science/math/cudnn/extension.nix index f1bdfb9836ed..265e09f436ab 100644 --- a/pkgs/development/libraries/science/math/cudnn/extension.nix +++ b/pkgs/development/libraries/science/math/cudnn/extension.nix @@ -85,7 +85,7 @@ final: prev: let rec { fileVersion = "10.2"; fullVersion = "8.3.2.44"; - hash = "sha256-mKh4TpKGLyABjSDCgbMNSgzZUfk2lPZDPM9K6cUCumo="; + hash = "sha256-1vVu+cqM+PketzIQumw9ykm6REbBZhv6/lXB7EC2aaw="; url = "${urlPrefix}/v${majorMinorPatch fullVersion}/local_installers/${fileVersion}/cudnn-linux-x86_64-${fullVersion}_cuda${fileVersion}-archive.tar.xz"; supportedCudaVersions = [ "10.2" ]; }