push sheeet

pkgs/development/cuda-modules/_cuda/fixups/cuda_compat.nix

@@ -0,0 +1,12 @@
+{ flags, lib }:
+prevAttrs: {
+  autoPatchelfIgnoreMissingDeps = prevAttrs.autoPatchelfIgnoreMissingDeps or [ ] ++ [
+    "libnvrm_gpu.so"
+    "libnvrm_mem.so"
+    "libnvdla_runtime.so"
+  ];
+  # `cuda_compat` only works on aarch64-linux, and only when building for Jetson devices.
+  badPlatformsConditions = prevAttrs.badPlatformsConditions or { } // {
+    "Trying to use cuda_compat on aarch64-linux targeting non-Jetson devices" = !flags.isJetsonBuild;
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/cuda_cudart.nix

@@ -0,0 +1,37 @@
+# TODO(@connorbaker): cuda_cudart.dev depends on crt/host_config.h, which is from
+# (getDev cuda_nvcc). It would be nice to be able to encode that.
+{ addDriverRunpath, lib }:
+prevAttrs: {
+  # Remove once cuda-find-redist-features has a special case for libcuda
+  outputs =
+    prevAttrs.outputs or [ ]
+    ++ lib.lists.optionals (!(builtins.elem "stubs" prevAttrs.outputs)) [ "stubs" ];
+
+  allowFHSReferences = false;
+
+  # The libcuda stub's pkg-config doesn't follow the general pattern:
+  postPatch =
+    prevAttrs.postPatch or ""
+    + ''
+      while IFS= read -r -d $'\0' path; do
+        sed -i \
+          -e "s|^libdir\s*=.*/lib\$|libdir=''${!outputLib}/lib/stubs|" \
+          -e "s|^Libs\s*:\(.*\)\$|Libs: \1 -Wl,-rpath,${addDriverRunpath.driverLink}/lib|" \
+          "$path"
+      done < <(find -iname 'cuda-*.pc' -print0)
+    ''
+    # Namelink may not be enough, add a soname.
+    # Cf. https://gitlab.kitware.com/cmake/cmake/-/issues/25536
+    + ''
+      if [[ -f lib/stubs/libcuda.so && ! -f lib/stubs/libcuda.so.1 ]]; then
+        ln -s libcuda.so lib/stubs/libcuda.so.1
+      fi
+    '';
+
+  postFixup = prevAttrs.postFixup or "" + ''
+    mv "''${!outputDev}/share" "''${!outputDev}/lib"
+    moveToOutput lib/stubs "$stubs"
+    ln -s "$stubs"/lib/stubs/* "$stubs"/lib/
+    ln -s "$stubs"/lib/stubs "''${!outputLib}/lib/stubs"
+  '';
+}

pkgs/development/cuda-modules/_cuda/fixups/cuda_demo_suite.nix

@@ -0,0 +1,18 @@
+{
+  libglut,
+  libcufft,
+  libcurand,
+  libGLU,
+  libglvnd,
+  libgbm,
+}:
+prevAttrs: {
+  buildInputs = prevAttrs.buildInputs or [ ] ++ [
+    libglut
+    libcufft
+    libcurand
+    libGLU
+    libglvnd
+    libgbm
+  ];
+}

pkgs/development/cuda-modules/_cuda/fixups/cuda_gdb.nix

@@ -0,0 +1,35 @@
+{
+  cudaAtLeast,
+  gmp,
+  expat,
+  libxcrypt-legacy,
+  ncurses6,
+  python310,
+  python311,
+  python312,
+  stdenv,
+  lib,
+}:
+prevAttrs: {
+  buildInputs =
+    prevAttrs.buildInputs or [ ]
+    ++ [
+      gmp
+      libxcrypt-legacy
+      ncurses6
+      python310
+      python311
+      python312
+    ]
+    # aarch64,sbsa needs expat
+    ++ lib.lists.optionals (stdenv.hostPlatform.isAarch64) [ expat ];
+
+  installPhase =
+    prevAttrs.installPhase or ""
+    # Python 3.8 is not in nixpkgs anymore, delete Python 3.8 cuda-gdb support
+    # to avoid autopatchelf failing to find libpython3.8.so.
+    + ''
+      find $bin -name '*python3.8*' -delete
+      find $bin -name '*python3.9*' -delete
+    '';
+}

pkgs/development/cuda-modules/_cuda/fixups/cuda_nvcc.nix

@@ -0,0 +1,62 @@
+{
+  lib,
+  backendStdenv,
+  setupCudaHook,
+}:
+prevAttrs: {
+  # Merge "bin" and "dev" into "out" to avoid circular references
+  outputs = builtins.filter (
+    x:
+    !(builtins.elem x [
+      "dev"
+      "bin"
+    ])
+  ) prevAttrs.outputs or [ ];
+
+  # Patch the nvcc.profile.
+  # Syntax:
+  # - `=` for assignment,
+  # - `?=` for conditional assignment,
+  # - `+=` to "prepend",
+  # - `=+` to "append".
+
+  # Cf. https://web.archive.org/web/20230308044351/https://arcb.csc.ncsu.edu/~mueller/cluster/nvidia/2.0/nvcc_2.0.pdf
+
+  # We set all variables with the lowest priority (=+), but we do force
+  # nvcc to use the fixed backend toolchain. Cf. comments in
+  # backend-stdenv.nix
+
+  postPatch =
+    prevAttrs.postPatch or ""
+    + ''
+      substituteInPlace bin/nvcc.profile \
+        --replace-fail \
+          '$(TOP)/$(_TARGET_DIR_)/include' \
+          "''${!outputDev}/include"
+    ''
+    + ''
+      cat << EOF >> bin/nvcc.profile
+
+      # Fix a compatible backend compiler
+      PATH += "${backendStdenv.cc}/bin":
+
+      # Expose the split-out nvvm
+      LIBRARIES =+ "-L''${!outputBin}/nvvm/lib"
+      INCLUDES =+ "-I''${!outputBin}/nvvm/include"
+      EOF
+    '';
+
+  # Entries here will be in nativeBuildInputs when cuda_nvcc is in nativeBuildInputs.
+  propagatedBuildInputs = prevAttrs.propagatedBuildInputs or [ ] ++ [ setupCudaHook ];
+
+  postInstall = prevAttrs.postInstall or "" + ''
+    moveToOutput "nvvm" "''${!outputBin}"
+  '';
+
+  # The nvcc and cicc binaries contain hard-coded references to /usr
+  allowFHSReferences = true;
+
+  meta = prevAttrs.meta or { } // {
+    mainProgram = "nvcc";
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/cuda_nvprof.nix

@@ -0,0 +1 @@
+{ cuda_cupti }: prevAttrs: { buildInputs = prevAttrs.buildInputs or [ ] ++ [ cuda_cupti ]; }

pkgs/development/cuda-modules/_cuda/fixups/cuda_sanitizer_api.nix

@@ -0,0 +1 @@
+_: _: { outputs = [ "out" ]; }

pkgs/development/cuda-modules/_cuda/fixups/cudnn.nix

@@ -0,0 +1,75 @@
+{
+  cudaOlder,
+  cudaMajorMinorVersion,
+  fetchurl,
+  lib,
+  libcublas,
+  patchelf,
+  zlib,
+}:
+let
+  inherit (lib)
+    attrsets
+    maintainers
+    meta
+    strings
+    ;
+in
+finalAttrs: prevAttrs: {
+  src = fetchurl { inherit (finalAttrs.passthru.redistribRelease) hash url; };
+
+  # Useful for inspecting why something went wrong.
+  badPlatformsConditions =
+    let
+      cudaTooOld = cudaOlder finalAttrs.passthru.featureRelease.minCudaVersion;
+      cudaTooNew =
+        (finalAttrs.passthru.featureRelease.maxCudaVersion != null)
+        && strings.versionOlder finalAttrs.passthru.featureRelease.maxCudaVersion cudaMajorMinorVersion;
+    in
+    prevAttrs.badPlatformsConditions or { }
+    // {
+      "CUDA version is too old" = cudaTooOld;
+      "CUDA version is too new" = cudaTooNew;
+    };
+
+  buildInputs = prevAttrs.buildInputs or [ ] ++ [
+    zlib
+    (attrsets.getLib libcublas)
+  ];
+
+  # Tell autoPatchelf about runtime dependencies. *_infer* libraries only
+  # exist in CuDNN 8.
+  # NOTE: Versions from CUDNN releases have four components.
+  postFixup =
+    prevAttrs.postFixup or ""
+    +
+      strings.optionalString
+        (
+          strings.versionAtLeast finalAttrs.version "8.0.5.0"
+          && strings.versionOlder finalAttrs.version "9.0.0.0"
+        )
+        ''
+          ${meta.getExe patchelf} $lib/lib/libcudnn.so --add-needed libcudnn_cnn_infer.so
+          ${meta.getExe patchelf} $lib/lib/libcudnn_ops_infer.so --add-needed libcublas.so --add-needed libcublasLt.so
+        '';
+
+  meta = prevAttrs.meta or { } // {
+    homepage = "https://developer.nvidia.com/cudnn";
+    maintainers =
+      prevAttrs.meta.maintainers or [ ]
+      ++ (with maintainers; [
+        mdaiter
+        samuela
+        connorbaker
+      ]);
+    # TODO(@connorbaker): Temporary workaround to avoid changing the derivation hash since introducing more
+    # brokenConditions would change the derivation as they're top-level and __structuredAttrs is set.
+    teams = prevAttrs.meta.teams or [ ];
+    license = {
+      shortName = "cuDNN EULA";
+      fullName = "NVIDIA cuDNN Software License Agreement (EULA)";
+      url = "https://docs.nvidia.com/deeplearning/sdk/cudnn-sla/index.html#supplement";
+      free = false;
+    };
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/default.nix

@@ -0,0 +1,11 @@
+{ lib }:
+lib.concatMapAttrs (
+  fileName: _type:
+  let
+    # Fixup is in `./${attrName}.nix` or in `./${fileName}/default.nix`:
+    attrName = lib.removeSuffix ".nix" fileName;
+    fixup = import (./. + "/${fileName}");
+    isFixup = fileName != "default.nix";
+  in
+  lib.optionalAttrs isFixup { ${attrName} = fixup; }
+) (builtins.readDir ./.)

pkgs/development/cuda-modules/_cuda/fixups/driver_assistant.nix

@@ -0,0 +1,5 @@
+_: prevAttrs: {
+  badPlatformsConditions = prevAttrs.badPlatformsConditions or { } // {
+    "Package is not supported; use drivers from linuxPackages" = true;
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/fabricmanager.nix

@@ -0,0 +1 @@
+{ zlib }: prevAttrs: { buildInputs = prevAttrs.buildInputs or [ ] ++ [ zlib ]; }

pkgs/development/cuda-modules/_cuda/fixups/imex.nix

@@ -0,0 +1 @@
+{ zlib }: prevAttrs: { buildInputs = prevAttrs.buildInputs or [ ] ++ [ zlib ]; }

pkgs/development/cuda-modules/_cuda/fixups/libcufile.nix

@@ -0,0 +1,12 @@
+{
+  libcublas,
+  numactl,
+  rdma-core,
+}:
+prevAttrs: {
+  buildInputs = prevAttrs.buildInputs or [ ] ++ [
+    libcublas
+    numactl
+    rdma-core
+  ];
+}

pkgs/development/cuda-modules/_cuda/fixups/libcusolver.nix

@@ -0,0 +1,19 @@
+{
+  cudaAtLeast,
+  lib,
+  libcublas,
+  libcusparse ? null,
+  libnvjitlink ? null,
+}:
+prevAttrs: {
+  buildInputs = prevAttrs.buildInputs or [ ] ++ [
+    libcublas
+    libnvjitlink
+    libcusparse
+  ];
+
+  brokenConditions = prevAttrs.brokenConditions or { } // {
+    "libnvjitlink missing (CUDA >= 12.0)" = libnvjitlink == null;
+    "libcusparse missing (CUDA >= 12.1)" = libcusparse == null;
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/libcusparse.nix

@@ -0,0 +1,12 @@
+{
+  cudaAtLeast,
+  lib,
+  libnvjitlink ? null,
+}:
+prevAttrs: {
+  buildInputs = prevAttrs.buildInputs or [ ] ++ [ libnvjitlink ];
+
+  brokenConditions = prevAttrs.brokenConditions or { } // {
+    "libnvjitlink missing (CUDA >= 12.0)" = libnvjitlink == null;
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/libcusparse_lt.nix

@@ -0,0 +1,23 @@
+{
+  cuda_cudart,
+  lib,
+  libcublas,
+}:
+finalAttrs: prevAttrs: {
+  buildInputs =
+    prevAttrs.buildInputs or [ ]
+    ++ [ (lib.getLib libcublas) ]
+    # For some reason, the 1.4.x release of cusparselt requires the cudart library.
+    ++ lib.optionals (lib.hasPrefix "1.4" finalAttrs.version) [ (lib.getLib cuda_cudart) ];
+  meta = prevAttrs.meta or { } // {
+    description = "cuSPARSELt: A High-Performance CUDA Library for Sparse Matrix-Matrix Multiplication";
+    homepage = "https://developer.nvidia.com/cusparselt-downloads";
+    maintainers = prevAttrs.meta.maintainers or [ ] ++ [ lib.maintainers.sepiabrown ];
+    teams = prevAttrs.meta.teams or [ ];
+    license = lib.licenses.unfreeRedistributable // {
+      shortName = "cuSPARSELt EULA";
+      fullName = "cuSPARSELt SUPPLEMENT TO SOFTWARE LICENSE AGREEMENT FOR NVIDIA SOFTWARE DEVELOPMENT KITS";
+      url = "https://docs.nvidia.com/cuda/cusparselt/license.html";
+    };
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/libcutensor.nix

@@ -0,0 +1,23 @@
+{
+  cuda_cudart,
+  lib,
+  libcublas,
+}:
+finalAttrs: prevAttrs: {
+  buildInputs =
+    prevAttrs.buildInputs or [ ]
+    ++ [ (lib.getLib libcublas) ]
+    # For some reason, the 1.4.x release of cuTENSOR requires the cudart library.
+    ++ lib.optionals (lib.hasPrefix "1.4" finalAttrs.version) [ (lib.getLib cuda_cudart) ];
+  meta = prevAttrs.meta or { } // {
+    description = "cuTENSOR: A High-Performance CUDA Library For Tensor Primitives";
+    homepage = "https://developer.nvidia.com/cutensor";
+    maintainers = prevAttrs.meta.maintainers or [ ] ++ [ lib.maintainers.obsidian-systems-maintenance ];
+    teams = prevAttrs.meta.teams;
+    license = lib.licenses.unfreeRedistributable // {
+      shortName = "cuTENSOR EULA";
+      fullName = "cuTENSOR SUPPLEMENT TO SOFTWARE LICENSE AGREEMENT FOR NVIDIA SOFTWARE DEVELOPMENT KITS";
+      url = "https://docs.nvidia.com/cuda/cutensor/license.html";
+    };
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/nsight_compute.nix

@@ -0,0 +1,86 @@
+{
+  cudaAtLeast,
+  cudaMajorMinorVersion,
+  cudaOlder,
+  e2fsprogs,
+  elfutils,
+  flags,
+  gst_all_1,
+  lib,
+  libjpeg8,
+  qt6,
+  rdma-core,
+  stdenv,
+  ucx,
+}:
+prevAttrs:
+let
+  qtwayland = lib.getLib qt6.qtwayland;
+  inherit (qt6) wrapQtAppsHook qtwebview;
+  archDir =
+    {
+      aarch64-linux = "linux-" + (if flags.isJetsonBuild then "v4l_l4t" else "desktop") + "-t210-a64";
+      x86_64-linux = "linux-desktop-glibc_2_11_3-x64";
+    }
+    .${stdenv.hostPlatform.system} or (throw "Unsupported system: ${stdenv.hostPlatform.system}");
+in
+{
+  outputs = [ "out" ]; # NOTE(@connorbaker): Force a single output so relative lookups work.
+  nativeBuildInputs = prevAttrs.nativeBuildInputs or [ ] ++ [ wrapQtAppsHook ];
+  buildInputs =
+    prevAttrs.buildInputs or [ ]
+    ++ [
+      qtwayland
+      qtwebview
+      (qt6.qtwebengine or qt6.full)
+      rdma-core
+    ]
+    ++ lib.optionals (cudaOlder "12.7") [
+      e2fsprogs
+      ucx
+    ]
+    ++ lib.optionals (cudaMajorMinorVersion == "12.9") [
+      elfutils
+    ];
+  dontWrapQtApps = true;
+  preInstall = prevAttrs.preInstall or "" + ''
+    if [[ -d nsight-compute ]]; then
+      nixLog "Lifting components of Nsight Compute to the top level"
+      mv -v nsight-compute/*/* .
+      nixLog "Removing empty directories"
+      rmdir -pv nsight-compute/*
+    fi
+
+    rm -rf host/${archDir}/Mesa/
+  '';
+  postInstall =
+    prevAttrs.postInstall or ""
+    + ''
+      moveToOutput 'ncu' "''${!outputBin}/bin"
+      moveToOutput 'ncu-ui' "''${!outputBin}/bin"
+      moveToOutput 'host/${archDir}' "''${!outputBin}/bin"
+      moveToOutput 'target/${archDir}' "''${!outputBin}/bin"
+      wrapQtApp "''${!outputBin}/bin/host/${archDir}/ncu-ui.bin"
+    ''
+    # NOTE(@connorbaker): No idea what this platform is or how to patchelf for it.
+    + lib.optionalString (flags.isJetsonBuild && cudaOlder "12.9") ''
+      nixLog "Removing QNX 700 target directory for Jetson builds"
+      rm -rfv "''${!outputBin}/target/qnx-700-t210-a64"
+    ''
+    + lib.optionalString (flags.isJetsonBuild && cudaAtLeast "12.8") ''
+      nixLog "Removing QNX 800 target directory for Jetson builds"
+      rm -rfv "''${!outputBin}/target/qnx-800-tegra-a64"
+    '';
+  # lib needs libtiff.so.5, but nixpkgs provides libtiff.so.6
+  preFixup = prevAttrs.preFixup or "" + ''
+    patchelf --replace-needed libtiff.so.5 libtiff.so "''${!outputBin}/bin/host/${archDir}/Plugins/imageformats/libqtiff.so"
+  '';
+  autoPatchelfIgnoreMissingDeps = prevAttrs.autoPatchelfIgnoreMissingDeps or [ ] ++ [
+    "libnvidia-ml.so.1"
+  ];
+  # NOTE(@connorbaker): It might be a problem that when nsight_compute contains hosts and targets of different
+  # architectures, that we patchelf just the binaries matching the builder's platform; autoPatchelfHook prints
+  # messages like
+  #   skipping [$out]/host/linux-desktop-glibc_2_11_3-x64/libQt6Core.so.6 because its architecture (x64) differs from
+  #   target (AArch64)
+}

pkgs/development/cuda-modules/_cuda/fixups/nsight_systems.nix

@@ -0,0 +1,136 @@
+{
+  boost178,
+  cuda_cudart,
+  cudaAtLeast,
+  e2fsprogs,
+  gst_all_1,
+  lib,
+  nss,
+  numactl,
+  pulseaudio,
+  qt6,
+  rdma-core,
+  stdenv,
+  ucx,
+  wayland,
+  xorg,
+}:
+prevAttrs:
+let
+  qtwayland = lib.getLib qt6.qtwayland;
+  qtWaylandPlugins = "${qtwayland}/${qt6.qtbase.qtPluginPrefix}";
+  # NOTE(@connorbaker): nsight_systems doesn't support Jetson, so no need for case splitting on aarch64-linux.
+  hostDir =
+    {
+      aarch64-linux = "host-linux-armv8";
+      x86_64-linux = "host-linux-x64";
+    }
+    .${stdenv.hostPlatform.system} or (throw "Unsupported system: ${stdenv.hostPlatform.system}");
+  targetDir =
+    {
+      aarch64-linux = "target-linux-sbsa-armv8";
+      x86_64-linux = "target-linux-x64";
+    }
+    .${stdenv.hostPlatform.system} or (throw "Unsupported system: ${stdenv.hostPlatform.system}");
+in
+{
+  outputs = [ "out" ]; # NOTE(@connorbaker): Force a single output so relative lookups work.
+
+  # An ad hoc replacement for
+  # https://github.com/ConnorBaker/cuda-redist-find-features/issues/11
+  env = prevAttrs.env or { } // {
+    rmPatterns =
+      prevAttrs.env.rmPatterns or ""
+      + toString [
+        "${hostDir}/lib{arrow,jpeg}*"
+        "${hostDir}/lib{ssl,ssh,crypto}*"
+        "${hostDir}/libboost*"
+        "${hostDir}/libexec"
+        "${hostDir}/libstdc*"
+        "${hostDir}/python/bin/python"
+        "${hostDir}/Mesa"
+      ];
+  };
+
+  # NOTE(@connorbaker): nsight-exporter and nsight-sys are deprecated scripts wrapping nsys, it's fine to remove them.
+  prePatch = prevAttrs.prePatch or "" + ''
+    if [[ -d bin ]]; then
+      nixLog "Removing bin wrapper scripts"
+      for knownWrapper in bin/{nsys{,-ui},nsight-{exporter,sys}}; do
+        [[ -e $knownWrapper ]] && rm -v "$knownWrapper"
+      done
+      unset -v knownWrapper
+
+      nixLog "Removing empty bin directory"
+      rmdir -v bin
+    fi
+
+    if [[ -d nsight-systems ]]; then
+      nixLog "Lifting components of Nsight System to the top level"
+      mv -v nsight-systems/*/* .
+      nixLog "Removing empty nsight-systems directory"
+      rmdir -pv nsight-systems/*
+    fi
+  '';
+
+  postPatch = prevAttrs.postPatch or "" + ''
+    for path in $rmPatterns; do
+      rm -r "$path"
+    done
+    patchShebangs nsight-systems
+  '';
+
+  nativeBuildInputs = prevAttrs.nativeBuildInputs or [ ] ++ [ qt6.wrapQtAppsHook ];
+
+  dontWrapQtApps = true;
+
+  buildInputs =
+    prevAttrs.buildInputs or [ ]
+    ++ [
+      (qt6.qtdeclarative or qt6.full)
+      (qt6.qtsvg or qt6.full)
+      (qt6.qtimageformats or qt6.full)
+      (qt6.qtpositioning or qt6.full)
+      (qt6.qtscxml or qt6.full)
+      (qt6.qttools or qt6.full)
+      (qt6.qtwebengine or qt6.full)
+      (qt6.qtwayland or qt6.full)
+      boost178
+      cuda_cudart.stubs
+      e2fsprogs
+      gst_all_1.gst-plugins-base
+      gst_all_1.gstreamer
+      nss
+      numactl
+      pulseaudio
+      qt6.qtbase
+      qtWaylandPlugins
+      rdma-core
+      ucx
+      wayland
+      xorg.libXcursor
+      xorg.libXdamage
+      xorg.libXrandr
+      xorg.libXtst
+    ]
+    # NOTE(@connorbaker): Seems to be required only for aarch64-linux.
+    ++ lib.optionals stdenv.hostPlatform.isAarch64 [
+      gst_all_1.gst-plugins-bad
+    ];
+
+  postInstall = prevAttrs.postInstall or "" + ''
+    moveToOutput '${hostDir}' "''${!outputBin}"
+    moveToOutput '${targetDir}' "''${!outputBin}"
+    moveToOutput 'bin' "''${!outputBin}"
+    wrapQtApp "''${!outputBin}/${hostDir}/nsys-ui.bin"
+  '';
+
+  # lib needs libtiff.so.5, but nixpkgs provides libtiff.so.6
+  preFixup = prevAttrs.preFixup or "" + ''
+    patchelf --replace-needed libtiff.so.5 libtiff.so "''${!outputBin}/${hostDir}/Plugins/imageformats/libqtiff.so"
+  '';
+
+  autoPatchelfIgnoreMissingDeps = prevAttrs.autoPatchelfIgnoreMissingDeps or [ ] ++ [
+    "libnvidia-ml.so.1"
+  ];
+}

pkgs/development/cuda-modules/_cuda/fixups/nvidia_driver.nix

@@ -0,0 +1,5 @@
+_: prevAttrs: {
+  badPlatformsConditions = prevAttrs.badPlatformsConditions or { } // {
+    "Package is not supported; use drivers from linuxPackages" = true;
+  };
+}

pkgs/development/cuda-modules/_cuda/fixups/tensorrt.nix

@@ -0,0 +1,128 @@
+{
+  _cuda,
+  cudaOlder,
+  cudaPackages,
+  cudaMajorMinorVersion,
+  lib,
+  patchelf,
+  requireFile,
+  stdenv,
+}:
+let
+  inherit (lib)
+    attrsets
+    maintainers
+    meta
+    strings
+    versions
+    ;
+  inherit (stdenv) hostPlatform;
+  # targetArch :: String
+  targetArch = attrsets.attrByPath [ hostPlatform.system ] "unsupported" {
+    x86_64-linux = "x86_64-linux-gnu";
+    aarch64-linux = "aarch64-linux-gnu";
+  };
+in
+finalAttrs: prevAttrs: {
+  # Useful for inspecting why something went wrong.
+  brokenConditions =
+    let
+      cudaTooOld = cudaOlder finalAttrs.passthru.featureRelease.minCudaVersion;
+      cudaTooNew =
+        (finalAttrs.passthru.featureRelease.maxCudaVersion != null)
+        && strings.versionOlder finalAttrs.passthru.featureRelease.maxCudaVersion cudaMajorMinorVersion;
+      cudnnVersionIsSpecified = finalAttrs.passthru.featureRelease.cudnnVersion != null;
+      cudnnVersionSpecified = versions.majorMinor finalAttrs.passthru.featureRelease.cudnnVersion;
+      cudnnVersionProvided = versions.majorMinor finalAttrs.passthru.cudnn.version;
+      cudnnTooOld =
+        cudnnVersionIsSpecified && (strings.versionOlder cudnnVersionProvided cudnnVersionSpecified);
+      cudnnTooNew =
+        cudnnVersionIsSpecified && (strings.versionOlder cudnnVersionSpecified cudnnVersionProvided);
+    in
+    prevAttrs.brokenConditions or { }
+    // {
+      "CUDA version is too old" = cudaTooOld;
+      "CUDA version is too new" = cudaTooNew;
+      "CUDNN version is too old" = cudnnTooOld;
+      "CUDNN version is too new" = cudnnTooNew;
+    };
+
+  src = requireFile {
+    name = finalAttrs.passthru.redistribRelease.filename;
+    inherit (finalAttrs.passthru.redistribRelease) hash;
+    message = ''
+      To use the TensorRT derivation, you must join the NVIDIA Developer Program and
+      download the ${finalAttrs.version} TAR package for CUDA ${cudaMajorMinorVersion} from
+      ${finalAttrs.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 ${finalAttrs.passthru.redistribRelease.filename}
+    '';
+  };
+
+  # We need to look inside the extracted output to get the files we need.
+  sourceRoot = "TensorRT-${finalAttrs.version}";
+
+  buildInputs = prevAttrs.buildInputs or [ ] ++ [ (finalAttrs.passthru.cudnn.lib or null) ];
+
+  preInstall =
+    prevAttrs.preInstall or ""
+    + strings.optionalString (targetArch != "unsupported") ''
+      # Replace symlinks to bin and lib with the actual directories from targets.
+      for dir in bin lib; do
+        rm "$dir"
+        mv "targets/${targetArch}/$dir" "$dir"
+      done
+
+      # Remove broken symlinks
+      for dir in include samples; do
+        rm "targets/${targetArch}/$dir" || :
+      done
+    '';
+
+  # Tell autoPatchelf about runtime dependencies.
+  postFixup =
+    let
+      versionTriple = "${versions.majorMinor finalAttrs.version}.${versions.patch finalAttrs.version}";
+    in
+    prevAttrs.postFixup or ""
+    + ''
+      ${meta.getExe' patchelf "patchelf"} --add-needed libnvinfer.so \
+        "$lib/lib/libnvinfer.so.${versionTriple}" \
+        "$lib/lib/libnvinfer_plugin.so.${versionTriple}" \
+        "$lib/lib/libnvinfer_builder_resource.so.${versionTriple}"
+    '';
+
+  passthru = prevAttrs.passthru or { } // {
+    # The CUDNN used with TensorRT.
+    # If null, the default cudnn derivation will be used.
+    # If a version is specified, the cudnn derivation with that version will be used,
+    # unless it is not available, in which case the default cudnn derivation will be used.
+    cudnn =
+      let
+        desiredName = _cuda.lib.mkVersionedName "cudnn" (
+          lib.versions.majorMinor finalAttrs.passthru.featureRelease.cudnnVersion
+        );
+      in
+      if finalAttrs.passthru.featureRelease.cudnnVersion == null || (cudaPackages ? desiredName) then
+        cudaPackages.cudnn
+      else
+        cudaPackages.${desiredName};
+  };
+
+  meta = prevAttrs.meta or { } // {
+    badPlatforms =
+      prevAttrs.meta.badPlatforms or [ ]
+      ++ lib.optionals (targetArch == "unsupported") [ hostPlatform.system ];
+    homepage = "https://developer.nvidia.com/tensorrt";
+    maintainers = prevAttrs.meta.maintainers or [ ] ++ [ maintainers.aidalgol ];
+    teams = prevAttrs.meta.teams or [ ];
+
+    # Building TensorRT on Hydra is impossible because of the non-redistributable
+    # license and because the source needs to be manually downloaded from the
+    # NVIDIA Developer Program (see requireFile above).
+    hydraPlatforms = lib.platforms.none;
+  };
+}