Merge pull request #178397 from aidalgol/tensorrt (12a73604) · Commits · nix / nixpkgs

pkgs/development/libraries/science/math/cudnn/extension.nix

+1 −1

Original line number	Diff line number	Diff line
		@@ -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" ];
		}

pkgs/development/libraries/science/math/tensorrt/extension.nix

0 → 100644

+63 −0

Original line number	Diff line number	Diff line
		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

pkgs/development/libraries/science/math/tensorrt/generic.nix

0 → 100644

+90 −0

Original line number	Diff line number	Diff line
		{ lib
		, stdenv
		, requireFile
		, autoPatchelfHook
		, autoAddOpenGLRunpathHook
		, cudaVersion
		, cudatoolkit
		, cudnn
		}:

		{ 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 = fullVersion;
		src = requireFile rec {
		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 for CUDA ${cudaVersion} 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; {
		# 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;
		platforms = [ "x86_64-linux" ];
		maintainers = with maintainers; [ aidalgol ];
		};
		}

pkgs/development/python-modules/tensorrt/default.nix

0 → 100644

+52 −0

Original line number	Diff line number	Diff line
		{ 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 ];
		};
		}

pkgs/top-level/cuda-packages.nix

+3 −2

Original line number	Diff line number	Diff line
		@@ -58,15 +58,16 @@ let

		};

		composedExtension = composeManyExtensions [
		composedExtension = composeManyExtensions ([
		extraPackagesExtension
		(import ../development/compilers/cudatoolkit/extension.nix)
		(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
		];
		]);

		in (scope.overrideScope' composedExtension)

Admin message