cudaPackages.tensorrt-samples: init

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>
This commit is contained in:
Connor Baker
2026-01-10 03:19:15 +00:00
parent 9c17880452
commit 6b45e4200a
14 changed files with 571 additions and 0 deletions
@@ -0,0 +1,193 @@
{
_cuda,
backendStdenv,
cmake,
cuda_cudart,
cuda_nvcc,
cuda_profiler_api,
cudaNamePrefix,
fetchFromGitHub,
fetchzip,
flags,
lib,
runCommand,
stdenvNoCC,
tensorrt,
writeShellApplication,
}:
let
inherit (_cuda.lib) majorMinorPatch;
inherit (lib)
cmakeBool
cmakeFeature
getAttr
getInclude
licenses
maintainers
optionalString
replaceStrings
teams
;
in
backendStdenv.mkDerivation (finalAttrs: {
__structuredAttrs = true;
strictDeps = true;
name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
pname = "tensorrt-samples";
version = majorMinorPatch tensorrt.version;
src = fetchFromGitHub (
{
owner = "NVIDIA";
repo = "TensorRT";
}
// getAttr finalAttrs.version {
"10.0.0" = {
tag = "v10.0.0";
hash = "sha256-k0FqEURPCtcPgowORHme/lhQ5SN63d0lYQvTvFXS6vw=";
};
"10.0.1" = {
tag = "v10.0.1";
hash = "sha256-lSEw0GM0eW2BHNBq/wTQA8v3aNueE3FT+k9F5nH1OgA=";
};
"10.1.0" = {
tag = "v10.1.0";
hash = "sha256-A3QwrQaI0EgRspgXEKcna/z6p7abOq3M7KDQMPQftvE=";
};
"10.2.0" = {
tag = "v10.2.0";
hash = "sha256-Euo9VD4VTpx8XJV97IMETTAx/YkPGXiNdA39Wjp3UMU=";
};
"10.3.0" = {
tag = "v10.3.0";
hash = "sha256-odSrsfOa8PlbJiPrkvWFm2hHc+ud0qSpLQanx9/M/90=";
};
"10.4.0" = {
tag = "v10.4.0";
hash = "sha256-GAu/VdHrC3UQw9okPexVItLPrRb1m3ZMpCkHNcfzRkE=";
};
"10.5.0" = {
tag = "v10.5.0";
hash = "sha256-No0JKfvi6ETXrnebMX+tAVhz7fuuCwYAp/WNUN73XzY=";
};
"10.6.0" = {
tag = "v10.6.0";
hash = "sha256-nnzicyCjVqpAonIhx3u9yNnoJkZ0XXjJ8oxQH+wfrtE=";
};
"10.7.0" = {
tag = "v10.7.0";
hash = "sha256-sbp61GverIWrHKvJV+oO9TctFTO4WUmH0oInZIwqF/s=";
};
"10.8.0" = {
tag = "v10.8.0";
hash = "sha256-SDlTZf8EQBq8vDCH3YFJCROHbf47RB5WUu4esLNpYuA=";
};
"10.9.0" = {
tag = "v10.9.0";
hash = "sha256-J8K9RjeGIem5ZxXyU+Rne8uBbul54ie6P/Y1In2mQ0g=";
};
"10.10.0" = {
tag = "v10.10.0";
hash = "sha256-/vkGmH+WKXMXsUizGfjBKRHOp5IpS236eUdCQ8tr1u8=";
};
"10.11.0" = {
tag = "v10.11";
hash = "sha256-OXI6mR2X+kF/0EO5RSBnnaGjMKD6AkuQMfl0OMzayxc=";
};
"10.12.0" = {
tag = "v10.12.0";
hash = "sha256-3pFiqDzWcMAk3GfnSOzzInddEfoGX0Fwqb+vEYr9eBw=";
};
"10.13.0" = {
tag = "v10.13.0";
hash = "sha256-hjl9fKFIE8p05/lmius2vuil6evPbNEjTT9yJyC44FI=";
};
"10.13.2" = {
tag = "v10.13.2";
hash = "sha256-1t4TyQKGTVPyPPNA3dlVDoBSHXKGcTms8AUejbvtVfw=";
};
"10.13.3" = {
tag = "v10.13.3";
hash = "sha256-d14R0UmSLT63wlmpCMi9ZvHZjottr8LJfig7EcqxLEY=";
};
"10.14.1" = {
tag = "v10.14";
hash = "sha256-pWvXpXiUriLDYHqro3HWAmO/9wbGznyUrc9qxq/t0/U=";
};
}
);
nativeBuildInputs = [
cmake
cuda_nvcc
];
postPatch = ''
nixLog "patching $PWD/CMakeLists.txt to avoid manually setting CMAKE_CXX_COMPILER"
substituteInPlace "$PWD"/CMakeLists.txt \
--replace-fail \
'find_program(CMAKE_CXX_COMPILER NAMES $ENV{CXX} g++)' \
'# find_program(CMAKE_CXX_COMPILER NAMES $ENV{CXX} g++)'
nixLog "patching $PWD/CMakeLists.txt to use find_package(CUDAToolkit) instead of find_package(CUDA)"
substituteInPlace "$PWD"/CMakeLists.txt \
--replace-fail \
'find_package(CUDA ''${CUDA_VERSION} REQUIRED)' \
'find_package(CUDAToolkit REQUIRED)'
''
# CMakeLists.txt only started using CMAKE_CUDA_ARCHITECTURES in 10.9, and this bug was fixed by 10.12.
+
optionalString
(lib.versionAtLeast finalAttrs.version "10.9" && lib.versionOlder finalAttrs.version "10.12")
''
nixLog "patching $PWD/CMakeLists.txt to fix CMake logic error"
substituteInPlace "$PWD"/CMakeLists.txt \
--replace-fail \
'list(APPEND CMAKE_CUDA_ARCHITECTURES SM)' \
'list(APPEND CMAKE_CUDA_ARCHITECTURES "''${SM}")'
'';
cmakeFlags = [
# Use tensorrt for these components; we only really want the samples.
(cmakeBool "BUILD_PARSERS" false)
(cmakeBool "BUILD_PLUGINS" false)
(cmakeBool "BUILD_SAMPLES" true)
# Build configuration
(cmakeFeature "GPU_ARCHS" (replaceStrings [ ";" ] [ " " ] flags.cmakeCudaArchitecturesString))
];
buildInputs = [
(getInclude cuda_nvcc)
cuda_cudart
cuda_profiler_api
tensorrt
];
passthru = import ./passthru.nix {
inherit
backendStdenv
fetchzip
finalAttrs
lib
runCommand
stdenvNoCC
writeShellApplication
;
};
meta = {
description = "Open Source Software (OSS) components of NVIDIA TensorRT";
homepage = "https://github.com/NVIDIA/TensorRT";
license = licenses.asl20;
platforms = [
"aarch64-linux"
"x86_64-linux"
];
teams = [ teams.cuda ];
maintainers = with maintainers; [ connorbaker ];
};
})
@@ -0,0 +1,85 @@
{
backendStdenv,
fetchzip,
finalAttrs,
lib,
runCommand,
stdenvNoCC,
writeShellApplication,
}:
let
older = lib.versionOlder finalAttrs.version;
atLeast = lib.versionAtLeast finalAttrs.version;
in
{
sample-data =
let
# Releases prior to 10.14.1 don't have any sample data available to them, so just use the 10.14.1 release's
# sample data.
sample-data_10_14_1 = {
url = "https://github.com/NVIDIA/TensorRT/releases/download/v10.14/tensorrt_sample_data_20251106.zip";
hash = "sha256-IA1pH8idtk/7FD1Tf0hKtyP7A5SW/2ugezyBRluG8yk=";
};
in
fetchzip (
if older "10.14.1" then
sample-data_10_14_1
else
lib.getAttr finalAttrs.version {
"10.14.1" = sample-data_10_14_1;
}
);
# TODO(@connorbaker): A number of the tests fail with 10.2:
# API Usage Error (Unable to load library: libnvinfer_builder_resource_win.so.10.2.0:
# libnvinfer_builder_resource_win.so.10.2.0: cannot open shared object file: No such file or directory)
# TODO(@connorbaker): Add tests for trtexec.
testers =
let
mkTester =
name: cmdArgs:
writeShellApplication {
name = finalAttrs.name + "-tester-" + name;
runtimeInputs = [ finalAttrs.finalPackage ];
text = ''
${lib.toShellVar "cmdArgs" cmdArgs}
echo "running ''${cmdArgs[*]@Q}"
"''${cmdArgs[@]}"
'';
};
in
lib.packagesFromDirectoryRecursive {
callPackage =
path: _:
import path {
inherit (finalAttrs.passthru) sample-data;
inherit
atLeast
backendStdenv
finalAttrs
lib
mkTester
older
;
};
directory = ./testers;
};
# Wrap each of the derivations in testers in a runCommand.
tests = lib.mapAttrsRecursiveCond (as: !(lib.isDerivation as)) (
path: drv:
runCommand (lib.replaceString "-tester-" "-test-" drv.name)
{
nativeBuildInputs = [ drv ];
requiredSystemFeatures = [ "cuda" ];
}
''
set -euo pipefail
mkdir -p "$out"
"${lib.getExe drv}" | tee -a "$out/test.log" || {
nixErrorLog "command failed with exit code $?"
exit 1
}
''
) finalAttrs.passthru.testers;
}
@@ -0,0 +1,44 @@
{
backendStdenv,
lib,
mkTester,
older,
sample-data,
...
}:
# While available on both 10.8 and 10.9, NVIDIA removed this sample from their CMakeLists.txt, so
# we cannot build it.
lib.optionalAttrs (older "10.8") (
{
default = mkTester "sample_algorithm_selector" [
"sample_algorithm_selector"
"--datadir=${sample-data.outPath + "/mnist"}"
];
int8 = mkTester "sample_algorithm_selector-int8" [
"sample_algorithm_selector"
"--datadir=${sample-data.outPath + "/mnist"}"
"--int8"
];
fp16 = mkTester "sample_algorithm_selector-fp16" [
"sample_algorithm_selector"
"--datadir=${sample-data.outPath + "/mnist"}"
"--fp16"
];
bf16 = mkTester "sample_algorithm_selector-bf16" [
"sample_algorithm_selector"
"--datadir=${sample-data.outPath + "/mnist"}"
"--bf16"
];
}
# Only Orin has a DLA
// lib.optionalAttrs (lib.elem "8.7" backendStdenv.cudaCapabilities) {
dla = mkTester "sample_algorithm_selector-dla" [
"sample_algorithm_selector"
"--datadir=${sample-data.outPath + "/mnist"}"
"--useDLACore=0"
];
}
)
@@ -0,0 +1,11 @@
{
mkTester,
sample-data,
...
}:
{
default = mkTester "sample_char_rnn" [
"sample_char_rnn"
"--datadir=${sample-data.outPath + "/char-rnn"}"
];
}
@@ -0,0 +1,30 @@
{
mkTester,
sample-data,
...
}:
{
default = mkTester "sample_dynamic_reshape" [
"sample_dynamic_reshape"
"--datadir=${sample-data.outPath + "/mnist"}"
];
# TODO(@connorbaker): Neither the sample data nor the sample sources provide train-images-idx3-ubyte, so we can't run the int8 test.
# int8 = mkTester "sample_dynamic_reshape-int8" [
# "sample_dynamic_reshape"
# "--datadir=${sample-data.outPath + "/mnist"}"
# "--int8"
# ];
fp16 = mkTester "sample_dynamic_reshape-fp16" [
"sample_dynamic_reshape"
"--datadir=${sample-data.outPath + "/mnist"}"
"--fp16"
];
bf16 = mkTester "sample_dynamic_reshape-bf16" [
"sample_dynamic_reshape"
"--datadir=${sample-data.outPath + "/mnist"}"
"--bf16"
];
}
@@ -0,0 +1,11 @@
{
atLeast,
lib,
mkTester,
...
}:
lib.optionalAttrs (atLeast "10.8") {
default = mkTester "sample_editable_timing_cache" [
"sample_editable_timing_cache"
];
}
@@ -0,0 +1,23 @@
{
backendStdenv,
lib,
mkTester,
sample-data,
...
}:
{
default = mkTester "sample_int8_api" [
"sample_int8_api"
"--model=${sample-data.outPath + "/resnet50/ResNet50.onnx"}"
"--data=${sample-data.outPath + "/int8_api"}"
];
}
# Only Orin has a DLA
// lib.optionalAttrs (lib.elem "8.7" backendStdenv.cudaCapabilities) {
dla = mkTester "sample_int8_api-dla" [
"sample_int8_api"
"--model=${sample-data.outPath + "/resnet50/ResNet50.onnx"}"
"--data=${sample-data.outPath + "/int8_api"}"
"--useDLACore=0"
];
}
@@ -0,0 +1,21 @@
{
backendStdenv,
lib,
mkTester,
sample-data,
...
}:
{
default = mkTester "sample_io_formats" [
"sample_io_formats"
"--datadir=${sample-data.outPath + "/mnist"}"
];
}
# Only Orin has a DLA
// lib.optionalAttrs (lib.elem "8.7" backendStdenv.cudaCapabilities) {
dla = mkTester "sample_io_formats-dla" [
"sample_io_formats"
"--datadir=${sample-data.outPath + "/mnist"}"
"--useDLACore=0"
];
}
@@ -0,0 +1,11 @@
{
finalAttrs,
mkTester,
...
}:
{
default = mkTester "sample_named_dimensions" [
"sample_named_dimensions"
"--datadir=${finalAttrs.src.outPath + "/samples/sampleNamedDimensions"}"
];
}
@@ -0,0 +1,27 @@
{
atLeast,
finalAttrs,
lib,
mkTester,
sample-data,
...
}:
{
default = mkTester "sample_non_zero_plugin" [
"sample_non_zero_plugin"
"--datadir=${sample-data.outPath + "/mnist"}"
];
fp16 = mkTester "sample_non_zero_plugin-fp16" [
"sample_non_zero_plugin"
"--datadir=${sample-data.outPath + "/mnist"}"
"--fp16"
];
}
// lib.optionalAttrs (atLeast "10.0.1") {
columnOrder = mkTester "sample_non_zero_plugin-columnOrder" [
"sample_non_zero_plugin"
"--datadir=${sample-data.outPath + "/mnist"}"
"--columnOrder"
];
}
@@ -0,0 +1,39 @@
{
backendStdenv,
lib,
mkTester,
sample-data,
...
}:
{
default = mkTester "sample_onnx_mnist" [
"sample_onnx_mnist"
"--datadir=${sample-data.outPath + "/mnist"}"
];
int8 = mkTester "sample_onnx_mnist-int8" [
"sample_onnx_mnist"
"--datadir=${sample-data.outPath + "/mnist"}"
"--int8"
];
fp16 = mkTester "sample_onnx_mnist-fp16" [
"sample_onnx_mnist"
"--datadir=${sample-data.outPath + "/mnist"}"
"--fp16"
];
bf16 = mkTester "sample_onnx_mnist-bf16" [
"sample_onnx_mnist"
"--datadir=${sample-data.outPath + "/mnist"}"
"--bf16"
];
}
# Only Orin has a DLA
// lib.optionalAttrs (lib.elem "8.7" backendStdenv.cudaCapabilities) {
dla = mkTester "sample_onnx_mnist-dla" [
"sample_onnx_mnist"
"--datadir=${sample-data.outPath + "/mnist"}"
"--useDLACore=0"
];
}
@@ -0,0 +1,36 @@
{
backendStdenv,
lib,
mkTester,
sample-data,
...
}:
# TODO(@connorbaker): Neither the sample data nor the sample sources provide mnist_with_coordconv.onnx, so we can't run these tests.
lib.optionalAttrs false (
{
default = mkTester "sample_onnx_mnist_coord_conv_ac" [
"sample_onnx_mnist_coord_conv_ac"
"--datadir=${sample-data.outPath + "/mnist"}"
];
int8 = mkTester "sample_onnx_mnist_coord_conv_ac-int8" [
"sample_onnx_mnist_coord_conv_ac"
"--datadir=${sample-data.outPath + "/mnist"}"
"--int8"
];
fp16 = mkTester "sample_onnx_mnist_coord_conv_ac-fp16" [
"sample_onnx_mnist_coord_conv_ac"
"--datadir=${sample-data.outPath + "/mnist"}"
"--fp16"
];
}
# Only Orin has a DLA
// lib.optionalAttrs (lib.elem "8.7" backendStdenv.cudaCapabilities) {
dla = mkTester "sample_onnx_mnist_coord_conv_ac-dla" [
"sample_onnx_mnist_coord_conv_ac"
"--datadir=${sample-data.outPath + "/mnist"}"
"--useDLACore=0"
];
}
)
@@ -0,0 +1,36 @@
{
backendStdenv,
lib,
mkTester,
sample-data,
...
}:
# NOTE: Disabled because it generates way too much output.
lib.optionalAttrs false (
{
default = mkTester "sample_progress_monitor" [
"sample_progress_monitor"
"--datadir=${sample-data.outPath + "/mnist"}"
];
int8 = mkTester "sample_progress_monitor-int8" [
"sample_progress_monitor"
"--datadir=${sample-data.outPath + "/mnist"}"
"--int8"
];
fp16 = mkTester "sample_progress_monitor-fp16" [
"sample_progress_monitor"
"--datadir=${sample-data.outPath + "/mnist"}"
"--fp16"
];
}
# Only Orin has a DLA
// lib.optionalAttrs (lib.elem "8.7" backendStdenv.cudaCapabilities) {
dla = mkTester "sample_progress_monitor-dla" [
"sample_progress_monitor"
"--datadir=${sample-data.outPath + "/mnist"}"
"--useDLACore=0"
];
}
)
@@ -0,0 +1,4 @@
{ ... }:
{
}