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.

1 files changed, 2 insertions, 12 deletions

diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix
index af8beb9b58c37..70d57672a0f56 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)