diff --git a/CHANGELOG.md b/CHANGELOG.md index 13cb4e30..256b382e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,11 +1,27 @@ # TensorRT OSS Release Changelog +## 10.14 GA - 2025-11-7 +- Sample changes + - Replace all pycuda usages with cuda-python APIs + - Removed the efficientnet samples + - Deprecated tensorflow_object_detection and efficientdet samples + - Samples will no longer be released with the packages. The TensorRT GitHub repository will be the single source. + + +- Parsers: + - Added support for the `Attention` operator + - Improved refit for `ConstantOfShape` nodes + +- Demos + - demoDiffusion: + - Added support for the Cosmos-Predict2 text2image and video2world pipelines + + ## 10.13.3 GA - 2025-9-8 - Added support for TensorRT API Capture and Replay feature, see the [developer guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html) for more information. - Demo changes - Added support for Flux Kontext pipeline. - ## 10.13.2 GA - 2025-8-18 - Added support for CUDA 13.0, dropped support for CUDA 11.X - Dropped support for Ubuntu 20.04 @@ -24,7 +40,6 @@ - Added `loadModelProto`, `loadInitializer` and `refitModelProto` APIs for IParserRefitter. These APIs are meant to be used to load user initializers when refitting ONNX models. - Deprecated `IParser::parseWithWeightDescriptors`. - ## 10.12.0 GA - 2025-6-10 - Plugin changes - Migrated `IPluginV2`-descendent version 1 of `cropAndResizeDynamic`, to version 2, which implements `IPluginV3`. diff --git a/CMakeLists.txt b/CMakeLists.txt index b5e12ce2..6030e6c0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -67,6 +67,9 @@ endif() set(CMAKE_SKIP_BUILD_RPATH True) # CUDA targets +set(DEFAULT_CUDA_VERSION 13.0.0) +set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION}) +message(STATUS "CUDA version set to ${CUDA_VERSION}") if (DEFINED GPU_ARCHS) message(STATUS "GPU_ARCHS defined as ${GPU_ARCHS}. Generating CUDA code for SM ${GPU_ARCHS}") @@ -81,6 +84,9 @@ else() list(APPEND CMAKE_CUDA_ARCHITECTURES 100 120) endif() + if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0) + list(APPEND CMAKE_CUDA_ARCHITECTURES 110) + endif() message(STATUS "GPU_ARCHS is not defined. Generating CUDA code for default SMs: ${CMAKE_CUDA_ARCHITECTURES}") endif() set(BERT_GENCODES) @@ -105,6 +111,10 @@ project(TensorRT DESCRIPTION "TensorRT is a C++ library that facilitates high-performance inference on NVIDIA GPUs and deep learning accelerators." HOMEPAGE_URL "https://github.com/NVIDIA/TensorRT") +if (WIN32) + enable_language(C) +endif() + if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT) set(CMAKE_INSTALL_PREFIX ${TRT_LIB_DIR}/../ CACHE PATH "TensorRT installation" FORCE) endif(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT) @@ -113,7 +123,7 @@ option(BUILD_PLUGINS "Build TensorRT plugin" ON) option(BUILD_PARSERS "Build TensorRT parsers" ON) option(BUILD_SAMPLES "Build TensorRT samples" ON) -# C++14 +# C++17 set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_EXTENSIONS OFF) @@ -141,14 +151,10 @@ endif() ############################################################################################ # Dependencies - -set(DEFAULT_CUDA_VERSION 12.2.0) set(DEFAULT_CUDNN_VERSION 8.9) -set(DEFAULT_PROTOBUF_VERSION 3.20.1) +set(DEFAULT_PROTOBUF_VERSION 3.20.3) # Dependency Version Resolution -set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION}) -message(STATUS "CUDA version set to ${CUDA_VERSION}") set_ifndef(CUDNN_VERSION ${DEFAULT_CUDNN_VERSION}) message(STATUS "cuDNN version set to ${CUDNN_VERSION}") set_ifndef(PROTOBUF_VERSION ${DEFAULT_PROTOBUF_VERSION}) @@ -221,16 +227,18 @@ endif() ############################################################################################ # TensorRT +set(HINT_PATHS "${TRT_OUT_DIR}" "${TRT_LIB_DIR}") + if(BUILD_PLUGINS) add_subdirectory(plugin) else() - find_library_create_target(nvinfer_plugin ${nvinfer_plugin_lib_name} SHARED "${TRT_OUT_DIR}" "${TRT_LIB_DIR}") + find_library_create_target(${nvinfer_plugin_lib_name} ${nvinfer_plugin_lib_name} SHARED "${HINT_PATHS}") endif() if(BUILD_PARSERS) add_subdirectory(parsers) else() - find_library_create_target(nvonnxparser ${nvonnxparser_lib_name} SHARED "${TRT_OUT_DIR}" "${TRT_LIB_DIR}") + find_library_create_target(${nvonnxparser_lib_name} ${nvonnxparser_lib_name} SHARED "${HINT_PATHS}") endif() if(BUILD_SAMPLES) diff --git a/README.md b/README.md index 2890afca..0e208840 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ To build the TensorRT-OSS components, you will first need the following software **TensorRT GA build** -- TensorRT v10.13.3.9 +- TensorRT v10.14.1.48 - Available from direct download links listed below **System Packages** @@ -86,24 +86,24 @@ To build the TensorRT-OSS components, you will first need the following software Else download and extract the TensorRT GA build from [NVIDIA Developer Zone](https://developer.nvidia.com) with the direct links below: - - [TensorRT 10.13.3.9 for CUDA 13.0, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz) - - [TensorRT 10.13.3.9 for CUDA 12.9, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz) - - [TensorRT 10.13.3.9 for CUDA 13.0, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/zip/TensorRT-10.13.3.9.Windows.win10.cuda-13.0.zip) - - [TensorRT 10.13.3.9 for CUDA 12.9, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/zip/TensorRT-10.13.3.9.Windows.win10.cuda-12.9.zip) + - [TensorRT 10.14.1.48 for CUDA 13.0, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz) + - [TensorRT 10.14.1.48 for CUDA 12.9, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz) + - [TensorRT 10.14.1.48 for CUDA 13.0, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/zip/TensorRT-10.14.1.48.Windows.win10.cuda-13.0.zip) + - [TensorRT 10.14.1.48 for CUDA 12.9, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/zip/TensorRT-10.14.1.48.Windows.win10.cuda-12.9.zip) **Example: Ubuntu 22.04 on x86-64 with cuda-13.0** ```bash cd ~/Downloads - tar -xvzf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz - export TRT_LIBPATH=`pwd`/TensorRT-10.13.3.9 + tar -xvzf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz + export TRT_LIBPATH=`pwd`/TensorRT-10.14.1.48 ``` **Example: Windows on x86-64 with cuda-12.9** ```powershell - Expand-Archive -Path TensorRT-10.13.3.9.Windows.win10.cuda-12.9.zip - $env:TRT_LIBPATH="$pwd\TensorRT-10.13.3.9\lib" + Expand-Archive -Path TensorRT-10.14.1.48.Windows.win10.cuda-12.9.zip + $env:TRT_LIBPATH="$pwd\TensorRT-10.14.1.48\lib" ``` ## Setting Up The Build Environment @@ -112,16 +112,16 @@ For Linux platforms, we recommend that you generate a docker container for build 1. #### Generate the TensorRT-OSS build container. - **Example: Ubuntu 22.04 on x86-64 with cuda-13.0 (default)** + **Example: Ubuntu 24.04 on x86-64 with cuda-13.0 (default)** ```bash - ./docker/build.sh --file docker/ubuntu-22.04.Dockerfile --tag tensorrt-ubuntu22.04-cuda13.0 + ./docker/build.sh --file docker/ubuntu-24.04.Dockerfile --tag tensorrt-ubuntu24.04-cuda13.0 ``` - **Example: Rockylinux8 on x86-64 with cuda-12.9** + **Example: Rockylinux8 on x86-64 with cuda-13.0** ```bash - ./docker/build.sh --file docker/rockylinux8.Dockerfile --tag tensorrt-rockylinux8-cuda12.9 + ./docker/build.sh --file docker/rockylinux8.Dockerfile --tag tensorrt-rockylinux8-cuda13.0 ``` **Example: Ubuntu 24.04 cross-compile for Jetson (aarch64) with cuda-13.0 (JetPack SDK)** @@ -137,9 +137,9 @@ For Linux platforms, we recommend that you generate a docker container for build ``` 2. #### Launch the TensorRT-OSS build container. - **Example: Ubuntu 22.04 build container** + **Example: Ubuntu 24.04 build container** ```bash - ./docker/launch.sh --tag tensorrt-ubuntu22.04-cuda13.0 --gpus all + ./docker/launch.sh --tag tensorrt-ubuntu24.04-cuda13.0 --gpus all ``` > NOTE: >
1. Use the `--tag` corresponding to build container generated in Step 1. @@ -175,7 +175,7 @@ For Linux platforms, we recommend that you generate a docker container for build ```bash cd $TRT_OSSPATH mkdir -p build && cd build - cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64 -DGPU_ARCHS=110 + cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64 CC=/usr/bin/gcc make -j$(nproc) ``` @@ -186,7 +186,16 @@ For Linux platforms, we recommend that you generate a docker container for build ```bash cd $TRT_OSSPATH mkdir -p build && cd build - cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_cross.toolchain -DGPU_ARCHS=110 + cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_cross.toolchain + make -j$(nproc) + ``` + + **Example: Ubuntu 24.04 Cross-Compile for DriveOS (aarch64) with cuda-13.0** + + ```bash + cd $TRT_OSSPATH + mkdir -p build && cd build + cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_dos_cross.toolchain make -j$(nproc) ``` @@ -196,7 +205,7 @@ For Linux platforms, we recommend that you generate a docker container for build cd $TRT_OSSPATH mkdir -p build cd -p build - cmake .. -DTRT_LIB_DIR="$env:TRT_LIBPATH" -DCUDNN_ROOT_DIR="$env:CUDNN_PATH" -DTRT_OUT_DIR="$pwd\\out" + cmake .. -DTRT_LIB_DIR="$env:TRT_LIBPATH" -DTRT_OUT_DIR="$pwd\\out" msbuild TensorRT.sln /property:Configuration=Release -m:$env:NUMBER_OF_PROCESSORS ``` diff --git a/VERSION b/VERSION index a863225f..553dcf2d 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -10.13.3.9 +10.14.1.48 diff --git a/cmake/modules/BundleLibraries.cmake b/cmake/modules/BundleLibraries.cmake new file mode 100644 index 00000000..f982d321 --- /dev/null +++ b/cmake/modules/BundleLibraries.cmake @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# This file handles logic relating to building "fat" static libraries, that is, ones which contain other static libraries. +# Idiomatically, CMake would prefer we ship all static libraries as part of our release and ship a file which specifies the linkages. +# However, for various reasons, this is both inadvisable and annoying. Instead, we would prefer to ship one static library containing all dependencies. +# +# To do that, we need rules to bundle static libraries into other static libraries. Hence, this class. + +define_property(TARGET + PROPERTY BUNDLED_LIBRARY_TEMPLATE_PATH + BRIEF_DOCS "File path to the template script which will be written to when calling target_bundle_libraries." +) + +define_property(TARGET + PROPERTY BUNDLE_LIBRARIES + BRIEF_DOCS "The list of libraries that have been bundled into this target by target_bundle_libraries." +) + +define_property(TARGET + PROPERTY BUNDLE_LIBRARY_KNOWN_TYPE + BRIEF_DOCS "Fallback type used when a target's TYPE is UNKNOWN_LIBRARY." +) + +# Internal helper to prefix all messages with "[target_bundle_libraries]: ". +# +# \param mode The message mode to be passed to message(...) +# \param argn The message contents. +macro(__bundleMessage mode) + message(${mode} "[target_bundle_libraries]: " ${ARGN}) +endmacro() + +# Illegal genex magicâ„¢ +# Given a string containing a generator expression (var), edits the variable in-place to escape any generator expression literals. +# The escaped generator expression is then suitable for use within a LIST:TRANSFORM replacement block. +# This more or less allows for mapping lists to generator expressions, which can be recursively evaluated. +function(escape_generator_expression var) + string(REPLACE ">" "__ANGLE_R__" ${var} "${${var}}") + string(REPLACE "$" "$<1:$>" ${var} "${${var}}") + string(REPLACE "," "$" ${var} "${${var}}") + string(REPLACE "__ANGLE_R__" "$" ${var} "${${var}}") + return(PROPAGATE ${var}) +endfunction() + +# Recursively unwraps alias targets until finding the real target. Non-targets are returned verbatim. +# We need to ensure that the generated .mri script contains a deduplicated list of bundled targets +# so we need to resolve aliases, otherwise we may end up with multiple entries for the same target. +# +# \param target_name The name of the target to unwrap. +# \param result_var The variable to store the unwrapped target name in. +function(unwrapAlias target_name result_var) + # First, try to unwrap common generator expressions that may wrap the target name. + string(REGEX MATCH "\\$" _ ${target_name}) + if(TARGET ${CMAKE_MATCH_1}) + set(target_name ${CMAKE_MATCH_1}) + endif() + + string(REGEX MATCH "\\$" _ ${target_name}) + if(TARGET ${CMAKE_MATCH_1}) + set(target_name ${CMAKE_MATCH_1}) + endif() + + if(TARGET ${target_name}) + get_target_property(aliased_target ${target_name} ALIASED_TARGET) + if(aliased_target) + # Recursively unwrap in case there are multiple levels + unwrapAlias(${aliased_target} unwrapped) + set(${result_var} ${unwrapped} PARENT_SCOPE) + else() + # Not an alias, return the original name + set(${result_var} ${target_name} PARENT_SCOPE) + endif() + else() + # Not a target at all, return the original name + set(${result_var} ${target_name} PARENT_SCOPE) + endif() +endfunction() + +# Internal function to retrieve the type of library for a given target. +# This function will fallback to the value of BUNDLE_LIBRARY_KNOWN_TYPE if it encounters an UNKNOWN_LIBRARY. +# +# \param lib A target to evaluate the type for. +# \param outVar The output variable to store the type name in. +function(__get_lib_type lib outVar) + get_target_property(libType ${lib} TYPE) + if (${libType} STREQUAL UNKNOWN_LIBRARY) + get_target_property(knownType ${lib} BUNDLE_LIBRARY_KNOWN_TYPE) + if (NOT ${knownType} STREQUAL "knownType-NOTFOUND") + set(libType ${knownType}) + endif() + __bundleMessage(DEBUG "Using known type of unknown library ${lib}: ${knownType}") + endif() + set(${outVar} ${libType} PARENT_SCOPE) +endfunction() + +# This is an internal-only function called the first time that target_bundle_libraries is called. +# It "registers" a target for bundling by creating the base template file and populating the target property BUNDLED_LIBRARY_TEMPLATE_PATH. +# Additionally, it registers the file generation logic for making the "final" script, as well as the custom command for running it after the build. +# +# \param lib The mainLib from target_bundle_libraries. +# \param templatePath The file path to the template file that will be created. +function(__registerTargetForBundling lib templatePath) + if(MSVC) + set(scriptPath $/archive-${lib}.bat) + + set(template "/OUT:\"$\" \"$\"\n") + + # Windows-syntax version of the same logic from the linux build below. + # Main differences is that windows does not have `addlib`, and uses `\n` instead of `\n`. Additionally, the values need to be quoted to account for spaces in file paths. + set(replaceExpr "\"$,$,\\1>\"") + escape_generator_expression(replaceExpr) + string(APPEND template "$,REPLACE,(.+),${replaceExpr}>,\n>>\n") + + file(WRITE ${templatePath} ${template}) + file(GENERATE + OUTPUT ${scriptPath} + INPUT ${templatePath} + ) + add_custom_command(TARGET ${lib} POST_BUILD + COMMAND ${CMAKE_AR} /NOLOGO @\"${scriptPath}\" + COMMAND ${CMAKE_COMMAND} -E echo "Bundled $> static libraries into target ${lib}. Script: ${scriptPath}" + WORKING_DIRECTORY $ + ) + else() + set(scriptPath $/archive-${lib}.mri) + + set(template "create $\n") + string(APPEND template "addlib $\n") + # Expand BUNDLE_LIBRARIES into the appropriate chain of addlib commands needed. + # BUNDLE_LIBRARIES will contain either (a) target names or (b) absolute file paths to libraries to include. + # This first part will disambiguate between (a) and (b) by evaluating (a) to `addlib $` and (b) to `addlib [[filepath]]` + set(replaceExpr "addlib $,$,\\1>") + escape_generator_expression(replaceExpr) + + # The second part maps every element in BUNDLE_LIBRARIES to `replaceExpr` and evaluates the resulting replacement, which produces the final .mri file. + string(APPEND template "$,REPLACE,(.+),${replaceExpr}>,\n>>\n") + string(APPEND template "save\n") + string(APPEND template "end\n") + + file(WRITE ${templatePath} ${template}) + file(GENERATE + OUTPUT ${scriptPath} + INPUT ${templatePath} + ) + + add_custom_command(TARGET ${lib} POST_BUILD + COMMAND ${CMAKE_AR} -M < ${scriptPath} + COMMAND ${CMAKE_RANLIB} $ + COMMAND ${CMAKE_COMMAND} -E echo "Bundled $> static libraries into target ${lib}. Script: ${scriptPath}" + WORKING_DIRECTORY $ + ) + endif() + + set_target_properties(${lib} + PROPERTIES BUNDLED_LIBRARY_TEMPLATE_PATH ${templatePath} + ) +endfunction() + +# Subcomponent of target_bundle_libraries which is responsible for walking the provided depLibs and recursively calling target_bundle_libraries. +# This macro must only be used within target_bundle_libraries. +# +# \param mainLib The current main library from target_bundle_libraries +# \param linkVis The link visibility from target_bundle_libraries +# \param argn The dependencies to walk. Usually the INTERFACE_LINK_LIBRARIES of a target currently being bundled. +function(__bundleRecursiveDeps mainLib linkVis) + if(ARGN) + get_target_property(bundledLibs ${mainLib} BUNDLE_LIBRARIES) + + foreach(dep IN LISTS ARGN) + unwrapAlias(${dep} dep) + if(${dep} IN_LIST bundledLibs) + continue() # Skip bundling of already-bundled libs to avoid many invocations of the same warnings. + endif() + + if(TARGET ${dep}) + __get_lib_type(${dep} depType) + if (${depType} STREQUAL STATIC_LIBRARY) + target_bundle_libraries(${mainLib} ${linkVis} ${dep}) + elseif(${depType} STREQUAL INTERFACE_LIBRARY) + # For interface libraries, we want to add all of the static libraries they may be pointing to, without the library itself (since it is not a static). + get_target_property(interfaceLibs ${dep} INTERFACE_LINK_LIBRARIES) + __bundleRecursiveDeps(${mainLib} ${linkVis} ${interfaceLibs}) + elseif(${depType} STREQUAL SHARED_LIBRARY) + # Skip SO's, since static libraries at the SO-boundary should not be bundled into the target mainLib. + else() + __bundleMessage(DEBUG "Skipping unhandled dependency ${dep} (a dependency of ${bundledLib}) with type ${depType} in ${mainLib}") + endif() + else() + __bundleMessage(DEBUG "Failed to recursively bundle dependency ${dep} (a dependency of ${bundledLib}) in ${mainLib}") + endif() + endforEach() + endif() +endfunction() + +# This function acts as a replacement target_link_libraries, instead linking +# one or more "bundled" static libraries into a "main" static library. +# The bundled libs are embedded inside the main lib during the post-build step. +# CMake interface properties for definitions, etc. are propagated using the specified visibility. +# +# This function is recursive. Any static dependencies of any bundled library will also be bundled into the mainLib. +# +# \param mainLib The main library that other libraries are to be bundled into. +# \param linkVis The link visiblity for interface properties. One of PRIVATE or PUBLIC. +# Using PRIVATE hides all interface properties of bundled libraries from consumers of this library. +# Using PUBLIC will share interface properties with dependents. +# \param argn Variadic arguments - The list of targets to be linked into the mainLib. +function(target_bundle_libraries mainLib linkVis) + if (NOT ${linkVis} STREQUAL PUBLIC AND NOT ${linkVis} STREQUAL PRIVATE) + __bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries with unknown visibility ${linkVis}. Must be either \"PUBLIC\" or \"PRIVATE\".") + endif() + + # TODO: Allow direct insertion of absolute file paths when CMAKE_LINK_LIBRARIES_ONLY_TARGETS is off, instead of always forcing targets. + if (NOT TARGET ${mainLib}) + __bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries for target ${mainLib}, but no target with that name is known.") + endif() + + __get_lib_type(${mainLib} mainLibType) + if (NOT mainLibType STREQUAL STATIC_LIBRARY) + __bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries for target ${mainLib}, but it is not a static library.") + endif() + + get_target_property(isMainLibImported ${mainLib} IMPORTED) + if(isMainLibImported) + __bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries for target ${mainLib}, but it is an imported target.") + endif() + + if(NOT ARGN) + __bundleMessage(WARNING "Called target_bundle_libraries with no libraries for target ${mainLib}") + endif() + + __bundleMessage(DEBUG "Bundling ${ARGN} into ${mainLib}") + + get_target_property(templatePath ${mainLib} BUNDLED_LIBRARY_TEMPLATE_PATH) + if(NOT EXISTS ${templatePath}) + if(MSVC) + set(templatePath ${PROJECT_BINARY_DIR}/archive-${mainLib}.bat.template) + else() + set(templatePath ${PROJECT_BINARY_DIR}/archive-${mainLib}.mri.template) + endif() + __bundleMessage(STATUS "Registering default script path ${templatePath} for target ${mainLib}") + __registerTargetForBundling(${mainLib} ${templatePath}) + endif() + + get_target_property(bundledLibs ${mainLib} BUNDLE_LIBRARIES) + if (NOT bundledLibs) + set(bundledLibs "") + endif() + + foreach(bundledLib IN LISTS ARGN) + unwrapAlias(${bundledLib} bundledLib) + __get_lib_type(${bundledLib} bundledLibType) + + if(${bundledLibType} STREQUAL INTERFACE_LIBRARY) + # Interface libraries are not bundled, since they do not contain any static libraries. + # Their dependencies will be bundled recursively in the next step. + continue() + endif() + + if (NOT ${bundledLibType} STREQUAL STATIC_LIBRARY AND NOT ${bundledLibType} STREQUAL OBJECT_LIBRARY) + __bundleMessage(FATAL_ERROR "Attempted to bundle ${bundledLibType} library ${bundledLib} into target ${mainLib} (only static and object libraries may be bundled)") + endif() + + if (${bundledLibType} STREQUAL STATIC_LIBRARY) + list(APPEND bundledLibs ${bundledLib}) + else() + # Exclude object libs from the BUNDLE_LIBRARIES property as they get added into the static lib using normal means. + endif() + endforeach() + + list(REMOVE_DUPLICATES bundledLibs) + set_target_properties(${mainLib} + PROPERTIES BUNDLE_LIBRARIES "${bundledLibs}" + ) + + foreach(bundledLib IN LISTS ARGN) + unwrapAlias(${bundledLib} bundledLib) + __get_lib_type(${bundledLib} bundledLibType) + + # Recursively bundle all static dependencies of each lib to be bundled. + get_target_property(depLibs ${bundledLib} LINK_LIBRARIES) + __bundleRecursiveDeps(${mainLib} ${linkVis} ${depLibs}) + + # Since we want the main library to be linkable standalone, we need both the LINK_LIBRARIES and INTERFACE_LINK_LIBRARIES bundled in. + # Otherwise, private static dependencies may be lost. + get_target_property(depLibs ${bundledLib} INTERFACE_LINK_LIBRARIES) + __bundleRecursiveDeps(${mainLib} ${linkVis} ${depLibs}) + + # Use `target_link_libraries` to propagate INTERFACE definitions from bundled libs + # BUILD_LOCAL_INTERFACE prevents clients from seeing this internal link relationship + # COMPILE_ONLY prevents the bundled lib from appearing in the link command redundantly + if (${bundledLibType} STREQUAL STATIC_LIBRARY OR ${bundledLibType} STREQUAL INTERFACE_LIBRARY) + target_link_libraries(${mainLib} ${linkVis} + $> + ) + else() + # Include Object Libraries as full libraries, since they do not get added by the bundling stage. + # To do this without breaking the link dependency logic, we need to steal the target objects and link the lib as local + compile only. + target_link_libraries(${mainLib} ${linkVis} + $> + ) + target_sources(${mainLib} PRIVATE $) + endif() + + # require that bundled libs are built before we try to bundle them + add_dependencies(${mainLib} ${bundledLib}) + endforeach() + + if(NOT EXISTS ${templatePath}) + __bundleMessage(FATAL_ERROR "Template file ${templatePath} for target ${mainLib} does not exist.") + endif() +endfunction() diff --git a/cmake/modules/ImportDL.cmake b/cmake/modules/ImportDL.cmake new file mode 100644 index 00000000..aef16c43 --- /dev/null +++ b/cmake/modules/ImportDL.cmake @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +include_guard() + +if(NOT TARGET dl) + # libdl is included in the system library on Windows and QNX. + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + find_library(DL_LIB_PATH + NAMES ${CMAKE_DL_LIBS} + REQUIRED + ) + + message(STATUS "Creating imported target 'dl' for ${DL_LIB_PATH}") + add_library(dl SHARED IMPORTED) + set_target_properties(dl PROPERTIES IMPORTED_LOCATION "${DL_LIB_PATH}") + else() + message(STATUS "Creating no-op target 'dl' since libdl is not available on this platform.") + add_library(dl INTERFACE) # Add a fake dl target so we can still call target_link_libraries without error, even though it's a no-op. + endif() +endif() diff --git a/cmake/modules/InstallUtils.cmake b/cmake/modules/InstallUtils.cmake index b6062e8a..7ca3552f 100644 --- a/cmake/modules/InstallUtils.cmake +++ b/cmake/modules/InstallUtils.cmake @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + include_guard() include(GNUInstallDirs) @@ -26,7 +27,7 @@ include(GNUInstallDirs) function(installLibraries) cmake_parse_arguments( ARG # Prefix for parsed args - "OPTIONAL" # Options (flags) + "OPTIONAL;RUNTIME_ONLY" # Options (flags) "COMPONENT" # Single value args "TARGETS;CONFIGURATIONS" # Multi-value args ${ARGN} @@ -50,12 +51,22 @@ function(installLibraries) set(optional_arg OPTIONAL) endif() + # When RUNTIME_ONLY is passed, we only want to install .dll files. + # Instead of also installing the import library (.lib) files. + # This is only relevant on Windows since Linux doesn't have this distinction. + if(ARG_RUNTIME_ONLY AND WIN32) + set(runtime_only_arg + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) + endif() + # Install the libraries install( TARGETS ${ARG_TARGETS} ${optional_arg} ${component_arg} ${config_arg} + ${runtime_only_arg} ) # Install PDB files for MSVC builds diff --git a/cmake/modules/Platforms.cmake b/cmake/modules/Platforms.cmake new file mode 100644 index 00000000..d5ff7082 --- /dev/null +++ b/cmake/modules/Platforms.cmake @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# Contains constants for the various platform names TRT supports. + +set(TRT_PLATFORM_X86 + "x86_64" + CACHE INTERNAL "Linux") +set(TRT_PLATFORM_AARCH64 + "aarch64" + CACHE INTERNAL "ARM Linux") +set(TRT_PLATFORM_QNX + "qnx" + CACHE INTERNAL "QNX") +set(TRT_PLATFORM_QNX_SAFE + "qnx-safe" + CACHE INTERNAL "QNX Safe") +set(TRT_PLATFORM_WIN10 + "win10" + CACHE INTERNAL "Windows 10") + + +# Checks if the current build platform matches any of the passed (ARGN) platforms. +# +# \param outVar The output variable name. +# \param argn The list of platforms to check against. +# \returns TRUE if TRT_BUILD_PLATFORM matches any of the platforms, FALSE otherwise. +function(checkPlatform outVar) + if(NOT DEFINED TRT_BUILD_PLATFORM) + message(FATAL_ERROR "checkPlatform was called before TRT_BUILD_PLATFORM was defined!") + endif() + + set(isPlatform FALSE) + foreach(platform IN LISTS ARGN) + if(${platform} STREQUAL ${TRT_BUILD_PLATFORM}) + set(isPlatform TRUE) + endif() + endforeach() + + set(${outVar} ${isPlatform} PARENT_SCOPE) +endfunction() diff --git a/cmake/modules/ShouldCompileKernel.cmake b/cmake/modules/ShouldCompileKernel.cmake index 12eadc02..7845c5a3 100644 --- a/cmake/modules/ShouldCompileKernel.cmake +++ b/cmake/modules/ShouldCompileKernel.cmake @@ -35,6 +35,43 @@ function(get_all_numeric_sms OUT_VAR) set(${OUT_VAR} ${ALL_NUMERIC_SMS} PARENT_SCOPE) endfunction() +# \brief Converts the list returned by get_all_numeric_sms into a list of arch values. +# \returns the list in the name specified by OUT_VAR for native platform and OUT_VAR_CROSS for cross OS support. e.g. ptx, sm75, sm80, sm86, sm89, sm100, sm120. +function(get_all_fatbin_archs OUT_VAR OUT_VAR_CROSS) + # Use get_all_numeric_sms to get SM values and convert them to sm-prefixed format + set(ARCH_LIST "") + set(ARCH_LIST_CROSS "") + get_all_numeric_sms(NUMERIC_SMS) + foreach(SM IN LISTS NUMERIC_SMS) + list(APPEND ARCH_LIST "sm${SM}") + endforeach() + + # Note: sm89 it is missing in NUMERIC_SMS since TRT treats sm89 as sm86. + # We should add sm89 to the list to generate the builder resource for sm89. + # If only sm86 is in the list, it means this build only supports sm86, + # so no need to add sm89. + list(FIND ARCH_LIST "sm86" SM86_INDEX) + list(FIND ARCH_LIST "sm89" SM89_INDEX) + list(LENGTH ARCH_LIST ARCH_LIST_COUNT) + if(${SM86_INDEX} GREATER_EQUAL 0 AND ${SM89_INDEX} EQUAL -1 AND ${ARCH_LIST_COUNT} GREATER 1) + list(APPEND ARCH_LIST "sm89") + endif() + + + # There is also a klib which only contains PTX code. + list(APPEND ARCH_LIST "ptx") + + set(ARCH_LIST_CROSS ${ARCH_LIST}) + # Cask5 does not include sm100 cubins. Exclude sm100 for both + # cross-OS support and native Windows build. + list(FILTER ARCH_LIST_CROSS EXCLUDE REGEX "sm100") + if(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_WIN10}) + list(FILTER ARCH_LIST EXCLUDE REGEX "sm100") + endif() + set(${OUT_VAR} ${ARCH_LIST} PARENT_SCOPE) + set(${OUT_VAR_CROSS} ${ARCH_LIST_CROSS} PARENT_SCOPE) +endfunction() + # Certain cubins are binary compatible between different SM versions, so they are reused. # This function checks if a SM-named file should be compiled based on current SM enablement. # Specifically, the SM80 files are compiled if either 80, 86, or 89 are enabled. diff --git a/cmake/modules/find_library_create_target.cmake b/cmake/modules/find_library_create_target.cmake index 998568e2..fa39e25b 100644 --- a/cmake/modules/find_library_create_target.cmake +++ b/cmake/modules/find_library_create_target.cmake @@ -31,6 +31,10 @@ macro(find_library_create_target target_name lib libtype hints) find_library(${lib}_LIB_PATH ${lib}) message(STATUS "Library that was found ${${lib}_LIB_PATH}") add_library(${target_name} ${libtype} IMPORTED) - set_property(TARGET ${target_name} PROPERTY IMPORTED_LOCATION ${${lib}_LIB_PATH}) + if(MSVC) + set_property(TARGET ${target_name} PROPERTY IMPORTED_IMPLIB ${${lib}_LIB_PATH}) + else() + set_property(TARGET ${target_name} PROPERTY IMPORTED_LOCATION ${${lib}_LIB_PATH}) + endif() message(STATUS "==========================================================================================") endmacro() diff --git a/cmake/toolchains/cmake_aarch64_dos_cross.toolchain b/cmake/toolchains/cmake_aarch64_dos_cross.toolchain new file mode 100644 index 00000000..0b49bcc7 --- /dev/null +++ b/cmake/toolchains/cmake_aarch64_dos_cross.toolchain @@ -0,0 +1,60 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +set(CMAKE_SYSTEM_NAME Linux) +set(CMAKE_SYSTEM_PROCESSOR aarch64) + +set(CMAKE_FIND_ROOT_PATH /usr/aarch64-linux-gnu/) + +set(TRT_PLATFORM_ID "aarch64") + +set(CUDA_PLATFORM_ID "aarch64-linux") + +set(CMAKE_C_COMPILER /usr/bin/aarch64-linux-gnu-gcc) +set(CMAKE_CXX_COMPILER /usr/bin/aarch64-linux-gnu-g++) +set(CMAKE_C_COMPILER_ID "GNU") +set(CMAKE_CXX_COMPILER_ID "GNU") +set(CMAKE_CXX_COMPILE_FEATURES cxx_std_17) + +set(CMAKE_C_FLAGS "" CACHE STRING "" FORCE) +set(CMAKE_CXX_FLAGS "" CACHE STRING "" FORCE) + +set(CMAKE_C_COMPILER_TARGET aarch64-linux-gnu) +set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu) + +set(CMAKE_C_COMPILER_FORCED TRUE) +set(CMAKE_CXX_COMPILER_FORCED TRUE) + +set(CUDA_ROOT /usr/local/cuda/targets/${CUDA_PLATFORM_ID} CACHE STRING "CUDA ROOT dir") + +set(CUDNN_LIB /usr/lib/aarch64-linux-gnu/libcudnn.so) + +set(BUILD_LIBRARY_ONLY 1) + +set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT}) +set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include) + +set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) +set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE) +set(CMAKE_CUDA_FLAGS "-I${CUDA_INCLUDE_DIRS} -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE) +set(CMAKE_CUDA_COMPILER_FORCED TRUE) + +set(CUDA_LIBS -L${CUDA_ROOT}/lib) + +set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart -lstdc++ -lm) + +link_directories(${CUDA_ROOT}/lib) diff --git a/demo/BERT/README.md b/demo/BERT/README.md index eee0fe24..95a36a54 100755 --- a/demo/BERT/README.md +++ b/demo/BERT/README.md @@ -86,7 +86,7 @@ This demo BERT application can be run within the TensorRT OSS build container. I - [NGC CLI](https://ngc.nvidia.com/setup/installers/cli) - for downloading BERT checkpoints from NGC. - PyPI Packages: - - [pycuda](https://pypi.org/project/pycuda/) (tested v2019.1.2) + - [cuda-python](https://pypi.org/project/cuda-python/) (tested v13.0.1) - [onnx](https://pypi.org/project/onnx) (tested v1.12.0) - [tensorflow](https://pypi.org/project/tensorflow/) (tested v2.9.1) - [torch](https://pypi.org/project/torch/) (tested v1.11.0) @@ -162,7 +162,7 @@ Completing these steps should resolve the error you encountered and allow the co 6. Run inference with CUDA Graph support. - A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs, not pyCUDA. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option. + A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option. ```bash mkdir -p build; pushd build diff --git a/demo/BERT/builder.py b/demo/BERT/builder.py index bd5b2632..886a4a9c 100755 --- a/demo/BERT/builder.py +++ b/demo/BERT/builder.py @@ -26,8 +26,7 @@ import re import sys import time import onnx -import pycuda.autoinit - +from helpers.cuda_utils import getComputeCapacity # TensorRT import tensorrt as trt from helpers.calibrator import BertCalibrator as BertCalibrator @@ -669,7 +668,7 @@ def main(): args.batch_size = args.batch_size or [1] args.sequence_length = args.sequence_length or [128] - cc = pycuda.autoinit.device.compute_capability() + cc = getComputeCapacity() if cc[0] * 10 + cc[1] < 75 and args.force_int8_multihead: raise RuntimeError("--force-int8-multihead option is only supported on Turing+ GPU.") if cc[0] * 10 + cc[1] < 72 and args.force_int8_skipln: diff --git a/demo/BERT/builder_varseqlen.py b/demo/BERT/builder_varseqlen.py index b0ec3ef8..986e8b21 100755 --- a/demo/BERT/builder_varseqlen.py +++ b/demo/BERT/builder_varseqlen.py @@ -26,8 +26,7 @@ import re import sys import time import onnx -import pycuda.autoinit - +from helpers.cuda_utils import getComputeCapacity # TensorRT import tensorrt as trt from builder_utils import load_tf_weights, load_pytorch_weights_and_quant, load_onnx_weights_and_quant, load_megatron_pickle_weights @@ -352,7 +351,6 @@ def bert_model(config, init_dict, network, input_tensor, residual, mask_idx, cu_ squad_logits_out.name = "logits_out" network.mark_output(squad_logits_out) - def squad_output(prefix, config, init_dict, network, input_tensor): """ Create the squad output @@ -645,7 +643,7 @@ def main(): if args.verbose: TRT_LOGGER.min_severity = TRT_LOGGER.VERBOSE - cc = pycuda.autoinit.device.compute_capability() + cc = getComputeCapacity() if cc[0] * 10 + cc[1] < 72: raise RuntimeError("This variable-length BERT demo only support Xavier+ GPU.") diff --git a/demo/BERT/helpers/calibrator.py b/demo/BERT/helpers/calibrator.py index 09e6014b..d5291f1c 100644 --- a/demo/BERT/helpers/calibrator.py +++ b/demo/BERT/helpers/calibrator.py @@ -19,8 +19,8 @@ import tensorrt as trt import os -import pycuda.driver as cuda -import pycuda.autoinit +from helpers.cuda_utils import cuda_call, memcpy_host_to_device +from cuda.bindings import driver as cuda, runtime as cudart import numpy as np import helpers.tokenization as tokenization import helpers.data_processing as dp @@ -44,11 +44,12 @@ class BertCalibrator(trt.IInt8LegacyCalibrator): self.max_query_length = 64 # Allocate enough memory for a whole batch. - self.device_inputs = [cuda.mem_alloc(self.max_seq_length * trt.int32.itemsize * self.batch_size) for binding in range(3)] + self.device_inputs = [cuda_call(cudart.cudaMalloc(self.max_seq_length * trt.int32.itemsize * self.batch_size)) for binding in range(3)] def free(self): for dinput in self.device_inputs: - dinput.free() + # dinput is a device pointer (int) returned by cudaMalloc + cuda_call(cudart.cudaFree(dinput)) def get_batch_size(self): return self.batch_size @@ -80,9 +81,9 @@ class BertCalibrator(trt.IInt8LegacyCalibrator): segment_ids = features[0].segment_ids input_mask = features[0].input_mask - cuda.memcpy_htod(self.device_inputs[0], input_ids.ravel()) - cuda.memcpy_htod(self.device_inputs[1], segment_ids.ravel()) - cuda.memcpy_htod(self.device_inputs[2], input_mask.ravel()) + memcpy_host_to_device(self.device_inputs[0], input_ids.ravel()) + memcpy_host_to_device(self.device_inputs[1], segment_ids.ravel()) + memcpy_host_to_device(self.device_inputs[2], input_mask.ravel()) self.current_index += self.batch_size return self.device_inputs diff --git a/demo/BERT/helpers/cuda_utils.py b/demo/BERT/helpers/cuda_utils.py new file mode 100644 index 00000000..acab7ca3 --- /dev/null +++ b/demo/BERT/helpers/cuda_utils.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import logging +from cuda.bindings import driver as cuda, runtime as cudart + +class CudaStreamContext: + """CUDA stream lifecycle management with context manager support""" + def __init__(self): + """Initialize CUDA stream""" + self._stream = cuda_call(cudart.cudaStreamCreate()) + + def __enter__(self): + """Create CUDA stream when entering context (if not already created)""" + if self._stream is None: + self._stream = cuda_call(cudart.cudaStreamCreate()) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Destroy CUDA stream when exiting context""" + self.free() + + @property + def stream(self) -> cudart.cudaStream_t: + if self._stream is None: + raise RuntimeError("Stream not created. Use 'with' statement.") + return self._stream + + def synchronize(self): + """Synchronize the stream""" + if self._stream is None: + raise RuntimeError("Stream not created. Use 'with' statement.") + cuda_call(cudart.cudaStreamSynchronize(self._stream)) + + def free(self): + """Explicitly free the CUDA stream""" + if self._stream is not None: + try: + cuda_call(cudart.cudaStreamDestroy(self._stream)) + self._stream = None + except Exception as e: + logging.warning(f"Failed to destroy CUDA stream: {e}") + + def __del__(self): + """Cleanup stream on destruction""" + if hasattr(self, '_stream') and self._stream is not None: + self.free() + + def __str__(self): + return f"CudaStreamContext: {self._stream}" + + def __repr__(self): + return self.__str__() + +def cuda_call(call): + """Helper function to make CUDA calls and check for errors""" + def _cudaGetErrorEnum(error): + if isinstance(error, cuda.CUresult): + err, name = cuda.cuGetErrorName(error) + return name if err == cuda.CUresult.CUDA_SUCCESS else "" + elif isinstance(error, cudart.cudaError_t): + return cudart.cudaGetErrorName(error)[1] + else: + raise RuntimeError("Unknown error type: {}".format(error)) + + err, res = call[0], call[1:] + if err.value: + raise RuntimeError( + "CUDA error code={}({})".format( + err.value, _cudaGetErrorEnum(err) + ) + ) + if len(res) == 1: + return res[0] + elif len(res) == 0: + return None + else: + return res + +def getComputeCapacity(devID=0): + major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID)) + minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID)) + return (major, minor) + +def memcpy_host_to_device_async(device_ptr: int, host_arr: np.ndarray, stream): + """Wrapper for async host-to-device memory copy""" + cuda_call(cudart.cudaMemcpyAsync(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)) + + +def memcpy_device_to_host_async(host_arr: np.ndarray, device_ptr: int, stream): + """Wrapper for async device-to-host memory copy""" + cuda_call(cudart.cudaMemcpyAsync(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)) + +def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray): + """Wrapper for synchronous host-to-device memory copy""" + cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)) + + +def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int): + """Wrapper for synchronous device-to-host memory copy""" + cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) + +# Initialize CUDA +cuda_call(cudart.cudaFree(0)) diff --git a/demo/BERT/inference.ipynb b/demo/BERT/inference.ipynb index 2882e0b6..3b72f3e7 100644 --- a/demo/BERT/inference.ipynb +++ b/demo/BERT/inference.ipynb @@ -221,8 +221,8 @@ "metadata": {}, "outputs": [], "source": [ - "import pycuda.driver as cuda\n", - "import pycuda.autoinit\n", + "from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host\n", + "from cuda.bindings import driver as cuda, runtime as cudart\n", "import collections\n", "import numpy as np\n", "import time\n", @@ -238,9 +238,7 @@ " input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n", " \n", " # Allocate device memory for inputs.\n", - " d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)]\n", - " # Create a stream in which to copy inputs/outputs and run inference.\n", - " stream = cuda.Stream()\n", + " d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]\n", "\n", " # Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n", " # Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n", @@ -250,57 +248,59 @@ " assert context.all_binding_shapes_specified\n", "\n", " # Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n", - " h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n", - " d_output = cuda.mem_alloc(h_output.nbytes)\n", + " h_output = np.empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n", + " cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))\n", + " d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))\n", "\n", - " print(\"\\nRunning Inference...\")\n", + " with CudaStreamContext() as stream:\n", + " print(\"\\nRunning Inference...\")\n", "\n", - " _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n", - " \"NetworkOutput\",\n", - " [\"start_logits\", \"end_logits\", \"feature_index\"])\n", - " networkOutputs = []\n", + " _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n", + " \"NetworkOutput\",\n", + " [\"start_logits\", \"end_logits\", \"feature_index\"])\n", + " networkOutputs = []\n", "\n", - " eval_time_elapsed = 0\n", - " for feature_index, feature in enumerate(features):\n", - " # Copy inputs\n", - " input_ids = cuda.register_host_memory(np.ascontiguousarray(feature.input_ids.ravel()))\n", - " segment_ids = cuda.register_host_memory(np.ascontiguousarray(feature.segment_ids.ravel()))\n", - " input_mask = cuda.register_host_memory(np.ascontiguousarray(feature.input_mask.ravel()))\n", + " eval_time_elapsed = 0\n", + " for feature_index, feature in enumerate(features):\n", + " # Copy inputs\n", + " input_ids = np.ascontiguousarray(feature.input_ids.ravel())\n", + " segment_ids = np.ascontiguousarray(feature.segment_ids.ravel())\n", + " input_mask = np.ascontiguousarray(feature.input_mask.ravel())\n", "\n", - " eval_start_time = time.time()\n", - " cuda.memcpy_htod_async(d_inputs[0], input_ids, stream)\n", - " cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream)\n", - " cuda.memcpy_htod_async(d_inputs[2], input_mask, stream)\n", + " eval_start_time = time.time()\n", + " memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)\n", + " memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)\n", + " memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)\n", "\n", - " # Setup tensor address\n", - " bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n", + " # Setup tensor address\n", + " bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n", "\n", - " for i in range(engine.num_io_tensors):\n", - " context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n", + " for i in range(engine.num_io_tensors):\n", + " context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n", "\n", - " # Run inference\n", - " context.execute_async_v3(stream_handle=stream.handle)\n", - " # Synchronize the stream\n", - " stream.synchronize()\n", - " eval_time_elapsed += (time.time() - eval_start_time)\n", + " # Run inference\n", + " context.execute_async_v3(stream_handle=stream.stream)\n", + " # Synchronize the stream\n", + " stream.synchronize()\n", + " eval_time_elapsed += (time.time() - eval_start_time)\n", "\n", - " # Transfer predictions back from GPU\n", - " cuda.memcpy_dtoh_async(h_output, d_output, stream)\n", - " stream.synchronize()\n", + " # Transfer predictions back from GPU\n", + " memcpy_device_to_host_async(h_output, d_output, stream.stream)\n", + " stream.synchronize()\n", "\n", - " for index, batch in enumerate(h_output):\n", - " # Data Post-processing\n", - " networkOutputs.append(_NetworkOutput(\n", - " start_logits = np.array(batch.squeeze()[:, 0]),\n", - " end_logits = np.array(batch.squeeze()[:, 1]),\n", - " feature_index = feature_index\n", - " ))\n", + " for index, batch in enumerate(h_output):\n", + " # Data Post-processing\n", + " networkOutputs.append(_NetworkOutput(\n", + " start_logits = np.array(batch.squeeze()[:, 0]),\n", + " end_logits = np.array(batch.squeeze()[:, 1]),\n", + " feature_index = feature_index\n", + " ))\n", "\n", - " eval_time_elapsed /= len(features)\n", - " \n", - " print(\"-----------------------------\")\n", - " print(\"Running Inference at {:.3f} Sentences/Sec\".format(1.0/eval_time_elapsed))\n", - " print(\"-----------------------------\")\n" + " eval_time_elapsed /= len(features)\n", + " \n", + " print(\"-----------------------------\")\n", + " print(\"Running Inference at {:.3f} Sentences/Sec\".format(1.0/eval_time_elapsed))\n", + " print(\"-----------------------------\")\n" ] }, { diff --git a/demo/BERT/inference.py b/demo/BERT/inference.py index aa0d0dd7..2cdb1a80 100644 --- a/demo/BERT/inference.py +++ b/demo/BERT/inference.py @@ -30,8 +30,8 @@ import argparse import collections import numpy as np import tensorrt as trt -import pycuda.driver as cuda -import pycuda.autoinit +from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async +from cuda.bindings import driver as cuda, runtime as cudart import helpers.tokenization as tokenization import helpers.data_processing as dp @@ -143,126 +143,133 @@ if __name__ == '__main__': raise RuntimeError("Could not find any profile that can run batch size {}.".format(args.batch_size)) # Create a stream in which to copy inputs/outputs and run inference. - stream = cuda.Stream() + with CudaStreamContext() as stream: + context.set_optimization_profile_async(selected_profile, stream.stream) + binding_idx_offset = selected_profile * engine.num_io_tensors - context.set_optimization_profile_async(selected_profile, stream.handle) - binding_idx_offset = selected_profile * engine.num_io_tensors + # Specify input shapes. These must be within the min/max bounds of the active profile + # Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape. + input_shape = (args.batch_size, max_seq_length) + input_nbytes = trt.volume(input_shape) * trt.int32.itemsize + for name in ["input_ids", "segment_ids", "input_mask"]: + context.set_input_shape(name, input_shape) + assert len(context.infer_shapes()) == 0 - # Specify input shapes. These must be within the min/max bounds of the active profile - # Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape. - input_shape = (args.batch_size, max_seq_length) - input_nbytes = trt.volume(input_shape) * trt.int32.itemsize - for name in ["input_ids", "segment_ids", "input_mask"]: - context.set_input_shape(name, input_shape) - assert len(context.infer_shapes()) == 0 + # Allocate device memory for inputs. + d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)] - # Allocate device memory for inputs. - d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)] + # Allocate output buffer by querying the size from the context. This may be different for different input shapes. + h_output = np.empty(tuple(context.get_tensor_shape("logits_out")), dtype=np.float32) + cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0)) + # Pin the memory for faster transfers + d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes)) - # Allocate output buffer by querying the size from the context. This may be different for different input shapes. - h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape("logits_out")), dtype=np.float32) - d_output = cuda.mem_alloc(h_output.nbytes) + def inference(features, tokens): + global h_output - def inference(features, tokens): - global h_output + _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name + "NetworkOutput", + ["start_logits", "end_logits", "feature_index"]) + networkOutputs = [] - _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name - "NetworkOutput", - ["start_logits", "end_logits", "feature_index"]) - networkOutputs = [] + eval_time_elapsed = 0 + for feature_index, feature in enumerate(features): + # Copy inputs + input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), args.batch_size, axis=0) + segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), args.batch_size, axis=0) + input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), args.batch_size, axis=0) - eval_time_elapsed = 0 - for feature_index, feature in enumerate(features): - # Copy inputs - input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), args.batch_size, axis=0) - segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), args.batch_size, axis=0) - input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), args.batch_size, axis=0) + input_ids = np.ascontiguousarray(input_ids_batch.ravel()) + segment_ids = np.ascontiguousarray(segment_ids_batch.ravel()) + input_mask = np.ascontiguousarray(input_mask_batch.ravel()) - input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids_batch.ravel())) - segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids_batch.ravel())) - input_mask = cuda.register_host_memory(np.ascontiguousarray(input_mask_batch.ravel())) + eval_start_time = time.time() + memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream) + memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream) + memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream) - eval_start_time = time.time() - cuda.memcpy_htod_async(d_inputs[0], input_ids, stream) - cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream) - cuda.memcpy_htod_async(d_inputs[2], input_mask, stream) + bindings = [0 for _ in range(binding_idx_offset)] + [int(d_inp) for d_inp in d_inputs] + [int(d_output)] - bindings = [0 for _ in range(binding_idx_offset)] + [int(d_inp) for d_inp in d_inputs] + [int(d_output)] - - # allocate address for IO tensor - for i in range(engine.num_io_tensors): - context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset]) + # allocate address for IO tensor + for i in range(engine.num_io_tensors): + context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset]) - # Run inference - context.execute_async_v3(stream_handle=stream.handle) - # Synchronize the stream - stream.synchronize() - eval_time_elapsed += (time.time() - eval_start_time) + # Run inference + context.execute_async_v3(stream_handle=stream.stream) + # Synchronize the stream + stream.synchronize() + eval_time_elapsed += (time.time() - eval_start_time) - # Transfer predictions back from GPU - cuda.memcpy_dtoh_async(h_output, d_output, stream) - stream.synchronize() + # Transfer predictions back from GPU + memcpy_device_to_host_async(h_output, d_output, stream.stream) + stream.synchronize() - # Only retrieve and post-process the first batch - batch = h_output[0] - networkOutputs.append(_NetworkOutput( - start_logits = np.array(batch.squeeze()[:, 0]), - end_logits = np.array(batch.squeeze()[:, 1]), - feature_index = feature_index - )) + # Only retrieve and post-process the first batch + batch = h_output[0] + networkOutputs.append(_NetworkOutput( + start_logits = np.array(batch.squeeze()[:, 0]), + end_logits = np.array(batch.squeeze()[:, 1]), + feature_index = feature_index + )) - eval_time_elapsed /= len(features) + eval_time_elapsed /= len(features) - # Total number of n-best predictions to generate in the nbest_predictions.json output file - n_best_size = 20 + # Total number of n-best predictions to generate in the nbest_predictions.json output file + n_best_size = 20 - # The maximum length of an answer that can be generated. This is needed - # because the start and end predictions are not conditioned on one another - max_answer_length = 30 + # The maximum length of an answer that can be generated. This is needed + # because the start and end predictions are not conditioned on one another + max_answer_length = 30 - prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features, - networkOutputs, args.n_best_size, args.max_answer_length) + prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features, + networkOutputs, args.n_best_size, args.max_answer_length) - return eval_time_elapsed, prediction, nbest_json + return eval_time_elapsed, prediction, nbest_json - def print_single_query(eval_time_elapsed, prediction, nbest_json): - print("------------------------") - print("Running inference in {:.3f} Sentences/Sec".format(args.batch_size/eval_time_elapsed)) - print("------------------------") + def print_single_query(eval_time_elapsed, prediction, nbest_json): + print("------------------------") + print("Running inference in {:.3f} Sentences/Sec".format(args.batch_size/eval_time_elapsed)) + print("------------------------") - print("Answer: '{}'".format(prediction)) - print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0)) + print("Answer: '{}'".format(prediction)) + print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0)) - if squad_examples: - all_predictions = collections.OrderedDict() + if squad_examples: + all_predictions = collections.OrderedDict() - for example in squad_examples: - features = question_features(example.doc_tokens, example.question_text) - eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens) - all_predictions[example.id] = prediction - - with open(output_prediction_file, "w") as f: - f.write(json.dumps(all_predictions, indent=4)) - print("\nOutput dump to {}".format(output_prediction_file)) - else: - # Extract tokecs from the paragraph - doc_tokens = dp.convert_doc_tokens(paragraph_text) - - if question_text: - print("\nPassage: {}".format(paragraph_text)) - print("\nQuestion: {}".format(question_text)) - - features = question_features(doc_tokens, question_text) - eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens) - print_single_query(eval_time_elapsed, prediction, nbest_json) + for example in squad_examples: + features = question_features(example.doc_tokens, example.question_text) + eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens) + all_predictions[example.id] = prediction + with open(output_prediction_file, "w") as f: + f.write(json.dumps(all_predictions, indent=4)) + print("\nOutput dump to {}".format(output_prediction_file)) else: - # If no question text is provided, loop until the question is 'exit' - EXIT_CMDS = ["exit", "quit"] - question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS)) + # Extract tokecs from the paragraph + doc_tokens = dp.convert_doc_tokens(paragraph_text) + + if question_text: + print("\nPassage: {}".format(paragraph_text)) + print("\nQuestion: {}".format(question_text)) - while question_text.strip() not in EXIT_CMDS: features = question_features(doc_tokens, question_text) eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens) print_single_query(eval_time_elapsed, prediction, nbest_json) + + else: + # If no question text is provided, loop until the question is 'exit' + EXIT_CMDS = ["exit", "quit"] question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS)) + + while question_text.strip() not in EXIT_CMDS: + features = question_features(doc_tokens, question_text) + eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens) + print_single_query(eval_time_elapsed, prediction, nbest_json) + question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS)) + + # free allocated memory + for d_input in d_inputs: + cuda_call(cudart.cudaFree(d_input)) + cuda_call(cudart.cudaFree(d_output)) + cuda_call(cudart.cudaHostUnregister(h_output)) diff --git a/demo/BERT/inference_varseqlen.py b/demo/BERT/inference_varseqlen.py index 700ddcce..e8a3ae79 100644 --- a/demo/BERT/inference_varseqlen.py +++ b/demo/BERT/inference_varseqlen.py @@ -30,8 +30,9 @@ import argparse import collections import numpy as np import tensorrt as trt -import pycuda.driver as cuda -import pycuda.autoinit +from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async +from cuda.bindings import driver as cuda, runtime as cudart + import helpers.tokenization as tokenization import helpers.data_processing as dp @@ -131,127 +132,133 @@ if __name__ == '__main__': with open(args.engine, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, \ runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context: # Create a stream in which to copy inputs/outputs and run inference. - stream = cuda.Stream() + with CudaStreamContext() as stream: + # select engine profile + context.set_optimization_profile_async(0, stream.stream) - # select engine profile - context.set_optimization_profile_async(0, stream.handle) + input_nbytes = max_seq_length * trt.int32.itemsize - input_nbytes = max_seq_length * trt.int32.itemsize + # Allocate device memory for inputs. + d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(4)] - # Allocate device memory for inputs. - d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(4)] - - # Allocate output buffer by querying the size from the context. This may be different for different input shapes. - h_output = cuda.pagelocked_empty((2 * max_seq_length), dtype=np.float32) - d_output = cuda.mem_alloc(h_output.nbytes) + # Allocate output buffer by querying the size from the context. This may be different for different input shapes. + h_output = np.empty((2 * max_seq_length), dtype=np.float32) + cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0)) + d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes)) - def inference(features, tokens): - global h_output + def inference(features, tokens): + global h_output - _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name - "NetworkOutput", - ["start_logits", "end_logits", "feature_index"]) - networkOutputs = [] + _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name + "NetworkOutput", + ["start_logits", "end_logits", "feature_index"]) + networkOutputs = [] - eval_time_elapsed = 0 - for feature_index, feature in enumerate(features): - # Copy inputs - B = 1 - S = np.sum(feature.input_mask) - input_ids = feature.input_ids[0:S] - segment_ids = feature.segment_ids[0:S] - cu_seq_lens = np.array([0, S], dtype=np.int32); + eval_time_elapsed = 0 + for feature_index, feature in enumerate(features): + # Copy inputs + B = 1 + S = np.sum(feature.input_mask) + input_ids = feature.input_ids[0:S] + segment_ids = feature.segment_ids[0:S] + cu_seq_lens = np.array([0, S], dtype=np.int32); - input_dim0_shape = {"input_ids": S, "segment_ids": S, "cu_seqlens": 2, "max_seqlen": S} - for name, val in input_dim0_shape.items(): - if context.get_tensor_shape(name)[0] != val: - context.set_input_shape(name, (val,)) + input_dim0_shape = {"input_ids": S, "segment_ids": S, "cu_seqlens": 2, "max_seqlen": S} + for name, val in input_dim0_shape.items(): + if context.get_tensor_shape(name)[0] != val: + context.set_input_shape(name, (val,)) - h_input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids.ravel())) - h_segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids.ravel())) - h_cu_seq_lens = cuda.register_host_memory(np.ascontiguousarray(cu_seq_lens.ravel())) + h_input_ids = np.ascontiguousarray(input_ids.ravel()) + h_segment_ids = np.ascontiguousarray(segment_ids.ravel()) + h_cu_seq_lens = np.ascontiguousarray(cu_seq_lens.ravel()) - eval_start_time = time.time() - cuda.memcpy_htod_async(d_inputs[0], h_input_ids, stream) - cuda.memcpy_htod_async(d_inputs[1], h_segment_ids, stream) - cuda.memcpy_htod_async(d_inputs[2], h_cu_seq_lens, stream) + eval_start_time = time.time() + memcpy_host_to_device_async(d_inputs[0], h_input_ids, stream.stream) + memcpy_host_to_device_async(d_inputs[1], h_segment_ids, stream.stream) + memcpy_host_to_device_async(d_inputs[2], h_cu_seq_lens, stream.stream) - # Setup tensor address - bindings = [int(d_inputs[i]) for i in range(4)] + [int(d_output)] + # Setup tensor address + bindings = [int(d_inputs[i]) for i in range(4)] + [int(d_output)] - for i in range(engine.num_io_tensors): - context.set_tensor_address(engine.get_tensor_name(i), bindings[i]) + for i in range(engine.num_io_tensors): + context.set_tensor_address(engine.get_tensor_name(i), bindings[i]) - # Run inference - context.execute_async_v3(stream_handle=stream.handle) - # Synchronize the stream - stream.synchronize() - eval_time_elapsed += (time.time() - eval_start_time) + # Run inference + context.execute_async_v3(stream_handle=stream.stream) + # Synchronize the stream + stream.synchronize() + eval_time_elapsed += (time.time() - eval_start_time) - # Transfer predictions back from GPU - cuda.memcpy_dtoh_async(h_output, d_output, stream) - stream.synchronize() + # Transfer predictions back from GPU + memcpy_device_to_host_async(h_output, d_output, stream.stream) + stream.synchronize() - # Only retrieve and post-process the first batch - networkOutputs.append(_NetworkOutput( - start_logits = np.array(h_output[0:S]), - end_logits = np.array(h_output[S:S*2]), - feature_index = feature_index - )) + # Only retrieve and post-process the first batch + networkOutputs.append(_NetworkOutput( + start_logits = np.array(h_output[0:S]), + end_logits = np.array(h_output[S:S*2]), + feature_index = feature_index + )) - eval_time_elapsed /= len(features) + eval_time_elapsed /= len(features) - # Total number of n-best predictions to generate in the nbest_predictions.json output file - n_best_size = 20 + # Total number of n-best predictions to generate in the nbest_predictions.json output file + n_best_size = 20 - # The maximum length of an answer that can be generated. This is needed - # because the start and end predictions are not conditioned on one another - max_answer_length = 30 + # The maximum length of an answer that can be generated. This is needed + # because the start and end predictions are not conditioned on one another + max_answer_length = 30 - prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features, - networkOutputs, args.n_best_size, args.max_answer_length) + prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features, + networkOutputs, args.n_best_size, args.max_answer_length) - return eval_time_elapsed, prediction, nbest_json + return eval_time_elapsed, prediction, nbest_json - def print_single_query(eval_time_elapsed, prediction, nbest_json): - print("------------------------") - print("Running inference in {:.3f} Sentences/Sec".format(1/eval_time_elapsed)) - print("------------------------") + def print_single_query(eval_time_elapsed, prediction, nbest_json): + print("------------------------") + print("Running inference in {:.3f} Sentences/Sec".format(1/eval_time_elapsed)) + print("------------------------") - print("Answer: '{}'".format(prediction)) - print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0)) + print("Answer: '{}'".format(prediction)) + print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0)) - if squad_examples: - all_predictions = collections.OrderedDict() + if squad_examples: + all_predictions = collections.OrderedDict() - for example in squad_examples: - features = question_features(example.doc_tokens, example.question_text) - eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens) - all_predictions[example.id] = prediction - - with open(output_prediction_file, "w") as f: - f.write(json.dumps(all_predictions, indent=4)) - print("\nOutput dump to {}".format(output_prediction_file)) - else: - # Extract tokecs from the paragraph - doc_tokens = dp.convert_doc_tokens(paragraph_text) - - if question_text: - print("\nPassage: {}".format(paragraph_text)) - print("\nQuestion: {}".format(question_text)) - - features = question_features(doc_tokens, question_text) - eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens) - print_single_query(eval_time_elapsed, prediction, nbest_json) + for example in squad_examples: + features = question_features(example.doc_tokens, example.question_text) + eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens) + all_predictions[example.id] = prediction + with open(output_prediction_file, "w") as f: + f.write(json.dumps(all_predictions, indent=4)) + print("\nOutput dump to {}".format(output_prediction_file)) else: - # If no question text is provided, loop until the question is 'exit' - EXIT_CMDS = ["exit", "quit"] - question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS)) + # Extract tokecs from the paragraph + doc_tokens = dp.convert_doc_tokens(paragraph_text) + + if question_text: + print("\nPassage: {}".format(paragraph_text)) + print("\nQuestion: {}".format(question_text)) - while question_text.strip() not in EXIT_CMDS: features = question_features(doc_tokens, question_text) eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens) print_single_query(eval_time_elapsed, prediction, nbest_json) + + else: + # If no question text is provided, loop until the question is 'exit' + EXIT_CMDS = ["exit", "quit"] question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS)) + + while question_text.strip() not in EXIT_CMDS: + features = question_features(doc_tokens, question_text) + eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens) + print_single_query(eval_time_elapsed, prediction, nbest_json) + question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS)) + + # free allocated memory + for d_input in d_inputs: + cuda_call(cudart.cudaFree(d_input)) + cuda_call(cudart.cudaFree(d_output)) + cuda_call(cudart.cudaHostUnregister(h_output)) diff --git a/demo/BERT/notebooks/Q-and-A.ipynb b/demo/BERT/notebooks/Q-and-A.ipynb index 9c82199a..84584077 100755 --- a/demo/BERT/notebooks/Q-and-A.ipynb +++ b/demo/BERT/notebooks/Q-and-A.ipynb @@ -58,6 +58,8 @@ "\n", "import sys\n", "sys.path.append('/workspace/TensorRT/demo/BERT')\n", + "import os\n", + "sys.path.append(os.path.abspath('..'))\n", "\n", "import ipywidgets as widgets\n", "import tensorrt as trt;\n", @@ -71,8 +73,8 @@ "import collections\n", "import numpy as np\n", "import tensorrt as trt\n", - "import pycuda.driver as cuda\n", - "import pycuda.autoinit\n", + "from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async\n", + "from cuda.bindings import driver as cuda, runtime as cudart\n", "\n", "from helpers import tokenization as tokenization\n", "from helpers import data_processing as dp\n", @@ -109,44 +111,45 @@ " networkOutputs = []\n", "\n", " eval_time_elapsed = 0\n", - " for feature_index, feature in enumerate(features):\n", - " # Copy inputs\n", - " input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), 1, axis=0)\n", - " segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), 1, axis=0)\n", - " input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), 1, axis=0)\n", + " with CudaStreamContext() as stream:\n", + " for feature_index, feature in enumerate(features):\n", + " # Copy inputs\n", + " input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), 1, axis=0)\n", + " segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), 1, axis=0)\n", + " input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), 1, axis=0)\n", "\n", - " input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids_batch.ravel()))\n", - " segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids_batch.ravel()))\n", - " input_mask = cuda.register_host_memory(np.ascontiguousarray(input_mask_batch.ravel()))\n", + " input_ids = np.ascontiguousarray(input_ids_batch.ravel())\n", + " segment_ids = np.ascontiguousarray(segment_ids_batch.ravel())\n", + " input_mask = np.ascontiguousarray(input_mask_batch.ravel())\n", "\n", - " eval_start_time = time.time()\n", - " cuda.memcpy_htod_async(d_inputs[0], input_ids, stream)\n", - " cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream)\n", - " cuda.memcpy_htod_async(d_inputs[2], input_mask, stream)\n", + " eval_start_time = time.time()\n", + " memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)\n", + " memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)\n", + " memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)\n", "\n", - " # Setup tensor address\n", - " bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n", + " # Setup tensor address\n", + " bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n", "\n", - " for i in range(engine.num_io_tensors):\n", - " context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n", + " for i in range(engine.num_io_tensors):\n", + " context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n", "\n", - " # Run inference\n", - " trt_context.execute_async_v3(stream_handle=stream.handle)\n", - " # Synchronize the stream\n", - " stream.synchronize()\n", - " eval_time_elapsed += (time.time() - eval_start_time)\n", + " # Run inference\n", + " trt_context.execute_async_v3(stream_handle=stream.stream)\n", + " # Synchronize the stream\n", + " stream.synchronize()\n", + " eval_time_elapsed += (time.time() - eval_start_time)\n", "\n", - " # Transfer predictions back from GPU\n", - " cuda.memcpy_dtoh_async(h_output, d_output, stream)\n", - " stream.synchronize()\n", + " # Transfer predictions back from GPU\n", + " memcpy_device_to_host_async(h_output, d_output, stream.stream)\n", + " stream.synchronize()\n", "\n", - " for index, batch in enumerate(h_output):\n", - " # Data Post-processing\n", - " networkOutputs.append(_NetworkOutput(\n", - " start_logits = np.array(batch.squeeze()[:, 0]),\n", - " end_logits = np.array(batch.squeeze()[:, 1]),\n", - " feature_index = feature_index\n", - " ))\n", + " for index, batch in enumerate(h_output):\n", + " # Data Post-processing\n", + " networkOutputs.append(_NetworkOutput(\n", + " start_logits = np.array(batch.squeeze()[:, 0]),\n", + " end_logits = np.array(batch.squeeze()[:, 1]),\n", + " feature_index = feature_index\n", + " ))\n", "\n", " eval_time_elapsed /= len(features)\n", "\n", @@ -172,59 +175,60 @@ " networkOutputs = []\n", "\n", " eval_time_elapsed = 0\n", - " for feature_index, feature in enumerate(features):\n", - " # Copy inputs\n", - " B = 1\n", - " S = np.sum(feature.input_mask)\n", - " input_ids = feature.input_ids[0:S]\n", - " segment_ids = feature.segment_ids[0:S]\n", - " cu_seq_lens = np.array([0, S], dtype=np.int32)\n", + " with CudaStreamContext() as stream:\n", + " for feature_index, feature in enumerate(features):\n", + " # Copy inputs\n", + " B = 1\n", + " S = np.sum(feature.input_mask)\n", + " input_ids = feature.input_ids[0:S]\n", + " segment_ids = feature.segment_ids[0:S]\n", + " cu_seq_lens = np.array([0, S], dtype=np.int32)\n", "\n", - " first_tensor_name = engine.get_tensor_name(0)\n", - " second_tensor_name = engine.get_tensor_name(1)\n", - " third_tensor_name = engine.get_tensor_name(2)\n", - " forth_tensor_name = engine.get_tensor_name(3)\n", + " first_tensor_name = engine.get_tensor_name(0)\n", + " second_tensor_name = engine.get_tensor_name(1)\n", + " third_tensor_name = engine.get_tensor_name(2)\n", + " forth_tensor_name = engine.get_tensor_name(3)\n", "\n", - " if context.get_tensor_shape(first_tensor_name)[0] != S:\n", - " context.set_input_shape(first_tensor_name, (S,))\n", - " if context.get_tensor_shape(second_tensor_name)[0] != S:\n", - " context.set_input_shape(second_tensor_name, (S,))\n", - " if context.get_tensor_shape(third_tensor_name)[0] != 2:\n", - " context.set_input_shape(third_tensor_name, (2,))\n", - " if context.get_tensor_shape(forth_tensor_name)[0] != S:\n", - " context.set_input_shapee(forth_tensor_name, (S,))\n", + " if context.get_tensor_shape(first_tensor_name)[0] != S:\n", + " context.set_input_shape(first_tensor_name, (S,))\n", + " if context.get_tensor_shape(second_tensor_name)[0] != S:\n", + " context.set_input_shape(second_tensor_name, (S,))\n", + " if context.get_tensor_shape(third_tensor_name)[0] != 2:\n", + " context.set_input_shape(third_tensor_name, (2,))\n", + " if context.get_tensor_shape(forth_tensor_name)[0] != S:\n", + " context.set_input_shape(forth_tensor_name, (S,))\n", "\n", - " h_input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids.ravel()))\n", - " h_segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids.ravel()))\n", - " h_cu_seq_lens = cuda.register_host_memory(np.ascontiguousarray(cu_seq_lens.ravel()))\n", + " h_input_ids = np.ascontiguousarray(input_ids.ravel())\n", + " h_segment_ids = np.ascontiguousarray(segment_ids.ravel())\n", + " h_cu_seq_lens = np.ascontiguousarray(cu_seq_lens.ravel())\n", "\n", - " eval_start_time = time.time()\n", - " cuda.memcpy_htod_async(d_inputs[0], h_input_ids, INT8_stream)\n", - " cuda.memcpy_htod_async(d_inputs[1], h_segment_ids, INT8_stream)\n", - " cuda.memcpy_htod_async(d_inputs[2], h_cu_seq_lens, INT8_stream)\n", + " eval_start_time = time.time()\n", + " memcpy_host_to_device_async(d_inputs[0], h_input_ids, stream.stream)\n", + " memcpy_host_to_device_async(d_inputs[1], h_segment_ids, stream.stream)\n", + " memcpy_host_to_device_async(d_inputs[2], h_cu_seq_lens, stream.stream)\n", "\n", - " # Setup tensor address\n", - " bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n", + " # Setup tensor address\n", + " bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n", "\n", - " for i in range(engine.num_io_tensors):\n", - " context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n", + " for i in range(engine.num_io_tensors):\n", + " context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n", "\n", - " # Run inference\n", - " trt_context.execute_async_v3(stream_handle=INT8_stream.handle)\n", - " # Synchronize the stream\n", - " INT8_stream.synchronize()\n", - " eval_time_elapsed += (time.time() - eval_start_time)\n", + " # Run inference\n", + " trt_context.execute_async_v3(stream_handle=stream.stream)\n", + " # Synchronize the stream\n", + " stream.synchronize()\n", + " eval_time_elapsed += (time.time() - eval_start_time)\n", "\n", - " # Transfer predictions back from GPU\n", - " cuda.memcpy_dtoh_async(h_output, d_output, INT8_stream)\n", - " INT8_stream.synchronize()\n", + " # Transfer predictions back from GPU\n", + " memcpy_device_to_host_async(h_output, d_output, stream.stream)\n", + " stream.synchronize()\n", "\n", - " # Only retrieve and post-process the first batch\n", - " networkOutputs.append(_NetworkOutput(\n", - " start_logits = np.array(h_output[0:S]),\n", - " end_logits = np.array(h_output[S:S*2]),\n", - " feature_index = feature_index\n", - " ))\n", + " # Only retrieve and post-process the first batch\n", + " networkOutputs.append(_NetworkOutput(\n", + " start_logits = np.array(h_output[0:S]),\n", + " end_logits = np.array(h_output[S:S*2]),\n", + " feature_index = feature_index\n", + " ))\n", "\n", " eval_time_elapsed /= len(features)\n", "\n", @@ -268,7 +272,7 @@ "input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n", "\n", "# Allocate device memory for inputs.\n", - "d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)]\n", + "d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]\n", "\n", "# Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n", "# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n", @@ -278,34 +282,35 @@ "assert context.all_binding_shapes_specified\n", "\n", "# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n", - "h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n", - "d_output = cuda.mem_alloc(h_output.nbytes)\n", + "h_output = np.empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n", + "cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))\n", + "d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))\n", "\n", "# Create a stream in which to copy inputs/outputs and run inference.\n", - "stream = cuda.Stream()\n", + "with CudaStreamContext() as stream:\n", + " ### INT8 TRT model\n", + " engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n", + " max_seq_length = 384\n", "\n", - "### INT8 TRT model\n", - "engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n", - "max_seq_length = 384\n", + " INT8_runtime = trt.Runtime(TRT_LOGGER)\n", + " INT8_engine = INT8_runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n", + " INT8_context = INT8_engine.create_execution_context()\n", "\n", - "INT8_runtime = trt.Runtime(TRT_LOGGER)\n", - "INT8_engine = INT8_runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n", - "INT8_context = INT8_engine.create_execution_context()\n", + " # select engine profile\n", + " INT8_context.set_optimization_profile_async(0, stream.stream)\n", "\n", - "# select engine profile\n", - "INT8_context.set_optimization_profile_async(0, stream.handle)\n", + " input_nbytes = max_seq_length * trt.int32.itemsize\n", "\n", - "input_nbytes = max_seq_length * trt.int32.itemsize\n", + " # Allocate device memory for inputs.\n", + " INT8_d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(4)]\n", "\n", - "# Allocate device memory for inputs.\n", - "INT8_d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(4)]\n", + " # Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n", + " INT8_h_output = np.empty((2 * max_seq_length), dtype=np.float32)\n", + " cuda_call(cudart.cudaHostRegister(INT8_h_output, INT8_h_output.nbytes, 0))\n", + " INT8_d_output = cuda_call(cudart.cudaMalloc(INT8_h_output.nbytes))\n", "\n", - "# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n", - "INT8_h_output = cuda.pagelocked_empty((2 * max_seq_length), dtype=np.float32)\n", - "INT8_d_output = cuda.mem_alloc(INT8_h_output.nbytes)\n", - "\n", - "# Create a stream in which to copy inputs/outputs and run inference.\n", - "INT8_stream = cuda.Stream()\n" + " # No separate INT8 stream; using context manager stream\n", + " INT8_stream = stream.stream\n" ] }, { @@ -439,7 +444,19 @@ "id": "musical-right", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "\n", + "# free allocated memory\n", + "for d_input in d_inputs:\n", + " cuda_call(cudart.cudaFree(d_input))\n", + "cuda_call(cudart.cudaFree(d_output))\n", + "cuda_call(cudart.cudaHostUnregister(h_output))\n", + "\n", + "for INT8_d_input in INT8_d_inputs:\n", + " cuda_call(cudart.cudaFree(INT8_d_input))\n", + "cuda_call(cudart.cudaFree(INT8_d_output))\n", + "cuda_call(cudart.cudaHostUnregister(INT8_h_output))" + ] } ], "metadata": { diff --git a/demo/BERT/notebooks/benchmark.ipynb b/demo/BERT/notebooks/benchmark.ipynb index 442b28f5..ddc9157c 100755 --- a/demo/BERT/notebooks/benchmark.ipynb +++ b/demo/BERT/notebooks/benchmark.ipynb @@ -41,7 +41,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "c89592ab-e50b-47c9-9b5b-6b5067d0b22a", "metadata": { "jupyter": { @@ -56,6 +56,8 @@ "\n", "import sys\n", "sys.path.append('/workspace/TensorRT/demo/BERT')\n", + "import os\n", + "sys.path.append(os.path.abspath('..'))\n", "\n", "import tensorrt as trt;\n", "TRT_VERSION = trt.__version__\n", @@ -65,8 +67,8 @@ "import ctypes\n", "import numpy as np\n", "import tensorrt as trt\n", - "import pycuda.driver as cuda\n", - "import pycuda.autoinit\n", + "from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host\n", + "from cuda.bindings import driver as cuda, runtime as cudart\n", "import numpy as np\n", "from tqdm import tqdm\n", "\n", @@ -80,15 +82,16 @@ "\n", "TRT_LOGGER = trt.Logger(trt.Logger.ERROR)\n", "\n", + "\n", "class DeviceBuffer(object):\n", " def __init__(self, shape, dtype=trt.int32):\n", - " self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize)\n", + " self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))\n", "\n", " def binding(self):\n", " return int(self.buf)\n", "\n", " def free(self):\n", - " self.buf.free()\n", + " cuda_call(cudart.cudaFree(self.buf))\n", "\n", "doc_stride = 128\n", "max_query_length = 64\n", @@ -138,56 +141,57 @@ " test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32)\n", "\n", " # Copy input h2d\n", - " cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())\n", - " cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())\n", - " cuda.memcpy_htod(buffers[2].buf, test_cu_seq_lens.ravel())\n", + " memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())\n", + " memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())\n", + " memcpy_host_to_device(buffers[2].buf, test_cu_seq_lens.ravel())\n", "\n", " bench_times = {}\n", - " stream = cuda.Stream()\n", + " with CudaStreamContext() as stream:\n", "\n", - " tensor_name = engine.get_tensor_name(engine.num_io_tensors-1)\n", - " for idx, batch_size in enumerate(sorted(args.batch_size)):\n", - " for idx in range(engine.num_optimization_profiles):\n", - " profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n", - " if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n", - " context.set_optimization_profile_async(idx, stream.handle)\n", - " binding_idx_offset = idx * engine.num_io_tensors\n", - " break\n", + " tensor_name = engine.get_tensor_name(engine.num_io_tensors-1)\n", + " for idx, batch_size in enumerate(sorted(args.batch_size)):\n", + " for idx in range(engine.num_optimization_profiles):\n", + " profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n", + " if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n", + " context.set_optimization_profile_async(idx, stream.stream)\n", + " binding_idx_offset = idx * engine.num_io_tensors\n", + " break\n", "\n", - " # Each profile has unique bindings\n", - " bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n", - " input_shape = (batch_size, args.sequence_length)\n", - " for binding in range(3):\n", - " tensor_name = engine.get_tensor_name(binding)\n", - " context.set_input_shape(tensor_name, input_shape)\n", - " assert context.all_binding_shapes_specified\n", + " # Each profile has unique bindings\n", + " bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n", + " input_shape = (batch_size, args.sequence_length)\n", + " for binding in range(3):\n", + " tensor_name = engine.get_tensor_name(binding)\n", + " context.set_input_shape(tensor_name, input_shape)\n", + " assert context.all_binding_shapes_specified\n", "\n", - " for i in range(engine.num_io_tensors):\n", - " context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n", + " for i in range(engine.num_io_tensors):\n", + " context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n", "\n", - " # Inference\n", - " total_time = 0\n", - " start = cuda.Event()\n", - " end = cuda.Event()\n", + " # Inference\n", + " total_time = 0\n", + " start = cuda_call(cudart.cudaEventCreate())\n", + " end = cuda_call(cudart.cudaEventCreate())\n", "\n", - " # Warmup\n", - " for _ in range(args.warm_up_runs):\n", - " context.execute_async_v3(stream_handle=stream.handle)\n", - " stream.synchronize()\n", + " # Warmup\n", + " for _ in range(args.warm_up_runs):\n", + " context.execute_async_v3(stream_handle=stream.stream)\n", + " stream.synchronize()\n", "\n", - " # Timing loop\n", - " times = []\n", - " progress_bar.value = 0\n", - " for _ in range(iteration_selector.value):\n", - " start.record(stream)\n", - " context.execute_async_v3(stream_handle=stream.handle)\n", - " end.record(stream)\n", - " stream.synchronize()\n", - " times.append(end.time_since(start))\n", - " progress_bar.value +=1\n", + " # Timing loop\n", + " times = []\n", + " progress_bar.value = 0\n", + " for _ in range(iteration_selector.value):\n", + " cuda_call(cudart.cudaEventRecord(start, stream.stream))\n", + " context.execute_async_v3(stream_handle=stream.stream)\n", + " cuda_call(cudart.cudaEventRecord(end, stream.stream))\n", + " stream.synchronize()\n", + " elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))\n", + " times.append(elapsed_time)\n", + " progress_bar.value +=1\n", "\n", - " # Compute average time, 95th percentile time and 99th percentile time.\n", - " bench_times[batch_size] = times\n", + " # Compute average time, 95th percentile time and 99th percentile time.\n", + " bench_times[batch_size] = times\n", "\n", " [b.free() for b in buffers]\n", "\n", @@ -227,57 +231,56 @@ " test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32)\n", "\n", " # Copy input h2d\n", - " cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())\n", - " cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())\n", - " cuda.memcpy_htod(buffers[2].buf, test_input_mask.ravel())\n", + " memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())\n", + " memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())\n", + " memcpy_host_to_device(buffers[2].buf, test_input_mask.ravel())\n", "\n", " bench_times = {}\n", - " stream = cuda.Stream()\n", + " with CudaStreamContext() as stream:\n", "\n", - " tensor_name = engine.get_tensor_name(engine.num_io_tensors - 1)\n", - " for idx, batch_size in enumerate(sorted(args.batch_size)):\n", - " for idx in range(engine.num_optimization_profiles):\n", - " profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n", - " if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n", - " context.set_optimization_profile_async(idx, stream.handle)\n", - " binding_idx_offset = idx * engine.num_io_tensors\n", - " break\n", + " tensor_name = engine.get_tensor_name(engine.num_io_tensors - 1)\n", + " for idx, batch_size in enumerate(sorted(args.batch_size)):\n", + " for idx in range(engine.num_optimization_profiles):\n", + " profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n", + " if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n", + " context.set_optimization_profile_async(idx, stream.stream)\n", + " binding_idx_offset = idx * engine.num_io_tensors\n", + " break\n", "\n", - " # Each profile has unique bindings\n", - " bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n", - " input_shape = (batch_size, args.sequence_length)\n", - " for binding in range(3):\n", - " tensor_name = engine.get_tensor_name(binding)\n", - " context.set_input_shape(tensor_name, input_shape)\n", - " assert context.all_binding_shapes_specified\n", + " # Each profile has unique bindings\n", + " bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n", + " input_shape = (batch_size, args.sequence_length)\n", + " for binding in range(3):\n", + " tensor_name = engine.get_tensor_name(binding)\n", + " context.set_input_shape(tensor_name, input_shape)\n", + " assert context.all_binding_shapes_specified\n", "\n", - " for i in range(engine.num_io_tensors):\n", - " context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n", + " for i in range(engine.num_io_tensors):\n", + " context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n", "\n", - " # Inference\n", - " total_time = 0\n", - " start = cuda.Event()\n", - " end = cuda.Event()\n", - " stream = cuda.Stream()\n", + " # Inference\n", + " total_time = 0\n", + " start = cuda_call(cudart.cudaEventCreate())\n", + " end = cuda_call(cudart.cudaEventCreate())\n", "\n", - " # Warmup\n", - " for _ in range(args.warm_up_runs):\n", - " context.execute_async_v3(stream_handle=stream.handle)\n", - " stream.synchronize()\n", + " # Warmup\n", + " for _ in range(args.warm_up_runs):\n", + " context.execute_async_v3(stream_handle=stream.stream)\n", + " stream.synchronize()\n", "\n", - " # Timing loop\n", - " times = []\n", - " progress_bar.value = 0\n", - " for _ in range(iteration_selector.value):\n", - " start.record(stream)\n", - " context.execute_async_v3(stream_handle=stream.handle)\n", - " end.record(stream)\n", - " stream.synchronize()\n", - " times.append(end.time_since(start))\n", - " progress_bar.value +=1\n", + " # Timing loop\n", + " times = []\n", + " progress_bar.value = 0\n", + " for _ in range(iteration_selector.value):\n", + " start.record(stream.stream)\n", + " context.execute_async_v3(stream_handle=stream.stream)\n", + " end.record(stream.stream)\n", + " stream.synchronize()\n", + " times.append(end.time_since(start))\n", + " progress_bar.value +=1\n", "\n", - " # Compute average time, 95th percentile time and 99th percentile time.\n", - " bench_times[batch_size] = times\n", + " # Compute average time, 95th percentile time and 99th percentile time.\n", + " bench_times[batch_size] = times\n", "\n", " [b.free() for b in buffers]\n", "\n", diff --git a/demo/BERT/perf.py b/demo/BERT/perf.py index f3d2ab74..f2e9fe99 100644 --- a/demo/BERT/perf.py +++ b/demo/BERT/perf.py @@ -18,24 +18,22 @@ import argparse import ctypes import time -import numpy as np import tensorrt as trt -import pycuda.driver as cuda -import pycuda.autoinit - +from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host +from cuda.bindings import driver as cuda, runtime as cudart import numpy as np TRT_LOGGER = trt.Logger(trt.Logger.ERROR) class DeviceBuffer(object): def __init__(self, shape, dtype=trt.int32): - self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize) + self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize)) def binding(self): return int(self.buf) def free(self): - self.buf.free() + cuda_call(cudart.cudaFree(self.buf)) def main(): @@ -73,63 +71,64 @@ def main(): test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32) # Copy input h2d - cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel()) - cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel()) - cuda.memcpy_htod(buffers[2].buf, test_input_mask.ravel()) + memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel()) + memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel()) + memcpy_host_to_device(buffers[2].buf, test_input_mask.ravel()) bench_times = {} - stream = cuda.Stream() - for batch_size in sorted(args.batch_size): - # Select engine profile - selected_profile = -1 - for idx in range(engine.num_optimization_profiles): - profile_shape = engine.get_tensor_profile_shape(name = "input_ids", profile_index = idx) - if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size and profile_shape[0][1] <= args.sequence_length and profile_shape[2][1] >= args.sequence_length: - selected_profile = idx - break - if selected_profile == -1: - raise RuntimeError("None of the dynamic shape profiles meets the requirement batch = {} and sequence = {}.".format(batch_size, args.sequence_length)) - context.set_optimization_profile_async(selected_profile, stream.handle) + with CudaStreamContext() as stream: + for batch_size in sorted(args.batch_size): + # Select engine profile + selected_profile = -1 + for idx in range(engine.num_optimization_profiles): + profile_shape = engine.get_tensor_profile_shape(name = "input_ids", profile_index = idx) + if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size and profile_shape[0][1] <= args.sequence_length and profile_shape[2][1] >= args.sequence_length: + selected_profile = idx + break + if selected_profile == -1: + raise RuntimeError("None of the dynamic shape profiles meets the requirement batch = {} and sequence = {}.".format(batch_size, args.sequence_length)) + context.set_optimization_profile_async(selected_profile, stream.stream) - # Each profile has unique bindings - binding_idx_offset = selected_profile * engine.num_io_tensors - bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers] + # Each profile has unique bindings + binding_idx_offset = selected_profile * engine.num_io_tensors + bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers] - input_shape = (batch_size, args.sequence_length) - for name in ["input_ids", "segment_ids", "input_mask"]: - context.set_input_shape(name, input_shape) - assert len(context.infer_shapes()) == 0 + input_shape = (batch_size, args.sequence_length) + for name in ["input_ids", "segment_ids", "input_mask"]: + context.set_input_shape(name, input_shape) + assert len(context.infer_shapes()) == 0 - for i in range(engine.num_io_tensors): - context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset]) + for i in range(engine.num_io_tensors): + context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset]) - # Inference - total_time = 0 - start = cuda.Event() - end = cuda.Event() + # Inference + total_time = 0 + start = cuda_call(cudart.cudaEventCreate()) + end = cuda_call(cudart.cudaEventCreate()) - # Warmup - for _ in range(args.warm_up_runs): - context.execute_async_v3(stream_handle=stream.handle) - stream.synchronize() + # Warmup + for _ in range(args.warm_up_runs): + context.execute_async_v3(stream_handle=stream.stream) + stream.synchronize() - # Timing loop - times = [] - actual_iterations = 0 - start_time = time.time() - while actual_iterations < args.iterations or (time.time() - start_time) < args.duration: - start.record(stream) - context.execute_async_v3(stream_handle=stream.handle) - end.record(stream) - stream.synchronize() - times.append(end.time_since(start)) - actual_iterations += 1 + # Timing loop + times = [] + actual_iterations = 0 + start_time = time.time() + while actual_iterations < args.iterations or (time.time() - start_time) < args.duration: + cuda_call(cudart.cudaEventRecord(start, stream.stream)) + context.execute_async_v3(stream_handle=stream.stream) + cuda_call(cudart.cudaEventRecord(end, stream.stream)) + stream.synchronize() + elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end)) + times.append(elapsed_time) + actual_iterations += 1 - # Compute average time, 95th percentile time and 99th percentile time. - bench_times[batch_size] = times + # Compute average time, 95th percentile time and 99th percentile time. + bench_times[batch_size] = times - [b.free() for b in buffers] + [b.free() for b in buffers] for batch_size, times in bench_times.items(): total_time = sum(times) diff --git a/demo/BERT/perf_varseqlen.py b/demo/BERT/perf_varseqlen.py index 6708f989..38223950 100644 --- a/demo/BERT/perf_varseqlen.py +++ b/demo/BERT/perf_varseqlen.py @@ -20,8 +20,8 @@ import ctypes import time import numpy as np import tensorrt as trt -import pycuda.driver as cuda -import pycuda.autoinit +from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host +from cuda.bindings import driver as cuda, runtime as cudart import numpy as np @@ -29,13 +29,13 @@ TRT_LOGGER = trt.Logger(trt.Logger.ERROR) class DeviceBuffer(object): def __init__(self, shape, dtype=trt.int32): - self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize) + self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize)) def binding(self): return int(self.buf) def free(self): - self.buf.free() + cuda_call(cudart.cudaFree(self.buf)) def main(): @@ -74,57 +74,58 @@ def main(): test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32) # Copy input h2d - cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel()) - cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel()) - cuda.memcpy_htod(buffers[2].buf, test_cu_seq_lens.ravel()) + memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel()) + memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel()) + memcpy_host_to_device(buffers[2].buf, test_cu_seq_lens.ravel()) bench_times = {} for idx, batch_size in enumerate(sorted(args.batch_size)): - stream = cuda.Stream() - context.set_optimization_profile_async(0, stream.handle) + with CudaStreamContext() as stream: + context.set_optimization_profile_async(0, stream.stream) - # Each profile has unique bindings - bindings = [buf.binding() for buf in buffers] + # Each profile has unique bindings + bindings = [buf.binding() for buf in buffers] - shapes = { - "input_ids": (args.sequence_length * batch_size, ), - "segment_ids": (args.sequence_length * batch_size, ), - "cu_seqlens": (batch_size + 1, ), - "max_seqlen": (args.sequence_length, ), - } + shapes = { + "input_ids": (args.sequence_length * batch_size, ), + "segment_ids": (args.sequence_length * batch_size, ), + "cu_seqlens": (batch_size + 1, ), + "max_seqlen": (args.sequence_length, ), + } - for binding, shape in shapes.items(): - context.set_input_shape(binding, shape) - assert len(context.infer_shapes()) == 0 + for binding, shape in shapes.items(): + context.set_input_shape(binding, shape) + assert len(context.infer_shapes()) == 0 - for i in range(engine.num_io_tensors): - context.set_tensor_address(engine.get_tensor_name(i), bindings[i]) + for i in range(engine.num_io_tensors): + context.set_tensor_address(engine.get_tensor_name(i), bindings[i]) - # Inference - total_time = 0 - start = cuda.Event() - end = cuda.Event() + # Inference + total_time = 0 + start = cuda_call(cudart.cudaEventCreate()) + end = cuda_call(cudart.cudaEventCreate()) - # Warmup - for _ in range(args.warm_up_runs): - context.execute_async_v3(stream_handle=stream.handle) - stream.synchronize() + # Warmup + for _ in range(args.warm_up_runs): + context.execute_async_v3(stream_handle=stream.stream) + stream.synchronize() - # Timing loop - times = [] - actual_iterations = 0 - start_time = time.time() - while actual_iterations < args.iterations or (time.time() - start_time) < args.duration: - start.record(stream) - context.execute_async_v3(stream_handle=stream.handle) - end.record(stream) - stream.synchronize() - times.append(end.time_since(start)) - actual_iterations += 1 + # Timing loop + times = [] + actual_iterations = 0 + start_time = time.time() + while actual_iterations < args.iterations or (time.time() - start_time) < args.duration: + cuda_call(cudart.cudaEventRecord(start, stream.stream)) + context.execute_async_v3(stream_handle=stream.stream) + cuda_call(cudart.cudaEventRecord(end, stream.stream)) + stream.synchronize() + elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end)) + times.append(elapsed_time) + actual_iterations += 1 - # Compute average time, 95th percentile time and 99th percentile time. - bench_times[batch_size] = times + # Compute average time, 95th percentile time and 99th percentile time. + bench_times[batch_size] = times [b.free() for b in buffers] diff --git a/demo/DeBERTa/cuda_utils.py b/demo/DeBERTa/cuda_utils.py new file mode 100644 index 00000000..5ffe769a --- /dev/null +++ b/demo/DeBERTa/cuda_utils.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import logging +from cuda.bindings import driver as cuda, runtime as cudart + +class CudaStreamContext: + """CUDA stream lifecycle management with context manager support""" + def __init__(self): + """Initialize CUDA stream""" + self._stream = cuda_call(cudart.cudaStreamCreate()) + + def __enter__(self): + """Create CUDA stream when entering context (if not already created)""" + if self._stream is None: + self._stream = cuda_call(cudart.cudaStreamCreate()) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Destroy CUDA stream when exiting context""" + self.free() + + @property + def stream(self) -> cudart.cudaStream_t: + if self._stream is None: + raise RuntimeError("Stream not created. Use 'with' statement.") + return self._stream + + def synchronize(self): + """Synchronize the stream""" + if self._stream is None: + raise RuntimeError("Stream not created. Use 'with' statement.") + cuda_call(cudart.cudaStreamSynchronize(self._stream)) + + def free(self): + """Explicitly free the CUDA stream""" + if self._stream is not None: + try: + cuda_call(cudart.cudaStreamDestroy(self._stream)) + self._stream = None + except Exception as e: + logging.warning(f"Failed to destroy CUDA stream: {e}") + + def __del__(self): + """Cleanup stream on destruction""" + if hasattr(self, '_stream') and self._stream is not None: + self.free() + + def __str__(self): + return f"CudaStreamContext: {self._stream}" + + def __repr__(self): + return self.__str__() + +def cuda_call(call): + """Helper function to make CUDA calls and check for errors""" + def _cudaGetErrorEnum(error): + if isinstance(error, cuda.CUresult): + err, name = cuda.cuGetErrorName(error) + return name if err == cuda.CUresult.CUDA_SUCCESS else "" + elif isinstance(error, cudart.cudaError_t): + return cudart.cudaGetErrorName(error)[1] + else: + raise RuntimeError("Unknown error type: {}".format(error)) + + err, res = call[0], call[1:] + if err.value: + raise RuntimeError( + "CUDA error code={}({})".format( + err.value, _cudaGetErrorEnum(err) + ) + ) + if len(res) == 1: + return res[0] + elif len(res) == 0: + return None + else: + return res + +def getComputeCapacity(devID=0): + major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID)) + minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID)) + return (major, minor) + +def memcpy_host_to_device_async(device_ptr: int, host_arr: np.ndarray, stream): + """Wrapper for async host-to-device memory copy""" + cuda_call(cudart.cudaMemcpyAsync(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)) + + +def memcpy_device_to_host_async(host_arr: np.ndarray, device_ptr: int, stream): + """Wrapper for async device-to-host memory copy""" + cuda_call(cudart.cudaMemcpyAsync(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)) + +def memcpy_device_to_device_async(dst_device_ptr: int, src_device_ptr: int, nbytes: int, stream): + """Wrapper for async device-to-device memory copy""" + cuda_call(cudart.cudaMemcpyAsync(dst_device_ptr, src_device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice, stream)) + +def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray): + """Wrapper for synchronous host-to-device memory copy""" + cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)) + + +def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int): + """Wrapper for synchronous device-to-host memory copy""" + cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) + +def memcpy_device_to_device(dst_device_ptr: int, src_device_ptr: int, nbytes: int): + """Wrapper for synchronous device-to-device memory copy""" + cuda_call(cudart.cudaMemcpy(dst_device_ptr, src_device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice)) + +# Initialize CUDA +cuda_call(cudart.cudaFree(0)) diff --git a/demo/DeBERTa/deberta_tensorrt_inference.py b/demo/DeBERTa/deberta_tensorrt_inference.py index 355ad7cf..753ca9a4 100644 --- a/demo/DeBERTa/deberta_tensorrt_inference.py +++ b/demo/DeBERTa/deberta_tensorrt_inference.py @@ -40,9 +40,17 @@ import torch import tensorrt as trt import os, sys, argparse import numpy as np -import pycuda.driver as cuda -import pycuda.autoinit # without this, "LogicError: explicit_context_dependent failed: invalid device context - no currently active context?" from time import time +from cuda.bindings import driver as cuda, runtime as cudart +from cuda_utils import ( + cuda_call, + CudaStreamContext, + memcpy_host_to_device_async, + memcpy_device_to_host_async, + memcpy_host_to_device, + memcpy_device_to_device_async, + memcpy_device_to_device, +) TRT_VERSION = int(trt.__version__[:3].replace('.','')) # e.g., version 8.4.1.5 becomes 84 @@ -167,7 +175,7 @@ class TRTModel: inputs = [] outputs = [] bindings = [] - stream = cuda.Stream() + stream = CudaStreamContext() for i in range(engine.num_io_tensors): tensor_name = engine.get_tensor_name(i) @@ -175,10 +183,10 @@ class TRTModel: dtype = trt.nptype(engine.get_tensor_dtype(tensor_name)) # Allocate host and device buffers - host_mem = cuda.pagelocked_empty(size, dtype) # page-locked memory buffer (won't swapped to disk) - device_mem = cuda.mem_alloc(host_mem.nbytes) + host_mem = np.empty(size, dtype) + device_mem = cuda_call(cudart.cudaMalloc(host_mem.nbytes)) - # Append the device buffer address to device bindings. When cast to int, it's a linear index into the context's memory (like memory address). See https://documen.tician.de/pycuda/driver.html#pycuda.driver.DeviceAllocation + # Append the device buffer address to device bindings. When cast to int, it's a linear index into the context's memory (like memory address). bindings.append(int(device_mem)) # Append to the appropriate input/output list. @@ -226,18 +234,19 @@ class TRTModel: # fill host memory with flattened input data np.copyto(self.inputs[i].host, model_input.ravel()) elif TORCH: + nbytes = model_input.element_size() * model_input.nelement() if timing: - cuda.memcpy_dtod(self.inputs[i].device, model_input.data_ptr(), model_input.element_size() * model_input.nelement()) + memcpy_device_to_device(self.inputs[i].device, model_input.data_ptr(), nbytes) else: # for Torch GPU tensor it's easier, can just do Device to Device copy - cuda.memcpy_dtod_async(self.inputs[i].device, model_input.data_ptr(), model_input.element_size() * model_input.nelement(), self.stream) # dtod need size in bytes + memcpy_device_to_device_async(self.inputs[i].device, model_input.data_ptr(), nbytes, self.stream.stream) if NUMPY: if timing: - [cuda.memcpy_htod(inp.device, inp.host) for inp in self.inputs] + [memcpy_host_to_device(inp.device, inp.host) for inp in self.inputs] else: # input, Host to Device - [cuda.memcpy_htod_async(inp.device, inp.host, self.stream) for inp in self.inputs] + [memcpy_host_to_device_async(inp.device, inp.host, self.stream.stream) for inp in self.inputs] for i in range(self.engine.num_io_tensors): self.context.set_tensor_address(self.engine.get_tensor_name(i), self.bindings[i]) @@ -250,13 +259,13 @@ class TRTModel: duration = end_time - start_time else: # run inference - self.context.execute_async_v3(stream_handle=self.stream.handle) + self.context.execute_async_v3(stream_handle=self.stream.stream) if timing: - [cuda.memcpy_dtoh(out.host, out.device) for out in self.outputs] + [cuda_call(cudart.cudaMemcpy(out.host.ctypes.data, out.device, out.host.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) for out in self.outputs] else: # output, Device to Host - [cuda.memcpy_dtoh_async(out.host, out.device, self.stream) for out in self.outputs] + [memcpy_device_to_host_async(out.host, out.device, self.stream.stream) for out in self.outputs] if not timing: # synchronize to ensure completion of async calls diff --git a/demo/Diffusion/README.md b/demo/Diffusion/README.md index 2c342814..ea3214b8 100755 --- a/demo/Diffusion/README.md +++ b/demo/Diffusion/README.md @@ -7,7 +7,7 @@ This demo application ("demoDiffusion") showcases the acceleration of Stable Dif ### Clone the TensorRT OSS repository ```bash -git clone git@github.com:NVIDIA/TensorRT.git -b release/10.13.3 --single-branch +git clone git@github.com:NVIDIA/TensorRT.git -b release/10.14 --single-branch cd TensorRT ``` @@ -23,10 +23,18 @@ NOTE: The demo supports CUDA>=12.0 ### Install the required packages +To install dependencies for modern pipelines (SDXL and later): + ```bash source setup.sh ``` +To install dependencies for legacy pipelines (SD 1.5/2.1): + +```bash +REQUIREMENTS_FILE=requirements_legacy.txt source setup.sh +``` + Check your installed version using: `python3 -c 'import tensorrt;print(tensorrt.__version__)'` @@ -41,7 +49,7 @@ onnx 1.18.0 onnx-graphsurgeon 0.5.2 onnxruntime 1.19.2 polygraphy 0.49.22 -tensorrt 10.13.3.9 +tensorrt 10.14.1.48 tokenizers 0.13.3 torch 2.8.0a0+5228986c39.nv25.6 transformers 4.52.4 @@ -225,12 +233,18 @@ Note that a denosing-percentage is applied to the number of denoising-steps when ### Generate an image with Stable Diffusion v3.5-large with ControlNet guided by an image and a text prompt ```bash -# Depth +# Depth BF16 python3 demo_controlnet_sd35.py "a photo of a man" --controlnet-type depth --hf-token=$HF_TOKEN --denoising-steps 40 --guidance-scale 4.5 --bf16 --download-onnx-models -# Canny +# Depth FP8 +python3 demo_controlnet_sd35.py "a photo of a man" --version=3.5-large --fp8 --controlnet-type depth --download-onnx-models --denoising-steps=40 --guidance-scale 4.5 --hf-token=$HF_TOKEN + +# Canny BF16 python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models +# Canny FP8 +python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --version=3.5-large --fp8 --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --download-onnx-models + # Blur python3 demo_controlnet_sd35.py "generated ai art, a tiny, lost rubber ducky in an action shot close-up, surfing the humongous waves, inside the tube, in the style of Kelly Slater" --controlnet-type blur --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models ``` @@ -455,7 +469,37 @@ Memory usage captured below excludes the ONNX export step, and assumes use of th NOTE: The FP8 and FP4 Pipelines are supported on Hopper/Ada/Blackwell devices only. The FP4 pipeline is most performant on Blackwell devices. -### Specify Custom Paths for ONNX models and TensorRT engines (FLUX only) + +### Run Cosmos2 World Foundation Models + +Select the prompts and export them as below + +```bash +export PROMPT="A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess." + +export NEGATIVE_PROMPT="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality." +``` + +#### 1. Generate an Image from a Text Prompt + +##### Run Cosmos-Predict2-2B-Text2Image + +```bash +# BF16 +python3 demo_txt2image_cosmos.py "$PROMPT" --negative-prompt="$NEGATIVE_PROMPT" --hf-token=$HF_TOKEN +``` + +#### 2. Generate a Video guided by an Initial Video Conditioning and a Text Prompt + +##### Run Cosmos-Predict2-2B-Video2World (only PyTorch backend enabled) + +```bash +# BF16 +python3 demo_vid2world_cosmos.py "$PROMPT" --negative-prompt="$NEGATIVE_PROMPT" --hf-token=$HF_TOKEN +``` + + +### Specify Custom Paths for ONNX models and TensorRT engines (FLUX, Stable Diffusion 3.5 and Cosmos only) Custom override paths to pre-exported ONNX model files can be provided using `--custom-onnx-paths`. These ONNX models are directly used to build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list of : pairs. For example: `--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx`. Call .get_model_names(...) for the list of supported model names. @@ -467,3 +511,4 @@ Custom override paths to pre-built engine files can be provided using `--custom- - To accelerate engine building time use `--timing-cache `. The cache file will be created if it does not already exist. Note that performance may degrade if cache files are used across multiple GPU targets. It is recommended to use timing caches only during development. To achieve the best perfromance in deployment, please build engines without timing cache. - Specify new directories for storing onnx and engine files when switching between versions, LoRAs, ControlNets, etc. This can be done using `--onnx-dir ` and `--engine-dir `. - Inference performance can be improved by enabling [CUDA graphs](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-graphs) using `--use-cuda-graph`. Enabling CUDA graphs requires fixed input shapes, so this flag must be combined with `--build-static-batch` and cannot be combined with `--build-dynamic-shape`. + diff --git a/demo/Diffusion/demo_diffusion/dd_argparse.py b/demo/Diffusion/demo_diffusion/dd_argparse.py index df5f90a6..362fc36f 100644 --- a/demo/Diffusion/demo_diffusion/dd_argparse.py +++ b/demo/Diffusion/demo_diffusion/dd_argparse.py @@ -79,6 +79,10 @@ def add_arguments(parser): "flux.1-dev-canny", "flux.1-dev-depth", "flux.1-kontext-dev", + "cosmos-predict2-2b-text2image", + "cosmos-predict2-14b-text2image", + "cosmos-predict2-2b-video2world", + "cosmos-predict2-14b-video2world", ), help="Version of Stable Diffusion", ) @@ -141,7 +145,7 @@ def add_arguments(parser): "--custom-onnx-paths", type=parse_key_value_pairs, help=( - "[FLUX only] Custom override paths to pre-exported ONNX model files. These ONNX models are directly used to " + "[FLUX, Stable Diffusion 3.5-large, Cosmos only] Custom override paths to pre-exported ONNX model files. These ONNX models are directly used to " "build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list " "of : pairs. For example: " "--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx. Call " @@ -278,6 +282,7 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic is_flux = args.version.startswith("flux") is_sd35 = args.version.startswith("3.5") + is_cosmos = args.version.startswith("cosmos") if args.height % 8 != 0 or args.width % 8 != 0: raise ValueError( @@ -322,9 +327,11 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic ) # Check controlnet compatibility - if hasattr(args, "controlnet_type") and args.version != "xl-1.0": - raise ValueError("fp8 controlnet quantization is only supported for SDXL.") - + if getattr(args, "controlnet_type", None) is not None: + if args.version not in ("xl-1.0", "3.5-large"): + raise ValueError("fp8 controlnet quantization is only supported for SDXL and SD3.5-large.") + if args.version == "3.5-large" and args.controlnet_type == "blur": + raise ValueError("Blur controlnet type is not supported for SD3.5.") # Check for conflicting quantization if args.int8: raise ValueError("Cannot apply both int8 and fp8 quantization, please choose only one.") @@ -386,7 +393,9 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic # Torch-fallback and Torch-inference if args.torch_fallback and not args.torch_inference: - assert is_flux or is_sd35, "PyTorch Fallback is only supported for Flux and Stable Diffusion 3.5 pipelines." + assert ( + is_flux or is_sd35 or is_cosmos + ), "PyTorch Fallback is only supported for Flux, Stable Diffusion 3.5 and Cosmos pipelines." args.torch_fallback = args.torch_fallback.split(",") if args.torch_fallback and args.torch_inference: @@ -397,7 +406,9 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic # low-vram if args.low_vram: - assert is_flux or is_sd35, "low-vram mode is only supported for Flux and Stable Diffusion 3.5 pipelines." + assert ( + is_flux or is_sd35 or is_cosmos + ), "low-vram mode is only supported for Flux, Stable Diffusion 3.5 and Cosmos pipelines." # Pack arguments kwargs_init_pipeline = { diff --git a/demo/Diffusion/demo_diffusion/model/__init__.py b/demo/Diffusion/demo_diffusion/model/__init__.py index a74e274f..699cb3d5 100644 --- a/demo/Diffusion/demo_diffusion/model/__init__.py +++ b/demo/Diffusion/demo_diffusion/model/__init__.py @@ -26,12 +26,13 @@ from demo_diffusion.model.clip import ( SD3_T5XXLModel, get_clip_embedding_dim, ) +from demo_diffusion.model.controlnet import SD3ControlNet from demo_diffusion.model.diffusion_transformer import ( + CosmosTransformerModel, FluxTransformerModel, SD3_MMDiTModel, SD3TransformerModel, ) -from demo_diffusion.model.controlnet import SD3ControlNet from demo_diffusion.model.gan import VQGANModel from demo_diffusion.model.load import unload_torch_model from demo_diffusion.model.lora import FLUXLoraLoader, SDLoraLoader, merge_loras @@ -47,6 +48,8 @@ from demo_diffusion.model.unet import ( UNetXLModelControlNet, ) from demo_diffusion.model.vae import ( + AutoencoderKLWanEncoderModel, + AutoencoderKLWanModel, SD3_VAEDecoderModel, SD3_VAEEncoderModel, TorchVAEEncoder, @@ -67,6 +70,7 @@ __all__ = [ "SD3_T5XXLModel", "CLIPVisionWithProjModel", "CLIPImageProcessorModel", + "CosmosTransformerModel", # diffusion_transformer "SD3_MMDiTModel", "FluxTransformerModel", @@ -98,6 +102,8 @@ __all__ = [ "TorchVAEEncoder", "VAEEncoderModel", "SD3_VAEEncoderModel", + "AutoencoderKLWanModel", + "AutoencoderKLWanEncoderModel", # load "unload_torch_model", ] diff --git a/demo/Diffusion/demo_diffusion/model/base_model.py b/demo/Diffusion/demo_diffusion/model/base_model.py index 2c755562..684e9775 100644 --- a/demo/Diffusion/demo_diffusion/model/base_model.py +++ b/demo/Diffusion/demo_diffusion/model/base_model.py @@ -19,13 +19,13 @@ import json import os import numpy as np +import onnx import torch from diffusers import DiffusionPipeline +from onnx import numpy_helper -import onnx from demo_diffusion.model import load, optimizer from demo_diffusion.model.lora import merge_loras -from onnx import numpy_helper class BaseModel: @@ -50,6 +50,7 @@ class BaseModel: ): self.name = self.__class__.__name__ + self.pipeline_type = pipeline self.pipeline = pipeline.name self.version = version self.path = load.get_path(version, pipeline) @@ -70,7 +71,7 @@ class BaseModel: self.min_batch = 1 self.max_batch = max_batch_size self.min_image_shape = 256 # min image resolution: 256x256 - self.max_image_shape = 1344 # max image resolution: 1344x1344 + self.max_image_shape = 1360 # max image resolution: 1360x1360 self.min_latent_shape = self.min_image_shape // self.compression_factor self.max_latent_shape = self.max_image_shape // self.compression_factor @@ -123,6 +124,7 @@ class BaseModel: enable_lora_merge=False, static_shape=False, lora_loader=None, + dynamo=False, ): onnx_opt_graph = None # Export optimized ONNX model (if missing) @@ -135,6 +137,11 @@ class BaseModel: assert lora_loader is not None model = merge_loras(model, lora_loader) + export_kwargs = {} + if dynamo: + export_kwargs["dynamic_shapes"] = self.get_dynamic_axes() + else: + export_kwargs["dynamic_axes"] = self.get_dynamic_axes() inputs = self.get_sample_input(1, opt_image_height, opt_image_width, static_shape) torch.onnx.export( model, @@ -145,8 +152,9 @@ class BaseModel: do_constant_folding=self.do_constant_folding, input_names=self.get_input_names(), output_names=self.get_output_names(), - dynamic_axes=self.get_dynamic_axes(), verbose=False, + dynamo=dynamo, + **export_kwargs, ) if custom_model: @@ -223,7 +231,7 @@ class BaseModel: print(f"[I] Found cached weights map: {weights_map_path} ") def optimize(self, onnx_graph, return_onnx=True, **kwargs): - opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose) + opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version) opt.info(self.name + ": original") opt.cleanup() opt.info(self.name + ": cleanup") diff --git a/demo/Diffusion/demo_diffusion/model/clip.py b/demo/Diffusion/demo_diffusion/model/clip.py index f9449d2c..4db97eb9 100644 --- a/demo/Diffusion/demo_diffusion/model/clip.py +++ b/demo/Diffusion/demo_diffusion/model/clip.py @@ -174,7 +174,7 @@ class CLIPModel(base_model.BaseModel): return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device) def optimize(self, onnx_graph): - opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose) + opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version) opt.info(self.name + ": original") keep_outputs = [0, 1] if self.keep_pooled_output else [0] opt.select_outputs(keep_outputs) @@ -356,7 +356,7 @@ class SD3_CLIPGModel(CLIPModel): return output def optimize(self, onnx_graph): - opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose) + opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version) opt.info(self.name + ": original") opt.select_outputs([0, 1]) opt.cleanup() diff --git a/demo/Diffusion/demo_diffusion/model/diffusion_transformer.py b/demo/Diffusion/demo_diffusion/model/diffusion_transformer.py index 97638bbe..df0e235c 100644 --- a/demo/Diffusion/demo_diffusion/model/diffusion_transformer.py +++ b/demo/Diffusion/demo_diffusion/model/diffusion_transformer.py @@ -27,7 +27,7 @@ from demo_diffusion.utils_sd3.other_impls import load_into from demo_diffusion.utils_sd3.sd3_impls import BaseModel as BaseModelSD3 # List of models to import from diffusers.models -models_to_import = ["FluxTransformer2DModel", "SD3Transformer2DModel"] +models_to_import = ["FluxTransformer2DModel", "SD3Transformer2DModel", "CosmosTransformer3DModel"] for model in models_to_import: globals()[model] = import_from_diffusers(model, "diffusers.models") @@ -475,9 +475,8 @@ class SD3TransformerModel(base_model.BaseModel): "encoder_hidden_states", "pooled_projections", "timestep", + "block_controlnet_hidden_states" ] - if not self.fp8: - input_names.append("block_controlnet_hidden_states") return input_names def get_output_names(self): @@ -491,9 +490,8 @@ class SD3TransformerModel(base_model.BaseModel): "pooled_projections": {0: xB}, "timestep": {0: xB}, "latent": {0: xB, 2: "H", 3: "W"}, + "block_controlnet_hidden_states": {1: xB, 2: "latent_dim"} } - if not self.fp8: - dynamic_axes["block_controlnet_hidden_states"] = {1: xB, 2: "latent_dim"} return dynamic_axes def get_input_profile( @@ -535,9 +533,7 @@ class SD3TransformerModel(base_model.BaseModel): (self.xB * max_batch, self.config["pooled_projection_dim"]), ], "timestep": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)], - } - if not self.fp8: - input_profile["block_controlnet_hidden_states"] = [ + "block_controlnet_hidden_states": [ ( self.num_controlnet_layers, self.xB * min_batch, @@ -557,6 +553,7 @@ class SD3TransformerModel(base_model.BaseModel): self.config["num_attention_heads"] * self.config["attention_head_dim"], ), ] + } return input_profile @@ -568,14 +565,13 @@ class SD3TransformerModel(base_model.BaseModel): "pooled_projections": (self.xB * batch_size, self.config["pooled_projection_dim"]), "timestep": (self.xB * batch_size,), "latent": (self.xB * batch_size, self.out_channels, latent_height, latent_width), - } - if not self.fp8: - shape_dict["block_controlnet_hidden_states"] = ( + "block_controlnet_hidden_states": ( self.num_controlnet_layers, self.xB * batch_size, latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"], self.config["num_attention_heads"] * self.config["attention_head_dim"], ) + } return shape_dict def get_sample_input(self, batch_size, image_height, image_width, static_shape): @@ -600,19 +596,225 @@ class SD3TransformerModel(base_model.BaseModel): ), torch.randn(self.xB * batch_size, self.config["pooled_projection_dim"], dtype=dtype, device=self.device), torch.randn(self.xB * batch_size, dtype=torch.float32, device=self.device), + { + "block_controlnet_hidden_states": torch.randn( + self.num_controlnet_layers, + self.xB * batch_size, + latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"], + self.config["num_attention_heads"] * self.config["attention_head_dim"], + dtype=dtype, + device=self.device, + ), + } ) - if not self.fp8: - sample_input += ( - { - "block_controlnet_hidden_states": torch.randn( - self.num_controlnet_layers, - self.xB * batch_size, - latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"], - self.config["num_attention_heads"] * self.config["attention_head_dim"], - dtype=dtype, - device=self.device, - ), - } + + return sample_input + + +class CosmosTransformerModel(base_model.BaseModel): + def __init__( + self, + version, + pipeline, + device, + hf_token, + verbose, + framework_model_dir, + fp16=False, + tf32=False, + int8=False, + fp8=False, + bf16=False, + max_batch_size=16, + text_maxlen=77, + build_strongly_typed=False, + weight_streaming=False, + weight_streaming_budget_percentage=None, + ): + super(CosmosTransformerModel, self).__init__( + version, + pipeline, + device=device, + hf_token=hf_token, + verbose=verbose, + framework_model_dir=framework_model_dir, + fp16=fp16, + tf32=tf32, + int8=int8, + fp8=fp8, + bf16=bf16, + max_batch_size=max_batch_size, + text_maxlen=text_maxlen, + ) + self.subfolder = "transformer" + self.transformer_model_dir = load.get_checkpoint_dir( + self.framework_model_dir, self.version, self.pipeline, self.subfolder + ) + if not os.path.exists(self.transformer_model_dir): + self.config = CosmosTransformer3DModel.load_config(self.path, subfolder=self.subfolder, token=self.hf_token) + else: + print(f"[I] Load CosmosTransformer3DModel config from: {self.transformer_model_dir}") + self.config = CosmosTransformer3DModel.load_config(self.transformer_model_dir) + self.build_strongly_typed = build_strongly_typed + self.weight_streaming = weight_streaming + self.weight_streaming_budget_percentage = weight_streaming_budget_percentage + + def get_model(self, torch_inference=""): + model_opts = ( + {"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {} + ) + if not load.is_model_cached(self.transformer_model_dir, model_opts, self.hf_safetensor): + model = CosmosTransformer3DModel.from_pretrained( + self.path, + subfolder=self.subfolder, + use_safetensors=self.hf_safetensor, + token=self.hf_token, + **model_opts, + ).to(self.device) + model.save_pretrained(self.transformer_model_dir, **model_opts) + else: + print(f"[I] Load CosmosTransformer3DModel model from: {self.transformer_model_dir}") + model = CosmosTransformer3DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device) + if torch_inference: + model.to(memory_format=torch.channels_last) + if self.fp16: + model.transformer_blocks[6].attn1.norm_q.float().to(self.device) + + model = optimizer.optimize_checkpoint(model, torch_inference) + return model.to(self.device) + + def get_input_names(self): + input_names = [ + "hidden_states", + "timestep", + "encoder_hidden_states", + "padding_mask", + ] + if self.pipeline_type.is_video2world(): + input_names.append("fps") + input_names.append("condition_mask") + return input_names + + def get_output_names(self): + return ["latent"] + + def get_dynamic_axes(self): + dynamic_axes = { + "hidden_states": {0: "B", 2: "latent_frames", 3: "latent_H", 4: "latent_W"}, + "timestep": {0: "B"}, + "encoder_hidden_states": {0: "B"}, + "padding_mask": {0: "B", 2: "H", 3: "W"}, + } + if self.pipeline_type.is_video2world(): + dynamic_axes["fps"] = {0: "B"} + dynamic_axes["condition_mask"] = {0: "B", 2: "latent_frames", 3: "latent_H", 4: "latent_W"} + + return dynamic_axes + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + ( + min_batch, + max_batch, + min_image_height, + max_image_height, + min_image_width, + max_image_width, + min_latent_height, + max_latent_height, + min_latent_width, + max_latent_width, + ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) + latent_frames = 24 if self.pipeline_type.is_video2world() else 1 + latent_channels = ( + self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"] + ) + input_profile = { + "hidden_states": [ + (min_batch, latent_channels, latent_frames, min_latent_height, min_latent_width), + (batch_size, latent_channels, latent_frames, latent_height, latent_width), + (max_batch, latent_channels, latent_frames, max_latent_height, max_latent_width), + ], + "timestep": [(min_batch,), (batch_size,), (max_batch,)], + "encoder_hidden_states": [ + (min_batch, self.text_maxlen, self.config["text_embed_dim"]), + (batch_size, self.text_maxlen, self.config["text_embed_dim"]), + (max_batch, self.text_maxlen, self.config["text_embed_dim"]), + ], + "padding_mask": [ + (1, 1, min_image_height, min_image_width), + (1, 1, image_height, image_width), + (1, 1, max_image_height, max_image_width), + ], + } + if self.pipeline_type.is_video2world(): + input_profile["fps"] = [(min_batch,), (batch_size,), (max_batch,)] + input_profile["condition_mask"] = [ + (min_batch, 1, latent_frames, min_latent_height, min_latent_width), + (batch_size, 1, latent_frames, latent_height, latent_width), + (max_batch, 1, latent_frames, max_latent_height, max_latent_width), + ] + return input_profile + + def get_shape_dict(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + # TODO: get latent_frames from infer call + latent_frames = 24 if self.pipeline_type.is_video2world() else 1 + latent_channels = ( + self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"] + ) + shape_dict = { + "hidden_states": (batch_size, latent_channels, latent_frames, latent_height, latent_width), + "timestep": (batch_size,), + "encoder_hidden_states": (batch_size, self.text_maxlen, self.config["text_embed_dim"]), + "padding_mask": (1, 1, image_height, image_width), + "latent": (batch_size, self.config["in_channels"], latent_frames, latent_height, latent_width), + } + + if self.pipeline_type.is_video2world(): + shape_dict["fps"] = (batch_size,) + shape_dict["condition_mask"] = (batch_size, 1, latent_frames, latent_height, latent_width) + return shape_dict + + def get_sample_input(self, batch_size, image_height, image_width, static_shape): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + dtype = torch.float32 + assert not (self.fp16 and self.bf16), "fp16 and bf16 cannot be enabled simultaneously" + tensor_dtype = torch.bfloat16 if self.bf16 else (torch.float16 if self.fp16 else torch.float32) + latent_frames = 1 + latent_channels = ( + self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"] + ) + sample_input = ( + { + "hidden_states": torch.randn( + batch_size, + latent_channels, + latent_frames, + latent_height, + latent_width, + dtype=tensor_dtype, + device=self.device, + ), + "timestep": torch.tensor([1.0] * batch_size, dtype=tensor_dtype, device=self.device), + "encoder_hidden_states": torch.randn( + batch_size, self.text_maxlen, self.config["text_embed_dim"], dtype=tensor_dtype, device=self.device + ), + "padding_mask": torch.ones( + batch_size, 1, image_height, image_width, dtype=tensor_dtype, device=self.device + ), + }, + ) + if self.pipeline_type.is_video2world(): + sample_input[-1]["fps"] = torch.tensor([30] * batch_size, dtype=dtype, device=self.device) + sample_input[-1]["condition_mask"] = torch.randn( + batch_size, + 1, + latent_frames, + latent_height, + latent_width, + dtype=tensor_dtype, + device=self.device, ) return sample_input diff --git a/demo/Diffusion/demo_diffusion/model/load.py b/demo/Diffusion/demo_diffusion/model/load.py index 9ee6f40c..11cd98f3 100644 --- a/demo/Diffusion/demo_diffusion/model/load.py +++ b/demo/Diffusion/demo_diffusion/model/load.py @@ -97,6 +97,14 @@ def get_path(version: str, pipeline: "pipeline.DiffusionPipeline", controlnets: return "black-forest-labs/FLUX.1-Depth-dev" elif version == "flux.1-kontext-dev": return "black-forest-labs/FLUX.1-Kontext-dev" + elif version == "cosmos-predict2-2b-text2image": + return "nvidia/Cosmos-Predict2-2B-Text2Image" + elif version == "cosmos-predict2-14b-text2image": + return "nvidia/Cosmos-Predict2-14B-Text2Image" + elif version == "cosmos-predict2-2b-video2world": + return "nvidia/Cosmos-Predict2-2B-Video2World" + elif version == "cosmos-predict2-14b-video2world": + return "nvidia/Cosmos-Predict2-14B-Video2World" else: raise ValueError(f"Unsupported version {version} + pipeline {pipeline.name}") diff --git a/demo/Diffusion/demo_diffusion/model/optimizer.py b/demo/Diffusion/demo_diffusion/model/optimizer.py index 52e4066b..35fffbbe 100644 --- a/demo/Diffusion/demo_diffusion/model/optimizer.py +++ b/demo/Diffusion/demo_diffusion/model/optimizer.py @@ -23,6 +23,7 @@ import onnx import onnx_graphsurgeon as gs import torch from onnx import shape_inference +from onnxconverter_common.float16 import convert_float_to_float16 from polygraphy.backend.onnx.loader import fold_constants from demo_diffusion.model import load @@ -46,9 +47,11 @@ def optimize_checkpoint(model, torch_inference: str): class Optimizer: - def __init__(self, onnx_graph, verbose=False): + + def __init__(self, onnx_graph, verbose=False, version=None): self.graph = gs.import_onnx(onnx_graph) self.verbose = verbose + self.version = version def info(self, prefix): if self.verbose: @@ -186,6 +189,15 @@ class Optimizer: # Convert INT8 Zero to FP8. onnx_graph = convert_zp_fp8(onnx_graph) + # WAR for legacy SD pipelines + legacy_versions = ( + "1.4", + "1.5", + "2.1", + ) + if any(self.version.startswith(prefix) for prefix in legacy_versions): + onnx_graph = convert_float_to_float16(onnx_graph, keep_io_types=False, disable_shape_infer=True) + self.graph = gs.import_onnx(onnx_graph) # Add cast nodes to Resize I/O. cast_resize_io(self.graph) diff --git a/demo/Diffusion/demo_diffusion/model/t5.py b/demo/Diffusion/demo_diffusion/model/t5.py index 84f5dd6f..03d9f0b8 100644 --- a/demo/Diffusion/demo_diffusion/model/t5.py +++ b/demo/Diffusion/demo_diffusion/model/t5.py @@ -27,6 +27,7 @@ from demo_diffusion.model import base_model, load, optimizer class T5Model(base_model.BaseModel): + def __init__( self, version, @@ -44,6 +45,7 @@ class T5Model(base_model.BaseModel): build_strongly_typed=False, weight_streaming=False, weight_streaming_budget_percentage=None, + use_attention_mask=False, ): super(T5Model, self).__init__( version, @@ -70,6 +72,7 @@ class T5Model(base_model.BaseModel): self.build_strongly_typed = build_strongly_typed self.weight_streaming = weight_streaming self.weight_streaming_budget_percentage = weight_streaming_budget_percentage + self.use_attention_mask = use_attention_mask def get_model(self, torch_inference=""): model_opts = ( @@ -91,12 +94,16 @@ class T5Model(base_model.BaseModel): return model def get_input_names(self): + if self.use_attention_mask: + return ["input_ids", "attention_mask"] return ["input_ids"] def get_output_names(self): return ["text_embeddings"] def get_dynamic_axes(self): + if self.use_attention_mask: + return {"input_ids": {0: "B"}, "attention_mask": {0: "B"}, "text_embeddings": {0: "B"}} return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}} def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): @@ -104,9 +111,16 @@ class T5Model(base_model.BaseModel): min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims( batch_size, image_height, image_width, static_batch, static_shape ) - return { + profile = { "input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)] } + if self.use_attention_mask: + profile["attention_mask"] = [ + (min_batch, self.text_maxlen), + (batch_size, self.text_maxlen), + (max_batch, self.text_maxlen), + ] + return profile def get_shape_dict(self, batch_size, image_height, image_width): self.check_dims(batch_size, image_height, image_width) @@ -114,8 +128,13 @@ class T5Model(base_model.BaseModel): "input_ids": (batch_size, self.text_maxlen), "text_embeddings": (batch_size, self.text_maxlen, self.config.d_model), } + if self.use_attention_mask: + output["attention_mask"] = (batch_size, self.text_maxlen) return output def get_sample_input(self, batch_size, image_height, image_width, static_shape): self.check_dims(batch_size, image_height, image_width) - return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device) + inputs = {"input_ids": torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)} + if self.use_attention_mask: + inputs["attention_mask"] = torch.ones(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device) + return inputs diff --git a/demo/Diffusion/demo_diffusion/model/vae.py b/demo/Diffusion/demo_diffusion/model/vae.py index ed76c1e1..26ceaf82 100644 --- a/demo/Diffusion/demo_diffusion/model/vae.py +++ b/demo/Diffusion/demo_diffusion/model/vae.py @@ -27,10 +27,7 @@ from demo_diffusion.utils_sd3.other_impls import load_into from demo_diffusion.utils_sd3.sd3_impls import SDVAE # List of models to import from diffusers.models -models_to_import = [ - "AutoencoderKL", - "AutoencoderKLTemporalDecoder", -] +models_to_import = ["AutoencoderKL", "AutoencoderKLTemporalDecoder", "AutoencoderKLWan"] for model in models_to_import: globals()[model] = import_from_diffusers(model, "diffusers.models") @@ -507,3 +504,163 @@ class SD3_VAEEncoderModel(base_model.BaseModel): def get_sample_input(self, batch_size, image_height, image_width, static_shape): dtype = torch.float16 if self.fp16 else torch.float32 return torch.randn(batch_size, 3, image_height, image_width, dtype=dtype, device=self.device) + + +class AutoencoderKLWanModel(base_model.BaseModel): + + def __init__( + self, + version, + pipeline, + device, + hf_token, + verbose, + framework_model_dir, + fp16=False, + tf32=False, + bf16=False, + max_batch_size=16, + ): + super(AutoencoderKLWanModel, self).__init__( + version, + pipeline, + device=device, + hf_token=hf_token, + verbose=verbose, + framework_model_dir=framework_model_dir, + fp16=fp16, + tf32=tf32, + bf16=bf16, + max_batch_size=max_batch_size, + ) + self.subfolder = "vae" + self.vae_decoder_model_dir = load.get_checkpoint_dir( + self.framework_model_dir, self.version, self.pipeline, self.subfolder + ) + if not os.path.exists(self.vae_decoder_model_dir): + self.config = AutoencoderKLWan.load_config(self.path, subfolder=self.subfolder, token=self.hf_token) + else: + print(f"[I] Load AutoencoderKLWan (decoder) config from: {self.vae_decoder_model_dir}") + self.config = AutoencoderKLWan.load_config(self.vae_decoder_model_dir) + + def get_model(self, torch_inference=""): + model_opts = ( + {"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {} + ) + if not load.is_model_cached(self.vae_decoder_model_dir, model_opts, self.hf_safetensor): + model = AutoencoderKLWan.from_pretrained( + self.path, + subfolder=self.subfolder, + use_safetensors=self.hf_safetensor, + token=self.hf_token, + **model_opts, + ).to(self.device) + model.save_pretrained(self.vae_decoder_model_dir, **model_opts) + else: + print(f"[I] Load AutoencoderKLWan (decoder) model from: {self.vae_decoder_model_dir}") + model = AutoencoderKLWan.from_pretrained(self.vae_decoder_model_dir, **model_opts).to(self.device) + model.forward = model.decode + model = optimizer.optimize_checkpoint(model, torch_inference) + return model + + def get_input_names(self): + return ["latent"] + + def get_output_names(self): + return ["images"] + + def get_dynamic_axes(self): + return {"latent": {0: "B", 3: "H", 4: "W"}, "images": {0: "B", 3: "8H", 4: "8W"}} + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = ( + self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) + ) + return { + "latent": [ + (min_batch, self.config["z_dim"], 1, min_latent_height, min_latent_width), + (batch_size, self.config["z_dim"], 1, latent_height, latent_width), + (max_batch, self.config["z_dim"], 1, max_latent_height, max_latent_width), + ] + } + + def get_shape_dict(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + return { + "latent": (batch_size, self.config["z_dim"], 1, latent_height, latent_width), + "images": (batch_size, 3, 1, image_height, image_width), + } + + def get_sample_input(self, batch_size, image_height, image_width, static_shape): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32 + return torch.randn( + batch_size, self.config["z_dim"], 1, latent_height, latent_width, dtype=dtype, device=self.device + ) + + +class AutoencoderKLWanEncoderModelWrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, x): + return self.model.encode(x).latent_dist.sample() + + +class AutoencoderKLWanEncoderModel(base_model.BaseModel): + def __init__( + self, + version, + pipeline, + device, + hf_token, + verbose, + framework_model_dir, + fp16=False, + tf32=False, + bf16=False, + max_batch_size=16, + ): + super(AutoencoderKLWanEncoderModel, self).__init__( + version, + pipeline, + device=device, + hf_token=hf_token, + verbose=verbose, + framework_model_dir=framework_model_dir, + fp16=fp16, + tf32=tf32, + bf16=bf16, + max_batch_size=max_batch_size, + ) + self.subfolder = "vae" + self.vae_encoder_model_dir = load.get_checkpoint_dir( + self.framework_model_dir, self.version, self.pipeline, self.subfolder + ) + if not os.path.exists(self.vae_encoder_model_dir): + self.config = AutoencoderKLWan.load_config(self.path, subfolder=self.subfolder, token=self.hf_token) + else: + print(f"[I] Load AutoencoderKLWan (encoder) config from: {self.vae_encoder_model_dir}") + self.config = AutoencoderKLWan.load_config(self.vae_encoder_model_dir) + + def get_model(self, torch_inference=""): + model_opts = ( + {"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {} + ) + if not load.is_model_cached(self.vae_encoder_model_dir, model_opts, self.hf_safetensor): + model = AutoencoderKLWan.from_pretrained( + self.path, + subfolder=self.subfolder, + use_safetensors=self.hf_safetensor, + token=self.hf_token, + **model_opts, + ).to(self.device) + model.save_pretrained(self.vae_encoder_model_dir, **model_opts) + else: + print(f"[I] Load AutoencoderKLWan (encoder) model from: {self.vae_encoder_model_dir}") + model = AutoencoderKLWan.from_pretrained(self.vae_encoder_model_dir, **model_opts).to(self.device) + model = AutoencoderKLWanEncoderModelWrapper(model) + model = optimizer.optimize_checkpoint(model, torch_inference) + return model diff --git a/demo/Diffusion/demo_diffusion/pipeline/__init__.py b/demo/Diffusion/demo_diffusion/pipeline/__init__.py index 5fafb745..cba7b441 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/__init__.py +++ b/demo/Diffusion/demo_diffusion/pipeline/__init__.py @@ -14,19 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import importlib +from typing import TYPE_CHECKING -from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline -from demo_diffusion.pipeline.flux_pipeline import FluxKontextPipeline, FluxPipeline -from demo_diffusion.pipeline.stable_cascade_pipeline import StableCascadePipeline -from demo_diffusion.pipeline.stable_diffusion_3_pipeline import StableDiffusion3Pipeline -from demo_diffusion.pipeline.stable_diffusion_35_pipeline import ( - StableDiffusion35Pipeline, -) -from demo_diffusion.pipeline.stable_diffusion_pipeline import StableDiffusionPipeline -from demo_diffusion.pipeline.stable_video_diffusion_pipeline import ( - StableVideoDiffusionPipeline, -) -from demo_diffusion.pipeline.type import PIPELINE_TYPE +# Expose public API while avoiding importing optional dependencies at module import time. +# Each attribute is imported on first access via __getattr__. __all__ = [ "DiffusionPipeline", @@ -36,6 +28,73 @@ __all__ = [ "StableDiffusion3Pipeline", "StableDiffusion35Pipeline", "StableDiffusionPipeline", + "CosmosPipeline", "StableVideoDiffusionPipeline", "PIPELINE_TYPE", ] + +_LAZY_ATTRS = { + # Core/base + "DiffusionPipeline": ("demo_diffusion.pipeline.diffusion_pipeline", "DiffusionPipeline"), + "PIPELINE_TYPE": ("demo_diffusion.pipeline.type", "PIPELINE_TYPE"), + # Stable Diffusion family + "StableDiffusionPipeline": ("demo_diffusion.pipeline.stable_diffusion_pipeline", "StableDiffusionPipeline"), + "StableDiffusion3Pipeline": ("demo_diffusion.pipeline.stable_diffusion_3_pipeline", "StableDiffusion3Pipeline"), + "StableDiffusion35Pipeline": ("demo_diffusion.pipeline.stable_diffusion_35_pipeline", "StableDiffusion35Pipeline"), + # Stable Cascade + "StableCascadePipeline": ("demo_diffusion.pipeline.stable_cascade_pipeline", "StableCascadePipeline"), + # Stable Video Diffusion + "StableVideoDiffusionPipeline": ("demo_diffusion.pipeline.stable_video_diffusion_pipeline", "StableVideoDiffusionPipeline"), + # Flux family (optional dependency: `flux`) + "FluxPipeline": ("demo_diffusion.pipeline.flux_pipeline", "FluxPipeline"), + "FluxKontextPipeline": ("demo_diffusion.pipeline.flux_pipeline", "FluxKontextPipeline"), + # Cosmos (optional dependency: `flux`) + "CosmosPipeline": ("demo_diffusion.pipeline.cosmos_pipeline", "CosmosPipeline"), +} + +def __getattr__(name): + if name not in _LAZY_ATTRS: + raise AttributeError(f"module 'demo_diffusion.pipeline' has no attribute {name!r}") + + module_path, attr_name = _LAZY_ATTRS[name] + try: + module = importlib.import_module(module_path) + except ModuleNotFoundError as e: + missing_pkg = e.name or "" + raise ModuleNotFoundError( + f"Optional dependency '{missing_pkg}' is required for '{name}'. " + "Install the appropriate extras/requirements for the selected pipeline " + "(e.g., use the non-legacy requirements for Flux/Cosmos), or install the missing package." + ) from e + try: + return getattr(module, attr_name) + except AttributeError as e: + raise AttributeError( + f"'{module_path}' does not export attribute '{attr_name}' (while resolving '{name}')." + ) from e + + +if TYPE_CHECKING: + from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline as DiffusionPipeline + from demo_diffusion.pipeline.type import PIPELINE_TYPE as PIPELINE_TYPE + from demo_diffusion.pipeline.stable_diffusion_pipeline import ( + StableDiffusionPipeline as StableDiffusionPipeline, + ) + from demo_diffusion.pipeline.stable_diffusion_3_pipeline import ( + StableDiffusion3Pipeline as StableDiffusion3Pipeline, + ) + from demo_diffusion.pipeline.stable_diffusion_35_pipeline import ( + StableDiffusion35Pipeline as StableDiffusion35Pipeline, + ) + from demo_diffusion.pipeline.stable_cascade_pipeline import ( + StableCascadePipeline as StableCascadePipeline, + ) + from demo_diffusion.pipeline.stable_video_diffusion_pipeline import ( + StableVideoDiffusionPipeline as StableVideoDiffusionPipeline, + ) + # Optional pipelines + from demo_diffusion.pipeline.flux_pipeline import ( + FluxPipeline as FluxPipeline, + FluxKontextPipeline as FluxKontextPipeline, + ) + from demo_diffusion.pipeline.cosmos_pipeline import CosmosPipeline as CosmosPipeline diff --git a/demo/Diffusion/demo_diffusion/pipeline/cosmos_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/cosmos_pipeline.py new file mode 100644 index 00000000..aeca1bb8 --- /dev/null +++ b/demo/Diffusion/demo_diffusion/pipeline/cosmos_pipeline.py @@ -0,0 +1,1009 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from __future__ import annotations + +import argparse +import inspect +import os +import random +import time +import warnings +from typing import Any, List + +import numpy as np +import tensorrt as trt +import torch +from cuda.bindings import runtime as cudart +from diffusers.video_processor import VideoProcessor +from flux.content_filters import PixtralContentFilter +from tqdm import tqdm + +from demo_diffusion import path as path_module +from demo_diffusion.model import ( + AutoencoderKLWanEncoderModel, + AutoencoderKLWanModel, + CosmosTransformerModel, + T5Model, + make_tokenizer, +) +from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline +from demo_diffusion.pipeline.type import PIPELINE_TYPE + +TRT_LOGGER = trt.Logger(trt.Logger.ERROR) + + +class CosmosPipeline(DiffusionPipeline): + """ + Application showcasing the acceleration of Cosmos pipelines using Nvidia TensorRT. + """ + + def __init__( + self, + version="cosmos-predict2-2b", + pipeline_type=PIPELINE_TYPE.TXT2IMG, + guidance_scale=6.0, + max_sequence_length=512, + t5_weight_streaming_budget_percentage=None, + transformer_weight_streaming_budget_percentage=None, + **kwargs, + ): + """ + Initializes the Cosmos pipeline. + + Args: + version (`str`, defaults to `cosmos-1.0-7B`) + Version of the underlying Cosmos model. + guidance_scale (`float`, defaults to 3.5): + Guidance scale is enabled by setting as > 1. + Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality. + max_sequence_length (`int`, defaults to 512): + Maximum sequence length to use with the `prompt`. + t5_weight_streaming_budget_percentage (`int`, defaults to None): + Weight streaming budget as a percentage of the size of total streamable weights for the T5 model. + transformer_weight_streaming_budget_percentage (`int`, defaults to None): + Weight streaming budget as a percentage of the size of total streamable weights for the Transformer model. + """ + super().__init__( + version=version, + pipeline_type=pipeline_type, + text_encoder_weight_streaming_budget_percentage=t5_weight_streaming_budget_percentage, + denoiser_weight_streaming_budget_percentage=transformer_weight_streaming_budget_percentage, + **kwargs, + ) + self.guidance_scale = guidance_scale + self.max_sequence_length = max_sequence_length + self.do_classifier_free_guidance = self.guidance_scale > 1 + + # WAR ONNX export error: Exporting the operator 'aten::_upsample_nearest_exact2d' to ONNX opset version 19 is not supported + self.config["vae_torch_fallback"] = True + self.config["vae_encoder_torch_fallback"] = True + + @classmethod + def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> CosmosPipeline: + """Factory method to construct a `CosmosPipeline` object from parsed arguments. + + Overrides: + DiffusionPipeline.FromArgs + """ + MAX_BATCH_SIZE = 4 + DEVICE = "cuda" + DO_RETURN_LATENTS = False + + # Resolve all paths. + dd_path = path_module.resolve_path( + cls.get_model_names(pipeline_type), args, pipeline_type, cls._get_pipeline_uid(args.version) + ) + + return cls( + dd_path=dd_path, + version=args.version, + pipeline_type=pipeline_type, + guidance_scale=args.guidance_scale, + max_sequence_length=args.max_sequence_length, + bf16=args.bf16, + low_vram=args.low_vram, + torch_fallback=args.torch_fallback, + weight_streaming=args.ws, + t5_weight_streaming_budget_percentage=args.t5_ws_percentage, + transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage, + max_batch_size=MAX_BATCH_SIZE, + denoising_steps=args.denoising_steps, + scheduler=args.scheduler, + device=DEVICE, + output_dir=args.output_dir, + hf_token=args.hf_token, + verbose=args.verbose, + nvtx_profile=args.nvtx_profile, + use_cuda_graph=args.use_cuda_graph, + framework_model_dir=args.framework_model_dir, + return_latents=DO_RETURN_LATENTS, + torch_inference=args.torch_inference, + ) + + @classmethod + def get_model_names(cls, pipeline_type: PIPELINE_TYPE, controlnet_type: str = None) -> List[str]: + """Return a list of model names used by this pipeline. + + Overrides: + DiffusionPipeline.get_model_names + """ + if pipeline_type.is_video2world(): + return ["vae_encoder", "t5", "transformer", "vae"] + return ["t5", "transformer", "vae"] + + def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None: + raise ValueError("ONNX models download is not supported for the Cosmos Pipeline") + + def _initialize_models(self, framework_model_dir, int8, fp8, fp4): + # Load text tokenizer(s) + self.tokenizer = make_tokenizer( + self.version, self.pipeline_type, self.hf_token, framework_model_dir, tokenizer_type="t5" + ) + + # Load pipeline models + models_args = { + "version": self.version, + "pipeline": self.pipeline_type, + "device": self.device, + "hf_token": self.hf_token, + "verbose": self.verbose, + "framework_model_dir": framework_model_dir, + "max_batch_size": self.max_batch_size, + } + + self.fp16 = True if not self.bf16 else False + self.tf32 = True + if "t5" in self.stages: + # Known accuracy issues with FP16 + self.models["t5"] = T5Model( + **models_args, + fp16=self.fp16, + tf32=self.tf32, + bf16=self.bf16, + text_maxlen=self.max_sequence_length, + build_strongly_typed=True, + use_attention_mask=True, + ) + + if "transformer" in self.stages: + self.models["transformer"] = CosmosTransformerModel( + **models_args, + bf16=self.bf16, + fp16=self.fp16, + int8=int8, + fp8=fp8, + tf32=self.tf32, + text_maxlen=self.max_sequence_length, + build_strongly_typed=True, + weight_streaming=self.weight_streaming, + weight_streaming_budget_percentage=self.denoiser_weight_streaming_budget_percentage, + ) + + if "vae" in self.stages: + self.models["vae"] = AutoencoderKLWanModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16) + + if "vae_encoder" in self.stages: + self.models["vae_encoder"] = AutoencoderKLWanEncoderModel( + **models_args, fp16=False, tf32=self.tf32, bf16=self.bf16 + ) + + self.vae_scale_factor_temporal = ( + 2 ** sum(self.models["vae"].config["temperal_downsample"]) + if "vae" in self.stages and self.models["vae"] is not None + else 4 + ) + self.vae_scale_factor_spatial = ( + 2 ** len(self.models["vae"].config["temperal_downsample"]) + if "vae" in self.stages and self.models["vae"] is not None + else 8 + ) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + def encode_video(self, video): + self.profile_start("vae_encoder", color="red") + cast_to = ( + torch.float16 + if self.models["vae_encoder"].fp16 + else torch.bfloat16 if self.models["vae_encoder"].bf16 else torch.float32 + ) + video = video.to(dtype=cast_to) + if self.torch_inference: + image_latents = self.torch_models["vae_encoder"](video) + else: + image_latents = self.run_engine("vae_encoder", {"images": video})["latent"] + self.profile_stop("vae_encoder") + return image_latents + + def initialize_latents_text2image( + self, + batch_size, + num_channels_latents, + num_latent_frames, + latent_height, + latent_width, + latents_dtype=torch.float32, + ): + latents_shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) + latents = torch.randn( + latents_shape, + device=self.device, + dtype=latents_dtype, + generator=self.generator, + ) + + return latents * self.scheduler.config.sigma_max + + def initialize_latents_video2world( + self, + video, + batch_size, + num_channels_latents, + num_frames, + latent_height, + latent_width, + latents_dtype=torch.float32, + do_classifier_free_guidance=False, + ): + num_cond_frames = video.size(2) + if num_cond_frames >= num_frames: + # Take the last `num_frames` frames for conditioning + num_cond_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + video = video[:, :, -num_frames:] + else: + num_cond_latent_frames = (num_cond_frames - 1) // self.vae_scale_factor_temporal + 1 + num_padding_frames = num_frames - num_cond_frames + last_frame = video[:, :, -1:] + padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1) + video = torch.cat([video, padding], dim=2) + + # Encode video + with self.model_memory_manager(["vae_encoder"], low_vram=self.low_vram): + video_latents = self.encode_video( + video=video, + ) + + latents_mean = ( + torch.tensor(self.models["vae"].config["latents_mean"]) + .view(1, self.models["vae"].config["z_dim"], 1, 1, 1) + .to(self.device, latents_dtype) + ) + latents_std = ( + torch.tensor(self.models["vae"].config["latents_std"]) + .view(1, self.models["vae"].config["z_dim"], 1, 1, 1) + .to(self.device, latents_dtype) + ) + init_latents = (video_latents - latents_mean) / latents_std * self.scheduler.config.sigma_data + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) + + latents = torch.randn( + shape, + device=self.device, + dtype=latents_dtype, + generator=self.generator, + ) + + latents = latents * self.scheduler.config.sigma_max + + padding_shape = (batch_size, 1, num_latent_frames, latent_height, latent_width) + ones_padding = latents.new_ones(padding_shape) + zeros_padding = latents.new_zeros(padding_shape) + + cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) + cond_indicator[:, :, :num_cond_latent_frames] = 1.0 + cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding + + uncond_indicator = uncond_mask = None + if do_classifier_free_guidance: + uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) + uncond_indicator[:, :, :num_cond_latent_frames] = 1.0 + uncond_mask = uncond_indicator * ones_padding + (1 - uncond_indicator) * zeros_padding + + return latents, init_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask + + # Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_img2img.py#L416C1 + def get_timesteps(self, num_inference_steps, strength): + # get the original timestep using init_timestep + init_timestep = min(num_inference_steps * strength, num_inference_steps) + + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + if hasattr(self.scheduler, "set_begin_index"): + self.scheduler.set_begin_index(t_start * self.scheduler.order) + + return timesteps, num_inference_steps - t_start + + def _duplicate_text_embeddings(self, batch_size, text_embeddings, num_outputs_per_prompt): + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = text_embeddings.shape + text_embeddings = text_embeddings.repeat(1, num_outputs_per_prompt, 1) + text_embeddings = text_embeddings.view(batch_size * num_outputs_per_prompt, seq_len, -1) + return text_embeddings + + def _prepare_timesteps(self, num_inference_steps): + """Prepare timesteps for the scheduler.""" + sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype) + accept_sigmas = "sigmas" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + self.scheduler.set_timesteps(sigmas=sigmas, device=self.device) + timesteps = self.scheduler.timesteps + num_inference_steps = len(timesteps) + if self.scheduler.config.get("final_sigmas_type", "zero") == "sigma_min": + # Replace the last sigma (which is zero) with the minimum sigma value + self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2] + return timesteps, num_inference_steps + + def _encode_text_prompts(self, prompt, negative_prompt, batch_size, num_outputs_per_prompt): + """Encode text prompts using T5 encoder.""" + with self.model_memory_manager(["t5"], low_vram=self.low_vram): + text_embeddings = self.encode_prompt(prompt) + text_embeddings = self._duplicate_text_embeddings(batch_size, text_embeddings, num_outputs_per_prompt) + negative_text_embeddings = None + if self.do_classifier_free_guidance: + negative_text_embeddings = self.encode_prompt(negative_prompt) + negative_text_embeddings = self._duplicate_text_embeddings( + batch_size, negative_text_embeddings, num_outputs_per_prompt + ) + return text_embeddings, negative_text_embeddings + + def _get_latents_normalization_params(self, device, dtype): + """Get latents normalization parameters from VAE config.""" + latents_mean = ( + torch.tensor(self.models["vae"].config["latents_mean"]) + .view(1, self.models["vae"].config["z_dim"], 1, 1, 1) + .to(device, dtype) + ) + latents_std = ( + torch.tensor(self.models["vae"].config["latents_std"]) + .view(1, self.models["vae"].config["z_dim"], 1, 1, 1) + .to(device, dtype) + ) + return latents_mean, latents_std + + def _normalize_and_decode_latents(self, latents, is_video2world=False): + """Normalize latents and decode using VAE.""" + latents_mean, latents_std = self._get_latents_normalization_params(latents.device, latents.dtype) + + if is_video2world: + # For video2world: latents * std / sigma_data + mean + latents = latents * latents_std / self.scheduler.config.sigma_data + latents_mean + else: + # For text2image: latents / (1/std) / sigma_data + mean + latents_std_inv = 1.0 / latents_std + latents = latents / latents_std_inv / self.scheduler.config.sigma_data + latents_mean + + with self.model_memory_manager(["vae"], low_vram=self.low_vram): + video = self.decode_latent(latents) + + return video + + def encode_prompt(self, prompt, encoder="t5"): + self.profile_start(encoder, color="green") + + def tokenize(prompt): + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.max_sequence_length, + truncation=True, + return_overflowing_tokens=False, + return_length=False, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(self.device) + attention_mask = text_inputs.attention_mask.bool().to(self.device) + + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(self.device) + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.max_sequence_length - 1 : -1]) + warnings.warn( + "The following part of your input was truncated because `max_sequence_length` is set to " + f"{self.max_sequence_length} tokens: {removed_text}" + ) + + if self.torch_inference or self.torch_fallback[encoder]: + text_encoder_output = self.torch_models[encoder]( + text_input_ids, attention_mask=attention_mask + ).last_hidden_state + else: + # NOTE: output tensor for the encoder must be cloned because it will be overwritten when called again for prompt2 + text_encoder_output = self.run_engine( + encoder, {"input_ids": text_input_ids, "attention_mask": attention_mask} + )["text_embeddings"] + + lengths = attention_mask.sum(dim=1).cpu() + for i, length in enumerate(lengths): + text_encoder_output[i, length:] = 0 + return text_encoder_output + + # Tokenize prompt + text_encoder_output = tokenize(prompt) + + self.profile_stop(encoder) + return ( + text_encoder_output.to(torch.float16) + if self.fp16 + else text_encoder_output.to(torch.bfloat16) if self.bf16 else text_encoder_output.to(torch.float32) + ) + + def denoise_latent( + self, + latents, + timesteps, + text_embeddings, + negative_text_embeddings, + padding_mask, + denoiser="transformer", + ): + do_autocast = self.torch_inference != "" and self.models[denoiser].fp16 + with torch.autocast("cuda", enabled=do_autocast, dtype=torch.float32): + self.profile_start(denoiser, color="blue") + + for step_index, timestep in tqdm(enumerate(timesteps), total=len(timesteps), desc="Denoising"): + # Prepare latents + cast_to = ( + torch.float16 + if self.models[denoiser].fp16 + else torch.bfloat16 if self.models[denoiser].bf16 else torch.float32 + ) + current_sigma = self.scheduler.sigmas[step_index] + current_t = current_sigma / (current_sigma + 1) + c_in = 1 - current_t + c_skip = 1 - current_t + c_out = -current_t + timestep_inp = current_t.expand(latents.shape[0]).to(cast_to) # [B, 1, T, 1, 1] + latents_input = (latents * c_in).to(cast_to) + # prepare inputs + params = { + "hidden_states": latents_input, + "timestep": timestep_inp, + "encoder_hidden_states": text_embeddings, + "padding_mask": padding_mask, + } + + if self.torch_inference or self.torch_fallback[denoiser]: + noise_pred = self.torch_models[denoiser](**params)["sample"] + else: + noise_pred = self.run_engine(denoiser, params)["latent"].clone() + + noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(cast_to) + if self.do_classifier_free_guidance: + params = { + "hidden_states": latents_input, + "timestep": timestep_inp, + "encoder_hidden_states": negative_text_embeddings, + "padding_mask": padding_mask, + } + + # Predict the noise residual + if self.torch_inference or self.torch_fallback[denoiser]: + noise_pred_uncond = self.torch_models[denoiser](**params)["sample"] + else: + noise_pred_uncond = self.run_engine(denoiser, params)["latent"].clone() + + noise_pred_uncond = (c_skip * latents + c_out * noise_pred_uncond.float()).to(cast_to) + noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond) + + noise_pred = (latents - noise_pred) / current_sigma + latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False)[0] + + self.profile_stop(denoiser) + return latents.to(dtype=torch.bfloat16) if self.bf16 else latents.to(dtype=torch.float32) + + def denoise_latent_video2world( + self, + latents, + timesteps, + text_embeddings, + negative_text_embeddings, + padding_mask, + fps, + cond_mask, + uncond_mask, + t_conditioning, + cond_indicator, + conditioning_latents, + uncond_indicator, + unconditioning_latents, + denoiser="transformer", + ): + do_autocast = self.torch_inference != "" and self.models[denoiser].fp16 + with torch.autocast("cuda", enabled=do_autocast, dtype=torch.float32): + self.profile_start(denoiser, color="blue") + + for step_index, timestep in tqdm(enumerate(timesteps), total=len(timesteps), desc="Denoising"): + # Prepare latents + cast_to = ( + torch.float16 + if self.models[denoiser].fp16 + else torch.bfloat16 if self.models[denoiser].bf16 else torch.float32 + ) + current_sigma = self.scheduler.sigmas[step_index] + current_t = current_sigma / (current_sigma + 1) + c_in = 1 - current_t + c_skip = 1 - current_t + c_out = -current_t + timestep_inp = current_t.view(1, 1, 1, 1, 1).expand( + latents.size(0), -1, latents.size(2), -1, -1 + ) # [B, 1, T, 1, 1] + latents_input = latents * c_in + latents_input = (cond_indicator * conditioning_latents + (1 - cond_indicator) * latents_input).to( + cast_to + ) + timestep_inp = (cond_indicator * t_conditioning + (1 - cond_indicator) * timestep_inp).to(cast_to) + + # prepare inputs + params = { + "hidden_states": latents_input, + "timestep": timestep_inp, + "encoder_hidden_states": text_embeddings, + "padding_mask": padding_mask, + "fps": fps, + "condition_mask": cond_mask, + } + + if self.torch_inference or self.torch_fallback[denoiser]: + noise_pred = self.torch_models[denoiser](**params)["sample"] + else: + noise_pred = self.run_engine(denoiser, params)["latent"].clone() + + noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(cast_to) + noise_pred = cond_indicator * conditioning_latents + (1 - cond_indicator) * noise_pred + if self.do_classifier_free_guidance: + latents_input = latents * c_in + latents_input = ( + uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * latents_input + ).to(cast_to) + timestep_inp = (uncond_indicator * t_conditioning + (1 - uncond_indicator) * timestep_inp).to( + cast_to + ) + params = { + "hidden_states": latents_input, + "timestep": timestep_inp, + "encoder_hidden_states": negative_text_embeddings, + "padding_mask": padding_mask, + "fps": fps, + "condition_mask": uncond_mask, + } + + # Predict the noise residual + if self.torch_inference or self.torch_fallback[denoiser]: + noise_pred_uncond = self.torch_models[denoiser](**params)["sample"] + else: + noise_pred_uncond = self.run_engine(denoiser, params)["latent"].clone() + + noise_pred_uncond = (c_skip * latents + c_out * noise_pred_uncond.float()).to(cast_to) + noise_pred_uncond = ( + uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * noise_pred_uncond + ) + noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond) + + noise_pred = (latents - noise_pred) / current_sigma + latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False)[0] + + self.profile_stop(denoiser) + return latents.to(dtype=torch.bfloat16) if self.bf16 else latents.to(dtype=torch.float32) + + def decode_latent(self, latents, decoder="vae"): + self.profile_start(decoder, color="red") + cast_to = ( + torch.float16 + if self.models[decoder].fp16 + else torch.bfloat16 if self.models[decoder].bf16 else torch.float32 + ) + latents = latents.to(dtype=cast_to) + + if self.torch_inference or self.torch_fallback[decoder]: + video = self.torch_models[decoder](latents, return_dict=False)[0] + else: + video = self.run_engine(decoder, {"latent": latents})["frames"] + + self.profile_stop(decoder) + return video + + def post_process_video(self, video): + # Post-process video + video = self.video_processor.postprocess_video(video, output_type="np") + video = (video * 255).astype(np.uint8) + video_batch = [] + for vid in video: + # vid = self.safety_checker.check_video_safety(vid) + video_batch.append(vid) + video = np.stack(video_batch).astype(np.float32) / 255.0 * 2 - 1 + video = torch.from_numpy(video).permute(0, 4, 1, 2, 3) + video = self.video_processor.postprocess_video(video, output_type="pil") + + if self.pipeline_type.is_video2world(): + return video + + image = [batch[0] for batch in video] + if isinstance(video, torch.Tensor): + image = torch.stack(image) + elif isinstance(video, np.ndarray): + image = np.stack(image) + + return image + + def _finalize_generation(self, video, walltime_ms, num_inference_steps, batch_size, warmup, save_output): + """Handle post-processing, saving, and performance reporting.""" + if not warmup: + self.print_summary(num_inference_steps, walltime_ms, batch_size) + if not self.return_latents and save_output: + # post-process video + processed_output = self.post_process_video(video) + + # save output + if self.pipeline_type.is_video2world(): + return (processed_output[0], walltime_ms) + return (np.array(processed_output), walltime_ms) + + def _check_integrity(self, images): + integrity_checker = PixtralContentFilter(self.device) + for image in images: + image_ = np.array(image) / 255.0 + image_ = 2 * image_ - 1 + image_ = torch.from_numpy(image_).to(self.device, dtype=torch.float32).permute(0, 3, 1, 2) + if integrity_checker.test_image(image_): + raise ValueError("Your image has been flagged. Choose another prompt/image or try again.") + + def save_images(self, prompt, images, check_integrity=False): + if check_integrity: + self._check_integrity(images) + for image in images: + self.save_image(image, self.pipeline_type.name.lower(), prompt, self.seed) + + def save_video( + self, + prompt, + videos, + check_integrity=False, + ): + for frames in videos: + if check_integrity: + self._check_integrity([frames]) + prompt_prefix = "".join(set([prompt[i].replace(" ", "_")[:10] for i in range(len(prompt))])) + video_name_prefix = "-".join( + [self.pipeline_type.name.lower(), "fp16", str(self.seed), str(random.randint(1000, 9999))] + ) + video_name_suffix = "torch" if self.torch_inference else "trt" + video_path = prompt_prefix + "-" + video_name_prefix + "-" + video_name_suffix + ".gif" + print(f"Saving video to: {video_path}") + frames[0].save( + os.path.join(self.output_dir, video_path), + save_all=True, + optimize=False, + append_images=frames[1:], + loop=0, + ) + + def print_summary(self, denoising_steps, walltime_ms, batch_size): + print("|-----------------|--------------|") + print("| {:^15} | {:^12} |".format("Module", "Latency")) + print("|-----------------|--------------|") + for stage in self.stages: + print( + "| {:^15} | {:>9.2f} ms |".format( + stage + " x " + str(denoising_steps) if stage == "transformer" else stage, + cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1], + ) + ) + print("|-----------------|--------------|") + print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms)) + print("|-----------------|--------------|") + print("Throughput: {:.2f} image/s".format(batch_size * 1000.0 / walltime_ms)) + + def generate_image( + self, + prompt, + negative_prompt, + image_height, + image_width, + num_frames=1, + num_images_per_prompt=1, + save_image=True, + warmup=False, + ): + batch_size = len(prompt) + + # Spatial dimensions of latent tensor + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + latent_height = image_height // self.vae_scale_factor_spatial + latent_width = image_width // self.vae_scale_factor_spatial + + num_inference_steps = self.denoising_steps + + with torch.inference_mode(), trt.Runtime(TRT_LOGGER): + torch.cuda.synchronize() + e2e_tic = time.perf_counter() + + # Prepare timesteps + timesteps, num_inference_steps = self._prepare_timesteps(num_inference_steps) + + # T5 text encoder + text_embeddings, negative_text_embeddings = self._encode_text_prompts( + prompt, negative_prompt, batch_size, num_images_per_prompt + ) + + num_channels_latents = self.models["transformer"].config["in_channels"] + latents_dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32 + + # Initialize latents + latents = self.initialize_latents_text2image( + batch_size=batch_size, + num_channels_latents=num_channels_latents, + num_latent_frames=num_latent_frames, + latent_height=latent_height, + latent_width=latent_width, + latents_dtype=latents_dtype, + ) + padding_mask = latents.new_zeros(1, 1, image_height, image_width, dtype=latents_dtype) + + # denoiser + with self.model_memory_manager(["transformer"], low_vram=self.low_vram): + latents = self.denoise_latent( + latents, + timesteps, + text_embeddings, + negative_text_embeddings, + padding_mask, + ) + + # VAE decode latent + video = self._normalize_and_decode_latents(latents, is_video2world=False) + + torch.cuda.synchronize() + e2e_toc = time.perf_counter() + + walltime_ms = (e2e_toc - e2e_tic) * 1000.0 + return self._finalize_generation( + video, + walltime_ms, + num_inference_steps, + batch_size, + warmup, + save_image, + ) + + def generate_video( + self, + prompt, + negative_prompt, + image_height, + image_width, + input_image=None, + input_video=None, + num_frames=1, + fps=16, + num_videos_per_prompt=1, + sigma_conditioning=0.0001, + save_video=True, + warmup=False, + ): + batch_size = len(prompt) + + # Spatial dimensions of latent tensor + latent_height = image_height // self.vae_scale_factor_spatial + latent_width = image_width // self.vae_scale_factor_spatial + + num_inference_steps = self.denoising_steps + + with torch.inference_mode(), trt.Runtime(TRT_LOGGER): + torch.cuda.synchronize() + e2e_tic = time.perf_counter() + + # Prepare timesteps + timesteps, num_inference_steps = self._prepare_timesteps(num_inference_steps) + + # T5 text encoder + text_embeddings, negative_text_embeddings = self._encode_text_prompts( + prompt, negative_prompt, batch_size, num_videos_per_prompt + ) + + num_channels_latents = self.models["transformer"].config["in_channels"] - 1 + latents_dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32 + + # Process input conditioning + if input_image is not None: + video = ( + self.video_processor.preprocess(input_image, image_height, image_width) + .unsqueeze(2) + .to(device=self.device, dtype=latents_dtype) + ) + elif input_video is not None: + video = self.video_processor.preprocess_video(input_video, image_height, image_width).to( + device=self.device, dtype=latents_dtype + ) + else: + raise ValueError("Video2world pipeline requires either input_image or input_video to be provided") + + # Initialize latents + latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = ( + self.initialize_latents_video2world( + video, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + num_frames=num_frames, + latent_height=latent_height, + latent_width=latent_width, + latents_dtype=latents_dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + ) + ) + unconditioning_latents = None + + cond_mask = cond_mask.to(latents_dtype) + if self.do_classifier_free_guidance: + uncond_mask = uncond_mask.to(latents_dtype) + unconditioning_latents = conditioning_latents + + padding_mask = latents.new_zeros(1, 1, image_height, image_width, dtype=latents_dtype) + sigma_conditioning = torch.tensor(sigma_conditioning, dtype=torch.float32, device=self.device) + t_conditioning = sigma_conditioning / (sigma_conditioning + 1) + + # denoiser + with self.model_memory_manager(["transformer"], low_vram=self.low_vram): + latents = self.denoise_latent_video2world( + latents, + timesteps, + text_embeddings, + negative_text_embeddings, + padding_mask, + fps, + cond_mask, + uncond_mask, + t_conditioning, + cond_indicator, + conditioning_latents, + uncond_indicator, + unconditioning_latents, + ) + + # VAE decode latent + video = self._normalize_and_decode_latents(latents, is_video2world=True) + + torch.cuda.synchronize() + e2e_toc = time.perf_counter() + + walltime_ms = (e2e_toc - e2e_tic) * 1000.0 + return self._finalize_generation( + video, + walltime_ms, + num_inference_steps, + batch_size, + warmup, + save_video, + ) + + def infer( + self, + prompt, + negative_prompt, + image_height, + image_width, + input_image=None, + input_video=None, + num_frames=1, + fps=16, + num_images_per_prompt=1, + num_videos_per_prompt=1, + sigma_conditioning=0.0001, + warmup=False, + save_output=True, + ): + """ + Run the diffusion pipeline. + + Args: + prompt (str): + The text prompt to guide image generation. + negative_prompt (str): + The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + input_image (image): + Input image used to initialize the latents. + input_video (video): + Input video used to initialize the latents. + image_height (int): + Height (in pixels) of the image to be generated. Must be a multiple of 8. + image_width (int): + Width (in pixels) of the image to be generated. Must be a multiple of 8. + num_frames (int): + The number of frames in the generated video. + fps (int): + The frames per second of the generated video. + num_images_per_prompt (int): + The number of images to generate per prompt. + num_videos_per_prompt (int): + The number of videos to generate per prompt. + sigma_conditioning (`float`, defaults to `0.0001`): + The sigma value used for scaling conditioning latents. Ideally, it should not be changed or should be + set to a small value close to zero. + warmup (bool): + Indicate if this is a warmup run. + save_output (bool): + Save the generated image or video (if applicable) + """ + if self.pipeline_type.is_txt2img(): + return self.generate_image( + prompt, + negative_prompt, + image_height, + image_width, + num_frames, + num_images_per_prompt, + save_output, + warmup, + ) + elif self.pipeline_type.is_video2world(): + return self.generate_video( + prompt, + negative_prompt, + image_height, + image_width, + input_image, + input_video, + num_frames, + fps, + num_videos_per_prompt, + sigma_conditioning, + save_output, + warmup, + ) + else: + raise ValueError(f"Invalid pipeline type: {self.pipeline_type}") + + def run( + self, + prompt, + negative_prompt, + height, + width, + batch_count, + num_warmup_runs, + use_cuda_graph, + **kwargs, + ): + if self.low_vram and self.use_cuda_graph: + print("[W] Using low_vram, use_cuda_graph will be disabled") + self.use_cuda_graph = False + num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs + if num_warmup_runs > 0: + print("[I] Warming up ..") + for _ in range(num_warmup_runs): + self.infer(prompt, negative_prompt, height, width, warmup=True, **kwargs) + + outputs = [] + for _ in range(batch_count): + print("[I] Running Cosmos pipeline") + if self.nvtx_profile: + cudart.cudaProfilerStart() + output, _ = self.infer(prompt, negative_prompt, height, width, warmup=False, **kwargs) + outputs.append(output) + if self.nvtx_profile: + cudart.cudaProfilerStop() + + return outputs diff --git a/demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py index 3150c211..d0cfc2b9 100755 --- a/demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py @@ -100,6 +100,10 @@ class DiffusionPipeline(ABC): "flux.1-dev-depth", "flux.1-schnell", "flux.1-kontext-dev", + "cosmos-predict2-2b-text2image", + "cosmos-predict2-14b-text2image", + "cosmos-predict2-2b-video2world", + "cosmos-predict2-14b-video2world", ) SCHEDULER_DEFAULTS = { "1.4": "PNDM", @@ -120,6 +124,10 @@ class DiffusionPipeline(ABC): "flux.1-dev-depth": "FlowMatchEuler", "flux.1-schnell": "FlowMatchEuler", "flux.1-kontext-dev": "FlowMatchEuler", + "cosmos-predict2-2b-text2image": "FlowMatchEuler", + "cosmos-predict2-14b-text2image": "FlowMatchEuler", + "cosmos-predict2-2b-video2world": "FlowMatchEuler", + "cosmos-predict2-14b-video2world": "FlowMatchEuler", } def __init__( @@ -385,9 +393,10 @@ class DiffusionPipeline(ABC): ), "fp8 quantization only supported for SDXL, SD1.5, SD2.1, SD3.5-large and FLUX pipelines" if ( (self.pipeline_type.is_sd_xl() and model_name == "unetxl") + or (self.version.startswith("flux.1") and model_name == "transformer") or ( - (self.version.startswith("flux.1") or self.version.startswith("3.5-large")) - and model_name == "transformer" + self.version.startswith("3.5-large") + and ("transformer" in model_name or "controlnet" in model_name) ) or (model_name == "unet") ): @@ -622,6 +631,7 @@ class DiffusionPipeline(ABC): else: self.is_native_export_supported(model_config) + dynamo = True if self.pipeline_type.is_video2world() and model_name == "transformer" else False if do_export_onnx or do_export_weights_map: if not model_config['use_int8'] and not model_config['use_fp8']: obj.export_onnx( @@ -633,6 +643,7 @@ class DiffusionPipeline(ABC): enable_lora_merge=model_config["do_lora_merge"], static_shape=static_shape, lora_loader=self.lora_loader, + dynamo=dynamo, ) else: print(f"[I] Generating quantized ONNX model: {model_config['onnx_path']}") @@ -656,6 +667,7 @@ class DiffusionPipeline(ABC): opt_image_width, custom_model=quantized_model, static_shape=static_shape, + dynamo=dynamo, ) # FIXME do_export_weights_map needs ONNX graph diff --git a/demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py index 1f5209ac..37191c63 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py @@ -794,7 +794,7 @@ class FluxPipeline(DiffusionPipeline): print("|-----------------|--------------|") print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms)) print("|-----------------|--------------|") - print("Throughput: {:.2f} image/s".format(batch_size * 1000.0 / walltime_ms)) + print("Throughput: {:.5f} image/s".format(batch_size * 1000.0 / walltime_ms)) def _check_integrity(self, images): integrity_checker = PixtralContentFilter(self.device) diff --git a/demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py index 3e3f2e19..5abeab1c 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py @@ -93,7 +93,7 @@ class StableCascadePipeline(StableDiffusionPipeline): self.models['vqgan'] = VQGANModel(**models_args, fp16=self.fp16, bf16=self.bf16, latent_dim_scale = self.latent_dim_scale) def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False): - self.profile_start('clip', color='green') + self.profile_start(encoder, color='green') tokenizer = self.tokenizer @@ -140,7 +140,7 @@ class StableCascadePipeline(StableDiffusionPipeline): if output_hidden_states: text_embeddings = torch.cat([text_hidden_states, uncond_hidden_states]) if self.do_classifier_free_guidance else text_hidden_states - self.profile_stop('clip') + self.profile_stop(encoder) if pooled_outputs: return text_embeddings, pooled_output return text_embeddings @@ -157,7 +157,7 @@ class StableCascadePipeline(StableDiffusionPipeline): do_autocast = False with torch.autocast('cuda', enabled=do_autocast): - self.profile_start('denoise', color='blue') + self.profile_start(denoiser, color='blue') for step_index, timestep in enumerate(timesteps): # ratio input required for stable cascade prior timestep_ratio = timestep.expand(latents.size(0)).to(latents.dtype) @@ -197,31 +197,34 @@ class StableCascadePipeline(StableDiffusionPipeline): latents = latents.to(dtype=torch.bfloat16 if self.bf16 else torch.float32) - self.profile_stop('denoise') + self.profile_stop(denoiser) return latents - def decode_latent(self, latents): - self.profile_start('vqgan', color='red') - latents = self.models['vqgan'].scale_factor * latents + def decode_latent(self, latents, model_name='vqgan'): + self.profile_start(model_name, color='red') + latents = self.models[model_name].scale_factor * latents if self.torch_inference: - images = self.torch_models['vqgan'](latents)['sample'] + images = self.torch_models[model_name](latents)['sample'] else: - images = self.runEngine('vqgan', {'latent': latents})['images'] - self.profile_stop('vqgan') + images = self.runEngine(model_name, {'latent': latents})['images'] + self.profile_stop(model_name) return images def print_summary(self, denoising_steps, walltime_ms, batch_size): print('|-----------------|--------------|') print('| {:^15} | {:^12} |'.format('Module', 'Latency')) print('|-----------------|--------------|') - print('| {:^15} | {:>9.2f} ms |'.format('CLIP', cudart.cudaEventElapsedTime(self.events['clip'][0], self.events['clip'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('UNet'+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1])) - if 'vqgan' in self.stages: - print('| {:^15} | {:>9.2f} ms |'.format('VQGAN', cudart.cudaEventElapsedTime(self.events['vqgan'][0], self.events['vqgan'][1])[1])) + for stage in self.stages: + stage_name = stage + ' x ' + str(denoising_steps) if stage == 'unet' else stage + print( + "| {:^15} | {:>9.2f} ms |".format( + stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1], + ) + ) print('|-----------------|--------------|') print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms)) print('|-----------------|--------------|') - print('Throughput: {:.2f} image/s'.format(batch_size*1000./walltime_ms)) + print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms)) def infer( self, diff --git a/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py index 52daed9f..a09a3065 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py @@ -178,6 +178,12 @@ class StableDiffusion35Pipeline(DiffusionPipeline): elif "controlnet" in model_name: hf_download_path_cnet = hf_download_path.replace("large", "controlnets") dirname = f"controlnet_{self.controlnet}" + if "blur" in model_name: + pass + elif model_config["use_fp8"]: + dirname = os.path.join(dirname, "fp8") + elif self.bf16: + dirname = os.path.join(dirname, "bf16") elif model_name in self.stages: dirname = model_name else: @@ -241,7 +247,7 @@ class StableDiffusion35Pipeline(DiffusionPipeline): self.bf16 = True if int8 or fp8 or fp4 else self.bf16 self.fp16 = True if not self.bf16 else False - self.tf32=True + self.tf32 = True self.fp8 = fp8 self.int8 = int8 self.fp4 = fp4 @@ -330,44 +336,24 @@ class StableDiffusion35Pipeline(DiffusionPipeline): print("|-----------------|--------------|") print("| {:^15} | {:^12} |".format("Module", "Latency")) print("|-----------------|--------------|") - if "vae_encoder" in self.stages: + for stage in self.stages: + # controlnet is profiled in the denoising step + if "controlnet" in stage: + continue + stage_name = stage + if "transformer" in stage: + if f"controlnet_{self.controlnet}" in self.stages: + stage_name += '+cnet' + stage_name += ' x ' + str(denoising_steps) print( "| {:^15} | {:>9.2f} ms |".format( - "VAE Encoder", - cudart.cudaEventElapsedTime(self.events["vae_encoder"][0], self.events["vae_encoder"][1])[1], + stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1], ) ) - print( - "| {:^15} | {:>9.2f} ms |".format( - "CLIP-G", cudart.cudaEventElapsedTime(self.events["clip_g"][0], self.events["clip_g"][1])[1] - ) - ) - print( - "| {:^15} | {:>9.2f} ms |".format( - "CLIP-L", cudart.cudaEventElapsedTime(self.events["clip_l"][0], self.events["clip_l"][1])[1] - ) - ) - print( - "| {:^15} | {:>9.2f} ms |".format( - "T5", cudart.cudaEventElapsedTime(self.events["t5"][0], self.events["t5"][1])[1] - ) - ) - print( - "| {:^15} | {:>9.2f} ms |".format( - "MMDiT" + " x " + str(denoising_steps), - cudart.cudaEventElapsedTime(self.events["transformer"][0], self.events["transformer"][1])[1], - ) - ) - print( - "| {:^15} | {:>9.2f} ms |".format( - "VAE Decoder", - cudart.cudaEventElapsedTime(self.events["vae"][0], self.events["vae"][1])[1], - ) - ) print("|-----------------|--------------|") print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms)) print("|-----------------|--------------|") - print("Throughput: {:.2f} image/s".format(self.batch_size * 1000.0 / walltime_ms)) + print("Throughput: {:.5f} image/s".format(self.batch_size * 1000.0 / walltime_ms)) @staticmethod def _tokenize( @@ -711,9 +697,8 @@ class StableDiffusion35Pipeline(DiffusionPipeline): "timestep": timestep_inp, "encoder_hidden_states": prompt_embeds, "pooled_projections": pooled_prompt_embeds, + "block_controlnet_hidden_states": control_block_samples, } - if not self.fp8: - params["block_controlnet_hidden_states"] = control_block_samples # Predict the noise residual if self.torch_inference or self.torch_fallback[denoiser]: diff --git a/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py index f295937b..05353591 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py @@ -126,9 +126,9 @@ class StableDiffusion3Pipeline: # Pipeline type self.pipeline_type = pipeline_type - self.stages = ['clip_g', 'clip_l', 't5xxl', 'mmdit', 'vae_decoder'] + self.stages = ['clip_g', 'clip_l', 't5xxl', 'transformer', 'vae_decoder'] if input_image is not None: - self.stages += ['vae_encoder'] + self.stages = ['vae_encoder'] + self.stages self.config = {} self.config['clip_hidden_states'] = True @@ -161,7 +161,7 @@ class StableDiffusion3Pipeline: self.generator = torch.Generator(device="cuda").manual_seed(seed) # Create CUDA events and stream - for stage in ['clip_g', 'clip_l', 't5xxl', 'denoise', 'vae_encode', 'vae_decode']: + for stage in self.stages: self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]] self.stream = cudart.cudaStreamCreate()[1] @@ -261,9 +261,9 @@ class StableDiffusion3Pipeline: if 't5xxl' in self.stages: self.models['t5xxl'] = SD3_T5XXLModel(**models_args, fp16=True, embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type)) - # Load MMDiT model - if 'mmdit' in self.stages: - self.models['mmdit'] = SD3_MMDiTModel(**models_args, fp16=True, shift=self.shift) + # Load Transformer model + if 'transformer' in self.stages: + self.models['transformer'] = SD3_MMDiTModel(**models_args, fp16=True, shift=self.shift) # Load VAE Encoder model if 'vae_encoder' in self.stages: @@ -320,7 +320,7 @@ class StableDiffusion3Pipeline: # Load torch models for model_name, obj in self.models.items(): - if self.torch_fallback[model_name] or model_name == 'mmdit': + if self.torch_fallback[model_name] or model_name == 'transformer': self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference) def calculateMaxDeviceMemory(self): @@ -361,17 +361,19 @@ class StableDiffusion3Pipeline: print('|-----------------|--------------|') print('| {:^15} | {:^12} |'.format('Module', 'Latency')) print('|-----------------|--------------|') - if 'vae_encoder' in self.stages: - print('| {:^15} | {:>9.2f} ms |'.format('VAE Encoder', cudart.cudaEventElapsedTime(self.events['vae_encode'][0], self.events['vae_encode'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('CLIP-G', cudart.cudaEventElapsedTime(self.events['clip_g'][0], self.events['clip_g'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('CLIP-L', cudart.cudaEventElapsedTime(self.events['clip_l'][0], self.events['clip_l'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('T5XXL', cudart.cudaEventElapsedTime(self.events['t5xxl'][0], self.events['t5xxl'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('MMDiT'+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('VAE Decoder', cudart.cudaEventElapsedTime(self.events['vae_decode'][0], self.events['vae_decode'][1])[1])) + for stage in self.stages: + stage_name = stage + if "transformer" in stage: + stage_name += ' x ' + str(denoising_steps) + print( + "| {:^15} | {:>9.2f} ms |".format( + stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1], + ) + ) print('|-----------------|--------------|') print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms)) print('|-----------------|--------------|') - print('Throughput: {:.2f} image/s'.format(batch_size*1000./walltime_ms)) + print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms)) def save_image(self, images, pipeline, prompt, seed): # Save image @@ -414,7 +416,7 @@ class StableDiffusion3Pipeline: neg_conditioning = tokenize(negative_prompt[0]) return conditioning, neg_conditioning - def denoise_latent(self, latent, conditioning, neg_conditioning, model_name='mmdit'): + def denoise_latent(self, latent, conditioning, neg_conditioning, model_name='transformer'): def get_noise(latent): return torch.randn(latent.size(), dtype=torch.float32, layout=latent.layout, generator=self.generator, device="cuda").to(latent.dtype) @@ -456,7 +458,7 @@ class StableDiffusion3Pipeline: scaled = neg_out + (pos_out - neg_out) * cond_scale return scaled - self.profile_start('denoise', color='blue') + self.profile_start(model_name, color='blue') latent = latent.half().cuda() noise = get_noise(latent).cuda() @@ -470,32 +472,32 @@ class StableDiffusion3Pipeline: latent = sample_euler(cfg_denoiser, noise_scaled, sigmas, extra_args=extra_args) latent = SD3LatentFormat().process_out(latent) - self.profile_stop('denoise') + self.profile_stop(model_name) return latent - def encode_image(self): + def encode_image(self, model_name='vae_encoder'): self.input_image = self.input_image.to(self.device) - self.profile_start('vae_encode', color='orange') + self.profile_start(model_name, color='orange') if self.torch_inference: with torch.autocast("cuda", dtype=torch.float16): - latent = self.torch_models['vae_encoder'](self.input_image) + latent = self.torch_models[model_name](self.input_image) else: - latent = self.runEngine('vae_encoder', {'images': self.input_image})['latent'] + latent = self.runEngine(model_name, {'images': self.input_image})['latent'] latent = SD3LatentFormat().process_in(latent) - self.profile_stop('vae_encode') + self.profile_stop(model_name) return latent - def decode_latent(self, latent): - self.profile_start('vae_decode', color='red') + def decode_latent(self, latent, model_name='vae_decoder'): + self.profile_start(model_name, color='red') if self.torch_inference: with torch.autocast("cuda", dtype=torch.float16): - image = self.torch_models['vae_decoder'](latent) + image = self.torch_models[model_name](latent) else: - image = self.runEngine('vae_decoder', {'latent': latent})['images'] + image = self.runEngine(model_name, {'latent': latent})['images'] image = image.float() - self.profile_stop('vae_decode') + self.profile_stop(model_name) return image def infer( diff --git a/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py index 724f4caf..833b26d3 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py @@ -49,10 +49,10 @@ from demo_diffusion.model import ( CLIPModel, CLIPWithProjModel, SDLoraLoader, + UNet2DConditionControlNetModel, UNetModel, UNetXLModel, UNetXLModelControlNet, - UNet2DConditionControlNetModel, VAEEncoderModel, VAEModel, get_clip_embedding_dim, @@ -265,7 +265,7 @@ class StableDiffusionPipeline: self.generator = torch.Generator(device="cuda").manual_seed(seed) # Create CUDA events and stream - for stage in ['clip', 'denoise', 'vae', 'vae_encoder', 'vqgan']: + for stage in self.stages: self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]] self.stream = cudart.cudaStreamCreate()[1] @@ -713,7 +713,7 @@ class StableDiffusionPipeline: return images def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False): - self.profile_start('clip', color='green') + self.profile_start(encoder, color='green') tokenizer = self.tokenizer2 if encoder == 'clip2' else self.tokenizer @@ -756,7 +756,7 @@ class StableDiffusionPipeline: if output_hidden_states: text_embeddings = torch.cat([uncond_hidden_states, text_hidden_states]).to(dtype=torch.float16) if self.do_classifier_free_guidance else text_hidden_states - self.profile_stop('clip') + self.profile_stop(encoder) if pooled_outputs: return text_embeddings, pooled_output return text_embeddings @@ -818,7 +818,7 @@ class StableDiffusionPipeline: do_autocast = self.torch_inference != '' and self.models[denoiser].fp16 with torch.autocast('cuda', enabled=do_autocast): - self.profile_start('denoise', color='blue') + self.profile_start(denoiser, color='blue') for step_index, timestep in enumerate(timesteps): # Expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents @@ -870,7 +870,7 @@ class StableDiffusionPipeline: latents = 1. / self.vae_scaling_factor * latents latents = latents.to(dtype=torch.float32) - self.profile_stop('denoise') + self.profile_stop(denoiser) return latents def encode_image(self, input_image): @@ -901,15 +901,21 @@ class StableDiffusionPipeline: print('|-----------------|--------------|') print('| {:^15} | {:^12} |'.format('Module', 'Latency')) print('|-----------------|--------------|') - if 'vae_encoder' in self.stages: - print('| {:^15} | {:>9.2f} ms |'.format('VAE-Enc', cudart.cudaEventElapsedTime(self.events['vae_encoder'][0], self.events['vae_encoder'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('CLIP', cudart.cudaEventElapsedTime(self.events['clip'][0], self.events['clip'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('UNet'+('+CNet' if self.pipeline_type.is_controlnet() else '')+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1])) - print('| {:^15} | {:>9.2f} ms |'.format('VAE-Dec', cudart.cudaEventElapsedTime(self.events['vae'][0], self.events['vae'][1])[1])) + for stage in self.stages: + stage_name = stage + if "unet" in stage: + if self.pipeline_type.is_controlnet(): + stage_name += '+cnet' + stage_name += ' x ' + str(denoising_steps) + print( + "| {:^15} | {:>9.2f} ms |".format( + stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1], + ) + ) print('|-----------------|--------------|') print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms)) print('|-----------------|--------------|') - print('Throughput: {:.2f} image/s'.format(batch_size*1000./walltime_ms)) + print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms)) def save_image(self, images, pipeline, prompt, seed): # Save image diff --git a/demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py b/demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py index 7dbd3100..aa11542c 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py +++ b/demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py @@ -118,7 +118,11 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline): # TODO user configurable cuda_device_id cuda_device_id = 0 - vram_size = cudart.cudaGetDeviceProperties(cuda_device_id)[1].totalGlobalMem + properties = cudart.cudaGetDeviceProperties(cuda_device_id) + if properties[0] != 0: + total_device_count = cudart.cudaGetDeviceCount()[1] + raise ValueError(f"Failed to get device properties for device {cuda_device_id}, total device count: {total_device_count}") + vram_size = properties[1].totalGlobalMem self.low_vram = vram_size < _GiB(40) if self.low_vram: print(f"[W] WARNING low VRAM ({vram_size/_GiB(1):.2f} GB) mode selected. Certain optimizations may be skipped.") @@ -397,7 +401,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline): print('|-----------------|--------------|') print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms)) print('|-----------------|--------------|') - print('Throughput: {:.2f} videos/min ({} frames)'.format(batch_size*60000./walltime_ms, num_frames)) + print('Throughput: {:.5f} videos/min ({} frames)'.format(batch_size*60000./walltime_ms, num_frames)) def save_video(self, frames, pipeline, seed): video_name_prefix = '-'.join([pipeline, 'fp16', str(seed), str(random.randint(1000,9999))]) diff --git a/demo/Diffusion/demo_diffusion/pipeline/type.py b/demo/Diffusion/demo_diffusion/pipeline/type.py index aca8df76..a2d0ab0d 100644 --- a/demo/Diffusion/demo_diffusion/pipeline/type.py +++ b/demo/Diffusion/demo_diffusion/pipeline/type.py @@ -29,6 +29,7 @@ class PIPELINE_TYPE(enum.Enum): XL_REFINER = enum.auto() CASCADE_PRIOR = enum.auto() CASCADE_DECODER = enum.auto() + VIDEO2WORLD = enum.auto() def is_txt2img(self): return self in (self.TXT2IMG, self.CONTROLNET) @@ -62,3 +63,6 @@ class PIPELINE_TYPE(enum.Enum): def is_cascade(self): return self.is_cascade_prior() or self.is_cascade_decoder() + + def is_video2world(self): + return self == self.VIDEO2WORLD diff --git a/demo/Diffusion/demo_diffusion/utils_modelopt.py b/demo/Diffusion/demo_diffusion/utils_modelopt.py index d0f84da2..5e384643 100755 --- a/demo/Diffusion/demo_diffusion/utils_modelopt.py +++ b/demo/Diffusion/demo_diffusion/utils_modelopt.py @@ -23,7 +23,6 @@ from typing import Set import modelopt.torch.quantization as mtq import numpy as np -import onnx import onnx_graphsurgeon as gs import torch import torch.nn.functional as F @@ -38,6 +37,8 @@ from modelopt.torch.quantization.calib.max import MaxCalibrator from PIL import Image from torch.utils.data import Dataset, Sampler +import onnx + USE_PEFT = True try: from peft.tuners.lora.layer import Conv2d as PEFTLoRAConv2d diff --git a/demo/Diffusion/demo_txt2image_cosmos.py b/demo/Diffusion/demo_txt2image_cosmos.py new file mode 100644 index 00000000..4dc12421 --- /dev/null +++ b/demo/Diffusion/demo_txt2image_cosmos.py @@ -0,0 +1,147 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import argparse + +from cuda import cudart + +from demo_diffusion import dd_argparse +from demo_diffusion import pipeline as pipeline_module + + +def parse_args(): + parser = argparse.ArgumentParser(description="Options for Cosmos text2image Demo", conflict_handler="resolve") + parser = dd_argparse.add_arguments(parser) + parser.add_argument( + "--version", + type=str, + default="cosmos-predict2-2b-text2image", + choices=("cosmos-predict2-2b-text2image", "cosmos-predict2-14b-text2image"), + help="Version of Cosmos", + ) + parser.add_argument( + "--height", + type=int, + default=768, + help="Height of image to generate (must be multiple of 8)", + ) + parser.add_argument( + "--width", + type=int, + default=1360, + help="Width of image to generate (must be multiple of 8)", + ) + parser.add_argument("--denoising-steps", type=int, default=35, help="Number of denoising steps") + parser.add_argument( + "--guidance-scale", + type=float, + default=7.0, + help="Value of classifier-free guidance scale (must be greater than 1)", + ) + parser.add_argument( + "--num-images-per-prompt", type=int, default=1, help="The number of images to generate per prompt." + ) + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help="Maximum sequence length to use with the prompt.", + ) + parser.add_argument( + "--t5-ws-percentage", + type=int, + default=None, + help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ", + ) + parser.add_argument( + "--transformer-ws-percentage", + type=int, + default=None, + help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.", + ) + parser.add_argument( + "--bf16", + action="store_true", + default=True, + help="Use bfloat16 precision by default.", + ) + return parser.parse_args() + + +def process_demo_args(args): + batch_size = args.batch_size + prompt = args.prompt + negative_prompt = args.negative_prompt + # Process input args + if not isinstance(prompt, list): + raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}") + prompt = prompt * batch_size + if not isinstance(negative_prompt, list): + raise ValueError(f"`negative_prompt` must be of type `str` list, but is {type(negative_prompt)}") + negative_prompt = negative_prompt * batch_size + + kwargs_run_demo = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "height": args.height, + "width": args.width, + "batch_count": args.batch_count, + "num_warmup_runs": args.num_warmup_runs, + "use_cuda_graph": args.use_cuda_graph, + "num_images_per_prompt": args.num_images_per_prompt, + } + + return kwargs_run_demo + + +if __name__ == "__main__": + print("[I] Initializing Cosmos text2image demo using TensorRT") + args = parse_args() + + _, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args) + kwargs_run_demo = process_demo_args(args) + + # Initialize demo + demo = pipeline_module.CosmosPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG) + + # Load TensorRT engines and pytorch modules + demo.load_engines( + framework_model_dir=args.framework_model_dir, + **kwargs_load_engine, + ) + + if args.onnx_export_only: + print("[I] ONNX export completed. Exiting...") + demo.teardown() + exit(0) + + # In low-vram mode we allocate the required device memory individually before each model is run. + if demo.low_vram: + demo.device_memory_sizes = demo.get_device_memory_sizes() + else: + _, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory()) + demo.activate_engines(shared_device_memory) + + demo.load_resources(args.height, args.width, args.batch_size, args.seed) + + # Run inference + images = demo.run(**kwargs_run_demo) + + demo.teardown() + + # save images + demo.save_images(kwargs_run_demo["prompt"], images, check_integrity=True) diff --git a/demo/Diffusion/demo_vid2world_cosmos.py b/demo/Diffusion/demo_vid2world_cosmos.py new file mode 100644 index 00000000..a3fb34a5 --- /dev/null +++ b/demo/Diffusion/demo_vid2world_cosmos.py @@ -0,0 +1,170 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import argparse + +from cuda import cudart +from diffusers.utils import load_image, load_video + +from demo_diffusion import dd_argparse +from demo_diffusion import pipeline as pipeline_module + + +def parse_args(): + parser = argparse.ArgumentParser(description="Options for Cosmos video2world Demo", conflict_handler="resolve") + parser = dd_argparse.add_arguments(parser) + parser.add_argument( + "--version", + type=str, + default="cosmos-predict2-2b-video2world", + choices=("cosmos-predict2-2b-video2world", "cosmos-predict2-14b-video2world"), + help="Version of Cosmos", + ) + parser.add_argument('--input-image', type=str, default=None, help="Path to the input image") + parser.add_argument('--input-video', type=str, default=None, help="Path to the input video") + parser.add_argument( + "--height", + type=int, + default=704, + help="Height of image to generate (must be multiple of 8)", + ) + parser.add_argument( + "--width", + type=int, + default=1280, + help="Width of image to generate (must be multiple of 8)", + ) + parser.add_argument("--denoising-steps", type=int, default=35, help="Number of denoising steps") + parser.add_argument( + "--guidance-scale", + type=float, + default=7.0, + help="Value of classifier-free guidance scale (must be greater than 1)", + ) + parser.add_argument("--num-frames", type=int, default=93, help="The number of frames in the generated video.") + parser.add_argument("--fps", type=int, default=16, help="The frames per second of the generated video.") + parser.add_argument("--num-videos-per-prompt", type=int, default=1, help="The number of videos to generate per prompt.") + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help="Maximum sequence length to use with the prompt.", + ) + parser.add_argument( + "--t5-ws-percentage", + type=int, + default=None, + help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ", + ) + parser.add_argument( + "--transformer-ws-percentage", + type=int, + default=None, + help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.", + ) + parser.add_argument( + "--bf16", + action="store_true", + default=True, + help="Use bfloat16 precision by default.", + ) + return parser.parse_args() + + +def process_demo_args(args): + batch_size = args.batch_size + prompt = args.prompt + negative_prompt = args.negative_prompt + # Process input args + if not isinstance(prompt, list): + raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}") + prompt = prompt * batch_size + if not isinstance(negative_prompt, list): + raise ValueError(f"`negative_prompt` must be of type `str` list, but is {type(negative_prompt)}") + negative_prompt = negative_prompt * batch_size + + # process input image and input video + if args.input_image and args.input_video: + raise ValueError("Only one of --input-image or --input-video can be provided") + if args.input_image: + args.input_image = load_image(args.input_image) + elif args.input_video: + args.input_video = load_video(args.input_video) + else: + # load default image + args.input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yellow-scrubber.png") + + kwargs_run_demo = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "height": args.height, + "width": args.width, + "batch_count": args.batch_count, + "num_warmup_runs": args.num_warmup_runs, + "use_cuda_graph": args.use_cuda_graph, + "num_frames": args.num_frames, + "fps": args.fps, + "input_image": args.input_image, + "input_video": args.input_video, + "num_videos_per_prompt": args.num_videos_per_prompt, + } + + return kwargs_run_demo + + +if __name__ == "__main__": + print("[I] Initializing Cosmos video2world demo using TensorRT") + args = parse_args() + + # Enforce torch-inference is enabled + if not args.torch_inference: + print("[W] The video2world demo only supports the PyTorch backend. Enabling torch-inference with 'eager' mode.") + args.torch_inference = "eager" + + _, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args) + kwargs_run_demo = process_demo_args(args) + + # Initialize demo + demo = pipeline_module.CosmosPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.VIDEO2WORLD) + + # Load TensorRT engines and pytorch modules + demo.load_engines( + framework_model_dir=args.framework_model_dir, + **kwargs_load_engine, + ) + + if args.onnx_export_only: + print("[I] ONNX export completed. Exiting...") + demo.teardown() + exit(0) + + # In low-vram mode we allocate the required device memory individually before each model is run. + if demo.low_vram: + demo.device_memory_sizes = demo.get_device_memory_sizes() + else: + _, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory()) + demo.activate_engines(shared_device_memory) + + demo.load_resources(args.height, args.width, args.batch_size, args.seed) + + # Run inference + videos = demo.run(**kwargs_run_demo) + + demo.teardown() + + # save video + demo.save_video(kwargs_run_demo["prompt"], videos, check_integrity=True) diff --git a/demo/Diffusion/docs/support_matrix.md b/demo/Diffusion/docs/support_matrix.md index 9701254b..7ac257d5 100644 --- a/demo/Diffusion/docs/support_matrix.md +++ b/demo/Diffusion/docs/support_matrix.md @@ -6,32 +6,34 @@ This demo supports Diffusion models that are popular in the Generative AI commun ## Pipeline Support Matrix -| Pipeline | Version | Task | Supported Precisions | Additional features | Hub | Restrictions | -|------------------------|---------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------|------------------------|--------------------------|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| Stable Diffusion | 1.4 |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) | -| Stable Diffusion | [1.5](../README.md#generate-an-image-guided-by-a-text-prompt) |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [KiwiXR/stable-diffusion-v1-5](https://huggingface.co/KiwiXR/stable-diffusion-v1-5) | -| Stable Diffusion | 1.4, [1.5](../README.md#generate-an-inpainted-image-guided-by-an-image-mask-and-a-text-prompt) |
  • Inpainting
| FP16 | N/A | [benjamin-paine/stable-diffusion-v1-5-inpainting](https://huggingface.co/benjamin-paine/stable-diffusion-v1-5-inpainting) | -| Stable Diffusion | dreamshaper-7 |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [Lykon/dreamshaper-7](https://huggingface.co/Lykon/dreamshaper-7) | -| Stable Diffusion | 2.0-base |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) | -| Stable Diffusion | 2.0 |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) | -| Stable Diffusion | 2.0, 2.0-base |
  • Inpainting
| FP16 | N/A | [stabilityai/stable-diffusion-2-inpainting](https://huggingface.co/stabilityai/stable-diffusion-2-inpainting) | -| Stable Diffusion | 2.1-base |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) | -| Stable Diffusion | 2.1 |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) | -| Stable Diffusion | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | LoRA (FP16, BF16, FP8) | [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) | -| Stable Diffusion | [XL 1.0-refiner](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-xl-refiner-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0) | -| Stable Diffusion | [XL-Turbo](../README.md#faster-text-to-image-using-sdxl-turbo) |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/sdxl-turbo](https://huggingface.co/stabilityai/sdxl-turbo) | -| Stable Diffusion | [3](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) |
  • Text-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) | -| Stable Diffusion | [3.5-medium](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) |
  • Text-to-image
| FP16, BF16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) | -| Stable Diffusion | [3.5-large](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) |
  • Text-to-image
| FP16, BF16, FP8 | N/A | [stabilityai/stable-diffusion-3-large](https://huggingface.co/stabilityai/stable-diffusion-3-large) | -| ControlNet | [1.5](../README.md#generate-an-image-with-controlnet-guided-by-images-and-text-prompts) |
  • Image-to-image
| FP16 | N/A |
  • [lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny)
  • [lllyasviel/sd-controlnet-depth](https://huggingface.co/lllyasviel/sd-controlnet-depth)
  • [lllyasviel/sd-controlnet-hed](https://huggingface.co/lllyasviel/sd-controlnet-hed)
  • [lllyasviel/sd-controlnet-mlsd](https://huggingface.co/lllyasviel/sd-controlnet-mlsd)
  • [lllyasviel/sd-controlnet-normal](https://huggingface.co/lllyasviel/sd-controlnet-normal)
  • [lllyasviel/sd-controlnet_openpose](https://huggingface.co/lllyasviel/sd-controlnet-openpose)
  • [lllyasviel/sd-controlnet_scribble](https://huggingface.co/lllyasviel/sd-controlnet-scribble)
  • [lllyasviel/sd-controlnet_seg](https://huggingface.co/lllyasviel/sd-controlnet-seg)
| -| ControlNet | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) |
  • Image-to-image
| FP16, FP8 | N/A | [stabilityai/controlnet-canny-sdxl-1.0](https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0) | -| ControlNet | [3.5-large](../README.md#generate-an-image-with-stable-diffusion-v35-large-with-controlnet-guided-by-an-image-and-a-text-prompt) |
  • Image-to-image
| FP16, BF16 | N/A |
  • [stabilityai/stable-diffusion-3.5-large-controlnet-canny](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-canny)
  • [stabilityai/stable-diffusion-3.5-large-controlnet-depth](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-depth)
  • [stabilityai/stable-diffusion-3.5-large-controlnet-blur](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-blur)
| -| Stable Video Diffusion | [XT-1.1](../README.md#generate-a-video-guided-by-an-initial-image-using-stable-video-diffusion) |
  • Text-to-video
| FP16, FP8 | N/A | [stabilityai/stable-video-diffusion-img2vid-xt-1-1](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1) | -| Stable Cascade | [N/A](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-cascade) |
  • Text-to-image
| BF16 | N/A |
  • [stabilityai/stable-cascade-prior](https://huggingface.co/stabilityai/stable-cascade-prior)
  • [stabilityai/stable-cascade](https://huggingface.co/stabilityai/stable-cascade)
| -| Flux | [1-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) |
  • Text-to-image
  • Image-to-image
| FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | -| Flux | [1-Schnell](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) |
  • Text-to-image
  • Image-to-image
| FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) | -| Flux | [1-Canny-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) |
  • Image-to-image
| FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Canny-dev](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) | -| Flux | [1-Depth-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) |
  • Image-to-image
| FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) | -| Flux | [1-Kontext-Dev](../README.md#5-edit-an-image-using-flux-kontext) |
  • Image-to-image
| BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) | +| Pipeline | Version | Task | Supported Precisions | Additional features | Hub | Restrictions | +|------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------|-----------------------------------------------|--------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| Stable Diffusion | 1.4 |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) | +| Stable Diffusion | [1.5](../README.md#generate-an-image-guided-by-a-text-prompt) |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [KiwiXR/stable-diffusion-v1-5](https://huggingface.co/KiwiXR/stable-diffusion-v1-5) | +| Stable Diffusion | 1.4, [1.5](../README.md#generate-an-inpainted-image-guided-by-an-image-mask-and-a-text-prompt) |
  • Inpainting
| FP16 | N/A | [benjamin-paine/stable-diffusion-v1-5-inpainting](https://huggingface.co/benjamin-paine/stable-diffusion-v1-5-inpainting) | +| Stable Diffusion | dreamshaper-7 |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [Lykon/dreamshaper-7](https://huggingface.co/Lykon/dreamshaper-7) | +| Stable Diffusion | 2.0-base |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) | +| Stable Diffusion | 2.0 |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) | +| Stable Diffusion | 2.0, 2.0-base |
  • Inpainting
| FP16 | N/A | [stabilityai/stable-diffusion-2-inpainting](https://huggingface.co/stabilityai/stable-diffusion-2-inpainting) | +| Stable Diffusion | 2.1-base |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) | +| Stable Diffusion | 2.1 |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) | +| Stable Diffusion | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) |
  • Text-to-image
  • Image-to-image
| FP16, FP8, INT8 * | LoRA (FP16, BF16, FP8) | [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) | +| Stable Diffusion | [XL 1.0-refiner](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-xl-refiner-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0) | +| Stable Diffusion | [XL-Turbo](../README.md#faster-text-to-image-using-sdxl-turbo) |
  • Text-to-image
  • Image-to-image
| FP16 | N/A | [stabilityai/sdxl-turbo](https://huggingface.co/stabilityai/sdxl-turbo) | +| Stable Diffusion | [3](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) |
  • Text-to-image
| FP16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) | +| Stable Diffusion | [3.5-medium](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) |
  • Text-to-image
| FP16, BF16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) | +| Stable Diffusion | [3.5-large](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) |
  • Text-to-image
| FP16, BF16, FP8 | N/A | [stabilityai/stable-diffusion-3-large](https://huggingface.co/stabilityai/stable-diffusion-3-large) | +| ControlNet | [1.5](../README.md#generate-an-image-with-controlnet-guided-by-images-and-text-prompts) |
  • Image-to-image
| FP16 | N/A |
  • [lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny)
  • [lllyasviel/sd-controlnet-depth](https://huggingface.co/lllyasviel/sd-controlnet-depth)
  • [lllyasviel/sd-controlnet-hed](https://huggingface.co/lllyasviel/sd-controlnet-hed)
  • [lllyasviel/sd-controlnet-mlsd](https://huggingface.co/lllyasviel/sd-controlnet-mlsd)
  • [lllyasviel/sd-controlnet-normal](https://huggingface.co/lllyasviel/sd-controlnet-normal)
  • [lllyasviel/sd-controlnet_openpose](https://huggingface.co/lllyasviel/sd-controlnet-openpose)
  • [lllyasviel/sd-controlnet_scribble](https://huggingface.co/lllyasviel/sd-controlnet-scribble)
  • [lllyasviel/sd-controlnet_seg](https://huggingface.co/lllyasviel/sd-controlnet-seg)
| +| ControlNet | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) |
  • Image-to-image
| FP16, FP8 | N/A | [stabilityai/controlnet-canny-sdxl-1.0](https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0) | +| ControlNet | [3.5-large](../README.md#generate-an-image-with-stable-diffusion-v35-large-with-controlnet-guided-by-an-image-and-a-text-prompt) |
  • Image-to-image
| FP16, BF16, FP8 (canny and depth only) | N/A |
  • [stabilityai/stable-diffusion-3.5-large-controlnet-canny](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-canny)
  • [stabilityai/stable-diffusion-3.5-large-controlnet-depth](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-depth)
  • [stabilityai/stable-diffusion-3.5-large-controlnet-blur](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-blur)
| +| Stable Video Diffusion | [XT-1.1](../README.md#generate-a-video-guided-by-an-initial-image-using-stable-video-diffusion) |
  • Text-to-video
| FP16, FP8 | N/A | [stabilityai/stable-video-diffusion-img2vid-xt-1-1](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1) | +| Stable Cascade | [N/A](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-cascade) |
  • Text-to-image
| BF16 | N/A |
  • [stabilityai/stable-cascade-prior](https://huggingface.co/stabilityai/stable-cascade-prior)
  • [stabilityai/stable-cascade](https://huggingface.co/stabilityai/stable-cascade)
| +| Flux | [1-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) |
  • Text-to-image
  • Image-to-image
| FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | +| Flux | [1-Schnell](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) |
  • Text-to-image
  • Image-to-image
| FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) | +| Flux | [1-Canny-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) |
  • Image-to-image
| FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Canny-dev](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) | +| Flux | [1-Depth-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) |
  • Image-to-image
| FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) | +| Flux | [1-Kontext-Dev](../README.md#5-edit-an-image-using-flux-kontext) |
  • Image-to-image
| BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) | +| Cosmos | [cosmos-predict2-2b-text2image](../README.md#1-generate-an-image-from-a-text-prompt-1), cosmos-predict2-14b-text2image |
  • Text-to-image
| BF16 | N/A |
  • [nvidia/Cosmos-Predict2-2B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image)
  • [nvidia/Cosmos-Predict2-14B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Text2Image) | +| Cosmos | [cosmos-predict2-2b-video2world](../README.md#2-generate-a-video-guided-by-an-initial-video-conditioning-and-a-text-prompt), cosmos-predict2-14b-video2world |
    • Video-to-World
    | BF16 | N/A |
    • [nvidia/Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)
    • [nvidia/Cosmos-Predict2-14B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Video2World) | *Note: Only the text2image pipelines support FP4/FP8/INT8 quantization. The image2image pipelines don't support quantization. diff --git a/demo/Diffusion/requirements.txt b/demo/Diffusion/requirements.txt index f0ebf437..9e923476 100755 --- a/demo/Diffusion/requirements.txt +++ b/demo/Diffusion/requirements.txt @@ -9,7 +9,6 @@ ftfy matplotlib nvtx onnx==1.18.0 -onnxruntime==1.19.2 onnxscript==0.3.2 opencv-python-headless==4.8.0.74 scipy @@ -21,3 +20,5 @@ peft==0.17.0 polygraphy==0.49.22 sentencepiece numpy==1.26.4 +imageio-ffmpeg + diff --git a/demo/Diffusion/requirements_legacy.txt b/demo/Diffusion/requirements_legacy.txt new file mode 100644 index 00000000..38304af8 --- /dev/null +++ b/demo/Diffusion/requirements_legacy.txt @@ -0,0 +1,24 @@ +apex==0.9.10dev +accelerate==1.2.1 +colored==2.3.1 +controlnet_aux==0.0.6 +cuda-python==13.0.2 +diffusers==0.35.0 +ftfy==6.3.1 +matplotlib==3.10.7 +nvtx==0.2.13 +onnx==1.17.0 +onnxscript==0.3.2 +onnxconverter-common==1.14.0 +opencv-python-headless==4.8.0.74 +scipy==1.15.3 +transformers==4.52.4 +--extra-index-url https://pypi.nvidia.com +nvidia-modelopt[torch,onnx]==0.29.0 +onnx-graphsurgeon==0.5.2 +peft==0.17.0 +polygraphy==0.49.22 +sentencepiece==0.2.1 +numpy==1.26.4 +imageio-ffmpeg==0.6.0 + diff --git a/demo/Diffusion/setup.sh b/demo/Diffusion/setup.sh index f55a614d..8da259e9 100644 --- a/demo/Diffusion/setup.sh +++ b/demo/Diffusion/setup.sh @@ -1,10 +1,15 @@ #!/bin/bash + +# Minimal and backward-compatible: default to modern requirements, allow override via env. +REQ_FILE="${REQUIREMENTS_FILE:-requirements.txt}" +echo "[I] Using requirements file: ${REQ_FILE} (override with REQUIREMENTS_FILE=)" + # Upgrade pip and install TensorRT python3 -m pip install --upgrade pip pip3 install --pre tensorrt-cu12 # Install the required packages -PIP_CONSTRAINT= pip3 install -r requirements.txt +PIP_CONSTRAINT= pip3 install -r "${REQ_FILE}" # Install libgl1 # Check if apt-get is available diff --git a/docker/rockylinux8.Dockerfile b/docker/rockylinux8.Dockerfile index 0378b8de..5946448a 100644 --- a/docker/rockylinux8.Dockerfile +++ b/docker/rockylinux8.Dockerfile @@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0 FROM nvidia/cuda:${CUDA_VERSION}-devel-rockylinux8 LABEL maintainer="NVIDIA CORPORATION" -ENV TRT_VERSION 10.13.3.9 +ENV TRT_VERSION 10.14.1.48 SHELL ["/bin/bash", "-c"] # Setup user account @@ -51,15 +51,15 @@ RUN dnf install -y python38 python38-devel &&\ # Install TensorRT RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp38-none-linux_x86_64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp38-none-linux_x86_64.whl ;\ elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp38-none-linux_x86_64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp38-none-linux_x86_64.whl ;\ else \ echo "Invalid CUDA_VERSION"; \ exit 1; \ diff --git a/docker/rockylinux9.Dockerfile b/docker/rockylinux9.Dockerfile index 02d9e4eb..8fd38c68 100644 --- a/docker/rockylinux9.Dockerfile +++ b/docker/rockylinux9.Dockerfile @@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0 FROM nvidia/cuda:${CUDA_VERSION}-devel-rockylinux9 LABEL maintainer="NVIDIA CORPORATION" -ENV TRT_VERSION 10.13.3.9 +ENV TRT_VERSION 10.14.1.48 SHELL ["/bin/bash", "-c"] # Setup user account @@ -56,15 +56,15 @@ RUN dnf -y install \ # Install TensorRT RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp39-none-linux_x86_64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp39-none-linux_x86_64.whl ;\ elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp39-none-linux_x86_64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp39-none-linux_x86_64.whl ;\ else \ echo "Invalid CUDA_VERSION"; \ exit 1; \ diff --git a/docker/ubuntu-22.04.Dockerfile b/docker/ubuntu-22.04.Dockerfile index f960909f..8ad84cdd 100644 --- a/docker/ubuntu-22.04.Dockerfile +++ b/docker/ubuntu-22.04.Dockerfile @@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 LABEL maintainer="NVIDIA CORPORATION" -ENV TRT_VERSION 10.13.3.9 +ENV TRT_VERSION 10.14.1.48 SHELL ["/bin/bash", "-c"] # Setup user account @@ -70,15 +70,15 @@ RUN apt-get install -y --no-install-recommends \ # Install TensorRT RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp310-none-linux_x86_64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp310-none-linux_x86_64.whl ;\ elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp310-none-linux_x86_64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp310-none-linux_x86_64.whl ;\ else \ echo "Invalid CUDA_VERSION"; \ exit 1; \ diff --git a/docker/ubuntu-24.04-aarch64.Dockerfile b/docker/ubuntu-24.04-aarch64.Dockerfile index 2ed57c44..7ea2f389 100644 --- a/docker/ubuntu-24.04-aarch64.Dockerfile +++ b/docker/ubuntu-24.04-aarch64.Dockerfile @@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0 # Multi-arch container support available in non-cudnn containers. FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04 -ENV TRT_VERSION 10.13.3.9 +ENV TRT_VERSION 10.14.1.48 SHELL ["/bin/bash", "-c"] # Setup user account and edit default account @@ -60,7 +60,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ lintian \ fakeroot \ dh-make \ - build-essential + build-essential \ + libffi-dev # Install python3 RUN apt-get install -y --no-install-recommends \ @@ -79,15 +80,15 @@ ENV PATH="/opt/venv/bin:$PATH" # Install TensorRT RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib/aarch64-linux-gnu/ \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp312-none-linux_aarch64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/aarch64-linux-gnu/ \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_aarch64.whl ;\ elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \ - wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-12.9.tar.gz \ - && tar -xf TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-12.9.tar.gz \ - && cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib/aarch64-linux-gnu/ \ - && pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp312-none-linux_aarch64.whl ;\ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-12.9.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-12.9.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/aarch64-linux-gnu/ \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_aarch64.whl ;\ else \ echo "Invalid CUDA_VERSION"; \ exit 1; \ diff --git a/docker/ubuntu-24.04.Dockerfile b/docker/ubuntu-24.04.Dockerfile new file mode 100644 index 00000000..ae6bfa31 --- /dev/null +++ b/docker/ubuntu-24.04.Dockerfile @@ -0,0 +1,121 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +ARG CUDA_VERSION=13.0.0 + +# Multi-arch container support available in non-cudnn containers. +FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04 +LABEL maintainer="NVIDIA CORPORATION" + +ENV TRT_VERSION=10.14.1.48 +SHELL ["/bin/bash", "-c"] + +# Setup user account and edit default account +RUN if id "ubuntu" &>/dev/null; then \ + usermod -u 1234 ubuntu && \ + groupmod -g 1234 ubuntu; \ + fi +ARG uid=1000 +ARG gid=1000 +RUN groupadd -r -f -g ${gid} trtuser && useradd -o -r -l -u ${uid} -g ${gid} -ms /bin/bash trtuser +RUN usermod -aG sudo trtuser +RUN echo 'trtuser:nvidia' | chpasswd +RUN mkdir -p /workspace && chown trtuser /workspace + +ENV DEBIAN_FRONTEND=noninteractive + +# Update CUDA signing key +RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/3bf863cc.pub + +# Install requried libraries +RUN apt-get update && apt-get install -y software-properties-common +RUN add-apt-repository ppa:ubuntu-toolchain-r/test +RUN apt-get update && apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev \ + wget \ + git \ + pkg-config \ + sudo \ + ssh \ + libssl-dev \ + pbzip2 \ + pv \ + bzip2 \ + unzip \ + devscripts \ + lintian \ + fakeroot \ + dh-make \ + build-essential + +# Install python3 +RUN apt-get install -y --no-install-recommends \ + python3 \ + python3-pip \ + python3-dev \ + python3-wheel \ + python3-venv &&\ + cd /usr/local/bin &&\ + ln -s /usr/bin/python3 python &&\ + ln -s /usr/bin/pip3 pip; + +# Create python3 virtualenv +RUN python3 -m venv /opt/venv +ENV PATH="/opt/venv/bin:$PATH" + +# Install TensorRT +RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/x86_64-linux-gnu/ \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_x86_64.whl ;\ + elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \ + wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \ + && cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/x86_64-linux-gnu/ \ + && pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_x86_64.whl ;\ + else \ + echo "Invalid CUDA_VERSION"; \ + exit 1; \ + fi + + # Install PyPI packages +RUN pip3 install --upgrade pip +RUN pip3 install setuptools>=41.0.0 +RUN pip3 install jupyter jupyterlab +# Workaround to remove numpy installed with tensorflow +RUN pip3 install --upgrade numpy + +# Install Cmake +RUN cd /tmp && \ + wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \ + chmod +x cmake-3.27.9-Linux-x86_64.sh && \ + ./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \ + rm ./cmake-3.27.9-Linux-x86_64.sh + +# Download NGC client +RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set + +# Set environment and working directory +ENV TRT_LIBPATH=/usr/lib/x86_64-linux-gnu/ +ENV TRT_OSSPATH=/workspace/TensorRT +ENV PATH="/workspace/TensorRT/build/out:${PATH}:/usr/local/bin/ngc-cli" +ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${TRT_OSSPATH}/build/out:${TRT_LIBPATH}" +WORKDIR /workspace + +USER trtuser +RUN ["/bin/bash"] diff --git a/docker/ubuntu-cross-aarch64.Dockerfile b/docker/ubuntu-cross-aarch64.Dockerfile index 00eaacb0..c0a60575 100644 --- a/docker/ubuntu-cross-aarch64.Dockerfile +++ b/docker/ubuntu-cross-aarch64.Dockerfile @@ -21,7 +21,7 @@ ARG OS_VERSION=24.04 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${OS_VERSION} LABEL maintainer="NVIDIA CORPORATION" -ENV TRT_VERSION 10.13.3.9 +ENV TRT_VERSION 10.14.1.48 ENV DEBIAN_FRONTEND=noninteractive # Setup user account and edit default account @@ -84,9 +84,9 @@ RUN wget https://developer.download.nvidia.com/compute/cuda/13.0.0/local_install # Unpack libnvinfer. -RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz && \ - tar -xf TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz && \ - cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib/aarch64-linux-gnu +RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz && \ + tar -xf TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz && \ + cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/aarch64-linux-gnu # Link required library RUN cd /usr/aarch64-linux-gnu/lib && ln -sf librt.so.1 librt.so diff --git a/include/NvInfer.h b/include/NvInfer.h index d8cea896..b02d73c8 100644 --- a/include/NvInfer.h +++ b/include/NvInfer.h @@ -19,7 +19,7 @@ #define NV_INFER_H #include "NvInferLegacyDims.h" -#include "NvInferRuntime.h" +#include "NvInferRuntime.h" // IWYU pragma: export //! //! \mainpage @@ -106,7 +106,9 @@ enum class LayerType : int32_t kSQUEEZE = 47, //!< Squeeze Layer. kUNSQUEEZE = 48, //!< Unsqueeze Layer. kCUMULATIVE = 49, //!< Cumulative layer. - kDYNAMIC_QUANTIZE = 50, //!< Dynamic Quantize layer. + kDYNAMIC_QUANTIZE = 50, //!< Dynamic Quantize layer. + kATTENTION_INPUT = 51, //!< Attention Input. + kATTENTION_OUTPUT = 52, //!< Attention Output. }; //! @@ -117,7 +119,7 @@ enum class LayerType : int32_t template <> constexpr inline int32_t EnumMax() noexcept { - return 51; + return 53; } //! @@ -151,6 +153,7 @@ enum class ActivationType : int32_t kGELU_TANH = 13 //!< GELU tanh activation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (0.044715F * pow(x, 3) + x))) }; + namespace impl { //! @@ -3645,6 +3648,33 @@ public: //! using ILayer::setInput; + //! + //! \brief Set the indices type for the layer. + //! + //! \param type The DataType of the indices tensor. + //! + //! \return true if set successfully, false otherwise. + //! + //! Set the indices (the second output) type of the TopK layer. Valid values are DataType::kINT32 and + //! DataType::kINT64, otherwise an error occurs and the type is not updated. + //! + bool setIndicesType(DataType type) noexcept + { + return mImpl->setIndicesType(type); + } + + //! + //! \brief Return the TopK layer indices type. + //! + //! \return indices type set during layer creation or by setIndicesType(). + //! The return value is the indices type of the TopK layer. + //! The default value is DataType::kINT32. + //! + DataType getIndicesType() const noexcept + { + return mImpl->getIndicesType(); + } + protected: apiv::VTopKLayer* mImpl; virtual ~ITopKLayer() noexcept = default; @@ -3757,7 +3787,7 @@ protected: //! //! The input may have type kFLOAT, kHALF, kINT32, or kBOOL. //! -//! The output is a matrix of type kINT32. +//! The output is a matrix of type kINT32 or kINT64. //! For an input with dimensions [L1, L2, ..., Lm], the output has dimensions [m,n], //! where n is the number of non-zero elements. I.e., each column denotes a m-D position. //! @@ -3771,6 +3801,34 @@ protected: //! class INonZeroLayer : public ILayer { +public: + //! + //! \brief Set the indices type for the layer. + //! + //! \param type The DataType of the indices tensor. + //! + //! \return true if set successfully, false otherwise. + //! + //! Set the indices (the first output) type of the NonZero layer. Valid values are DataType::kINT32 and + //! DataType::kINT64, otherwise an error occurs and the type is not updated. + //! + bool setIndicesType(DataType type) noexcept + { + return mImpl->setIndicesType(type); + } + + //! + //! \brief Return the NonZero layer indices type. + //! + //! \return indices type set during layer creation or by setIndicesType(). + //! The return value is the indices type of the NonZero layer. + //! The default value is DataType::kINT32. + //! + DataType getIndicesType() const noexcept + { + return mImpl->getIndicesType(); + } + protected: virtual ~INonZeroLayer() noexcept = default; apiv::VNonZeroLayer* mImpl; @@ -6145,10 +6203,12 @@ constexpr inline int32_t EnumMax() noexcept //! intersection-over-union (IoU) with previously selected boxes is less than or equal to a given threshold. //! This layer implements NMS per batch item and per class. //! -//! Per batch item, boxes are initially sorted by their scores without regard to class. Only boxes up to a maximum of the TopK limit are considered for selection (per batch). -//! During selection, only overlapping boxes of the same class are compared, so that overlapping boxes of different classes do not suppress each other. +//! Per batch item, boxes are initially sorted by their scores without regard to class. Only boxes up to a maximum of +//! the TopK limit are considered for selection (per batch). During selection, only overlapping boxes of the same class +//! are compared, so that overlapping boxes of different classes do not suppress each other. //! -//! For each batch item, the ordering of candidate bounding boxes with the same score is unspecified, but the ordering will be consistent across different runs for the same inputs. +//! For each batch item, the ordering of candidate bounding boxes with the same score is unspecified, but the ordering +//! will be consistent across different runs for the same inputs. //! //! The layer has the following inputs, in order of input index: //! @@ -6161,19 +6221,21 @@ constexpr inline int32_t EnumMax() noexcept //! It is a scalar (0D tensor) of type kINT32. //! * IoUThreshold is the maximum IoU for selected boxes. It is a scalar (0D tensor) of type kFLOAT in the range //! [0.0f, 1.0f]. It is an optional input with default 0.0f. -//! * ScoreThreshold is the value that a box score must exceed in order to be selected. It is a scalar (0D tensor) of type kFLOAT. It is an optional +//! * ScoreThreshold is the value that a box score must exceed in order to be selected. It is a scalar (0D tensor) of +//! type kFLOAT. It is an optional //! input with default 0.0f. //! //! The layer has the following outputs, in order of output index: //! -//! * SelectedIndices contains the indices of the selected boxes. It is a linear tensor of type kINT32. It has shape +//! * SelectedIndices contains the indices of the selected boxes. It is a linear tensor of type kINT32 or kINT64. It has +//! shape //! [NumOutputBoxes, 3]. Each row contains a (batchIndex, classIndex, boxIndex) tuple. -//! The output boxes are sorted in order of increasing batchIndex and then in order of decreasing score within each batchIndex. -//! For each batchIndex, the ordering of output boxes with the same score is unspecified. -//! If MaxOutputBoxesPerClass is a constant input, the maximum number of output boxes is -//! batchSize * numClasses * min(numInputBoundingBoxes, MaxOutputBoxesPerClass). -//! Otherwise, the maximum number of output boxes is batchSize * numClasses * numInputBoundingBoxes. -//! The maximum number of output boxes is used to determine the upper-bound on allocated memory for this output tensor. +//! The output boxes are sorted in order of increasing batchIndex and then in order of decreasing score within each +//! batchIndex. For each batchIndex, the ordering of output boxes with the same score is unspecified. If +//! MaxOutputBoxesPerClass is a constant input, the maximum number of output boxes is batchSize * numClasses * +//! min(numInputBoundingBoxes, MaxOutputBoxesPerClass). Otherwise, the maximum number of output boxes is batchSize * +//! numClasses * numInputBoundingBoxes. The maximum number of output boxes is used to determine the upper-bound on +//! allocated memory for this output tensor. //! * NumOutputBoxes is the number of output boxes in SelectedIndices. It is a scalar (0D tensor) of type kINT32. //! //! \warning There is a hardware-dependent limit K such that only the K highest scoring boxes in each batch item @@ -6254,6 +6316,33 @@ public: //! using ILayer::setInput; + //! + //! \brief Set the indices type for the layer. + //! + //! \param type The DataType of the indices tensor. + //! + //! \return true if set successfully, false otherwise. + //! + //! Set the indices (the first output) type of the NMS layer. Valid values are DataType::kINT32 and + //! DataType::kINT64, otherwise an error occurs and the type is not updated. + //! + bool setIndicesType(DataType type) noexcept + { + return mImpl->setIndicesType(type); + } + + //! + //! \brief Return the NMS layer indices type. + //! + //! \return indices type set during layer creation or by setIndicesType(). + //! The return value is the indices type of the NMS layer. + //! The default value is DataType::kINT32. + //! + DataType getIndicesType() const noexcept + { + return mImpl->getIndicesType(); + } + protected: apiv::VNMSLayer* mImpl; virtual ~INMSLayer() noexcept = default; @@ -6661,6 +6750,409 @@ protected: virtual ~ICumulativeLayer() noexcept = default; }; +//! +//! \enum AttentionNormalizationOp +//! +//! \brief Enumerates the operations that may be performed by the normalization in the attention subgraph. +//! +enum class AttentionNormalizationOp : int32_t +{ + kNONE + = 0, //!< Apply no normalization on the attention scores. Must be used with decomposable=True on pre-Blackwell GPUs + kSOFTMAX = 1, //!< Apply softmax normalization on the attention scores on the `s_kv` dimension. +}; + +namespace impl +{ +//! +//! Maximum number of elements in AttentionNormalizationOp enum. +//! +//! \see AttentionNormalizationOp +//! +template <> +struct EnumMaxImpl +{ + static constexpr int32_t kVALUE = 2; +}; + +} // namespace impl + +//! +//! \class IAttentionBoundaryLayer +//! +//! \brief This is a base class for Attention boundary layers. +//! +//! Boundary layers are used to demarcate the boundaries of IAttention. +//! Typically client code does not deal directly with the boundary layers. +//! However, they are indirectly visible via method `INetworkDefinition::getLayer(int32_t index)`. +//! +class IAttentionBoundaryLayer : public ILayer +{ +public: + //! + //! \brief Get a pointer to the IAttention associated with this boundary layer. + //! + IAttention* getAttention() const noexcept + { + return mBoundary->getAttention(); + } + +protected: + virtual ~IAttentionBoundaryLayer() noexcept = default; + apiv::VAttentionBoundaryLayer* mBoundary; +}; + +//! +//! \class IAttentionInputLayer +//! +//! \brief This layer represents an input to an attention subgraph. +//! +//! This layer is automatically created when an `IAttention` is created. Clients typically do not +//! deal with the layer directly, but instead specify its input via `addAttention` or `IAttention::setInput`. +//! +//! An IAttentionInputLayer has three to four inputs and one output. +//! +class IAttentionInputLayer : public IAttentionBoundaryLayer +{ +public: + //! + //! \brief Append or replace an input of this layer with a specific tensor + //! + //! \param index the index of the input to modify. + //! \param tensor the new input tensor + //! + //! The indices are as follows: + //! + //! Input 0 is the input query tensor. + //! Input 1 is the input key tensor. + //! Input 2 is the input value tensor. + //! Input 3 is the optional mask tensor. setMask should be used instead of setInput + //! Input 4 is the optional normalizationQuantizeScale tensor. setNormalizationQuantizeScale should be used instead + //! of setInput + //! + using ILayer::setInput; + +protected: + virtual ~IAttentionInputLayer() noexcept = default; + apiv::VAttentionInputLayer* mImpl; +}; + +//! +//! \class IAttentionOutputLayer +//! +//! \brief This layer represents an output of an IAttention. +//! +//! This layer is automatically created when an `IAttention` is created. Clients typically do not +//! deal with the layer directly, but instead getting its output via `IAttention::getOutput`. +//! +//! An IAttentionOutputLayer has one input and one output. +//! +class IAttentionOutputLayer : public IAttentionBoundaryLayer +{ +public: +protected: + virtual ~IAttentionOutputLayer() noexcept = default; + apiv::VAttentionOutputLayer* mImpl; +}; + +//! +//! \class IAttention +//! +//! \brief Helper for constructing an attention that consumes query, key and value tensors. +//! +//! An attention subgraph implicitly includes three main components, two MatrixMultiply layers +//! known as BMM1 and BMM2, and one normalization operation which defaults to be a Softmax. +//! By default, IAttention is not decomposable and TensorRT will try to use a single fused kernel, which may be more +//! efficient than if the subgraph is expressed without IAttention. Setting the IAttention to decomposable=True can +//! allow IAttention to be decomposed to use multiple kernels if no fused kernel support found. +//! +//! Query Key Value Mask (optional) NormalizationQuantizeScale (optional) +//! | | | | | +//! | Transpose | | | +//! | | | | | +//! ----BMM1---- | | | +//! | | | | +//! *--------------------------- | +//! | | | +//! Normalization | | +//! | | | +//! *------------------------------------------------ +//! | | +//! -------BMM2------ +//! | +//! Output +//! +//! The attention has the following inputs, in order of input index: +//! +//! * Query contains the input query. It is a tensor of type kFLOAT, kHALF or kBF16 with +//! shape [batchSize, numHeadsQuery, sequenceLengthQuery, dimHead] +//! * Key contains the input key. It is a tensor of type kFLOAT, kHALF or kBF16 with +//! shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead] +//! * Value contains the input value. It is a tensor of type kFLOAT, kHALF or kBF16 with +//! shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead] +//! * Mask (optional) contains the mask value. It is a tensor of type kBOOL or the same data type of +//! BMM1 output with shape [batchSize, numHeadsQuery, sequenceLengthQuery, sequenceLengthKeyValue] +//! with batchSize and numHeadsQuery broadcastable. For a kBOOL mask, a True value indicates that the corresponding +//! position is allowed to attend. For other data types, the mask values will be added to the BMM1 output, known +//! as an add mask. +//! * NormalizationQuantizeScale (optional) contains the quantization scale for the attention normalization output. +//! It is a tensor of type kFLOAT, kHALF or kBF16 with dimension 0 or 1. +//! +//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI. +//! +class IAttention : public INoCopy +{ +public: + //! + //! \brief Set the normalization operation for the attention. + //! + //! \see getNormalizationOperation(), AttentionNormalizationOp + //! + //! \return True if the normalization operation is set successfully, false otherwise. + //! + bool setNormalizationOperation(AttentionNormalizationOp op) noexcept + { + return mImpl->setNormalizationOperation(op); + } + + //! + //! \brief Get the normalization operation for the attention. + //! + //! \see setNormalizationOperation(), AttentionNormalizationOp + //! + //! \return The normalization operation for the attention. Default is kSOFTMAX. + //! + AttentionNormalizationOp getNormalizationOperation() const noexcept + { + return mImpl->getNormalizationOperation(); + } + + //! + //! \brief Set whether a mask will be used for the normalization operation. + //! + //! \param mask the mask tensor of type kBOOL or the same data type of + //! BMM1 output with shape [batchSize, sequenceLengthQuery, sequenceLengthKeyValue]. For a kBOOL mask, a True value + //! indicates that the corresponding position is allowed to attend. For other data types, the mask values will + //! be added to the BMM1 output, known as an add mask. + //! + //! \see getMask + //! + //! \return True if the mask is set successfully, false otherwise. + //! + bool setMask(ITensor& mask) noexcept + { + return mImpl->setMask(mask); + } + + //! + //! \brief Get the optional mask in attention. + //! + //! \see setMask + //! + //! \return The optional mask in attention, nullptr if no mask is set. + //! + ITensor* getMask() noexcept + { + return mImpl->getMask(); + } + + //! + //! \brief Set whether the attention will run a causal inference. + //! Cannot be used together with setMask(). + //! + //! \see getCausal + //! + //! \return True if the causal inference is set successfully, false otherwise. + //! + bool setCausal(bool isCausal) noexcept + { + return mImpl->setCausal(isCausal); + } + + //! + //! \brief Get whether the attention will run a causal inference. + //! + //! \see setCausal + //! + //! \return True if the attention will run a causal inference, false otherwise. Default is false. + //! + bool getCausal() const noexcept + { + return mImpl->getCausal(); + } + + //! + //! \brief Set whether the attention can be decomposed to use multiple kernels if no fused kernel support found. + //! + //! \see getDecomposable + //! + //! \return True if the decomposable attention is set successfully, false otherwise. + //! + bool setDecomposable(bool decomposable) noexcept + { + return mImpl->setDecomposable(decomposable); + } + + //! + //! \brief Get whether the attention can be decomposed to use multiple kernels if no fused kernel support found. + //! + //! \return True if the attention can be decomposed to use multiple kernels by the compiler, + //! false otherwise. Default is false. + //! + //! \see setDecomposable + //! + bool getDecomposable() const noexcept + { + return mImpl->getDecomposable(); + } + + //! + //! \brief Append or replace an input of this layer with a specific tensor. + //! + //! \param index the index of the input to modify. + //! \param input the new input tensor. + //! + //! The indices are as follows: + //! + //! Input 0 is the input query tensor. + //! Input 1 is the input key tensor. + //! Input 2 is the input value tensor. + //! + //! \return True if the input tensor is set successfully, false otherwise. + //! + bool setInput(int32_t index, ITensor& input) noexcept + { + return mImpl->setInput(index, input); + } + + //! + //! \brief Get the number of inputs of IAttention. IAttention has three inputs. + //! + //! \return The number of inputs of IAttention. + int32_t getNbInputs() const noexcept + { + return mImpl->getNbInputs(); + } + + //! + //! \brief Get the IAttention input corresponding to the given index. + //! + //! \param index The index of the input tensor. + //! + //! \return The input tensor, or nullptr if the index is out of range. + //! + ITensor* getInput(int32_t index) const noexcept + { + return mImpl->getInput(index); + } + + //! + //! \brief Get the number of outputs of a layer. IAttention has one output. + //! + int32_t getNbOutputs() const noexcept + { + return mImpl->getNbOutputs(); + } + + //! + //! \brief Get the IAttention output corresponding to the given index. IAttention has only one output. + //! + //! \param index The index of the output tensor. + //! + //! \return The indexed output tensor, or nullptr if the index is out of range. + //! + ITensor* getOutput(int32_t index) const noexcept + { + return mImpl->getOutput(index); + } + + //! + //! \brief Set the name of the attention. + //! + //! The name is used in error diagnostics. + //! This method copies the name string. + //! + //! \warning The string name must be null-terminated, and be at most 4096 bytes including the terminator. + //! + //! \see getName() + //! + //! \return True if the name is set successfully, false otherwise. + //! + bool setName(char const* name) noexcept + { + return mImpl->setName(name); + } + + //! + //! \brief Return the name of the attention. + //! + //! \see setName() + //! + //! \return The name of the attention. + //! + char const* getName() const noexcept + { + return mImpl->getName(); + } + + //! + //! \brief Set the quantization scale for the attention normalization output. + //! + //! \param tensor for quantization scale. Data type must be DataType::kFLOAT, DataType::kHALF or DataType::kBF16. + //! Must be a 0-d or 1-d. + //! + //! \return True if the quantization scale is set successfully, false otherwise. + //! + //! \warning Must be used together with setNormalizationQuantizeToType to set normalization output datatype to + //! DataType::kFP8 or DataType::kINT8. + //! + bool setNormalizationQuantizeScale(ITensor& tensor) noexcept + { + return mImpl->setNormalizationQuantizeScale(tensor); + } + + //! + //! \brief Get the quantization scale for the attention normalization output. + //! + //! \return The quantization scale for the attention normalization output or nullptr if no quantization scale is + //! set. + //! + ITensor* getNormalizationQuantizeScale() const noexcept + { + return mImpl->getNormalizationQuantizeScale(); + } + + //! + //! \brief Set the datatype the attention normalization is quantized to. + //! + //! \param type the datatype the attention normalization is quantized to. Must be one of DataType::kFP8, + //! DataType::kINT8. + //! + //! \return True if the quantization to type is set successfully, false otherwise. + //! + bool setNormalizationQuantizeToType(DataType type) noexcept + { + return mImpl->setNormalizationQuantizeToType(type); + } + + //! + //! \brief Get the datatype the attention normalization is quantized to. + //! + //! \return The datatype the attention normalization is quantized to. + //! The default value is DataType::kFLOAT. + //! + //! \warning Must be used after normalization quantization to type is set by setNormalizationQuantizeToType. + DataType getNormalizationQuantizeToType() const noexcept + { + return mImpl->getNormalizationQuantizeToType(); + } + + +protected: + apiv::VAttention* mImpl; + virtual ~IAttention() noexcept = default; +}; + //! //! \class INetworkDefinition //! @@ -7112,6 +7604,8 @@ public: //! //! Currently only values of K up to 3840 are supported. //! + //! The default indices tensor (the second output) data type is DataType::kINT32. + //! //! \param input The input tensor to the layer. //! //! \param op Operation to perform. @@ -7122,19 +7616,53 @@ public: //! \param reduceAxes The reduction dimensions. //! The bit in position i of bitmask reduceAxes corresponds to explicit dimension i of the result. //! E.g., the least significant bit corresponds to the first explicit dimension and the next to least - //! significant bit corresponds to the second explicit dimension. - //! - //! Currently reduceAxes must specify exactly one dimension, and it must be one of the last four dimensions. + //! significant bit corresponds to the second explicit dimension. Currently reduceAxes must specify + //! exactly one dimension, and it must be one of the last four dimensions. //! //! \see ITopKLayer //! //! \return The new TopK layer, or nullptr if it could not be created. //! - ITopKLayer* addTopK(ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes) noexcept + //! \deprecated Deprecated in TensorRT 10.14. Superseded by five-argument addTopK. + //! + TRT_DEPRECATED ITopKLayer* addTopK(ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes) noexcept { return mImpl->addTopK(input, op, k, reduceAxes); } + //! + //! \brief Add a TopK layer to the network. + //! + //! The TopK layer has two outputs of the same dimensions. The first contains data values, + //! the second contains index positions for the values. Output values are sorted, largest first + //! for operation kMAX and smallest first for operation kMIN. + //! + //! Currently only values of K up to 3840 are supported. + //! + //! \param input The input tensor to the layer. + //! + //! \param op Operation to perform. + //! + //! \param k The number of elements to keep. For dynamic k, use the setInput() method to pass in k as a tensor + //! instead, which will override the static k value passed here in calculations. + //! + //! \param reduceAxes The reduction dimensions. + //! The bit in position i of bitmask reduceAxes corresponds to explicit dimension i of the result. + //! E.g., the least significant bit corresponds to the first explicit dimension and the next to least + //! significant bit corresponds to the second explicit dimension. Currently reduceAxes must specify + //! exactly one dimension, and it must be one of the last four dimensions. + //! + //! \param indicesType Indices tensor (the second output) data type, must be DataType::kINT32 or DataType::kINT64. + //! + //! \see ITopKLayer + //! + //! \return The new TopK layer, or nullptr if it could not be created. + //! + ITopKLayer* addTopK(ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes, DataType indicesType) noexcept + { + return mImpl->addTopKV2(input, op, k, reduceAxes, indicesType); + } + //! //! \brief Add gather with mode GatherMode::kDEFAULT and specified axis and nbElementWiseDims=0. //! @@ -7211,17 +7739,37 @@ public: //! //! \brief Add a nonzero layer to the network. //! + //! The default indices tensor (the first output) data type is DataType::kINT32. + //! //! \param input The input tensor to the layer. //! //! \see INonZeroLayer //! - //! \return The new nonzero layer, or nullptr if it could be created. + //! \return The new nonzero layer, or nullptr if it could not be created. //! - INonZeroLayer* addNonZero(ITensor& input) noexcept + //! \deprecated Deprecated in TensorRT 10.14. Superseded by two-argument addNonZero. + //! + TRT_DEPRECATED INonZeroLayer* addNonZero(ITensor& input) noexcept { return mImpl->addNonZero(input); } + //! + //! \brief Add a nonzero layer to the network. + //! + //! \param input The input tensor to the layer. + //! + //! \param indicesType Indices tensor (the first output) data type, must be DataType::kINT32 or DataType::kINT64. + //! + //! \see INonZeroLayer + //! + //! \return The new nonzero layer, or nullptr if it could not be created. + //! + INonZeroLayer* addNonZero(ITensor& input, DataType indicesType) noexcept + { + return mImpl->addNonZeroV2(input, indicesType); + } + //! //! \brief Add a constant layer to the network. //! @@ -7940,7 +8488,7 @@ public: //! //! This layer performs dynamic block quantization of its input tensor and outputs the //! quantized data and the computed block scale-factors. - //! The block size is currently limited to 16 and the size of the blocked axis must be divisible by 16. + //! The blocked axis dimension size must be divisible by the block size. //! //! \param input The input tensor to be quantized. Its data type must be one of DataType::kFLOAT, //! DataType::kHALF, or DataType::kBF16. Currently only 2D and 3D inputs are supported. @@ -7999,6 +8547,8 @@ public: //! //! \brief Add a non-maximum suppression layer to the network. //! + //! The default indices tensor (the first output) data type is DataType::kINT32. + //! //! \param boxes The input boxes tensor to the layer. //! //! \param scores The input scores tensor to the layer. @@ -8009,11 +8559,33 @@ public: //! //! \return The new NMS layer, or nullptr if it could not be created. //! - INMSLayer* addNMS(ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass) noexcept + //! \deprecated Deprecated in TensorRT 10.14. Superseded by four-argument addNMS. + //! + TRT_DEPRECATED INMSLayer* addNMS(ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass) noexcept { return mImpl->addNMS(boxes, scores, maxOutputBoxesPerClass); } + //! + //! \brief Add a non-maximum suppression layer to the network. + //! + //! \param boxes The input boxes tensor to the layer. + //! + //! \param scores The input scores tensor to the layer. + //! + //! \param maxOutputBoxesPerClass The input maxOutputBoxesPerClass tensor to the layer. + //! + //! \param indicesType Indices tensor (the first output) data type, must be DataType::kINT32 or DataType::kINT64. + //! + //! \see INMSLayer + //! + //! \return The new NMS layer, or nullptr if it could not be created. + //! + INMSLayer* addNMS(ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass, DataType indicesType) noexcept + { + return mImpl->addNMSV2(boxes, scores, maxOutputBoxesPerClass, indicesType); + } + //! //! \brief Add a ReverseSequence layer to the network. //! @@ -8079,6 +8651,34 @@ public: return mImpl->addCumulative(input, axis, operation, exclusive, reverse); } + //! + //! \brief Add an attention to the network. + //! + //! \param query A 4d input query tensor to the layer. + //! \param key A 4d input key tensor to the layer. + //! \param value A 4d input value tensor to the layer. + //! \param normOp The normalization operation to perform. + //! \param causal Use causual inference or not. + //! + //! query must have shape [batchSize, numHeadsQuery, sequenceLengthQuery, dimHead]. + //! key and value must have shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead]. + //! pastKey and pastValue must have shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead]. + //! normOp defaults to kSOFTMAX isCausal defaults to false. + //! + //! By default, IAttention is not decomposable and TensorRT will try to use a single fused kernel, which may be more + //! efficient than if the subgraph is expressed without IAttention. Setting the IAttention to decomposable=True can + //! allow IAttention to be to use multiple kernels if no fused kernel support found. + //! + //! \see IAttention + //! + //! \return The new attention, or nullptr if it could not be created. + //! + IAttention* addAttention( + ITensor& query, ITensor& key, ITensor& value, AttentionNormalizationOp normOp, bool causal) noexcept + { + return mImpl->addAttention(query, key, value, normOp, causal); + } + //! //! \brief Return the builder from which this INetworkDefinition was created. //! @@ -8812,6 +9412,8 @@ enum class RuntimePlatform : int32_t //! Designates the target platform for engine execution as Windows AMD64 system. Currently this flag can only be //! enabled when building engines on Linux AMD64 platforms. kWINDOWS_AMD64 = 1, + + }; namespace impl @@ -9018,6 +9620,8 @@ enum class BuilderFlag : int32_t //! For layers that perform einsum: //! Let n be the leftmost reduction axis. The axes to the left of n are distributive axes. kDISTRIBUTIVE_INDEPENDENCE = 28, + + }; //! diff --git a/include/NvInferImpl.h b/include/NvInferImpl.h index f166fca6..79a87165 100644 --- a/include/NvInferImpl.h +++ b/include/NvInferImpl.h @@ -68,6 +68,7 @@ class IAlgorithmIOInfo; class IAlgorithmVariant; #endif // !STRIP_TRT_RTX_INTERNAL_API class IAssertionLayer; +class IAttention; class IBuilder; class IBuilderConfig; class IConcatenationLayer; @@ -175,6 +176,7 @@ struct Permutation; class Weights; enum class ActivationType : int32_t; +enum class AttentionNormalizationOp : int32_t; enum class BoundingBoxFormat : int32_t; enum class BuilderFlag : int32_t; #if !STRIP_TRT_RTX_INTERNAL_API @@ -218,6 +220,7 @@ enum class HardwareCompatibilityLevel : int32_t; enum class ExecutionContextAllocationStrategy : int32_t; enum class RuntimePlatform : int32_t; enum class TilingOptimizationLevel : int32_t; +enum class EngineStat : int32_t; using TacticSources = uint32_t; @@ -420,6 +423,7 @@ public: TRT_NODISCARD virtual IExecutionContext* createExecutionContextWithRuntimeConfig( IRuntimeConfig* runtimeConfig) noexcept = 0; TRT_NODISCARD virtual IRuntimeConfig* createRuntimeConfig() noexcept = 0; + TRT_NODISCARD virtual int64_t getEngineStat(EngineStat stat) const noexcept = 0; }; class VExecutionContext : public VRoot @@ -776,6 +780,8 @@ public: virtual int32_t getK() const noexcept = 0; virtual void setReduceAxes(uint32_t reduceAxes) noexcept = 0; virtual uint32_t getReduceAxes() const noexcept = 0; + virtual bool setIndicesType(DataType type) noexcept = 0; + virtual DataType getIndicesType() const noexcept = 0; }; class VMatrixMultiplyLayer : public VRoot @@ -788,6 +794,8 @@ public: class VNonZeroLayer : public VRoot { public: + virtual bool setIndicesType(DataType type) noexcept = 0; + virtual DataType getIndicesType() const noexcept = 0; }; class VRaggedSoftMaxLayer : public VRoot @@ -917,6 +925,46 @@ public: virtual char const* getName() const noexcept = 0; }; +class VAttentionBoundaryLayer : public VRoot +{ +public: + virtual IAttention* getAttention() const noexcept = 0; +}; + +class VAttentionInputLayer : public VRoot +{ +public: +}; + +class VAttentionOutputLayer : public VRoot +{ +public: +}; + +class VAttention : public VRoot +{ +public: + TRT_NODISCARD virtual bool setInput(int32_t index, ITensor& input) noexcept = 0; + TRT_NODISCARD virtual int32_t getNbInputs() const noexcept = 0; + TRT_NODISCARD virtual ITensor* getInput(int32_t index) const noexcept = 0; + TRT_NODISCARD virtual int32_t getNbOutputs() const noexcept = 0; + TRT_NODISCARD virtual ITensor* getOutput(int32_t index) const noexcept = 0; + TRT_NODISCARD virtual bool setName(char const* name) noexcept = 0; + TRT_NODISCARD virtual char const* getName() const noexcept = 0; + TRT_NODISCARD virtual bool setNormalizationOperation(AttentionNormalizationOp op) noexcept = 0; + TRT_NODISCARD virtual AttentionNormalizationOp getNormalizationOperation() const noexcept = 0; + TRT_NODISCARD virtual bool setCausal(bool isCausal) noexcept = 0; + TRT_NODISCARD virtual bool getCausal() const noexcept = 0; + TRT_NODISCARD virtual bool setMask(ITensor& mask) noexcept = 0; + TRT_NODISCARD virtual ITensor* getMask() const noexcept = 0; + TRT_NODISCARD virtual bool setDecomposable(bool decomposable) noexcept = 0; + TRT_NODISCARD virtual bool getDecomposable() const noexcept = 0; + TRT_NODISCARD virtual bool setNormalizationQuantizeScale(ITensor& tensor) noexcept = 0; + TRT_NODISCARD virtual ITensor* getNormalizationQuantizeScale() const noexcept = 0; + TRT_NODISCARD virtual bool setNormalizationQuantizeToType(DataType type) noexcept = 0; + TRT_NODISCARD virtual DataType getNormalizationQuantizeToType() const noexcept = 0; +}; // class VAttention + class VSelectLayer : public VRoot { }; @@ -1020,6 +1068,8 @@ public: virtual BoundingBoxFormat getBoundingBoxFormat() const noexcept = 0; virtual void setTopKBoxLimit(int32_t limit) noexcept = 0; virtual int32_t getTopKBoxLimit() const noexcept = 0; + virtual bool setIndicesType(DataType type) noexcept = 0; + virtual DataType getIndicesType() const noexcept = 0; }; // class VNMSLayer class VReverseSequenceLayer : public VRoot @@ -1155,11 +1205,23 @@ public: virtual ISqueezeLayer* addSqueeze(ITensor& input, ITensor& axes) noexcept = 0; virtual IUnsqueezeLayer* addUnsqueeze(ITensor& input, ITensor& axes) noexcept = 0; virtual IDynamicQuantizeLayer* addDynamicQuantize( - ITensor& input, int32_t axis, int32_t blockSize, DataType toType, DataType scaleType) noexcept = 0; + ITensor& input, int32_t axis, int32_t blockSize, DataType toType, DataType scaleType) noexcept + = 0; virtual ICumulativeLayer* addCumulative( - ITensor& input, ITensor& axis, CumulativeOperation operation, bool exclusive, bool reverse) noexcept = 0; + ITensor& input, ITensor& axis, CumulativeOperation operation, bool exclusive, bool reverse) noexcept + = 0; virtual bool markUnfusedTensorsAsDebugTensors() noexcept = 0; virtual bool unmarkUnfusedTensorsAsDebugTensors() noexcept = 0; + virtual ITopKLayer* addTopKV2( + ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes, DataType indicesType) noexcept + = 0; + virtual INonZeroLayer* addNonZeroV2(ITensor& input, DataType indicesType) noexcept = 0; + virtual INMSLayer* addNMSV2( + ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass, DataType indicesType) noexcept + = 0; + virtual IAttention* addAttention( + ITensor& query, ITensor& key, ITensor& value, AttentionNormalizationOp normOp, bool isCausal) noexcept + = 0; }; #if !STRIP_TRT_RTX_INTERNAL_API diff --git a/include/NvInferRuntime.h b/include/NvInferRuntime.h index f56edeb3..8e7a0e6b 100644 --- a/include/NvInferRuntime.h +++ b/include/NvInferRuntime.h @@ -24,11 +24,11 @@ //! This is the top-level API file for TensorRT extended runtime library. //! -#include "NvInferImpl.h" +#include "NvInferImpl.h" // IWYU pragma: export #define NV_INFER_INTERNAL_INCLUDE 1 -#include "NvInferPluginBase.h" // IWYU pragma: exports +#include "NvInferPluginBase.h" // IWYU pragma: export #undef NV_INFER_INTERNAL_INCLUDE -#include "NvInferRuntimeCommon.h" +#include "NvInferRuntimeCommon.h" // IWYU pragma: export namespace nvinfer1 { @@ -1854,6 +1854,7 @@ protected: //! using IGpuAllocator = v_1_0::IGpuAllocator; + //! //! \class IRuntime //! @@ -2178,6 +2179,7 @@ public: return mImpl->getEngineHostCodeAllowed(); } + protected: apiv::VRuntime* mImpl; }; @@ -3000,13 +3002,14 @@ enum class SerializationFlag : int32_t { kEXCLUDE_WEIGHTS = 0, //!< Exclude the weights that can be refitted. kEXCLUDE_LEAN_RUNTIME = 1, //!< Exclude the lean runtime. + kINCLUDE_REFIT = 2, //!< Remain refittable if originally so. }; //! Maximum number of serialization flags in SerializationFlag enum. \see SerializationFlag template <> constexpr inline int32_t EnumMax() noexcept { - return 2; + return 3; } //! @@ -3155,6 +3158,34 @@ protected: apiv::VRuntimeConfig* mImpl; }; // class IRuntimeConfig +//! +//! \enum EngineStat +//! +//! \brief The kind of engine statistics that queried from the ICudaEngine. +//! +//! \see ICudaEngine::getEngineStat() +//! \see BuilderFlag::kSTRIP_PLAN +//! +enum class EngineStat : int32_t +{ + //! Return the total weight size in bytes. + kTOTAL_WEIGHTS_SIZE = 0, + + //! Return the stripped weight size in bytes for engines built with BuilderFlag::kSTRIP_PLAN. + kSTRIPPED_WEIGHTS_SIZE = 1, +}; + +//! +//! \brief Maximum number of engine statistic kinds in EngineStat enum. +//! +//! \see EngineStat +//! +template <> +constexpr inline int32_t EnumMax() noexcept +{ + return 2; +} + //! //! \class ICudaEngine //! @@ -3827,6 +3858,10 @@ public: //! Serializing plan file with SerializationFlag::kEXCLUDE_WEIGHTS requires building the engine with kREFIT, //! kREFIT_IDENTICAL or kREFIT_INDIVIDUAL. //! + //! The only applicable scenario for SerializationFlag::kINCLUDE_REFIT is when serializing weight-stripping + //! engines without kEXCLUDE_WEIGHTS. By default, the resulting serialized engine is unrefittable. Setting + //! SerializationFlag::kINCLUDE_REFIT ensures that the serialized engine remains refittable. + //! //! \see IRuntime::deserializeCudaEngine() //! IHostMemory* serializeWithConfig(ISerializationConfig& config) const noexcept @@ -4090,6 +4125,34 @@ public: return mImpl->getProfileTensorValuesV2(tensorName, profileIndex, select); } + //! + //! \brief Get engine statistics according to the given enum value. + //! + //! \param stat The kind of statistics to query. + //! + //! If stat is kTOTAL_WEIGHTS_SIZE, the return value is the total weights size in bytes in the engine. + //! If stat is kSTRIPPED_WEIGHTS_SIZE, the return value is the stripped weight size in bytes for engines + //! built with BuilderFlag::kSTRIP_PLAN. + //! + //! When the BuilderFlag::kWEIGHT_STREAMING flag is enabled, engine weights may not be fully copied to the device. + //! The reported total weight size reflects the sum of all weights utilized by the engine, + //! which does not necessarily correspond to the actual GPU memory allocated. + //! + //! \return The kind of statistics specified by EngineStat. + //! + //! \warning if kSTRIPPED_WEIGHTS_SIZE is passed to query a normal engine, this function will + //! return -1 to indicate invalid enum value. + //! + //! \see EngineStat + //! \see BuilderFlag::kWEIGHT_STREAMING + //! \see setWeightStreamingBudget() + //! \see getStreamableWeightsSize() + //! + int64_t getEngineStat(EngineStat stat) const noexcept + { + return mImpl->getEngineStat(stat); + } + protected: apiv::VCudaEngine* mImpl; }; diff --git a/include/NvInferRuntimeBase.h b/include/NvInferRuntimeBase.h index bd021865..a1e7e283 100644 --- a/include/NvInferRuntimeBase.h +++ b/include/NvInferRuntimeBase.h @@ -138,6 +138,8 @@ constexpr int32_t EnumMax() noexcept //! //! \enum DataType //! \brief The type of weights and tensors. +//! The datatypes other than kBOOL, kINT32, and kINT64 are "activation datatypes," +//! as they often represent values corresponding to inference results. //! enum class DataType : int32_t { diff --git a/include/NvInferVersion.h b/include/NvInferVersion.h index 3365e357..2ce6b509 100644 --- a/include/NvInferVersion.h +++ b/include/NvInferVersion.h @@ -24,9 +24,9 @@ #define NV_INFER_VERSION_H #define TRT_MAJOR_ENTERPRISE 10 -#define TRT_MINOR_ENTERPRISE 13 -#define TRT_PATCH_ENTERPRISE 3 -#define TRT_BUILD_ENTERPRISE 9 +#define TRT_MINOR_ENTERPRISE 14 +#define TRT_PATCH_ENTERPRISE 1 +#define TRT_BUILD_ENTERPRISE 48 #define NV_TENSORRT_MAJOR TRT_MAJOR_ENTERPRISE //!< TensorRT major version. #define NV_TENSORRT_MINOR TRT_MINOR_ENTERPRISE //!< TensorRT minor version. #define NV_TENSORRT_PATCH TRT_PATCH_ENTERPRISE //!< TensorRT patch version. diff --git a/include/NvOnnxParser.h b/include/NvOnnxParser.h new file mode 100644 index 00000000..7825c330 --- /dev/null +++ b/include/NvOnnxParser.h @@ -0,0 +1,640 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef NV_ONNX_PARSER_H +#define NV_ONNX_PARSER_H + +#include "NvInfer.h" +#include +#include +#include + +//! +//! \file NvOnnxParser.h +//! +//! This is the API for the ONNX Parser +//! + +#define NV_ONNX_PARSER_MAJOR 0 +#define NV_ONNX_PARSER_MINOR 1 +#define NV_ONNX_PARSER_PATCH 0 + +static constexpr int32_t NV_ONNX_PARSER_VERSION + = ((NV_ONNX_PARSER_MAJOR * 10000) + (NV_ONNX_PARSER_MINOR * 100) + NV_ONNX_PARSER_PATCH); + +//! +//! \typedef SubGraph_t +//! +//! \brief The data structure containing the parsing capability of +//! a set of nodes in an ONNX graph. +//! +typedef std::pair, bool> SubGraph_t; + +//! +//! \typedef SubGraphCollection_t +//! +//! \brief The data structure containing all SubGraph_t partitioned +//! out of an ONNX graph. +//! +typedef std::vector SubGraphCollection_t; + +//! +//! \namespace nvonnxparser +//! +//! \brief The TensorRT ONNX parser API namespace +//! +namespace nvonnxparser +{ + +template +constexpr inline int32_t EnumMax() noexcept; + +//! +//! \enum ErrorCode +//! +//! \brief The type of error that the parser or refitter may return +//! +enum class ErrorCode : int +{ + kSUCCESS = 0, + kINTERNAL_ERROR = 1, + kMEM_ALLOC_FAILED = 2, + kMODEL_DESERIALIZE_FAILED = 3, + kINVALID_VALUE = 4, + kINVALID_GRAPH = 5, + kINVALID_NODE = 6, + kUNSUPPORTED_GRAPH = 7, + kUNSUPPORTED_NODE = 8, + kUNSUPPORTED_NODE_ATTR = 9, + kUNSUPPORTED_NODE_INPUT = 10, + kUNSUPPORTED_NODE_DATATYPE = 11, + kUNSUPPORTED_NODE_DYNAMIC = 12, + kUNSUPPORTED_NODE_SHAPE = 13, + kREFIT_FAILED = 14 +}; + +//! +//! Maximum number of flags in the ErrorCode enum. +//! +//! \see ErrorCode +//! +template <> +constexpr inline int32_t EnumMax() noexcept +{ + return 14; +} + +//! +//! \brief Represents one or more OnnxParserFlag values using binary OR +//! operations, e.g., 1U << OnnxParserFlag::kNATIVE_INSTANCENORM +//! +//! \see IParser::setFlags() and IParser::getFlags() +//! +using OnnxParserFlags = uint32_t; + +enum class OnnxParserFlag : int32_t +{ + //! Parse the ONNX model into the INetworkDefinition with the intention of using TensorRT's native layer + //! implementation over the plugin implementation for InstanceNormalization nodes. + //! This flag is required when building version-compatible or hardware-compatible engines. + //! This flag is set to be ON by default. + kNATIVE_INSTANCENORM = 0, + //! Enable UINT8 as a quantization data type and asymmetric quantization with non-zero zero-point values + //! in Quantize and Dequantize nodes. This flag is set to be OFF by default. + //! The resulting engine must be built targeting DLA version >= 3.16. + kENABLE_UINT8_AND_ASYMMETRIC_QUANTIZATION_DLA = 1, +}; + +//! +//! Maximum number of flags in the OnnxParserFlag enum. +//! +//! \see OnnxParserFlag +//! +template <> +constexpr inline int32_t EnumMax() noexcept +{ + return 2; +} + +//! +//! \class IParserError +//! +//! \brief an object containing information about an error +//! +class IParserError +{ +public: + //! + //!\brief the error code. + //! + virtual ErrorCode code() const = 0; + //! + //!\brief description of the error. + //! + virtual char const* desc() const = 0; + //! + //!\brief source file in which the error occurred. + //! + virtual char const* file() const = 0; + //! + //!\brief source line at which the error occurred. + //! + virtual int line() const = 0; + //! + //!\brief source function in which the error occurred. + //! + virtual char const* func() const = 0; + //! + //!\brief index of the ONNX model node in which the error occurred. + //! + virtual int node() const = 0; + //! + //!\brief name of the node in which the error occurred. + //! + virtual char const* nodeName() const = 0; + //! + //!\brief name of the node operation in which the error occurred. + //! + virtual char const* nodeOperator() const = 0; + //! + //!\brief A list of the local function names, from the top level down, constituting the current + //! stack trace in which the error occurred. A top-level node that is not inside any + //! local function would return a nullptr. + //! + virtual char const* const* localFunctionStack() const = 0; + //! + //!\brief The size of the stack of local functions at the point where the error occurred. + //! A top-level node that is not inside any local function would correspond to + // a stack size of 0. + //! + virtual int32_t localFunctionStackSize() const = 0; + +protected: + virtual ~IParserError() {} +}; + +//! +//! \class IParser +//! +//! \brief an object for parsing ONNX models into a TensorRT network definition +//! +//! \warning If the ONNX model has a graph output with the same name as a graph input, +//! the output will be renamed by prepending "__". +//! +//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI. +//! +class IParser +{ +public: + //! + //! \brief Parse a serialized ONNX model into the TensorRT network. + //! This method has very limited diagnostics. If parsing the serialized model + //! fails for any reason (e.g. unsupported IR version, unsupported opset, etc.) + //! it the user responsibility to intercept and report the error. + //! To obtain a better diagnostic, use the parseFromFile method below. + //! + //! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns. + //! \param serialized_onnx_model_size Size of the serialized ONNX model + //! in bytes + //! \param model_path Absolute path to the model file for loading external weights if required + //! \return true if the model was parsed successfully + //! \see getNbErrors() getError() + //! + virtual bool parse( + void const* serialized_onnx_model, size_t serialized_onnx_model_size, const char* model_path = nullptr) noexcept + = 0; + + //! + //! \brief Parse an onnx model file, which can be a binary protobuf or a text onnx model + //! calls parse method inside. + //! + //! \param onnxModelFile name + //! \param verbosity Level + //! + //! \return true if the model was parsed successfully + //! + //! + virtual bool parseFromFile(const char* onnxModelFile, int verbosity) noexcept = 0; + + //! + //! [DEPRECATED] Deprecated in TensorRT 10.1. See supportsModelV2. + //! + //! \brief Check whether TensorRT supports a particular ONNX model. + //! If the function returns True, one can proceed to engine building + //! without having to call \p parse or \p parseFromFile. + //! + //! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns. + //! \param serialized_onnx_model_size Size of the serialized ONNX model + //! in bytes + //! \param sub_graph_collection Container to hold supported subgraphs + //! \param model_path Absolute path to the model file for loading external weights if required + //! \return true if the model is supported + //! + TRT_DEPRECATED virtual bool supportsModel(void const* serialized_onnx_model, size_t serialized_onnx_model_size, + SubGraphCollection_t& sub_graph_collection, const char* model_path = nullptr) noexcept = 0; + + //! + //! [DEPRECATED] Deprecated in TensorRT 10.13. See loadInitializer(). + //! + //!\brief Parse a serialized ONNX model into the TensorRT network + //! with consideration of user provided weights + //! + //! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns. + //! \param serialized_onnx_model_size Size of the serialized ONNX model + //! in bytes + //! \return true if the model was parsed successfully + //! \see getNbErrors() getError() + //! + TRT_DEPRECATED virtual bool parseWithWeightDescriptors( + void const* serialized_onnx_model, size_t serialized_onnx_model_size) noexcept = 0; + + //! + //!\brief Returns whether the specified operator may be supported by the + //! parser. + //! + //! Note that a result of true does not guarantee that the operator will be + //! supported in all cases (i.e., this function may return false-positives). + //! + //! \param op_name The name of the ONNX operator to check for support + //! + virtual bool supportsOperator(const char* op_name) const noexcept = 0; + + //! + //!\brief Get the number of errors that occurred during prior calls to + //! \p parse + //! + //! \see getError() clearErrors() IParserError + //! + virtual int getNbErrors() const noexcept = 0; + + //! + //!\brief Get an error that occurred during prior calls to \p parse + //! + //! \see getNbErrors() clearErrors() IParserError + //! + virtual IParserError const* getError(int index) const noexcept = 0; + + //! + //!\brief Clear errors from prior calls to \p parse + //! + //! \see getNbErrors() getError() IParserError + //! + virtual void clearErrors() noexcept = 0; + + virtual ~IParser() noexcept = default; + + //! + //! \brief Query the plugin libraries needed to implement operations used by the parser in a version-compatible + //! engine. + //! + //! This provides a list of plugin libraries on the filesystem needed to implement operations + //! in the parsed network. If you are building a version-compatible engine using this network, + //! provide this list to IBuilderConfig::setPluginsToSerialize to serialize these plugins along + //! with the version-compatible engine, or, if you want to ship these plugin libraries externally + //! to the engine, ensure that IPluginRegistry::loadLibrary is used to load these libraries in the + //! appropriate runtime before deserializing the corresponding engine. + //! + //! \param[out] nbPluginLibs Returns the number of plugin libraries in the array, or -1 if there was an error. + //! \return Array of `nbPluginLibs` C-strings describing plugin library paths on the filesystem if nbPluginLibs > 0, + //! or nullptr otherwise. This array is owned by the IParser, and the pointers in the array are only valid until + //! the next call to parse(), supportsModel(), parseFromFile(), or parseWithWeightDescriptors(). + //! + virtual char const* const* getUsedVCPluginLibraries(int64_t& nbPluginLibs) const noexcept = 0; + + //! + //! \brief Set the parser flags. + //! + //! The flags are listed in the OnnxParserFlag enum. + //! + //! \param OnnxParserFlags The flags used when parsing an ONNX model. + //! + //! \note This function will override the previous set flags, rather than bitwise ORing the new flag. + //! + //! \see getFlags() + //! + virtual void setFlags(OnnxParserFlags onnxParserFlags) noexcept = 0; + + //! + //! \brief Get the parser flags. Defaults to 0. + //! + //! \return The parser flags as a bitmask. + //! + //! \see setFlags() + //! + virtual OnnxParserFlags getFlags() const noexcept = 0; + + //! + //! \brief clear a parser flag. + //! + //! clears the parser flag from the enabled flags. + //! + //! \see setFlags() + //! + virtual void clearFlag(OnnxParserFlag onnxParserFlag) noexcept = 0; + + //! + //! \brief Set a single parser flag. + //! + //! Add the input parser flag to the already enabled flags. + //! + //! \see setFlags() + //! + virtual void setFlag(OnnxParserFlag onnxParserFlag) noexcept = 0; + + //! + //! \brief Returns true if the parser flag is set + //! + //! \see getFlags() + //! + //! \return True if flag is set, false if unset. + //! + virtual bool getFlag(OnnxParserFlag onnxParserFlag) const noexcept = 0; + + //! + //!\brief Return the i-th output ITensor object for the ONNX layer "name". + //! + //! Return the i-th output ITensor object for the ONNX layer "name". + //! If "name" is not found or i is out of range, return nullptr. + //! In the case of multiple nodes sharing the same name this function will return + //! the output tensors of the first instance of the node in the ONNX graph. + //! + //! \param name The name of the ONNX layer. + //! + //! \param i The index of the output. i must be in range [0, layer.num_outputs). + //! + virtual nvinfer1::ITensor const* getLayerOutputTensor(char const* name, int64_t i) noexcept = 0; + + //! + //! \brief Check whether TensorRT supports a particular ONNX model. + //! If the function returns True, one can proceed to engine building + //! without having to call \p parse or \p parseFromFile. + //! Results can be queried through \p getNbSubgraphs, \p isSubgraphSupported, + //! \p getSubgraphNodes. + //! + //! \param serializedOnnxModel Pointer to the serialized ONNX model. Can be freed after this function returns. + //! \param serializedOnnxModelSize Size of the serialized ONNX model in bytes + //! \param modelPath Absolute path to the model file for loading external weights if required + //! \return true if the model is supported + //! + virtual bool supportsModelV2( + void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0; + + //! + //! \brief Get the number of subgraphs. Calling this function before calling \p supportsModelV2 results in undefined + //! behavior. + //! + //! + //! \return Number of subgraphs. + //! + virtual int64_t getNbSubgraphs() noexcept = 0; + + //! + //! \brief Returns whether the subgraph is supported. Calling this function before calling \p supportsModelV2 + //! results in undefined behavior. + //! + //! + //! \param index Index of the subgraph. + //! \return Whether the subgraph is supported. + //! + virtual bool isSubgraphSupported(int64_t const index) noexcept = 0; + + //! + //! \brief Get the nodes of the specified subgraph. Calling this function before calling \p supportsModelV2 results + //! in undefined behavior. + //! + //! + //! \param index Index of the subgraph. + //! \param subgraphLength Returns the length of the subgraph as reference. + //! + //! \return Pointer to the subgraph nodes array. This pointer is owned by the Parser. + //! + virtual int64_t* getSubgraphNodes(int64_t const index, int64_t& subgraphLength) noexcept = 0; + + //! + //! \brief Load a serialized ONNX model into the parser. Unlike the parse(), parseFromFile(), or + //! parseWithWeightDescriptors() functions, this function does not immediately convert the model into a TensorRT + //! INetworkDefinition. Using this function allows users to provide their own initializers for the ONNX model + //! through the loadInitializer() function. + //! + //! Only one model can be loaded at a time. Subsequent calls to loadModelProto() will result in an error. + //! + //! To begin the conversion of the model into a TensorRT INetworkDefinition, use parseModelProto(). + //! + //! \param serializedOnnxModel Pointer to the serialized ONNX model. Can be freed after this function returns. + //! \param serializedOnnxModelSize Size of the serialized ONNX model in bytes. + //! \param modelPath Absolute path to the model file for loading external weights if required. + //! \return true if the model was loaded successfully + //! \see getNbErrors() getError() + //! + virtual bool loadModelProto( + void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0; + + //! + //! \brief Prompt the ONNX parser to load an initializer with user-provided binary data. + //! The lifetime of the data must exceed the lifetime of the parser. + //! + //! All user-provided initializers must be provided prior to calling refitModelProto(). + //! + //! This function can be called multiple times to specify the names of multiple initializers. + //! + //! Calling this function with an initializer previously specified will overwrite the previous instance. + //! + //! + //! This function will return false if initializer validation fails. Possible validation errors are: + //! * This function was called prior to loadModelProto(). + //! * The requested initializer was not found in the model. + //! * The size of the data provided is different from the corresponding initializer in the model. + //! + //! \param name Name of the initializer. + //! \param data Binary data containing the values of the initializer. + //! \param size Size of the initializer in bytes. + //! \return true if the initializer was loaded successfully + //! \see loadModelProto() + //! + virtual bool loadInitializer(char const* name, void const* data, size_t size) noexcept = 0; + + //! \brief Begin the parsing and conversion process of the loaded ONNX model into a TensorRT INetworkDefinition. + //! + //! \return true if conversion was successful + //! \see getNbErrors() getError() loadModelProto() loadModelProtoFromFile() + //! + virtual bool parseModelProto() noexcept = 0; +}; + +//! +//! \class IParserRefitter +//! +//! \brief An interface designed to refit weights from an ONNX model. +//! +//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI. +//! +class IParserRefitter +{ +public: + //! + //! \brief Load a serialized ONNX model from memory and perform weight refit. + //! + //! \param serializedOnnxModel Pointer to the serialized ONNX model + //! \param serializedOnnxModelSize Size of the serialized ONNX model + //! in bytes + //! \param modelPath Absolute path to the model file for loading external weights if required + //! \return true if all the weights in the engine were refit successfully. + //! + //! The serialized ONNX model must be identical to the one used to generate the engine + //! that will be refit. + //! + virtual bool refitFromBytes( + void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept + = 0; + + //! + //! \brief Load and parse a ONNX model from disk and perform weight refit. + //! + //! \param onnxModelFile Path to the ONNX model to load from disk. + //! + //! \return true if the model was loaded successfully, and if all the weights in the engine were refit successfully. + //! + //! The provided ONNX model must be identical to the one used to generate the engine + //! that will be refit. + //! + virtual bool refitFromFile(char const* onnxModelFile) noexcept = 0; + + //! + //!\brief Get the number of errors that occurred during prior calls to \p refitFromBytes or \p refitFromFile + //! + //! \see getError() IParserError + //! + virtual int32_t getNbErrors() const noexcept = 0; + + //! + //!\brief Get an error that occurred during prior calls to \p refitFromBytes or \p refitFromFile + //! + //! \see getNbErrors() IParserError + //! + virtual IParserError const* getError(int32_t index) const noexcept = 0; + + //! + //!\brief Clear errors from prior calls to \p refitFromBytes or \p refitFromFile + //! + //! \see getNbErrors() getError() IParserError + //! + virtual void clearErrors() = 0; + + virtual ~IParserRefitter() noexcept = default; + + //! + //! \brief Load a serialized ONNX model into the parser. Unlike the refit(), or refitFromFile() + //! functions, this function does not immediately begin the refit process. Using this function + //! allows users to provide their own initializers for the ONNX model through the loadInitializer() function. + //! + //! Only one model can be loaded at a time. Subsequent calls to loadModelProto() will result in an error. + //! + //! To begin the refit process, use refitModelProto(). + //! + //! \param serializedOnnxModel Pointer to the serialized ONNX model. Can be freed after this function returns. + //! \param serializedOnnxModelSize Size of the serialized ONNX model in bytes. + //! \param modelPath Absolute path to the model file for loading external weights if required. + //! \return true if the model was loaded successfully + //! \see getNbErrors() getError() + //! + virtual bool loadModelProto( + void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0; + + //! + //! \brief Prompt the ONNX refitter to load an initializer with user-provided binary data. + //! The lifetime of the data must exceed the lifetime of the refitter. + //! + //! All user-provided initializers must be provided prior to calling refitModelProto(). + //! + //! This function can be called multiple times to specify the names of multiple initializers. + //! + //! Calling this function with an initializer previously specified will overwrite the previous instance. + //! + //! This function will return false if initializer validation fails. Possible validation errors are: + //! * This function was called prior to loadModelProto() + //! * The requested initializer was not found in the model. + //! * The size of the data provided is different from the corresponding initializer in the model. + //! + //! \param name Name of the initializer. + //! \param data Binary data containing the values of the initializer. + //! \param size Size of the initializer in bytes. + //! \return true if the initializer was loaded successfully + //! \see loadModelProto() + //! + virtual bool loadInitializer(char const* name, void const* data, size_t size) noexcept = 0; + + //! \brief Begin the refit process from the loaded ONNX model. + //! + //! \return true if refit was successful + //! \see getNbErrors() getError() loadModelProto() + //! + virtual bool refitModelProto() noexcept = 0; +}; + +} // namespace nvonnxparser + +extern "C" TENSORRTAPI void* createNvOnnxParser_INTERNAL(void* network, void* logger, int version) noexcept; +extern "C" TENSORRTAPI void* createNvOnnxParserRefitter_INTERNAL( + void* refitter, void* logger, int32_t version) noexcept; +extern "C" TENSORRTAPI int getNvOnnxParserVersion() noexcept; + +namespace nvonnxparser +{ + +namespace +{ + +//! +//! \brief Create a new parser object +//! +//! \param network The network definition that the parser will write to +//! \param logger The logger to use +//! \return a new parser object or NULL if an error occurred +//! +//! Any input dimensions that are constant should not be changed after parsing, +//! because correctness of the translation may rely on those constants. +//! Changing a dynamic input dimension, i.e. one that translates to -1 in +//! TensorRT, to a constant is okay if the constant is consistent with the model. +//! Each instance of the parser is designed to only parse one ONNX model once. +//! +//! \see IParser +//! +inline IParser* createParser(nvinfer1::INetworkDefinition& network, nvinfer1::ILogger& logger) noexcept +{ + return static_cast(createNvOnnxParser_INTERNAL(&network, &logger, NV_ONNX_PARSER_VERSION)); +} + +//! +//! \brief Create a new ONNX refitter object +//! +//! \param refitter The Refitter object used to refit the model +//! \param logger The logger to use +//! \return a new ParserRefitter object or NULL if an error occurred +//! +//! \see IParserRefitter +//! +inline IParserRefitter* createParserRefitter(nvinfer1::IRefitter& refitter, nvinfer1::ILogger& logger) noexcept +{ + return static_cast( + createNvOnnxParserRefitter_INTERNAL(&refitter, &logger, NV_ONNX_PARSER_VERSION)); +} + +} // namespace + +} // namespace nvonnxparser + +#endif // NV_ONNX_PARSER_H diff --git a/python/include/impl/NvInferPythonPlugin.h b/include/impl/NvInferPythonPlugin.h similarity index 99% rename from python/include/impl/NvInferPythonPlugin.h rename to include/impl/NvInferPythonPlugin.h index d703ba52..fd90943f 100644 --- a/python/include/impl/NvInferPythonPlugin.h +++ b/include/impl/NvInferPythonPlugin.h @@ -115,7 +115,6 @@ protected: virtual ~ISymExprs() noexcept = default; }; - //! \enum QuickPluginCreationRequest //! \brief Communicates preference when a quickly deployable plugin is to be added to the network enum class QuickPluginCreationRequest : int32_t diff --git a/parsers/onnx b/parsers/onnx index 9a9f7883..c7272770 160000 --- a/parsers/onnx +++ b/parsers/onnx @@ -1 +1 @@ -Subproject commit 9a9f7883dd7b8cb0a718395bac2075fab6f97da8 +Subproject commit c72727708d4ed8a379c04c4ea235c157ce9c3870 diff --git a/plugin/CMakeLists.txt b/plugin/CMakeLists.txt index 290b577b..dcd2c255 100644 --- a/plugin/CMakeLists.txt +++ b/plugin/CMakeLists.txt @@ -73,6 +73,7 @@ set(TRT_PLUGIN_NAMES ) + if(${TRT_BUILD_INCLUDE_BERT_QKV_PLUGIN}) list(APPEND TRT_PLUGIN_NAMES bertQKVToContextPlugin @@ -117,20 +118,23 @@ endforeach() target_compile_options(trt_plugins PUBLIC $<$:--expt-relaxed-constexpr>) target_compile_options(trt_vc_plugins PUBLIC $<$:--expt-relaxed-constexpr>) +target_compile_definitions(trt_vc_plugins PRIVATE + COMPILE_VFC_PLUGIN=1 +) + # Create all the library targets, reusing the objects we've compiled in the first step. -add_library(tensorrt_plugins SHARED $) -add_library(tensorrt_plugins_internal SHARED $) -add_library(tensorrt_plugins_static STATIC $) -add_library(tensorrt_vc_plugins SHARED $) -add_library(tensorrt_vc_plugins_static STATIC $) +add_library(tensorrt_plugins SHARED) +add_library(tensorrt_plugins_internal SHARED) +add_library(tensorrt_plugins_static STATIC) +foreach(lib tensorrt_plugins tensorrt_plugins_internal tensorrt_plugins_static) + target_link_libraries(${lib} PRIVATE trt_plugins) +endforeach() -target_compile_definitions(tensorrt_vc_plugins PRIVATE - COMPILE_VFC_PLUGIN=1 -) - -target_compile_definitions(tensorrt_vc_plugins_static PRIVATE - COMPILE_VFC_PLUGIN=1 -) +add_library(tensorrt_vc_plugins SHARED) +add_library(tensorrt_vc_plugins_static STATIC) +foreach(lib tensorrt_vc_plugins tensorrt_vc_plugins_static) + target_link_libraries(${lib} PRIVATE trt_vc_plugins) +endforeach() if (NOT MSVC) set(trt_plugins_link_options @@ -168,8 +172,10 @@ set_target_properties( OUTPUT_NAME nvinfer_plugin VERSION ${TensorRT_VERSION} SOVERSION ${TRT_MAJOR} - LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports.map) + LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports.map +) +update_windows_output_name(tensorrt_plugins ${TRT_MAJOR} ${TRT_MINOR}) if (NOT MSVC) set(trt_plugins_internal_link_options @@ -198,8 +204,10 @@ set_target_properties( OUTPUT_NAME nvinfer_plugin_internal VERSION ${TensorRT_VERSION} SOVERSION ${TRT_MAJOR} - LINK_DEPENDS ${TensorRT_SOURCE_DIR}/Exports-plugin_internal.map) + LINK_DEPENDS ${TensorRT_SOURCE_DIR}/Exports-plugin_internal.map +) +update_windows_output_name(tensorrt_plugins_internal ${TRT_MAJOR} ${TRT_MINOR}) ### Static Plugin Setup set(trt_plugin_static_dependencies @@ -225,6 +233,8 @@ if(NOT ${TRT_BUILD_ENABLE_STATIC_LIBS}) ) endif() +smoke_test_static_lib(tensorrt_plugins_static) + ### VC Plugin Setup if (NOT MSVC) set(trt_vc_plugins_link_options @@ -242,7 +252,7 @@ endif() # Target properties for tensorrt_vc_plugins # This library includes a minimal subset of the plugins used for version compatibility. target_include_directories(tensorrt_vc_plugins PRIVATE ${trt_plugin_include_dirs}) -target_link_libraries(tensorrt_vc_plugins PRIVATE ${trt_plugin_dependencies}) +target_link_libraries(tensorrt_vc_plugins PRIVATE trt_global_definitions) target_link_options(tensorrt_vc_plugins PRIVATE ${trt_vc_plugins_link_options}) set_target_properties( @@ -252,7 +262,10 @@ set_target_properties( OUTPUT_NAME nvinfer_vc_plugin VERSION ${TensorRT_VERSION} SOVERSION ${TRT_MAJOR} - LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports-vfc_plugin.map) + LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports-vfc_plugin.map +) + +update_windows_output_name(tensorrt_vc_plugins ${TRT_MAJOR} ${TRT_MINOR}) ### VC Plugin Static Setup target_include_directories(tensorrt_vc_plugins_static PRIVATE ${trt_plugin_include_dirs}) @@ -273,16 +286,23 @@ if(NOT ${TRT_BUILD_ENABLE_STATIC_LIBS}) ) endif() +smoke_test_static_lib(tensorrt_vc_plugins_static) + +if(${TRT_BUILD_STUB_LIBS} AND NOT MSVC) + create_stub_lib(tensorrt_plugins) + create_stub_lib(tensorrt_vc_plugins) +endif() + installLibraries( TARGETS tensorrt_plugins tensorrt_plugins_static tensorrt_vc_plugins tensorrt_vc_plugins_static OPTIONAL - COMPONENT release + COMPONENT external ) installLibraries( TARGETS tensorrt_plugins_internal OPTIONAL - COMPONENT full + COMPONENT internal ) else() # TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW @@ -387,6 +407,7 @@ endif() include_directories(common common/kernels ${CMAKE_SOURCE_DIR}/third_party) + foreach(PLUGIN_ITER ${PLUGIN_LISTS}) include_directories(${PLUGIN_ITER}) add_subdirectory(${PLUGIN_ITER}) diff --git a/plugin/batchedNMSPlugin/gatherNMSOutputs.cu b/plugin/batchedNMSPlugin/gatherNMSOutputs.cu index a30b8f7a..45dfe86e 100644 --- a/plugin/batchedNMSPlugin/gatherNMSOutputs.cu +++ b/plugin/batchedNMSPlugin/gatherNMSOutputs.cu @@ -174,7 +174,7 @@ struct nmsOutLaunchConfig , function(function) { } - bool operator==(const nmsOutLaunchConfig& other) + bool operator==(nmsOutLaunchConfig const& other) const { return t_bbox == other.t_bbox && t_score == other.t_score; } diff --git a/plugin/bertQKVToContextPlugin/CMakeLists.txt b/plugin/bertQKVToContextPlugin/CMakeLists.txt index d9a6e6eb..1c84a336 100644 --- a/plugin/bertQKVToContextPlugin/CMakeLists.txt +++ b/plugin/bertQKVToContextPlugin/CMakeLists.txt @@ -49,3 +49,4 @@ if (NOT DEFINED ${TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW}) set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE) set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE) endif() + diff --git a/plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginConfig.yaml b/plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginConfig.yaml index e14df037..2aa8ae8b 100644 --- a/plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginConfig.yaml +++ b/plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -69,7 +69,7 @@ versions: - hidden_size - num_heads - has_mask - golden_io_path: "plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginGoldenIO.json" + golden_io_path: "plugin/CustomQKVToContextPluginDynamic_PluginGoldenIO.json" abs_tol: 1e-5 rel_tol: 1e-5 fp16_atol: 1e-2 diff --git a/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp b/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp index 789f9856..a69983fe 100644 --- a/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp +++ b/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp @@ -54,7 +54,7 @@ QKVToContextInterleavedPlugin::QKVToContextInterleavedPlugin(std::string const& , mQkvScale(qkvScale) , mCtxScale(ctxScale) { - mSM = getSMVersion(); + mSM = getSmVersion(); mUseInt8ScaleMax = static_cast(useInt8ScaleMax); mUseExplicitInt8 = static_cast(useExplicitInt8); // variable sequence length is only supported with the fused MHA kernels diff --git a/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPluginLegacy.cpp b/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPluginLegacy.cpp index faa70850..64df35a0 100644 --- a/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPluginLegacy.cpp +++ b/plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPluginLegacy.cpp @@ -57,7 +57,7 @@ QKVToContextInterleavedPluginLegacy::QKVToContextInterleavedPluginLegacy(std::st , mQkvScale(qkvScale) , mCtxScale(ctxScale) { - mSM = getSMVersion(); + mSM = getSmVersion(); // variable sequence length is only supported with the fused MHA kernels // we should not override mS! bool isSMSupported = elem(mSM, diff --git a/plugin/bertQKVToContextPlugin/qkvToContextPlugin.cpp b/plugin/bertQKVToContextPlugin/qkvToContextPlugin.cpp index 9d6746d6..ea98d1dd 100644 --- a/plugin/bertQKVToContextPlugin/qkvToContextPlugin.cpp +++ b/plugin/bertQKVToContextPlugin/qkvToContextPlugin.cpp @@ -66,7 +66,7 @@ QKVToContextPluginDynamic::QKVToContextPluginDynamic(const std::string name, con { mHasImask = static_cast(hasImask); - mSM = getSMVersion(); + mSM = getSmVersion(); } QKVToContextPluginDynamic::QKVToContextPluginDynamic(const std::string name, const DataType type, const int32_t S, @@ -772,7 +772,7 @@ QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name, , mDqProbs(dqProbs) , mHdim(HDIM) { - mSM = getSMVersion(); + mSM = getSmVersion(); mUseVarSeqlen = static_cast(varSeqlen); mUseInt8ScaleMax = static_cast(useInt8ScaleMax); mHasImask = static_cast(hasImask); @@ -802,7 +802,7 @@ QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name, , mDqProbs(dqProbs) , mHdim(HDIM) { - mSM = getSMVersion(); + mSM = getSmVersion(); mUseVarSeqlen = static_cast(varSeqlen); mUseInt8ScaleMax = static_cast(useInt8ScaleMax); mHasImask = static_cast(hasImask); @@ -825,7 +825,6 @@ QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name, mDispatcher->deserialize(runnerStateBuffer, length); } - IPluginCapability* QKVToContextVarSeqlenPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept { try diff --git a/plugin/bertQKVToContextPlugin/qkvToContextPluginLegacy.cpp b/plugin/bertQKVToContextPlugin/qkvToContextPluginLegacy.cpp index e01d4fcf..9fe6a27c 100644 --- a/plugin/bertQKVToContextPlugin/qkvToContextPluginLegacy.cpp +++ b/plugin/bertQKVToContextPlugin/qkvToContextPluginLegacy.cpp @@ -65,7 +65,7 @@ QKVToContextPluginDynamicLegacy::QKVToContextPluginDynamicLegacy(std::string con , mDqProbs(dqProbs) { - mSM = getSMVersion(); + mSM = getSmVersion(); } QKVToContextPluginDynamicLegacy::QKVToContextPluginDynamicLegacy( @@ -603,7 +603,7 @@ QKVToContextVarSeqlenPluginLegacy::QKVToContextVarSeqlenPluginLegacy(std::string , mUseVarSeqlen(varSeqlen) , mUseInt8ScaleMax(useInt8ScaleMax) { - mSM = getSMVersion(); + mSM = getSmVersion(); if (varSeqlen) { diff --git a/plugin/clipPlugin/CMakeLists.txt b/plugin/clipPlugin/CMakeLists.txt index 20ad1f81..297cd6c9 100644 --- a/plugin/clipPlugin/CMakeLists.txt +++ b/plugin/clipPlugin/CMakeLists.txt @@ -21,3 +21,4 @@ add_plugin_source( clipPlugin.cpp clipPlugin.h ) + diff --git a/plugin/common/CMakeLists.txt b/plugin/common/CMakeLists.txt index 60a961dc..3fd70841 100644 --- a/plugin/common/CMakeLists.txt +++ b/plugin/common/CMakeLists.txt @@ -14,7 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # -add_plugin_source( +# Base source files +set(PLUGIN_COMMON_SOURCES bboxUtils.h bertCommon.h checkMacrosPlugin.cpp @@ -37,10 +38,14 @@ add_plugin_source( plugin.cpp plugin.h reducedMathPlugin.cpp + scopedCudaStream.h serialize.hpp templates.h ) + +add_plugin_source(${PLUGIN_COMMON_SOURCES}) + add_subdirectory(kernels) # Promote PLUGIN_SOURCES and PLUGIN_CU_SOURCES added by `add_subdirectory` to this file's parent. diff --git a/plugin/common/bertCommon.h b/plugin/common/bertCommon.h index 692a0a14..06df9fbc 100644 --- a/plugin/common/bertCommon.h +++ b/plugin/common/bertCommon.h @@ -139,7 +139,7 @@ inline int32_t getMHAMaskPackedSize(int32_t smVersion, nvinfer1::DataType dataTy return packedSize; } -inline uint32_t getElementSize(nvinfer1::DataType t) noexcept +inline uint32_t getElementSize(nvinfer1::DataType t) { switch (t) { diff --git a/plugin/common/cudaDriverWrapper.cpp b/plugin/common/cudaDriverWrapper.cpp index e1267173..f81635a0 100644 --- a/plugin/common/cudaDriverWrapper.cpp +++ b/plugin/common/cudaDriverWrapper.cpp @@ -50,18 +50,45 @@ CUDADriverWrapper::CUDADriverWrapper() return ret; }; - *(void**) (&_cuGetErrorName) = load_sym(handle, "cuGetErrorName"); - *(void**) (&_cuFuncSetAttribute) = load_sym(handle, "cuFuncSetAttribute"); - *(void**) (&_cuLinkComplete) = load_sym(handle, "cuLinkComplete"); - *(void**) (&_cuModuleUnload) = load_sym(handle, "cuModuleUnload"); - *(void**) (&_cuLinkDestroy) = load_sym(handle, "cuLinkDestroy"); - *(void**) (&_cuModuleLoadData) = load_sym(handle, "cuModuleLoadData"); - *(void**) (&_cuLinkCreate) = load_sym(handle, "cuLinkCreate_v2"); - *(void**) (&_cuModuleGetFunction) = load_sym(handle, "cuModuleGetFunction"); - *(void**) (&_cuLinkAddFile) = load_sym(handle, "cuLinkAddFile_v2"); - *(void**) (&_cuLinkAddData) = load_sym(handle, "cuLinkAddData_v2"); - *(void**) (&_cuLaunchCooperativeKernel) = load_sym(handle, "cuLaunchCooperativeKernel"); - *(void**) (&_cuLaunchKernel) = load_sym(handle, "cuLaunchKernel"); + _cuGetErrorName = reinterpret_cast(load_sym(handle, "cuGetErrorName")); + _cuGetErrorString = reinterpret_cast(load_sym(handle, "cuGetErrorString")); + _cuFuncSetAttribute = reinterpret_cast( + load_sym(handle, "cuFuncSetAttribute")); + _cuLinkComplete = reinterpret_cast(load_sym(handle, "cuLinkComplete")); + _cuModuleUnload = reinterpret_cast(load_sym(handle, "cuModuleUnload")); + _cuLinkDestroy = reinterpret_cast(load_sym(handle, "cuLinkDestroy")); + _cuModuleLoadData = reinterpret_cast(load_sym(handle, "cuModuleLoadData")); + _cuLinkCreate = reinterpret_cast( + load_sym(handle, "cuLinkCreate_v2")); + _cuModuleGetFunction + = reinterpret_cast(load_sym(handle, "cuModuleGetFunction")); + _cuLinkAddFile + = reinterpret_cast( + load_sym(handle, "cuLinkAddFile_v2")); + _cuLinkAddData = reinterpret_cast(load_sym(handle, "cuLinkAddData_v2")); + _cuLaunchCooperativeKernel = reinterpret_cast(load_sym(handle, "cuLaunchCooperativeKernel")); + _cuLaunchKernel = reinterpret_cast(load_sym(handle, "cuLaunchKernel")); +#if CUDA_VERSION >= 11060 + _cuLaunchKernelEx + = reinterpret_cast( + dllGetSym(handle, "cuLaunchKernelEx")); +#endif +#if CUDA_VERSION >= 12000 + _cuTensorMapEncodeTiled + = reinterpret_cast(load_sym(handle, "cuTensorMapEncodeTiled")); +#endif + _cuMemcpyDtoH = reinterpret_cast(load_sym(handle, "cuMemcpyDtoH_v2")); + _cuDeviceGetAttribute = reinterpret_cast( + load_sym(handle, "cuDeviceGetAttribute")); +#if CUDA_VERSION >= 12000 + _cuOccupancyMaxActiveClusters = reinterpret_cast( + load_sym(handle, "cuOccupancyMaxActiveClusters")); +#endif } CUDADriverWrapper::~CUDADriverWrapper() @@ -74,6 +101,11 @@ CUresult CUDADriverWrapper::cuGetErrorName(CUresult error, char const** pStr) co return (*_cuGetErrorName)(error, pStr); } +CUresult CUDADriverWrapper::cuGetErrorString(CUresult error, char const** pStr) const +{ + return (*_cuGetErrorString)(error, pStr); +} + CUresult CUDADriverWrapper::cuFuncSetAttribute(CUfunction hfunc, CUfunction_attribute attrib, int32_t value) const { return (*_cuFuncSetAttribute)(hfunc, attrib, value); @@ -137,3 +169,43 @@ CUresult CUDADriverWrapper::cuLaunchKernel(CUfunction f, uint32_t gridDimX, uint return (*_cuLaunchKernel)( f, gridDimX, gridDimY, gridDimZ, blockDimX, blockDimY, blockDimZ, sharedMemBytes, hStream, kernelParams, extra); } + +#if CUDA_VERSION >= 11060 +CUresult CUDADriverWrapper::cuLaunchKernelEx( + CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra) const +{ + PLUGIN_ASSERT(_cuLaunchKernelEx != nullptr); + return (*_cuLaunchKernelEx)(config, f, kernelParams, extra); +} +#endif + +#if CUDA_VERSION >= 12000 +CUresult CUDADriverWrapper::cuTensorMapEncodeTiled(CUtensorMap* tensorMap, CUtensorMapDataType tensorDataType, + cuuint32_t tensorRank, void const* globalAddress, cuuint64_t const* globalDim, cuuint64_t const* globalStrides, + cuuint32_t const* boxDim, cuuint32_t const* elementStrides, CUtensorMapInterleave interleave, + CUtensorMapSwizzle swizzle, CUtensorMapL2promotion l2Promotion, CUtensorMapFloatOOBfill oobFill) const +{ + PLUGIN_ASSERT(_cuTensorMapEncodeTiled != nullptr); + return (*_cuTensorMapEncodeTiled)(tensorMap, tensorDataType, tensorRank, globalAddress, globalDim, globalStrides, + boxDim, elementStrides, interleave, swizzle, l2Promotion, oobFill); +} +#endif + +CUresult CUDADriverWrapper::cuMemcpyDtoH(void* dstHost, CUdeviceptr srcDevice, size_t ByteCount) const +{ + return (*_cuMemcpyDtoH)(dstHost, srcDevice, ByteCount); +} + +CUresult CUDADriverWrapper::cuDeviceGetAttribute(int32_t* pi, CUdevice_attribute attrib, CUdevice dev) const +{ + return (*_cuDeviceGetAttribute)(pi, attrib, dev); +} + +#if CUDA_VERSION >= 12000 +CUresult CUDADriverWrapper::cuOccupancyMaxActiveClusters( + int32_t* maxActiveClusters, CUfunction f, CUlaunchConfig const* config) const +{ + PLUGIN_ASSERT(_cuOccupancyMaxActiveClusters != nullptr); + return (*_cuOccupancyMaxActiveClusters)(maxActiveClusters, f, config); +} +#endif diff --git a/plugin/common/cudaDriverWrapper.h b/plugin/common/cudaDriverWrapper.h index 209ed3f8..387c8379 100644 --- a/plugin/common/cudaDriverWrapper.h +++ b/plugin/common/cudaDriverWrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,6 +21,8 @@ #include #include #include +#include +#include #define cuErrCheck(stat, wrap) \ { \ @@ -42,6 +44,8 @@ public: CUresult cuGetErrorName(CUresult error, char const** pStr) const; + CUresult cuGetErrorString(CUresult error, char const** pStr) const; + CUresult cuFuncSetAttribute(CUfunction hfunc, CUfunction_attribute attrib, int32_t value) const; CUresult cuLinkComplete(CUlinkState state, void** cubinOut, size_t* sizeOut) const; @@ -70,9 +74,29 @@ public: uint32_t blockDimY, uint32_t blockDimZ, uint32_t sharedMemBytes, CUstream hStream, void** kernelParams, void** extra) const; +#if CUDA_VERSION >= 11060 + CUresult cuLaunchKernelEx(CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra) const; +#endif + +#if CUDA_VERSION >= 12000 + CUresult cuTensorMapEncodeTiled(CUtensorMap* tensorMap, CUtensorMapDataType tensorDataType, cuuint32_t tensorRank, + void const* globalAddress, cuuint64_t const* globalDim, cuuint64_t const* globalStrides, + cuuint32_t const* boxDim, cuuint32_t const* elementStrides, CUtensorMapInterleave interleave, + CUtensorMapSwizzle swizzle, CUtensorMapL2promotion l2Promotion, CUtensorMapFloatOOBfill oobFill) const; +#endif + + CUresult cuMemcpyDtoH(void* dstHost, CUdeviceptr srcDevice, size_t ByteCount) const; + + CUresult cuDeviceGetAttribute(int32_t* pi, CUdevice_attribute attrib, CUdevice dev) const; + +#if CUDA_VERSION >= 12000 + CUresult cuOccupancyMaxActiveClusters(int32_t* maxActiveClusters, CUfunction f, CUlaunchConfig const* config) const; +#endif + private: void* handle; CUresult (*_cuGetErrorName)(CUresult, char const**); + CUresult (*_cuGetErrorString)(CUresult, char const**); CUresult (*_cuFuncSetAttribute)(CUfunction, CUfunction_attribute, int32_t); CUresult (*_cuLinkComplete)(CUlinkState, void**, size_t*); CUresult (*_cuModuleUnload)(CUmodule); @@ -88,6 +112,20 @@ private: CUresult (*_cuLaunchKernel)(CUfunction f, uint32_t gridDimX, uint32_t gridDimY, uint32_t gridDimZ, uint32_t blockDimX, uint32_t blockDimY, uint32_t blockDimZ, uint32_t sharedMemBytes, CUstream hStream, void** kernelParams, void** extra); +#if CUDA_VERSION >= 11060 + CUresult (*_cuLaunchKernelEx)(CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra); +#endif +#if CUDA_VERSION >= 12000 + CUresult (*_cuTensorMapEncodeTiled)(CUtensorMap* tensorMap, CUtensorMapDataType tensorDataType, + cuuint32_t tensorRank, void const* globalAddress, cuuint64_t const* globalDim, cuuint64_t const* globalStrides, + cuuint32_t const* boxDim, cuuint32_t const* elementStrides, CUtensorMapInterleave interleave, + CUtensorMapSwizzle swizzle, CUtensorMapL2promotion l2Promotion, CUtensorMapFloatOOBfill oobFill); +#endif + CUresult (*_cuMemcpyDtoH)(void* dstHost, CUdeviceptr srcDevice, size_t ByteCount); + CUresult (*_cuDeviceGetAttribute)(int32_t*, CUdevice_attribute attrib, CUdevice dev); +#if CUDA_VERSION >= 12000 + CUresult (*_cuOccupancyMaxActiveClusters)(int32_t*, CUfunction f, CUlaunchConfig const* config); +#endif }; inline void cuErrCheck_(CUresult stat, CUDADriverWrapper const& wrap, char const* file, int32_t line) @@ -106,6 +144,108 @@ constexpr int32_t getCudaLibVersionMaj() noexcept return CUDA_VERSION / 1000U; } +// RAII wrapper for CUDA module +// Automatically manages CUDA module lifecycle - loading in constructor, unloading in destructor +// Provides safe module management and prevents memory leaks +class CudaModule +{ +public: + // Default constructor - creates an uninitialized module + CudaModule() noexcept + : mModule(nullptr) + , mDriverWrapper(nullptr) + , mLoadResult(CUDA_ERROR_NOT_INITIALIZED) + { + } + + // Constructor that loads a CUDA module from data + CudaModule(CUDADriverWrapper const& driverWrapper, void const* image) + : mModule(nullptr) + , mDriverWrapper(&driverWrapper) + , mLoadResult(CUDA_ERROR_NOT_INITIALIZED) + { + mLoadResult = mDriverWrapper->cuModuleLoadData(&mModule, image); + } + + // Destructor - automatically unloads the module + ~CudaModule() + { + if (mModule != nullptr && mDriverWrapper != nullptr) + { + mDriverWrapper->cuModuleUnload(mModule); + } + } + + // Delete copy constructor and assignment + CudaModule(CudaModule const&) = delete; + CudaModule& operator=(CudaModule const&) = delete; + + // Move constructor + CudaModule(CudaModule&& other) noexcept + : mModule(std::exchange(other.mModule, nullptr)) + , mDriverWrapper(std::exchange(other.mDriverWrapper, nullptr)) + , mLoadResult(std::exchange(other.mLoadResult, CUDA_ERROR_NOT_INITIALIZED)) + { + } + + // Move assignment + CudaModule& operator=(CudaModule&& other) noexcept + { + CudaModule tmp{std::move(other)}; + std::swap(mModule, tmp.mModule); + std::swap(mDriverWrapper, tmp.mDriverWrapper); + std::swap(mLoadResult, tmp.mLoadResult); + return *this; + } + + //! Get the underlying CUDA module handle (`CUmodule`; a raw pointer). + [[nodiscard]] CUmodule get() noexcept + { + return mModule; + } + + //! Get the underlying const CUDA module handle (`CUmodule`; a raw pointer). + //! \note Since `CUmodule` is a raw pointer, we remove the pointer from the + //! type, add `const`, and re-add the pointer. + [[nodiscard]] std::remove_pointer_t const* get() const noexcept + { + return mModule; + } + + // Implicit conversion to CUmodule + operator CUmodule() const noexcept + { + return mModule; + } + + // Check if the module is valid + bool isValid() const noexcept + { + return mModule != nullptr && mLoadResult == CUDA_SUCCESS; + } + + // Get function from module + CUresult getFunction(CUfunction* hfunc, char const* name) const + { + if (mModule != nullptr && mDriverWrapper != nullptr) + { + return mDriverWrapper->cuModuleGetFunction(hfunc, mModule, name); + } + return CUDA_ERROR_INVALID_HANDLE; + } + +private: + CUmodule mModule; + CUDADriverWrapper const* mDriverWrapper; + CUresult mLoadResult; +}; + +// Helper function to create a unique_ptr to CudaModule +inline std::unique_ptr makeCudaModule(CUDADriverWrapper const& driverWrapper, void const* image) +{ + return std::make_unique(driverWrapper, image); +} + } // namespace nvinfer1 #endif // CUDA_DRIVER_WRAPPER_H diff --git a/plugin/common/cudnnWrapper.cpp b/plugin/common/cudnnWrapper.cpp index 1e300e89..a16270f9 100644 --- a/plugin/common/cudnnWrapper.cpp +++ b/plugin/common/cudnnWrapper.cpp @@ -91,8 +91,8 @@ void* CudnnWrapper::tryLoadingCudnn(char const* callerPluginName) static constexpr int32_t kSM_BLACKWELL_100 = 100; std::string errorMsgCudnnSupport - = "At least one plugin (" + std::string(callerPluginName) + ") that requires cuDNN is being used. TensorRT does not provide cuDNN support for Blackwell (compute capability: 10.0) and later architectures. Detected compute capability: " + std::to_string(nvinfer1::plugin::getSMVersion() / 10) + "." + std::to_string(nvinfer1::plugin::getSMVersion() % 10) + ". Please run on a platform with compute capability < 10.0, or use an alternative to " + std::string(callerPluginName) + "."; - PLUGIN_VALIDATE(nvinfer1::plugin::getSMVersion() < kSM_BLACKWELL_100, errorMsgCudnnSupport.c_str()); + = "At least one plugin (" + std::string(callerPluginName) + ") that requires cuDNN is being used. TensorRT does not provide cuDNN support for Blackwell (compute capability: 10.0) and later architectures. Detected compute capability: " + std::to_string(nvinfer1::plugin::getSmVersion() / 10) + "." + std::to_string(nvinfer1::plugin::getSmVersion() % 10) + ". Please run on a platform with compute capability < 10.0, or use an alternative to " + std::string(callerPluginName) + "."; + PLUGIN_VALIDATE(nvinfer1::plugin::getSmVersion() < kSM_BLACKWELL_100, errorMsgCudnnSupport.c_str()); #endif // CUDART_VERSION >= 12070 && CUDNN_MAJOR == 8 void* cudnnLib = dllOpen(kCUDNN_PLUGIN_LIBNAME.c_str()); std::string errorMsg = "Failed to load " + kCUDNN_PLUGIN_LIBNAME + "."; diff --git a/plugin/common/kernels/allClassNMS.cu b/plugin/common/kernels/allClassNMS.cu index 35ab989b..b80e00e2 100644 --- a/plugin/common/kernels/allClassNMS.cu +++ b/plugin/common/kernels/allClassNMS.cu @@ -360,7 +360,7 @@ struct nmsLaunchConfigSSD , function(function) { } - bool operator==(const nmsLaunchConfigSSD& other) + bool operator==(nmsLaunchConfigSSD const& other) const { return t_score == other.t_score && t_bbox == other.t_bbox; } diff --git a/plugin/common/kernels/bboxDeltas2Proposals.cu b/plugin/common/kernels/bboxDeltas2Proposals.cu index 0be5e90d..1b4fb1e5 100644 --- a/plugin/common/kernels/bboxDeltas2Proposals.cu +++ b/plugin/common/kernels/bboxDeltas2Proposals.cu @@ -236,7 +236,7 @@ struct bd2pLaunchConfig { } - bool operator==(const bd2pLaunchConfig& other) + bool operator==(bd2pLaunchConfig const& other) const { return t_deltas == other.t_deltas && l_deltas == other.l_deltas && t_proposals == other.t_proposals && l_proposals == other.l_proposals && t_scores == other.t_scores && l_scores == other.l_scores; } diff --git a/plugin/common/kernels/decodeBBoxes.cu b/plugin/common/kernels/decodeBBoxes.cu index 1c71d2ee..844322ea 100644 --- a/plugin/common/kernels/decodeBBoxes.cu +++ b/plugin/common/kernels/decodeBBoxes.cu @@ -355,7 +355,7 @@ struct dbbLaunchConfig , function(function) { } - bool operator==(const dbbLaunchConfig& other) + bool operator==(dbbLaunchConfig const& other) const { return t_bbox == other.t_bbox; } diff --git a/plugin/common/kernels/gatherTopDetections.cu b/plugin/common/kernels/gatherTopDetections.cu index e5655a1d..64551c5f 100644 --- a/plugin/common/kernels/gatherTopDetections.cu +++ b/plugin/common/kernels/gatherTopDetections.cu @@ -173,7 +173,7 @@ struct gtdLaunchConfig , function(function) { } - bool operator==(const gtdLaunchConfig& other) + bool operator==(gtdLaunchConfig const& other) const { return t_bbox == other.t_bbox && t_score == other.t_score; } diff --git a/plugin/common/kernels/nmsLayer.cu b/plugin/common/kernels/nmsLayer.cu index 8ce2a8f2..722e65ae 100644 --- a/plugin/common/kernels/nmsLayer.cu +++ b/plugin/common/kernels/nmsLayer.cu @@ -390,7 +390,7 @@ struct nmsLaunchConfig { } - bool operator==(const nmsLaunchConfig& other) + bool operator==(nmsLaunchConfig const& other) const { return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores) && (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals) && (t_rois == other.t_rois); } diff --git a/plugin/common/kernels/permuteData.cu b/plugin/common/kernels/permuteData.cu index 185e4c53..e9a53e8d 100644 --- a/plugin/common/kernels/permuteData.cu +++ b/plugin/common/kernels/permuteData.cu @@ -80,7 +80,7 @@ struct pdLaunchConfig , function(function) { } - bool operator==(const pdLaunchConfig& other) + bool operator==(pdLaunchConfig const& other) const { return t_data == other.t_data; } diff --git a/plugin/common/kernels/proposalKernel.cu b/plugin/common/kernels/proposalKernel.cu index 82f2db9b..52d97ea5 100644 --- a/plugin/common/kernels/proposalKernel.cu +++ b/plugin/common/kernels/proposalKernel.cu @@ -420,7 +420,7 @@ struct nmsLaunchConfig { } - bool operator==(const nmsLaunchConfig& other) + bool operator==(nmsLaunchConfig const& other) const { return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores) && (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals) diff --git a/plugin/common/kernels/roiPooling.cu b/plugin/common/kernels/roiPooling.cu index 353173cc..4b2f1109 100644 --- a/plugin/common/kernels/roiPooling.cu +++ b/plugin/common/kernels/roiPooling.cu @@ -285,7 +285,7 @@ struct roiFwdLaunchConfig { } - bool operator==(const roiFwdLaunchConfig& other) + bool operator==(roiFwdLaunchConfig const& other) const { return (t_rois == other.t_rois) && (t_featureMap == other.t_featureMap) diff --git a/plugin/common/kernels/sortScoresPerClass.cu b/plugin/common/kernels/sortScoresPerClass.cu index cd62df64..b2448a8c 100644 --- a/plugin/common/kernels/sortScoresPerClass.cu +++ b/plugin/common/kernels/sortScoresPerClass.cu @@ -218,7 +218,7 @@ struct sspcLaunchConfig , function(function) { } - bool operator==(const sspcLaunchConfig& other) + bool operator==(sspcLaunchConfig const& other) const { return t_score == other.t_score; } diff --git a/plugin/common/kernels/sortScoresPerImage.cu b/plugin/common/kernels/sortScoresPerImage.cu index 99749c53..802c02ec 100644 --- a/plugin/common/kernels/sortScoresPerImage.cu +++ b/plugin/common/kernels/sortScoresPerImage.cu @@ -79,7 +79,7 @@ struct sspiLaunchConfig , function(function) { } - bool operator==(const sspiLaunchConfig& other) + bool operator==(sspiLaunchConfig const& other) const { return t_score == other.t_score; } diff --git a/plugin/common/plugin.h b/plugin/common/plugin.h index 385fec20..eaed1de8 100644 --- a/plugin/common/plugin.h +++ b/plugin/common/plugin.h @@ -107,7 +107,7 @@ OutType read(BufferType const*& buffer) return val; } -inline int32_t getTrtSMVersionDec(int32_t majorVersion, int32_t minorVersion) +inline int32_t getTrtSmVersionDec(int32_t majorVersion, int32_t minorVersion) { return majorVersion * 10 + minorVersion; } @@ -132,12 +132,12 @@ struct DeviceComputeCapability } }; -inline int32_t getSMVersion() +inline int32_t getSmVersion() { int32_t device{-1}; PLUGIN_CHECK_CUDA(cudaGetDevice(&device)); auto const cc = DeviceComputeCapability::forDevice(device); - return getTrtSMVersionDec(cc.major, cc.minor); + return getTrtSmVersionDec(cc.major, cc.minor); } // Check that all required field names are present in the PluginFieldCollection. diff --git a/plugin/common/scopedCudaStream.h b/plugin/common/scopedCudaStream.h new file mode 100644 index 00000000..5e680abb --- /dev/null +++ b/plugin/common/scopedCudaStream.h @@ -0,0 +1,120 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TRT_SCOPED_CUDA_STREAM_H +#define TRT_SCOPED_CUDA_STREAM_H + +#include +#include +#include +#include + +namespace nvinfer1 +{ +namespace pluginInternal +{ + +// RAII wrapper for CUDA stream +// Automatically manages CUDA stream lifecycle - creation in constructor, destruction in destructor +// Provides safe stream management and prevents memory leaks +class ScopedCudaStream +{ +public: + // Constructor that creates a new CUDA stream with default flags + ScopedCudaStream() + : mStream(nullptr) + { + cudaError_t result = cudaStreamCreate(&mStream); + if (result != cudaSuccess) + { + throw std::runtime_error("Failed to create CUDA stream: " + std::string(cudaGetErrorString(result))); + } + } + + // Constructor that creates a new CUDA stream with custom flags + explicit ScopedCudaStream(uint32_t const flags) + : mStream(nullptr) + { + cudaError_t result = cudaStreamCreateWithFlags(&mStream, flags); + if (result != cudaSuccess) + { + mStream = nullptr; + } + } + + // Destructor - automatically destroys the stream + ~ScopedCudaStream() + { + if (mStream != nullptr) + { + cudaStreamDestroy(mStream); + } + } + + // Delete copy constructor and assignment + ScopedCudaStream(ScopedCudaStream const&) = delete; + ScopedCudaStream& operator=(ScopedCudaStream const&) = delete; + + // Move constructor + ScopedCudaStream(ScopedCudaStream&& other) noexcept + : mStream(std::exchange(other.mStream, nullptr)) + { + } + + // Move assignment + ScopedCudaStream& operator=(ScopedCudaStream&& other) noexcept + { + ScopedCudaStream tmp{std::move(other)}; + std::swap(mStream, tmp.mStream); + return *this; + } + + // Get the underlying CUDA stream handle + cudaStream_t get() const noexcept + { + return mStream; + } + + // Implicit conversion to cudaStream_t + operator cudaStream_t() const noexcept + { + return mStream; + } + + // Check if the stream is valid + bool isValid() const noexcept + { + return mStream != nullptr; + } + +private: + cudaStream_t mStream; +}; + +// Helper function to create a unique_ptr to ScopedCudaStream +inline std::unique_ptr makeScopedCudaStream(uint32_t const flags = cudaStreamDefault) +{ + if (flags == cudaStreamDefault) + { + return std::make_unique(); + } + return std::make_unique(flags); +} + +} // namespace pluginInternal +} // namespace nvinfer1 + +#endif // TRT_SCOPED_CUDA_STREAM_H diff --git a/plugin/cropAndResizePlugin/CMakeLists.txt b/plugin/cropAndResizePlugin/CMakeLists.txt index 9279bef3..38fdc9de 100644 --- a/plugin/cropAndResizePlugin/CMakeLists.txt +++ b/plugin/cropAndResizePlugin/CMakeLists.txt @@ -21,3 +21,4 @@ add_plugin_source( cropAndResizePluginLegacy.cpp cropAndResizePluginLegacy.h ) + diff --git a/plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginConfig.yaml b/plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginConfig.yaml index 63a170e4..b57ac7be 100644 --- a/plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginConfig.yaml +++ b/plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginConfig.yaml @@ -52,7 +52,7 @@ versions: attributes_required: - crop_height - crop_width - golden_io_path: "plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginGoldenIO.json" + golden_io_path: "plugin/CropAndResizeDynamic_PluginGoldenIO.json" abs_tol: 1e-6 rel_tol: 1e-6 fp16_atol: 1e-3 @@ -111,7 +111,7 @@ versions: attributes_required: - crop_height - crop_width - golden_io_path: "plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginGoldenIO.json" + golden_io_path: "plugin/CropAndResizeDynamic_PluginGoldenIO.json" abs_tol: 1e-6 rel_tol: 1e-6 fp16_atol: 1e-3 diff --git a/plugin/detectionLayerPlugin/CMakeLists.txt b/plugin/detectionLayerPlugin/CMakeLists.txt index 2199f51d..018b2301 100644 --- a/plugin/detectionLayerPlugin/CMakeLists.txt +++ b/plugin/detectionLayerPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( detectionLayerPlugin.cpp detectionLayerPlugin.h ) + diff --git a/plugin/disentangledAttentionPlugin/CMakeLists.txt b/plugin/disentangledAttentionPlugin/CMakeLists.txt index 8df5d3fd..fb88dc2e 100644 --- a/plugin/disentangledAttentionPlugin/CMakeLists.txt +++ b/plugin/disentangledAttentionPlugin/CMakeLists.txt @@ -23,3 +23,4 @@ add_plugin_source( disentangledAttentionPluginLegacy.h disentangledKernel.cu ) + diff --git a/plugin/disentangledAttentionPlugin/DisentangledAttentionPlugin_PluginConfig.yaml b/plugin/disentangledAttentionPlugin/DisentangledAttentionPlugin_PluginConfig.yaml index 3cc47c34..c432aff1 100644 --- a/plugin/disentangledAttentionPlugin/DisentangledAttentionPlugin_PluginConfig.yaml +++ b/plugin/disentangledAttentionPlugin/DisentangledAttentionPlugin_PluginConfig.yaml @@ -45,7 +45,7 @@ versions: attributes_required: - span - factor - golden_io_path: "plugin/disentangledAttentionPlugin/DisentangledAttention_PluginGoldenIO.json" + golden_io_path: "plugin/DisentangledAttention_PluginGoldenIO.json" abs_tol: 1e-5 rel_tol: 1e-5 fp16_atol: 1e-2 @@ -120,7 +120,7 @@ versions: attributes_required: - span - factor - golden_io_path: "plugin/disentangledAttentionPlugin/DisentangledAttention_PluginGoldenIO.json" + golden_io_path: "plugin/DisentangledAttention_PluginGoldenIO.json" abs_tol: 1e-5 rel_tol: 1e-5 fp16_atol: 1e-2 diff --git a/plugin/efficientNMSPlugin/CMakeLists.txt b/plugin/efficientNMSPlugin/CMakeLists.txt index b16adac1..81d4cae1 100644 --- a/plugin/efficientNMSPlugin/CMakeLists.txt +++ b/plugin/efficientNMSPlugin/CMakeLists.txt @@ -31,3 +31,4 @@ if (NOT DEFINED ${TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW}) set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE) set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE) endif() + diff --git a/plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginConfig.yaml b/plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginConfig.yaml index e40535cb..b74ae154 100644 --- a/plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginConfig.yaml +++ b/plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -81,7 +81,7 @@ versions: - background_class - score_activation - box_coding - golden_io_path: "plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginGoldenIO.json" + golden_io_path: "plugin/EfficientNMSPlugin_PluginGoldenIO.json" abs_tol: 1e-5 rel_tol: 1e-5 configs: diff --git a/plugin/embLayerNormPlugin/CMakeLists.txt b/plugin/embLayerNormPlugin/CMakeLists.txt index f19f180b..87fca8c6 100644 --- a/plugin/embLayerNormPlugin/CMakeLists.txt +++ b/plugin/embLayerNormPlugin/CMakeLists.txt @@ -9,8 +9,7 @@ # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # @@ -28,3 +27,4 @@ add_plugin_source( embLayerNormVarSeqlenPluginLegacy.cpp embLayerNormVarSeqlenPluginLegacy.h ) + diff --git a/plugin/embLayerNormPlugin/CustomEmbLayerNormPluginDynamic_PluginConfig.yaml b/plugin/embLayerNormPlugin/CustomEmbLayerNormPluginDynamic_PluginConfig.yaml index 62b95b35..79cb49d7 100644 --- a/plugin/embLayerNormPlugin/CustomEmbLayerNormPluginDynamic_PluginConfig.yaml +++ b/plugin/embLayerNormPlugin/CustomEmbLayerNormPluginDynamic_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -138,7 +138,7 @@ versions: - bert_embeddings_word_embeddings - bert_embeddings_token_type_embeddings - bert_embeddings_position_embeddings - golden_reference_script: "plugin/embLayerNormPlugin/CustomEmbLayerNormPluginDynamic_PluginReference.py" + golden_reference_script: "plugin/CustomEmbLayerNormPluginDynamic_PluginReference.py" abs_tol: 1e-5 rel_tol: 1e-5 configs: diff --git a/plugin/embLayerNormPlugin/embLayerNormPlugin.cpp b/plugin/embLayerNormPlugin/embLayerNormPlugin.cpp index 86d075fa..0df771fb 100644 --- a/plugin/embLayerNormPlugin/embLayerNormPlugin.cpp +++ b/plugin/embLayerNormPlugin/embLayerNormPlugin.cpp @@ -55,7 +55,7 @@ EmbLayerNormPluginDynamic::EmbLayerNormPluginDynamic(std::string const& name, Da mWordVocabSize = wordEmb.count / mLd; mPosVocabSize = posEmb.count / mLd; mTokVocabSize = tokEmb.count / mLd; - mSM = getSMVersion(); + mSM = getSmVersion(); mOutputFp16 = mType == DataType::kHALF ? 1 : 0; mUseFullMask = static_cast(useFullMask); // NOTE: mS is set during configure diff --git a/plugin/embLayerNormPlugin/embLayerNormPluginLegacy.cpp b/plugin/embLayerNormPlugin/embLayerNormPluginLegacy.cpp index 0d8d166e..62cb3644 100644 --- a/plugin/embLayerNormPlugin/embLayerNormPluginLegacy.cpp +++ b/plugin/embLayerNormPlugin/embLayerNormPluginLegacy.cpp @@ -56,7 +56,7 @@ EmbLayerNormPluginDynamicLegacy::EmbLayerNormPluginDynamicLegacy(std::string con mWordVocabSize = wordEmb.count / mLd; mPosVocabSize = posEmb.count / mLd; mTokVocabSize = tokEmb.count / mLd; - mSM = getSMVersion(); + mSM = getSmVersion(); // mS is set during configure mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT); diff --git a/plugin/fcPlugin/CMakeLists.txt b/plugin/fcPlugin/CMakeLists.txt index db08cd47..d089c872 100644 --- a/plugin/fcPlugin/CMakeLists.txt +++ b/plugin/fcPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( fcPlugin.cpp fcPlugin.h ) + diff --git a/plugin/fcPlugin/CustomFCPluginDynamic_PluginConfig.yaml b/plugin/fcPlugin/CustomFCPluginDynamic_PluginConfig.yaml index c29c1f84..1792671d 100644 --- a/plugin/fcPlugin/CustomFCPluginDynamic_PluginConfig.yaml +++ b/plugin/fcPlugin/CustomFCPluginDynamic_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -73,8 +73,8 @@ versions: - out_dims - type_id - W - golden_io_path: "plugin/fcPlugin/CustomFCPluginDynamic_PluginGoldenIO.json" - golden_reference_script: "plugin/fcPlugin/CustomFCPluginDynamic_PluginReference.py" + golden_io_path: "plugin/CustomFCPluginDynamic_PluginGoldenIO.json" + golden_reference_script: "plugin/CustomFCPluginDynamic_PluginReference.py" abs_tol: 1e-5 rel_tol: 1e-5 fp16_atol: 1e-3 diff --git a/plugin/flattenConcat/CMakeLists.txt b/plugin/flattenConcat/CMakeLists.txt index 6d9af595..de1d1085 100644 --- a/plugin/flattenConcat/CMakeLists.txt +++ b/plugin/flattenConcat/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( flattenConcat.cpp flattenConcat.h ) + diff --git a/plugin/geluPlugin/CMakeLists.txt b/plugin/geluPlugin/CMakeLists.txt index f8af3481..764e6d8f 100644 --- a/plugin/geluPlugin/CMakeLists.txt +++ b/plugin/geluPlugin/CMakeLists.txt @@ -20,3 +20,4 @@ add_plugin_source( geluPlugin.cpp geluPlugin.h ) + diff --git a/plugin/geluPlugin/CustomGeluPluginDynamic_PluginConfig.yaml b/plugin/geluPlugin/CustomGeluPluginDynamic_PluginConfig.yaml index b863f6ed..69d67e06 100644 --- a/plugin/geluPlugin/CustomGeluPluginDynamic_PluginConfig.yaml +++ b/plugin/geluPlugin/CustomGeluPluginDynamic_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -66,7 +66,7 @@ versions: - type_id abs_tol: 1e-2 rel_tol: 1e-2 - golden_reference_script: "plugin/geluPlugin/CustomGeluPluginDynamic_PluginReference.py" + golden_reference_script: "plugin/CustomGeluPluginDynamic_PluginReference.py" configs: config1: input_types: diff --git a/plugin/generateDetectionPlugin/CMakeLists.txt b/plugin/generateDetectionPlugin/CMakeLists.txt index a6e07663..71a69692 100644 --- a/plugin/generateDetectionPlugin/CMakeLists.txt +++ b/plugin/generateDetectionPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( generateDetectionPlugin.cpp generateDetectionPlugin.h ) + diff --git a/plugin/gridAnchorPlugin/CMakeLists.txt b/plugin/gridAnchorPlugin/CMakeLists.txt index bb6e041b..1b80e278 100644 --- a/plugin/gridAnchorPlugin/CMakeLists.txt +++ b/plugin/gridAnchorPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( gridAnchorPlugin.cpp gridAnchorPlugin.h ) + diff --git a/plugin/gridAnchorPlugin/GridAnchor_TRT_PluginConfig.yaml b/plugin/gridAnchorPlugin/GridAnchor_TRT_PluginConfig.yaml index 331916b5..c3912701 100644 --- a/plugin/gridAnchorPlugin/GridAnchor_TRT_PluginConfig.yaml +++ b/plugin/gridAnchorPlugin/GridAnchor_TRT_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -69,7 +69,7 @@ versions: - variance - aspectRatios - featureMapShapes - golden_io_path: "plugin/gridAnchorPlugin/GridAnchor_TRT_PluginGoldenIO.json" + golden_io_path: "plugin/GridAnchor_TRT_PluginGoldenIO.json" abs_tol: 1e-2 rel_tol: 1e-2 configs: diff --git a/plugin/groupNormalizationPlugin/CMakeLists.txt b/plugin/groupNormalizationPlugin/CMakeLists.txt index 2628e537..d593655d 100644 --- a/plugin/groupNormalizationPlugin/CMakeLists.txt +++ b/plugin/groupNormalizationPlugin/CMakeLists.txt @@ -20,3 +20,4 @@ add_plugin_source( groupNormalizationPlugin.cpp groupNormalizationPlugin.h ) + diff --git a/plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginConfig.yaml b/plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginConfig.yaml index ac5a6401..bb3887f6 100644 --- a/plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginConfig.yaml +++ b/plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -76,7 +76,7 @@ versions: attributes_required: [] abs_tol: 1e-2 rel_tol: 1e-2 - golden_reference_script: "plugin/groupNormalizationPlugin/GroupNormalizationPlugin_PluginReference.py" + golden_reference_script: "plugin/GroupNormalizationPlugin_PluginReference.py" configs: config1: input_types: diff --git a/plugin/instanceNormalizationPlugin/instanceNormFwd.h b/plugin/instanceNormalizationPlugin/instanceNormFwd.h index 1836eb41..e09096b3 100644 --- a/plugin/instanceNormalizationPlugin/instanceNormFwd.h +++ b/plugin/instanceNormalizationPlugin/instanceNormFwd.h @@ -70,54 +70,34 @@ template struct Instance_norm_kernel_params { - enum - { - USE_ONLINE_APPROACH = 1 - }; - enum - { - THREADS_PER_CTA = THREADS_PER_CTA_ - }; - enum - { - THREADS_PER_PIXEL = THREADS_PER_PIXEL_ - }; // 8 or 16 - enum - { - SM = SM_ - }; + static constexpr int32_t USE_ONLINE_APPROACH = 1; + + static constexpr int32_t THREADS_PER_CTA = THREADS_PER_CTA_; + + //! 8 or 16 + static constexpr int32_t THREADS_PER_PIXEL = THREADS_PER_PIXEL_; + + static constexpr int32_t SM = SM_; typedef Input_Data_Type_ Input_Data_Type; typedef Output_Data_Type_ Output_Data_Type; typedef StorageType_ StorageType; - enum - { - PIXELS_PER_THREAD_IN_REGISTERS = getPixelsPerThreadInRegisters() - }; - enum - { - PIXELS_PER_THREAD_IN_SMEM = getPixelsPerThreadInSmem() - }; - enum - { - C_ELEMENTS_PER_CTA = C_ELEMENTS_PER_CTA_ - }; // 64; - enum - { - ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL - }; // 4 default + static constexpr int32_t PIXELS_PER_THREAD_IN_REGISTERS = getPixelsPerThreadInRegisters(); + + static constexpr int32_t PIXELS_PER_THREAD_IN_SMEM = getPixelsPerThreadInSmem(); + + //! 64 + static constexpr int32_t C_ELEMENTS_PER_CTA = C_ELEMENTS_PER_CTA_; + + //! 4 default + static constexpr int32_t ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL; // Derived params. - enum - { - PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL - }; - enum - { - MIN_PIXELS_PER_CTA = PIXELS_PER_LDG * PIXELS_PER_THREAD_IN_REGISTERS - }; + static constexpr int32_t PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL; + + static constexpr int32_t MIN_PIXELS_PER_CTA = PIXELS_PER_LDG * PIXELS_PER_THREAD_IN_REGISTERS; }; struct InstanceNormFwdContext diff --git a/plugin/leakyReluPlugin/CMakeLists.txt b/plugin/leakyReluPlugin/CMakeLists.txt index c6fc6c55..f0b23b26 100644 --- a/plugin/leakyReluPlugin/CMakeLists.txt +++ b/plugin/leakyReluPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( lReluPlugin.cpp lReluPlugin.h ) + diff --git a/plugin/modulatedDeformConvPlugin/CMakeLists.txt b/plugin/modulatedDeformConvPlugin/CMakeLists.txt index 268a8070..06235b59 100644 --- a/plugin/modulatedDeformConvPlugin/CMakeLists.txt +++ b/plugin/modulatedDeformConvPlugin/CMakeLists.txt @@ -26,3 +26,4 @@ add_plugin_source( modulatedDeformConvPluginLegacy.cpp modulatedDeformConvPluginLegacy.h ) + diff --git a/plugin/modulatedDeformConvPlugin/CustomModulatedDeformConv2d_PluginConfig.yaml b/plugin/modulatedDeformConvPlugin/CustomModulatedDeformConv2d_PluginConfig.yaml index 198a3948..794da0c2 100644 --- a/plugin/modulatedDeformConvPlugin/CustomModulatedDeformConv2d_PluginConfig.yaml +++ b/plugin/modulatedDeformConvPlugin/CustomModulatedDeformConv2d_PluginConfig.yaml @@ -132,7 +132,7 @@ versions: - dilation - group - deformable_group - golden_io_path: "plugin/modulatedDeformConvPlugin/CustomModulatedDeformConv2d_PluginGoldenIO.json" + golden_io_path: "plugin/CustomModulatedDeformConv2d_PluginGoldenIO.json" abs_tol: 1e-5 rel_tol: 1e-5 configs: @@ -261,7 +261,7 @@ versions: - dilation - group - deformable_group - golden_io_path: "plugin/modulatedDeformConvPlugin/CustomModulatedDeformConv2d_PluginGoldenIO.json" + golden_io_path: "plugin/CustomModulatedDeformConv2d_PluginGoldenIO.json" abs_tol: 1e-5 rel_tol: 1e-5 configs: diff --git a/plugin/multilevelCropAndResizePlugin/CMakeLists.txt b/plugin/multilevelCropAndResizePlugin/CMakeLists.txt index 08428b45..736ea97f 100644 --- a/plugin/multilevelCropAndResizePlugin/CMakeLists.txt +++ b/plugin/multilevelCropAndResizePlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( multilevelCropAndResizePlugin.cpp multilevelCropAndResizePlugin.h ) + diff --git a/plugin/multilevelProposeROI/CMakeLists.txt b/plugin/multilevelProposeROI/CMakeLists.txt index 36ba44f3..a10ffbc6 100644 --- a/plugin/multilevelProposeROI/CMakeLists.txt +++ b/plugin/multilevelProposeROI/CMakeLists.txt @@ -20,3 +20,4 @@ add_plugin_source( multilevelProposeROIPlugin.h tlt_mrcnn_config.h ) + diff --git a/plugin/pillarScatterPlugin/CMakeLists.txt b/plugin/pillarScatterPlugin/CMakeLists.txt index 58c04748..c0823846 100644 --- a/plugin/pillarScatterPlugin/CMakeLists.txt +++ b/plugin/pillarScatterPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( pillarScatter.cpp pillarScatter.h ) + diff --git a/plugin/proposalLayerPlugin/CMakeLists.txt b/plugin/proposalLayerPlugin/CMakeLists.txt index 85908ebf..f917e5d7 100644 --- a/plugin/proposalLayerPlugin/CMakeLists.txt +++ b/plugin/proposalLayerPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( proposalLayerPlugin.cpp proposalLayerPlugin.h ) + diff --git a/plugin/reorgPlugin/CMakeLists.txt b/plugin/reorgPlugin/CMakeLists.txt index f1e1a23e..b31d1b2b 100644 --- a/plugin/reorgPlugin/CMakeLists.txt +++ b/plugin/reorgPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( reorgPlugin.cpp reorgPlugin.h ) + diff --git a/plugin/resizeNearestPlugin/CMakeLists.txt b/plugin/resizeNearestPlugin/CMakeLists.txt index 089c0f6a..8c20350f 100644 --- a/plugin/resizeNearestPlugin/CMakeLists.txt +++ b/plugin/resizeNearestPlugin/CMakeLists.txt @@ -19,3 +19,4 @@ add_plugin_source( resizeNearestPlugin.cpp resizeNearestPlugin.h ) + diff --git a/plugin/roiAlignPlugin/CMakeLists.txt b/plugin/roiAlignPlugin/CMakeLists.txt index fefd0b6b..778dce47 100644 --- a/plugin/roiAlignPlugin/CMakeLists.txt +++ b/plugin/roiAlignPlugin/CMakeLists.txt @@ -9,8 +9,7 @@ # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # @@ -32,3 +31,4 @@ add_vc_plugin_source( roiAlignPluginLegacy.cpp roiAlignPluginLegacy.h ) + diff --git a/plugin/roiAlignPlugin/ROIAlign_PluginConfig.yaml b/plugin/roiAlignPlugin/ROIAlign_PluginConfig.yaml index 3bed92e4..c8e00ecf 100644 --- a/plugin/roiAlignPlugin/ROIAlign_PluginConfig.yaml +++ b/plugin/roiAlignPlugin/ROIAlign_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -77,7 +77,7 @@ versions: attributes_required: [] abs_tol: 1e-4 rel_tol: 1e-4 - golden_io_path: "plugin/roiAlignPlugin/ROIAlign_PluginGoldenIO.json" + golden_io_path: "plugin/ROIAlign_PluginGoldenIO.json" configs: aligned_false: input_types: diff --git a/plugin/scatterElementsPlugin/CMakeLists.txt b/plugin/scatterElementsPlugin/CMakeLists.txt index 169b813f..95a536f1 100644 --- a/plugin/scatterElementsPlugin/CMakeLists.txt +++ b/plugin/scatterElementsPlugin/CMakeLists.txt @@ -27,3 +27,4 @@ add_plugin_source( scatterElementsPluginLegacy.h TensorInfo.cuh ) + diff --git a/plugin/scatterElementsPlugin/ScatterElementsPlugin_PluginConfig.yaml b/plugin/scatterElementsPlugin/ScatterElementsPlugin_PluginConfig.yaml index 12337677..bb499756 100644 --- a/plugin/scatterElementsPlugin/ScatterElementsPlugin_PluginConfig.yaml +++ b/plugin/scatterElementsPlugin/ScatterElementsPlugin_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -143,7 +143,7 @@ versions: - "max" attributes_required: - reduction - golden_io_path: "plugin/scatterElementsPlugin/ScatterElementsPlugin_PluginGoldenIO.json" + golden_io_path: "plugin/ScatterElementsPlugin_PluginGoldenIO.json" abs_tol: 1e-2 rel_tol: 1e-2 diff --git a/plugin/scatterElementsPlugin/scatterElementsPluginKernel.cu b/plugin/scatterElementsPlugin/scatterElementsPluginKernel.cu index c5dab332..0f2c2f50 100644 --- a/plugin/scatterElementsPlugin/scatterElementsPluginKernel.cu +++ b/plugin/scatterElementsPlugin/scatterElementsPluginKernel.cu @@ -67,7 +67,7 @@ bool hasBfloat16AtomicAdd() return deviceProp.major >= 8; } -inline uint32_t getElementSize(nvinfer1::DataType t) noexcept +inline uint32_t getElementSize(nvinfer1::DataType t) { switch (t) { diff --git a/plugin/skipLayerNormPlugin/CMakeLists.txt b/plugin/skipLayerNormPlugin/CMakeLists.txt index 59449856..d6251b7b 100644 --- a/plugin/skipLayerNormPlugin/CMakeLists.txt +++ b/plugin/skipLayerNormPlugin/CMakeLists.txt @@ -28,3 +28,4 @@ add_plugin_source( skipLayerNormPluginLegacy.cpp skipLayerNormPluginLegacy.h ) + diff --git a/plugin/skipLayerNormPlugin/CustomSkipLayerNormPluginDynamic_PluginConfig.yaml b/plugin/skipLayerNormPlugin/CustomSkipLayerNormPluginDynamic_PluginConfig.yaml index 117fcbf1..ba4f5aaf 100644 --- a/plugin/skipLayerNormPlugin/CustomSkipLayerNormPluginDynamic_PluginConfig.yaml +++ b/plugin/skipLayerNormPlugin/CustomSkipLayerNormPluginDynamic_PluginConfig.yaml @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -104,7 +104,7 @@ versions: - ld - beta - gamma - golden_reference_script: "plugin/skipLayerNormPlugin/CustomSkipLayerNormPluginDynamic_PluginReference.py" + golden_reference_script: "plugin/CustomSkipLayerNormPluginDynamic_PluginReference.py" abs_tol: 1e-2 rel_tol: 1e-2 configs: @@ -214,7 +214,7 @@ versions: - type_id - beta - gamma - golden_reference_script: "plugin/skipLayerNormPlugin/CustomSkipLayerNormPluginDynamic_PluginReference.py" + golden_reference_script: "plugin/CustomSkipLayerNormPluginDynamic_PluginReference.py" abs_tol: 1e-2 rel_tol: 1e-2 configs: diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 6666511d..dafd35fa 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -63,7 +63,9 @@ set(PYBIND11_NOPYTHON ON CACHE INTERNAL "") add_subdirectory(${TRT_BUILD_PYTHON_EXTERNALS_PATH}/pybind11 ${CMAKE_CURRENT_BINARY_DIR}/externals/pybind11) # Pybind11 would normally enable this by default, but does not do so under NOPYTHON mode, so we do it manually. -set(CMAKE_INTERPROCEDURAL_OPTIMIZATION ON) +if(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_X86}) + set(CMAKE_INTERPROCEDURAL_OPTIMIZATION ON) +endif() add_custom_target(tensorrt_python_bindings) @@ -171,9 +173,6 @@ function(createBindingLibrary moduleName pyVersion) else() message(FATAL_ERROR "Unknown TensorRT module " ${moduleName}) endif() - message(STATUS "TRT LIBS: ${TRT_LIBS}") - - find_package(CUDAToolkit REQUIRED) target_link_libraries(${libName} PRIVATE ${TRT_LIBS} @@ -224,44 +223,6 @@ function(createBindingLibrary moduleName pyVersion) endif() endfunction() -# Processes one or more wheel templates file, replacing any markers with concrete information and -# copying the result into the per-python per-module build dir. -# -# \param moduleName The module name to create the bindings for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". -# \param pyVersion The python version to create bindings for, i.e. "3.12". -# \param ARGN A list of file paths, relative to './packaging/bindings_wheel/tensorrt', of the file(s) to copy. -# \returns generatedWheelFiles A list containing paths to all generated files, which can be used to create a custom target. -function(processWheelTemplates moduleName pyVersion) - foreach(filePath IN LISTS ARGN) - set(outputDir ${CMAKE_CURRENT_BINARY_DIR}/${moduleName}_bindings-py${pyVersion}/$) - set(outputFile ${outputDir}/${filePath}) - - get_target_property(TRT_OUTPUT_NAME tensorrt OUTPUT_NAME) - get_target_property(PARSER_OUTPUT_NAME nvonnxparser OUTPUT_NAME) - - add_custom_command( - OUTPUT ${outputFile} - COMMAND - ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/scripts/process_wheel_template.py - --src-dir ${CMAKE_CURRENT_LIST_DIR}/packaging/bindings_wheel/tensorrt - --dst-dir ${outputDir} - --filepath ${filePath} - --trt-module ${moduleName} - --trt-py-version ${TensorRT_PACKAGE_VERSION} - --cuda-version ${TRT_CUDA_VERSION} - --trt-version ${TensorRT_VERSION} - --trt-nvinfer-name ${TRT_OUTPUT_NAME} - --trt-onnxparser-name ${PARSER_OUTPUT_NAME} - DEPENDS - scripts/process_wheel_template.py - ${CMAKE_CURRENT_LIST_DIR}/packaging/bindings_wheel/tensorrt/${filePath} - VERBATIM) - - list(APPEND generatedFiles ${outputFile}) - endforeach() - set(generatedWheelFiles ${generatedFiles} PARENT_SCOPE) -endfunction() - # Enumerate all the combinations and create the per-python per-module targets. foreach(moduleName IN LISTS TRT_PYTHON_MODULE_NAMES) foreach(pyVersion IN LISTS TRT_BUILD_PYTHON_PY_VERSIONS) @@ -272,14 +233,6 @@ endforeach() # Enter the packaging subdir to actually build the wheels. add_subdirectory(packaging) -if(${TRT_BUILD_PLUGINS}) - install( - FILES include/impl/NvInferPythonPlugin.h - TYPE INCLUDE - COMPONENT release - ) -endif() - else() # TRT_BUILD_ENABLE_NEW_PYTHON_FLOW - old flow is below this line set(TRT_BUILD_WINML OFF) diff --git a/python/docstrings/infer/pyAlgorithmSelectorDoc.h b/python/docstrings/infer/pyAlgorithmSelectorDoc.h index e43325a1..2390f9b6 100644 --- a/python/docstrings/infer/pyAlgorithmSelectorDoc.h +++ b/python/docstrings/infer/pyAlgorithmSelectorDoc.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/python/docstrings/infer/pyCoreDoc.h b/python/docstrings/infer/pyCoreDoc.h index 8845032e..648469d7 100644 --- a/python/docstrings/infer/pyCoreDoc.h +++ b/python/docstrings/infer/pyCoreDoc.h @@ -616,6 +616,8 @@ constexpr char const* get_runtime_config = R"trtdoc( :returns: The runtime configuration. )trtdoc"; + + } // namespace IExecutionContextDoc namespace IDebugListenerDoc @@ -724,6 +726,14 @@ constexpr char const* get_execution_context_allocation_strategy = R"trtdoc( } // namespace IRuntimeConfigDoc +namespace EngineStatDoc +{ +constexpr char const* descr = R"trtdoc(The kind of engine statistics that queried from the ICudaEngine.)trtdoc"; +constexpr char const* TOTAL_WEIGHTS_SIZE = R"trtdoc(The total weights size in bytes in the engine.)trtdoc"; +constexpr char const* STRIPPED_WEIGHTS_SIZE + = R"trtdoc(The stripped weight size in bytes for engines built with BuilderFlag::kSTRIP_PLAN.)trtdoc"; +} // namespace EngineStatDoc + namespace ICudaEngineDoc { constexpr char const* descr = R"trtdoc( @@ -924,6 +934,14 @@ constexpr char const* is_debug_tensor = R"trtdoc( :arg name: The tensor name. )trtdoc"; + +constexpr char const* get_engine_stat = R"trtdoc( + Return the engine statistics specified by the given enum value. + If STRIPPED_WEIGHTS_SIZE is passed to query a normal engine, this function will + return -1 to indicate invalid enum value. + + :arg stat: The engine statistic kind to get. +)trtdoc"; } // namespace ICudaEngineDoc namespace OutputAllocatorDoc @@ -1816,6 +1834,7 @@ namespace SerializationFlagDoc constexpr char const* descr = R"trtdoc(Valid flags that can be use to creating binary file from engine.)trtdoc"; constexpr char const* EXCLUDE_WEIGHTS = R"trtdoc(Exclude weights that can be refitted.)trtdoc"; constexpr char const* EXCLUDE_LEAN_RUNTIME = R"trtdoc(Exclude lean runtime from the plan.)trtdoc"; +constexpr char const* INCLUDE_REFIT = R"trtdoc(Remain refittable if originally so.)trtdoc"; } // namespace SerializationFlagDoc namespace ExecutionContextAllocationStrategyDoc diff --git a/python/docstrings/infer/pyGraphDoc.h b/python/docstrings/infer/pyGraphDoc.h index 70bd8e5a..2518bac4 100644 --- a/python/docstrings/infer/pyGraphDoc.h +++ b/python/docstrings/infer/pyGraphDoc.h @@ -76,6 +76,8 @@ constexpr const char* SQUEEZE = R"trtdoc(Squeeze layer)trtdoc"; constexpr const char* UNSQUEEZE = R"trtdoc(Unsqueeze layer)trtdoc"; constexpr const char* CUMULATIVE = R"trtdoc(Cumulative layer)trtdoc"; constexpr const char* DYNAMIC_QUANTIZE = R"trtdoc(DynamicQuantize layer)trtdoc"; +constexpr const char* ATTENTION_INPUT = R"trtdoc(Attention input layer)trtdoc"; +constexpr const char* ATTENTION_OUTPUT = R"trtdoc(Attention output layer)trtdoc"; constexpr const char* SPLIT_TO_RAGGED = R"trtdoc(SplitToRagged layer)trtdoc"; constexpr const char* CONCAT_FROM_RAGGED = R"trtdoc(ConcatFromRagged layer)trtdoc"; } // namespace LayerTypeDoc @@ -905,6 +907,8 @@ constexpr const char* descr = R"trtdoc( :ivar k: :class:`TopKOperation` the k value for the layer. Currently only values up to 3840 are supported. Use the set_input() method with index 1 to pass in dynamic k as a tensor. :ivar axes: :class:`TopKOperation` The axes along which to reduce. + :ivar indices_type: :class:`DataType` The specified data type of the output indices tensor. Must be tensorrt.int32 or tensorrt.int64. + )trtdoc"; constexpr const char* set_input = R"trtdoc( @@ -959,6 +963,25 @@ constexpr const char* descr = R"trtdoc( )trtdoc"; } // namespace IMatrixMultiplyLayerDoc +namespace CollectiveOperationDoc +{ +constexpr const char* descr + = R"trtdoc(The collective operations that may be performed by a DistCollective layer)trtdoc"; + +constexpr const char* ALL_REDUCE = R"trtdoc(All reduce collective operation)trtdoc"; +constexpr const char* ALL_GATHER = R"trtdoc(All gather collective operation)trtdoc"; +constexpr const char* BROADCAST = R"trtdoc(Broadcast collective operation)trtdoc"; +constexpr const char* REDUCE = R"trtdoc(Reduce collective operation)trtdoc"; +constexpr const char* REDUCE_SCATTER = R"trtdoc(Reduce scatter collective operation)trtdoc"; +} // namespace CollectiveOperationDoc + +namespace IDistCollectiveLayerDoc +{ +constexpr const char* descr = R"trtdoc( + A dist collective layer in an :class:`INetworkDefinition` . +)trtdoc"; +} // namespace IDistCollectiveLayerDoc + namespace IRaggedSoftMaxLayerDoc { constexpr const char* descr = R"trtdoc( @@ -1410,7 +1433,7 @@ constexpr const char* descr = R"trtdoc( Use :func:`set_input` to add this optional tensor. The SelectedIndices output tensor contains the indices of the selected boxes. - It is a linear tensor of type ``int32``. It has shape [NumOutputBoxes, 3].] + It is a linear tensor of type ``int32`` or ``int64``. It has shape [NumOutputBoxes, 3].] Each row contains a (batchIndex, classIndex, boxIndex) tuple. The output boxes are sorted in order of increasing batchIndex and then in order of decreasing score within each batchIndex. For each batchIndex, the ordering of output boxes with the same score is unspecified. @@ -1432,6 +1455,7 @@ constexpr const char* descr = R"trtdoc( :ivar bounding_box_format: :class:`BoundingBoxFormat` The bounding box format used by the layer. Default is CORNER_PAIRS. :ivar topk_box_limit: :class:`int` The maximum number of filtered boxes considered for selection. Default is 2000 for SM 5.3 and 6.2 devices, and 5000 otherwise. The TopK box limit must be less than or equal to {2000 for SM 5.3 and 6.2 devices, 5000 otherwise}. + :ivar indices_type: :class:`DataType` The specified data type of the output indices tensor. Must be tensorrt.int32 or tensorrt.int64. )trtdoc"; constexpr const char* set_input = R"trtdoc( @@ -1832,6 +1856,8 @@ constexpr char const* descr = R"trtdoc( Computes the indices of the input tensor where the value is non-zero. The returned indices are in row-major order. The output shape is always `{D, C}`, where `D` is the number of dimensions of the input and `C` is the number of non-zero values. + + :ivar indices_type: :class:`DataType` The specified data type of the output indices tensor. Must be tensorrt.int32 or tensorrt.int64. )trtdoc"; } // namespace INonZeroLayerDoc @@ -1977,6 +2003,94 @@ constexpr const char* descr = R"trtdoc( )trtdoc"; } // namespace ICumulativeLayerDoc + +namespace AttentionNormalizationOpDoc +{ +constexpr const char* descr = R"trtdoc(The normalization operations that may be performed by an Attention layer)trtdoc"; +constexpr const char* NONE = R"trtdoc()trtdoc"; +constexpr const char* SOFTMAX = R"trtdoc()trtdoc"; +} // namespace AttentionNormalizationOpDoc + +namespace IAttentionBoundaryLayerDoc +{ +constexpr const char* descr = R"trtdoc( + :ivar attention: :class:`IAttention` associated with this boundary layer. +)trtdoc"; +} // namespace IAttentionBoundaryLayerDoc + +namespace IAttentionInputLayerDoc +{ +constexpr const char* descr = R"trtdoc( + Marks input boundary to an :class:`IAttention` scope +)trtdoc"; +} // namespace IAttentionInputLayerDoc + +namespace IAttentionOutputLayerDoc +{ +constexpr const char* descr = R"trtdoc( + Marks output boundary to an :class:`IAttention` scope +)trtdoc"; +} // namespace IAttentionOutputLayerDoc + +namespace IAttentionDoc +{ +constexpr const char* descr = R"trtdoc( + An attention in a :class:`INetworkDefinition` . + + :ivar mask: :class:`ITensor` The mask tensor for attention. Cannot be set together with causal attention. + :ivar norm_op: :class:`AttentionNormalizationOp` The normalization operation for the attention layer. Default to AttentionNormalizationOp::kSOFTMAX. + :ivar decomposable: :class:`bool` Specifies whether decomposition into primitive ops is allowed when no attention fusion is supported. Default to False. + :ivar causal: :class:`bool` Specifies whether the attention will run a causal inference. Cannot be used together with mask. + :ivar name: :class:`str` The name of the attention. + :ivar normalization_quantize_scale: :class:`ITensor` The quantization scale for the attention normalization output. + :ivar normalization_quantize_to_type: :class:`DataType` The datatype the attention normalization is quantized to. + :ivar num_inputs: :class:`int` The number of inputs of the attention. + :ivar num_outputs: :class:`int` The number of outputs of the attention. +)trtdoc"; + +constexpr char const* init = R"trtdoc( + :arg query: The input query tensor. + :arg key: The input key tensor. + :arg value: The input value tensor. + :arg norm_op: The normalization operation for the attention. + :arg casual: The boolean specifies whether the attention will run a causal inference. +)trtdoc"; + +constexpr const char* set_input = R"trtdoc( + Set the input tensor specified by the given index. + + The indices are as follows: + + ===== ================================================================================== + Index Description + ===== ================================================================================== + 0 query. + 1 key. + 2 value. + ===== ================================================================================== + + :arg index: The index of the input tensor. query:0, key:1, value:2 + :arg tensor: The input tensor. +)trtdoc"; + +constexpr const char* get_input = R"trtdoc( + Get the input tensor specified by the given index. + + :arg index: The index of the input tensor. + + :returns: The tensor, or :class:`None` if it is out of range. +)trtdoc"; + +constexpr const char* get_output = R"trtdoc( + Get the output tensor specified by the given index. + + :arg index: The index of the output tensor. + + :returns: The tensor, or :class:`None` if it is out of range. +)trtdoc"; + +} // namespace IAttentionDoc + namespace INetworkDefinitionDoc { constexpr const char* descr = R"trtdoc( @@ -1985,6 +2099,7 @@ constexpr const char* descr = R"trtdoc( :ivar num_layers: :class:`int` The number of layers in the network. :ivar num_inputs: :class:`int` The number of inputs of the network. :ivar num_outputs: :class:`int` The number of outputs of the network. + :ivar num_ranks: :class:`int` The number of ranks to use for multi-device execution. :ivar name: :class:`str` The name of the network. This is used so that it can be associated with a built engine. The name must be at most 128 characters in length. TensorRT makes no use of this string except storing it as part of the engine so that it may be retrieved at runtime. A name unique to the builder will be generated by default. :ivar has_implicit_batch_dimension: :class:`bool` [DEPRECATED] Deprecated in TensorRT 10.0. Always flase since the implicit batch dimensions support has been removed. :ivar error_recorder: :class:`IErrorRecorder` Application-implemented error reporting interface for TensorRT objects. @@ -2063,7 +2178,7 @@ constexpr const char* mark_unfused_tensors_as_debug_tensors = R"trtdoc( Tensors marked this way will not prevent fusion like mark_debug() does, thus preserving performance. Tensors marked this way cannot be detected by is_debug_tensor(). - DebugListener can only get internal tensor names instead of the original tensor names in the NetworkDefinition for tensors marked this way. + DebugListener can only get internal tensor names instead of the original tensor names in the NetworkDefinition for tensors marked this way. But the names correspond to the names obtained by IEngineInspector. There is no guarantee that all unfused tensors are marked. @@ -2288,6 +2403,8 @@ constexpr const char* add_topk = R"trtdoc( significant bit corresponds to the second explicit dimension. Currently axes must specify exactly one dimension, and it must be one of the last four dimensions. + :arg indices_type: The datatype of the output indices tensor. Specifying indices_type is optional (default value tensorrt.int32). + :returns: The new TopK layer, or :class:`None` if it could not be created. )trtdoc"; @@ -2466,6 +2583,7 @@ constexpr const char* add_nms = R"trtdoc( :arg max_output_boxes_per_class: The maxOutputBoxesPerClass tensor to the layer. :ivar bounding_box_format: :class:`BoundingBoxFormat` The bounding box format used by the layer. Default is CORNER_PAIRS. :ivar topk_box_limit: :class:`int` The maximum number of filtered boxes considered for selection per batch item. Default is 2000 for SM 5.3 and 6.2 devices, and 5000 otherwise. The TopK box limit must be less than or equal to {2000 for SM 5.3 and 6.2 devices, 5000 otherwise}. + :arg indices_type: The datatype of the output indices tensor. Specifying indices_type is optional (default value tensorrt.int32). :returns: The new NMS layer, or :class:`None` if it could not be created. )trtdoc"; @@ -2688,6 +2806,8 @@ constexpr char const* add_non_zero = R"trtdoc( :arg input: The input tensor to the layer. + :arg indices_type: The datatype of the output indices tensor. Specifying indices_type is optional (default value tensorrt.int32). + :returns: the new NonZero layer, or :class:`None` if it could not be created. )trtdoc"; @@ -2750,6 +2870,19 @@ constexpr const char* add_cumulative = R"trtdoc( :returns: The new cumulative layer, or :class:`None` if it could not be created. )trtdoc"; +constexpr const char* add_attention = R"trtdoc( + Add an attention to the network. + See :class:`IAttention` for more information. + + :arg query: The 4d query input tensor to the attention. + :arg key: The 4d key input tensor to the attention. + :arg value: The 4d value input tensor to the attention. + :arg normOp: The normalization operation to perform. + :arg causal: The boolean that specifies whether an attention will run casual inference. + + :returns: The new Attention, or :class:`None` if it could not be created. +)trtdoc"; + } // namespace INetworkDefinitionDoc } // namespace tensorrt diff --git a/python/include/utils.h b/python/include/utils.h index 97ab53c7..2aa351cb 100644 --- a/python/include/utils.h +++ b/python/include/utils.h @@ -169,6 +169,18 @@ constexpr auto deprecateMember(RetVal (Cls::*func)(Args...), const char* useInst return DeprecatedMemberFunc{func, useInstead}; } +template +constexpr auto deprecateInTrtRtxOnly(T&& func, const char* /*unused*/) -> T&& +{ + return std::forward(func); +} + +template +constexpr auto deprecateMemberInTrtRtxOnly(T&& func, const char* /*unused*/) -> T&& +{ + return std::forward(func); +} + template void doNothingDel(const T& self) { diff --git a/python/packaging/CMakeLists.txt b/python/packaging/CMakeLists.txt index 5b0191be..8916c3ef 100644 --- a/python/packaging/CMakeLists.txt +++ b/python/packaging/CMakeLists.txt @@ -32,8 +32,14 @@ function(processWheelTemplates wheelType moduleName pyVersion) string(REGEX REPLACE "^tensorrt/" "" adjustedFilePath ${filePath}) set(__srcDir ${CMAKE_CURRENT_LIST_DIR}/tensorrt) set(__srcFilePath ${adjustedFilePath}) - set(__outDir ${outputDir}/${moduleName}) - set(__outputFile ${outputDir}/${moduleName}/${adjustedFilePath}) + # Bit of a hack branch for standalone binding wheels + if(wheelType STREQUAL "binding_standalone") + set(__outDir ${outputDir}/${moduleName}_bindings) + set(__outputFile ${outputDir}/${moduleName}_bindings/${adjustedFilePath}) + else() + set(__outDir ${outputDir}/${moduleName}) + set(__outputFile ${outputDir}/${moduleName}/${adjustedFilePath}) + endif() elseif(filePath MATCHES "^tensorrt_libs/") # Standalone lib wheels use _libs/ instead of / string(REGEX REPLACE "^tensorrt_libs/" "" adjustedFilePath ${filePath}) set(__srcDir ${CMAKE_CURRENT_LIST_DIR}/tensorrt_libs) @@ -50,7 +56,9 @@ function(processWheelTemplates wheelType moduleName pyVersion) flagToInt(TRT_BUILD_WINML) get_target_property(TRT_OUTPUT_NAME tensorrt OUTPUT_NAME) + string(REGEX MATCH "[a-zA-Z_]+[a-zA-Z]" TRT_OUTPUT_NAME ${TRT_OUTPUT_NAME}) get_target_property(PARSER_OUTPUT_NAME nvonnxparser OUTPUT_NAME) + string(REGEX MATCH "[a-zA-Z_]+[a-zA-Z]" PARSER_OUTPUT_NAME ${PARSER_OUTPUT_NAME}) add_custom_command( OUTPUT ${__outputFile} @@ -69,7 +77,9 @@ function(processWheelTemplates wheelType moduleName pyVersion) DEPENDS ${TensorRT_SOURCE_DIR}/python/scripts/process_wheel_template.py ${CMAKE_CURRENT_LIST_DIR}/${filePath} - VERBATIM) + COMMENT "Expanding wheel template ${filePath} for ${moduleName}" + VERBATIM + ) list(APPEND generatedFiles ${__outputFile}) endforeach() @@ -98,3 +108,5 @@ add_custom_target(trt_packaging_requirements_installed DEPENDS package_install_c add_subdirectory(bindings_wheel) add_subdirectory(libs_wheel) +add_subdirectory(frontend_sdist) +add_subdirectory(metapackage) diff --git a/python/packaging/bindings_wheel/CMakeLists.txt b/python/packaging/bindings_wheel/CMakeLists.txt index 8ea13b73..7ec67acb 100644 --- a/python/packaging/bindings_wheel/CMakeLists.txt +++ b/python/packaging/bindings_wheel/CMakeLists.txt @@ -15,16 +15,44 @@ add_custom_target(tensorrt_bindings_wheels ALL) add_custom_target(trt_bindings_wheel_files) +define_property(TARGET + PROPERTY TRT_WHEEL_STAGING_DIR + BRIEF_DOCS "The directory containing all the files that were packaged into this wheel." +) + + # \brief Creates a target named tensorrt_bindings_wheel_${moduleName}_${pyVersion} which will build the Bindings Wheel for that combination. -# +# # \details The wheel is created by expanding all template files (from this directory) into the per-module per-python build directory. # Then, the binding library (tensorrt.so) is copied into the same directory as the generated files. # Finally, the wheel is built by running setup.py with the appropriate arguments in the binary directory. # -# \param moduleName The module name to create the bindings for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". -# \param pyVersion The python version to create bindings for, i.e. "3.12". -function(buildBindingsWheel moduleName pyVersion) - set(filesTarget trt_wheel_files_binding_${moduleName}_${pyVersion}) +# \param moduleName The module name to create the bindings for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". +# \param pyVersion The python version to create bindings for, i.e. "3.12". +# \param isStandalone If we are building the wheels for distribution with the standalone libs wheels. +# If true, the binding wheels' top-level folder will be _bindings instead of +function(buildBindingsWheel moduleName pyVersion isStandalone) + if (isStandalone) + set(standaloneMarker "_standalone") + set(bindingSuffix "_bindings") + else() + set(standaloneMarker "") + set(bindingSuffix "") + endif() + + set(filesTarget trt_wheel_files_binding_${moduleName}_${pyVersion}${standaloneMarker}) + + string(REPLACE "." "" pyVerStr ${pyVersion}) + + if (MSVC) + if(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_WIN10}) + set(wheelPlatform win_amd64) + else() + message(FATAL_ERROR "Unable to determine python binding wheel platform name for TRT Platform: ${TRT_BUILD_PLATFORM}") + endif() + else() + set(wheelPlatform linux_${TRT_CONFIG_ARCH}) + endif() set(wheelTemplateFiles tensorrt/__init__.py @@ -51,45 +79,54 @@ function(buildBindingsWheel moduleName pyVersion) # Expands all template files for the bindings for the target module and python version. # File paths starting with "tensorrt/" are expanded into "${moduleName}/". - processWheelTemplates(binding ${moduleName} ${pyVersion} ${wheelTemplateFiles}) + processWheelTemplates(binding${standaloneMarker} ${moduleName} ${pyVersion} ${wheelTemplateFiles}) + + # We need to copy the binding library into the directory that the wheel is created from. + set(copiedBindingLibrary ${generatedFileOutDir}/${moduleName}${bindingSuffix}/${moduleName}${CMAKE_SHARED_LIBRARY_SUFFIX}) + + # I bet you thought you were free from more windows exceptions, I certainly did. + # But here we are again, with yet another naming conventionâ„¢ + if(MSVC) + set(copiedBindingLibrary ${generatedFileOutDir}/${moduleName}${bindingSuffix}/${moduleName}.cp${pyVerStr}-${wheelPlatform}.pyd) + endif() # Creates a new custom target, and makes trt_bindings_wheel_files depend on the new target. - add_custom_target(${filesTarget} DEPENDS ${generatedWheelFiles}) + add_custom_target(${filesTarget} DEPENDS ${generatedWheelFiles} ${copiedBindingLibrary}) add_dependencies(trt_bindings_wheel_files ${filesTarget}) # Copies the binding library (tensorrt.so) into the same directory as the generated files. add_custom_command( - TARGET trt_wheel_files_binding_${moduleName}_${pyVersion} - POST_BUILD + OUTPUT ${copiedBindingLibrary} COMMAND ${CMAKE_COMMAND} -E copy $ - ${generatedFileOutDir}/${moduleName} - COMMENT "Copying bindings library to wheel directory" + ${copiedBindingLibrary} + DEPENDS tensorrt_bindings_${moduleName}_${pyVersion} + COMMENT "Copying bindings library for ${moduleName} to ${copiedBindingLibrary}" VERBATIM ) - string(REPLACE "." "" pyVerStr ${pyVersion}) - - if (MSVC) - set(wheelPlatform win_${TRT_CONFIG_ARCH}) - else() - set(wheelPlatform linux_${TRT_CONFIG_ARCH}) - endif() - # Define the output directory for the wheel - set(wheelOutDir ${TRT_WHEEL_OUTPUT_DIR}/bindings) - set(wheelOutputFile ${wheelOutDir}/${moduleName}-${TensorRT_PACKAGE_VERSION}-cp${pyVerStr}-none-${wheelPlatform}.whl) + set(wheelOutDir ${TRT_WHEEL_OUTPUT_DIR}/bindings${standaloneMarker}) + + # The standalone binding wheels' setup.py requires that STANDALONE=1 be set in the environment to update the package directory. + if(isStandalone) + set(standalone_env_prefix ${CMAKE_COMMAND} -E env STANDALONE=1 --) + set(wheelOutputFile ${wheelOutDir}/${moduleName}_cu${CUDAToolkit_VERSION_MAJOR}_bindings-${TensorRT_PACKAGE_VERSION}-cp${pyVerStr}-none-${wheelPlatform}.whl) + else() + set(standalone_env_prefix "") + set(wheelOutputFile ${wheelOutDir}/${moduleName}-${TensorRT_PACKAGE_VERSION}-cp${pyVerStr}-none-${wheelPlatform}.whl) + endif() # Add a custom command to build the wheel add_custom_command( OUTPUT ${wheelOutputFile} - COMMAND ${Python3_EXECUTABLE} setup.py -q bdist_wheel --python-tag=cp${pyVerStr} --plat-name=${wheelPlatform} --dist-dir=${wheelOutDir} + COMMAND ${standalone_env_prefix} ${Python3_EXECUTABLE} setup.py -q bdist_wheel --python-tag=cp${pyVerStr} --plat-name=${wheelPlatform} --dist-dir=${wheelOutDir} WORKING_DIRECTORY ${generatedFileOutDir} - DEPENDS tensorrt_bindings_${moduleName}_${pyVersion} ${generatedWheelFiles} trt_packaging_requirements_installed + DEPENDS ${copiedBindingLibrary} ${generatedWheelFiles} trt_packaging_requirements_installed VERBATIM ) - set(wheelTarget tensorrt_bindings_wheel_${moduleName}_${pyVersion}) + set(wheelTarget tensorrt_bindings_wheel_${moduleName}_${pyVersion}${standaloneMarker}) # Add a custom target for the wheel add_custom_target( @@ -101,18 +138,39 @@ function(buildBindingsWheel moduleName pyVersion) add_dependencies(${wheelTarget} trt_bindings_wheel_files) add_dependencies(tensorrt_bindings_wheels ${wheelTarget}) + set_target_properties(${wheelTarget} + PROPERTIES TRT_WHEEL_STAGING_DIR "${generatedFileOutDir}" + ) + + if(isStandalone) + set(bindingsInstallComponent internal) + else() + set(bindingsInstallComponent external) + endif() + install(FILES ${wheelOutputFile} - DESTINATION wheels - COMPONENT release + DESTINATION python${standaloneMarker} + COMPONENT ${bindingsInstallComponent} OPTIONAL ) endfunction() foreach(moduleName IN LISTS TRT_PYTHON_MODULE_NAMES) foreach(pyVersion IN LISTS TRT_BUILD_PYTHON_PY_VERSIONS) - buildBindingsWheel(${moduleName} ${pyVersion}) + buildBindingsWheel(${moduleName} ${pyVersion} FALSE) endforeach() endforeach() add_dependencies(tensorrt_python_wheels tensorrt_bindings_wheels) + +# For the standalone wheels, we need to build an additional copy of the wheels that uses _bindings as the package directory. +# This prevents clashes with the metapackage, which uses the "tensorrt" module name for the standalone path. +# TODO TRT-11.0: Can we drop this and change the name of the metapackage? +if(${TRT_BUILD_PYTHON_STANDALONE_WHEELS}) + foreach(moduleName IN LISTS TRT_PYTHON_MODULE_NAMES) + foreach(pyVersion IN LISTS TRT_BUILD_PYTHON_PY_VERSIONS) + buildBindingsWheel(${moduleName} ${pyVersion} TRUE) + endforeach() + endforeach() +endif() diff --git a/python/packaging/frontend_sdist/CMakeLists.txt b/python/packaging/frontend_sdist/CMakeLists.txt new file mode 100644 index 00000000..dbb8a34b --- /dev/null +++ b/python/packaging/frontend_sdist/CMakeLists.txt @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Adding a custom target to the all target requires passing the "ALL" keyword to the initial call. +# We use this variable so we can conditionally add the keyword based on the build options. +if(${TRT_BUILD_PYTHON_STANDALONE_WHEELS}) + set(ENABLE_ALL "ALL") +else() + set(ENABLE_ALL "") +endif() + +add_custom_target(tensorrt_frontend_sdist ${ENABLE_ALL}) + +# \brief Creates a target named tensorrt_frontend_sdist_${moduleName} which will build the Frontend SDist for that module. +# +# \param moduleName The module name to create the frontend for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". +function(buildFrontendWheel moduleName) + set(filesTarget trt_wheel_files_frontend_${moduleName}) + + set(wheelTemplateFiles + tensorrt/__init__.py + tensorrt/plugin/__init__.py + LICENSE.txt + setup.cfg + setup.py + ) + + # Expands all template files for the frontend for the target module. The python version is unused (and set to 0). + # File paths starting with "tensorrt/" are expanded into "${moduleName}/". + processWheelTemplates(frontend ${moduleName} 0 ${wheelTemplateFiles}) + + # Creates a new custom target, and makes trt_standalone_wheel_files depend on the new target. + add_custom_target(${filesTarget} DEPENDS ${generatedWheelFiles}) + + # Define the output directory for the wheel + set(sdistOutDir ${TRT_WHEEL_OUTPUT_DIR}/frontend) + set(sdistOutputFile ${sdistOutDir}/${moduleName}_cu${CUDAToolkit_VERSION_MAJOR}-${TensorRT_PACKAGE_VERSION}.tar.gz) + + # Add a custom command to build the wheel + add_custom_command( + OUTPUT ${sdistOutputFile} + COMMAND ${Python3_EXECUTABLE} setup.py -q sdist --dist-dir=${sdistOutDir} + WORKING_DIRECTORY ${generatedFileOutDir} + DEPENDS ${filesTarget} trt_packaging_requirements_installed + VERBATIM + ) + + set(sdistTarget tensorrt_frontend_sdist_${moduleName}) + + # Add a custom target for the wheel + add_custom_target( + ${sdistTarget} + ${ENABLE_ALL} + DEPENDS ${sdistOutputFile} + ) + + add_dependencies(tensorrt_frontend_sdist ${sdistTarget}) + + # Standalone wheels are only published to the internal tarfile as they are released separately. + install(FILES + ${sdistOutputFile} + DESTINATION python_standalone + COMPONENT internal + OPTIONAL + ) +endfunction() + +foreach(moduleName IN LISTS TRT_PYTHON_MODULE_NAMES) + buildFrontendWheel(${moduleName}) +endforeach() + +add_dependencies(tensorrt_python_wheels tensorrt_frontend_sdist) diff --git a/python/packaging/frontend_sdist/setup.py b/python/packaging/frontend_sdist/setup.py index d6232764..1dd59ddb 100644 --- a/python/packaging/frontend_sdist/setup.py +++ b/python/packaging/frontend_sdist/setup.py @@ -15,6 +15,7 @@ # limitations under the License. # +import contextlib import os import platform import subprocess @@ -26,6 +27,7 @@ from setuptools.command.install import install distribution_package_name = "##TENSORRT_MODULE##_cu##CUDA_MAJOR##" import_package_name = "##TENSORRT_MODULE##" +plugin_import_package_name = f"{import_package_name}.plugin" tensorrt_version = "##TENSORRT_PYTHON_VERSION##" tensorrt_submodules = [ "{}_libs=={}".format(distribution_package_name, tensorrt_version), @@ -89,10 +91,9 @@ class InstallCommand(install): def pip_config_list(): """Get the current pip config (env vars, config file, etc).""" - try: + with contextlib.suppress(subprocess.CalledProcessError, OSError, UnicodeDecodeError): return run_pip_command(["config", "list"], subprocess.check_output).decode() - except: - return "" + return "" def parent_command_line(): @@ -100,17 +101,13 @@ def parent_command_line(): pid = os.getppid() # try retrieval using psutil - try: + with contextlib.suppress(ImportError, ModuleNotFoundError, Exception): import psutil - return " ".join(psutil.Process(pid).cmdline()) - except: - pass # fall back to shell - try: + with contextlib.suppress(subprocess.CalledProcessError, OSError, UnicodeDecodeError): return subprocess.check_output(["ps", "-p", str(pid), "-o", "command", "--no-headers"]).decode() - except: - return "" + return "" # use pip-inside-pip hack only if the nvidia index is not set in the environment @@ -145,7 +142,7 @@ pip install tensorrt "Intended Audience :: Developers", "Programming Language :: Python :: 3", ], - packages=[import_package_name], + packages=[import_package_name, plugin_import_package_name], install_requires=install_requires, setup_requires=["wheel", "pip"], python_requires=">=3.6", # ref https://pypi.nvidia.com/tensorrt-bindings/ diff --git a/python/packaging/frontend_sdist/tensorrt/plugin/__init__.py b/python/packaging/frontend_sdist/tensorrt/plugin/__init__.py new file mode 100644 index 00000000..520a0203 --- /dev/null +++ b/python/packaging/frontend_sdist/tensorrt/plugin/__init__.py @@ -0,0 +1,18 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from tensorrt_bindings.plugin import * diff --git a/python/packaging/libs_wheel/CMakeLists.txt b/python/packaging/libs_wheel/CMakeLists.txt index f99b419a..8b91f14d 100644 --- a/python/packaging/libs_wheel/CMakeLists.txt +++ b/python/packaging/libs_wheel/CMakeLists.txt @@ -28,7 +28,7 @@ add_custom_target(tensorrt_libs_wheels ${ENABLE_ALL}) # Then, the relevant TRT libraries (e.g. libnvinfer.so) are copied into the same directory as the generated files. # Finally, the wheel is built by running setup.py with the appropriate arguments in the binary directory. # -# \param moduleName The module name to create the bindings for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". +# \param moduleName The module name to create the libs wheel for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". function(buildLibsWheel moduleName) set(filesTarget trt_wheel_files_libs_${moduleName}) @@ -53,8 +53,15 @@ function(buildLibsWheel moduleName) endif() elseif(${moduleName} STREQUAL "tensorrt" OR ${moduleName} STREQUAL "tensorrt_rtx") set(moduleLibraryTargets tensorrt) + get_all_fatbin_archs(KLIB_ARCHS KLIB_ARCHS_CROSS) if(NOT ${TRT_BUILD_WINML}) - list(APPEND moduleLibraryTargets tensorrt_builder_resource) + if(NOT ${TRT_BUILD_SPLIT_KLIB}) + list(APPEND moduleLibraryTargets tensorrt_builder_resource) + else() + foreach(ARCH IN LISTS KLIB_ARCHS) + list(APPEND moduleLibraryTargets tensorrt_builder_resource_${ARCH}) + endforeach() + endif() endif() if(${TRT_BUILD_PLUGINS}) list(APPEND moduleLibraryTargets tensorrt_plugins) @@ -63,7 +70,9 @@ function(buildLibsWheel moduleName) list(APPEND moduleLibraryTargets nvonnxparser) endif() if(NOT ${TRT_BUILD_SKIP_WIN_BUILDER_RESOURCE}) - list(APPEND moduleLibraryTargets tensorrt_builder_resource_win) + foreach(ARCH IN LISTS KLIB_ARCHS_CROSS) + list(APPEND moduleLibraryTargets tensorrt_builder_resource_win_${ARCH}) + endforeach() endif() else() message(FATAL_ERROR "Unknown module name: ${moduleName}. Expected 'tensorrt', 'tensorrt_dispatch', or 'tensorrt_lean'.") @@ -115,8 +124,9 @@ function(buildLibsWheel moduleName) ${pageSizeArgs} --build-dir=${CMAKE_LIBRARY_OUTPUT_DIRECTORY} # Since we give absolute paths, this is mostly unused, but we'll keep it for consistency. --output=${generatedFileOutDir}/${moduleName}_libs/ + --trim-version ${moduleLibraryPaths} - COMMENT "Copying TRT library files to wheel directory" + COMMENT "Copying TRT libraries for ${moduleName} to ${generatedFileOutDir}/${moduleName}_libs/" DEPENDS ${moduleLibraryTargets} VERBATIM ) @@ -167,10 +177,11 @@ function(buildLibsWheel moduleName) add_dependencies(tensorrt_libs_wheels ${wheelTarget}) + # Standalone wheels are only published to the internal tarfile as they are released separately. install(FILES ${wheelOutputFile} - DESTINATION wheels - COMPONENT release + DESTINATION python_standalone + COMPONENT internal OPTIONAL ) endfunction() diff --git a/python/packaging/libs_wheel/setup.py b/python/packaging/libs_wheel/setup.py index 054c565b..a83f590e 100644 --- a/python/packaging/libs_wheel/setup.py +++ b/python/packaging/libs_wheel/setup.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,7 +23,7 @@ import_package_name = "##TENSORRT_MODULE##_libs" def get_requirements(): - reqs = ["nvidia-cuda-runtime-cu##CUDA_MAJOR##"] + reqs = [f"cuda-toolkit[cudart] >=##CUDA_MAJOR##,<{##CUDA_MAJOR## + 1}"] return reqs diff --git a/python/packaging/metapackage/CMakeLists.txt b/python/packaging/metapackage/CMakeLists.txt new file mode 100644 index 00000000..15b03967 --- /dev/null +++ b/python/packaging/metapackage/CMakeLists.txt @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Adding a custom target to the all target requires passing the "ALL" keyword to the initial call. +# We use this variable so we can conditionally add the keyword based on the build options. +if(${TRT_BUILD_PYTHON_STANDALONE_WHEELS}) + set(ENABLE_ALL "ALL") +else() + set(ENABLE_ALL "") +endif() + +add_custom_target(tensorrt_metapackage_sdist ${ENABLE_ALL}) + +# \brief Creates a target named tensorrt_metapackage_sdist_${moduleName} which will build the metapackage SDist for that module. +# +# \param moduleName The module name to create the metapackage for. One of "tensorrt", "tensorrt_dispatch", or "tensorrt_lean". +function(buildMetapackageWheel moduleName) + set(filesTarget trt_wheel_files_metapackage_${moduleName}) + + set(wheelTemplateFiles + LICENSE.txt + setup.cfg + setup.py + ) + + # Expands all template files for the metapackage for the target module. The python version is unused (and set to 0). + # File paths starting with "tensorrt/" are expanded into "${moduleName}/". + processWheelTemplates(metapackage ${moduleName} 0 ${wheelTemplateFiles}) + + # Creates a new custom target, and makes trt_standalone_wheel_files depend on the new target. + add_custom_target(${filesTarget} DEPENDS ${generatedWheelFiles}) + + # Define the output directory for the wheel + set(sdistOutDir ${TRT_WHEEL_OUTPUT_DIR}/metapackage) + set(sdistOutputFile ${sdistOutDir}/${moduleName}-${TensorRT_PACKAGE_VERSION}.tar.gz) + + # Add a custom command to build the wheel + add_custom_command( + OUTPUT ${sdistOutputFile} + COMMAND ${Python3_EXECUTABLE} setup.py -q sdist --dist-dir=${sdistOutDir} + WORKING_DIRECTORY ${generatedFileOutDir} + DEPENDS ${filesTarget} trt_packaging_requirements_installed + VERBATIM + ) + + set(sdistTarget tensorrt_metapackage_sdist_${moduleName}) + + # Add a custom target for the wheel + add_custom_target( + ${sdistTarget} + ${ENABLE_ALL} + DEPENDS ${sdistOutputFile} + ) + + add_dependencies(tensorrt_metapackage_sdist ${sdistTarget}) + + # Standalone wheels are only published to the internal tarfile as they are released separately. + install(FILES + ${sdistOutputFile} + DESTINATION python_standalone + COMPONENT internal + OPTIONAL + ) +endfunction() + +foreach(moduleName IN LISTS TRT_PYTHON_MODULE_NAMES) + buildMetapackageWheel(${moduleName}) +endforeach() + +add_dependencies(tensorrt_python_wheels tensorrt_metapackage_sdist) diff --git a/python/packaging/metapackage/setup.py b/python/packaging/metapackage/setup.py index 1ea1138e..e4db6126 100644 --- a/python/packaging/metapackage/setup.py +++ b/python/packaging/metapackage/setup.py @@ -19,7 +19,6 @@ from setuptools import setup distribution_package_name = "##TENSORRT_MODULE##" -plugin_import_package_name = f"{distribution_package_name}.plugin" DISABLE_INTERNAL_PIP_FLAG = "NVIDIA_TENSORRT_DISABLE_INTERNAL_PIP" @@ -46,7 +45,6 @@ pip install tensorrt "Intended Audience :: Developers", "Programming Language :: Python :: 3", ], - packages=[plugin_import_package_name], install_requires=["##TENSORRT_MODULE##_cu##CUDA_MAJOR##==##TENSORRT_PYTHON_VERSION##"], include_package_data=True, zip_safe=True, diff --git a/python/packaging/requirements.txt b/python/packaging/requirements.txt index 1c7e37ee..25ff62ae 100644 --- a/python/packaging/requirements.txt +++ b/python/packaging/requirements.txt @@ -1,5 +1,4 @@ # Required for building wheel files -wheel==0.37.1 -setuptools~=75.3.2; python_version<"3.10" -setuptools~=80.9.0; python_version>="3.10" - +wheel==0.45.1 +setuptools~=75.3.2; python_version<"3.12" +setuptools~=80.9.0; python_version>="3.12" diff --git a/python/src/infer/pyAlgorithmSelector.cpp b/python/src/infer/pyAlgorithmSelector.cpp index 1d0255c8..14a76b27 100644 --- a/python/src/infer/pyAlgorithmSelector.cpp +++ b/python/src/infer/pyAlgorithmSelector.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -162,8 +162,7 @@ void bindAlgorithm(py::module& m) .def_property_readonly("name", &IAlgorithmContext::getName) .def("get_shape", lambdas::get_shape, "index"_a, IAlgorithmContextDoc::get_shape) .def_property_readonly("num_inputs", &IAlgorithmContext::getNbInputs) - .def_property_readonly("num_outputs", &IAlgorithmContext::getNbOutputs) - ; + .def_property_readonly("num_outputs", &IAlgorithmContext::getNbOutputs); // IAlgorithm py::class_>( diff --git a/python/src/infer/pyCore.cpp b/python/src/infer/pyCore.cpp index f787b136..36501ee0 100644 --- a/python/src/infer/pyCore.cpp +++ b/python/src/infer/pyCore.cpp @@ -175,6 +175,7 @@ static auto const reader_v2_read = [](IStreamReaderV2& self, void* destination, }; + // For ICudaEngine // TODO: Add slicing support? static auto const engine_getitem = [](ICudaEngine& self, int32_t pyIndex) { @@ -1310,6 +1311,26 @@ void bindCore(py::module& m) "step"_a) .def("phase_finish", &IProgressMonitor::phaseFinish, IProgressMonitorDoc::phase_finish, "phase_name"_a); + py::enum_(m, "ExecutionContextAllocationStrategy", py::arithmetic{}, + ExecutionContextAllocationStrategyDoc::descr, py::module_local()) + .value("STATIC", ExecutionContextAllocationStrategy::kSTATIC, ExecutionContextAllocationStrategyDoc::STATIC) + .value("ON_PROFILE_CHANGE", ExecutionContextAllocationStrategy::kON_PROFILE_CHANGE, + ExecutionContextAllocationStrategyDoc::ON_PROFILE_CHANGE) + .value("USER_MANAGED", ExecutionContextAllocationStrategy::kUSER_MANAGED, + ExecutionContextAllocationStrategyDoc::USER_MANAGED); + + + py::class_(m, "IRuntimeConfig", IRuntimeConfigDoc::descr, py::module_local()) + .def("set_execution_context_allocation_strategy", &IRuntimeConfig::setExecutionContextAllocationStrategy, + IRuntimeConfigDoc::set_execution_context_allocation_strategy, + py::arg("strategy") = ExecutionContextAllocationStrategy::kSTATIC, py::keep_alive<0, 1>{}, + py::call_guard{}) + .def("get_execution_context_allocation_strategy", &IRuntimeConfig::getExecutionContextAllocationStrategy, + IRuntimeConfigDoc::get_execution_context_allocation_strategy, py::keep_alive<0, 1>{}, + py::call_guard{}) + ; + + py::class_(m, "IExecutionContext", IExecutionContextDoc::descr, py::module_local()) .def("execute_v2", lambdas::execute_v2, "bindings"_a, IExecutionContextDoc::execute_v2, py::call_guard{}) @@ -1381,22 +1402,15 @@ void bindCore(py::module& m) .def_property("unfused_tensors_debug_state", &IExecutionContext::getUnfusedTensorsDebugState, &IExecutionContext::setUnfusedTensorsDebugState) .def("get_runtime_config", &IExecutionContext::getRuntimeConfig, IExecutionContextDoc::get_runtime_config, - py::keep_alive<1, 0>{}, py::call_guard{}); - - py::enum_(m, "ExecutionContextAllocationStrategy", py::arithmetic{}, - ExecutionContextAllocationStrategyDoc::descr, py::module_local()) - .value("STATIC", ExecutionContextAllocationStrategy::kSTATIC, ExecutionContextAllocationStrategyDoc::STATIC) - .value("ON_PROFILE_CHANGE", ExecutionContextAllocationStrategy::kON_PROFILE_CHANGE, - ExecutionContextAllocationStrategyDoc::ON_PROFILE_CHANGE) - .value("USER_MANAGED", ExecutionContextAllocationStrategy::kUSER_MANAGED, - ExecutionContextAllocationStrategyDoc::USER_MANAGED); - + py::keep_alive<1, 0>{}, py::call_guard{}) + ; py::enum_( m, "SerializationFlag", py::arithmetic{}, SerializationFlagDoc::descr, py::module_local()) .value("EXCLUDE_WEIGHTS", SerializationFlag::kEXCLUDE_WEIGHTS, SerializationFlagDoc::EXCLUDE_WEIGHTS) .value("EXCLUDE_LEAN_RUNTIME", SerializationFlag::kEXCLUDE_LEAN_RUNTIME, - SerializationFlagDoc::EXCLUDE_LEAN_RUNTIME); + SerializationFlagDoc::EXCLUDE_LEAN_RUNTIME) + .value("INCLUDE_REFIT", SerializationFlag::kINCLUDE_REFIT, SerializationFlagDoc::INCLUDE_REFIT); py::class_(m, "ISerializationConfig", ISerializationConfigDoc::descr, py::module_local()) .def_property("flags", &ISerializationConfig::getFlags, &lambdas::serialization_config_set_flags) @@ -1436,16 +1450,9 @@ void bindCore(py::module& m) .value("INPUT", TensorIOMode::kINPUT, TensorIOModeDoc::INPUT) .value("OUTPUT", TensorIOMode::kOUTPUT, TensorIOModeDoc::OUTPUT); - py::class_(m, "IRuntimeConfig", IRuntimeConfigDoc::descr, py::module_local()) - .def("set_execution_context_allocation_strategy", &IRuntimeConfig::setExecutionContextAllocationStrategy, - IRuntimeConfigDoc::set_execution_context_allocation_strategy, - py::arg("strategy") = ExecutionContextAllocationStrategy::kSTATIC, py::keep_alive<0, 1>{}, - py::call_guard{}) - .def("get_execution_context_allocation_strategy", &IRuntimeConfig::getExecutionContextAllocationStrategy, - IRuntimeConfigDoc::get_execution_context_allocation_strategy, py::keep_alive<0, 1>{}, - py::call_guard{}) - ; - + py::enum_(m, "EngineStat", py::arithmetic{}, EngineStatDoc::descr, py::module_local()) + .value("TOTAL_WEIGHTS_SIZE", EngineStat::kTOTAL_WEIGHTS_SIZE, EngineStatDoc::TOTAL_WEIGHTS_SIZE) + .value("STRIPPED_WEIGHTS_SIZE", EngineStat::kSTRIPPED_WEIGHTS_SIZE, EngineStatDoc::STRIPPED_WEIGHTS_SIZE); py::class_(m, "ICudaEngine", ICudaEngineDoc::descr, py::module_local()) .def("__getitem__", lambdas::engine_getitem) @@ -1491,36 +1498,42 @@ void bindCore(py::module& m) .def( "get_tensor_bytes_per_component", - [](ICudaEngine& self, std::string const& name) -> int32_t - { return self.getTensorBytesPerComponent(name.c_str()); }, + [](ICudaEngine& self, std::string const& name) -> int32_t { + return self.getTensorBytesPerComponent(name.c_str()); + }, "name"_a, ICudaEngineDoc::get_tensor_bytes_per_component) .def( "get_tensor_bytes_per_component", - [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> int32_t - { return self.getTensorBytesPerComponent(name.c_str(), profileIndex); }, + [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> int32_t { + return self.getTensorBytesPerComponent(name.c_str(), profileIndex); + }, "name"_a, "profile_index"_a, ICudaEngineDoc::get_tensor_bytes_per_component) .def( "get_tensor_components_per_element", - [](ICudaEngine& self, std::string const& name) -> int32_t - { return self.getTensorComponentsPerElement(name.c_str()); }, + [](ICudaEngine& self, std::string const& name) -> int32_t { + return self.getTensorComponentsPerElement(name.c_str()); + }, "name"_a, ICudaEngineDoc::get_tensor_components_per_element) .def( "get_tensor_components_per_element", - [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> int32_t - { return self.getTensorComponentsPerElement(name.c_str(), profileIndex); }, + [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> int32_t { + return self.getTensorComponentsPerElement(name.c_str(), profileIndex); + }, "name"_a, "profile_index"_a, ICudaEngineDoc::get_tensor_components_per_element) .def( "get_tensor_format", - [](ICudaEngine& self, std::string const& name) -> TensorFormat - { return self.getTensorFormat(name.c_str()); }, + [](ICudaEngine& self, std::string const& name) -> TensorFormat { + return self.getTensorFormat(name.c_str()); + }, "name"_a, ICudaEngineDoc::get_tensor_format) .def( "get_tensor_format", - [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> TensorFormat - { return self.getTensorFormat(name.c_str(), profileIndex); }, + [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> TensorFormat { + return self.getTensorFormat(name.c_str(), profileIndex); + }, "name"_a, "profile_index"_a, ICudaEngineDoc::get_tensor_format) .def( @@ -1538,13 +1551,15 @@ void bindCore(py::module& m) .def( "get_tensor_vectorized_dim", - [](ICudaEngine& self, std::string const& name) -> int32_t - { return self.getTensorVectorizedDim(name.c_str()); }, + [](ICudaEngine& self, std::string const& name) -> int32_t { + return self.getTensorVectorizedDim(name.c_str()); + }, "name"_a, ICudaEngineDoc::get_tensor_vectorized_dim) .def( "get_tensor_vectorized_dim", - [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> int32_t - { return self.getTensorVectorizedDim(name.c_str(), profileIndex); }, + [](ICudaEngine& self, std::string const& name, int32_t profileIndex) -> int32_t { + return self.getTensorVectorizedDim(name.c_str(), profileIndex); + }, "name"_a, "profile_index"_a, ICudaEngineDoc::get_tensor_vectorized_dim) .def("get_tensor_profile_shape", lambdas::get_tensor_profile_shape, "name"_a, "profile_index"_a, @@ -1586,7 +1601,10 @@ void bindCore(py::module& m) ICudaEngineDoc::create_execution_context, py::arg("runtime_config") = nullptr, py::keep_alive<0, 1>{}, py::call_guard{}) .def("create_runtime_config", &ICudaEngine::createRuntimeConfig, ICudaEngineDoc::create_runtime_config, - py::keep_alive<0, 1>{}, py::call_guard{}) + py::call_guard{}) + .def("get_engine_stat", &ICudaEngine::getEngineStat, ICudaEngineDoc::get_engine_stat, + py::arg("stat") = EngineStat::kTOTAL_WEIGHTS_SIZE, py::keep_alive<0, 1>{}, + py::call_guard{}) .def("__del__", &utils::doNothingDel); @@ -1679,7 +1697,8 @@ void bindCore(py::module& m) .value("MONITOR_MEMORY", BuilderFlag::kMONITOR_MEMORY, BuilderFlagDoc::MONITOR_MEMORY) .value("FP4", BuilderFlag::kFP4, BuilderFlagDoc::FP4) .value("DISTRIBUTIVE_INDEPENDENCE", BuilderFlag::kDISTRIBUTIVE_INDEPENDENCE, - BuilderFlagDoc::DISTRIBUTIVE_INDEPENDENCE); + BuilderFlagDoc::DISTRIBUTIVE_INDEPENDENCE) + ; py::enum_(m, "MemoryPoolType", MemoryPoolTypeDoc::descr, py::module_local()) .value("WORKSPACE", MemoryPoolType::kWORKSPACE, MemoryPoolTypeDoc::WORKSPACE) @@ -1739,18 +1758,29 @@ void bindCore(py::module& m) .def_static("parse", &lambdas::parseTimingCacheKey, "text"_a, TimingCacheKeyDoc::parse) .def("__str__", &lambdas::convertTimingCacheKeyToString, TimingCacheKeyDoc::convertTimingCacheKeyToString); + const char* const timing_cache_deprecation_str + = "Deprecated in TensorRT-RTX 1.2. Timing cache operations are no-ops in TensorRT-RTX."; + py::class_(m, "TimingCacheValue", TimingCacheValueDoc::descr, py::module_local()) .def(py::init()) - .def_property("tacticHash", &lambdas::getTacticHash, &lambdas::setTacticHash) - .def_property("timingMSec", &lambdas::getTimingMSec, &lambdas::setTimingMSec); + .def_property("tacticHash", utils::deprecateInTrtRtxOnly(&lambdas::getTacticHash, timing_cache_deprecation_str), + utils::deprecateInTrtRtxOnly(&lambdas::setTacticHash, timing_cache_deprecation_str)) + .def_property("timingMSec", utils::deprecateInTrtRtxOnly(&lambdas::getTimingMSec, timing_cache_deprecation_str), + utils::deprecateInTrtRtxOnly(&lambdas::setTimingMSec, timing_cache_deprecation_str)); py::class_(m, "ITimingCache", ITimingCacheDoc::descr, py::module_local()) - .def("serialize", &ITimingCache::serialize, ITimingCacheDoc::serialize) - .def("combine", &ITimingCache::combine, "input_cache"_a, "ignore_mismatch"_a, ITimingCacheDoc::combine) - .def("reset", &ITimingCache::reset, ITimingCacheDoc::reset) - .def("queryKeys", &lambdas::queryTimingCacheKeys, ITimingCacheDoc::queryKeys) - .def("query", &ITimingCache::query, "key"_a, ITimingCacheDoc::query) - .def("update", &ITimingCache::update, "key"_a, "value"_a, ITimingCacheDoc::update); + .def("serialize", utils::deprecateMemberInTrtRtxOnly(&ITimingCache::serialize, timing_cache_deprecation_str), + ITimingCacheDoc::serialize) + .def("combine", utils::deprecateMemberInTrtRtxOnly(&ITimingCache::combine, timing_cache_deprecation_str), + "input_cache"_a, "ignore_mismatch"_a, ITimingCacheDoc::combine) + .def("reset", utils::deprecateMemberInTrtRtxOnly(&ITimingCache::reset, timing_cache_deprecation_str), + ITimingCacheDoc::reset) + .def("queryKeys", utils::deprecateInTrtRtxOnly(&lambdas::queryTimingCacheKeys, timing_cache_deprecation_str), + ITimingCacheDoc::queryKeys) + .def("query", utils::deprecateMemberInTrtRtxOnly(&ITimingCache::query, timing_cache_deprecation_str), "key"_a, + ITimingCacheDoc::query) + .def("update", utils::deprecateMemberInTrtRtxOnly(&ITimingCache::update, timing_cache_deprecation_str), "key"_a, + "value"_a, ITimingCacheDoc::update); py::enum_( m, "TilingOptimizationLevel", TilingOptimizationLevelDoc::descr, py::module_local()) @@ -1821,9 +1851,12 @@ void bindCore(py::module& m) .def("get_tactic_sources", &IBuilderConfig::getTacticSources, IBuilderConfigDoc::get_tactic_sources) .def("create_timing_cache", lambdas::netconfig_create_timing_cache, "serialized_timing_cache"_a, IBuilderConfigDoc::create_timing_cache) - .def("set_timing_cache", &IBuilderConfig::setTimingCache, "cache"_a, "ignore_mismatch"_a, - IBuilderConfigDoc::set_timing_cache, py::keep_alive<1, 2>{}) - .def("get_timing_cache", &IBuilderConfig::getTimingCache, IBuilderConfigDoc::get_timing_cache) + .def("set_timing_cache", + utils::deprecateMemberInTrtRtxOnly(&IBuilderConfig::setTimingCache, timing_cache_deprecation_str), + "cache"_a, "ignore_mismatch"_a, IBuilderConfigDoc::set_timing_cache, py::keep_alive<1, 2>{}) + .def("get_timing_cache", + utils::deprecateMemberInTrtRtxOnly(&IBuilderConfig::getTimingCache, timing_cache_deprecation_str), + IBuilderConfigDoc::get_timing_cache) .def("set_preview_feature", &IBuilderConfig::setPreviewFeature, "feature"_a, "enable"_a, IBuilderConfigDoc::set_preview_feature) .def("get_preview_feature", &IBuilderConfig::getPreviewFeature, "feature"_a, diff --git a/python/src/infer/pyGraph.cpp b/python/src/infer/pyGraph.cpp index 582a2144..844f1ae2 100644 --- a/python/src/infer/pyGraph.cpp +++ b/python/src/infer/pyGraph.cpp @@ -35,13 +35,13 @@ namespace tensorrt // Long lambda functions should go here rather than being inlined into the bindings (1 liners are OK). namespace lambdas { - Weights optionalWeights(Weights* weights) + Weights optionalWeights(Weights* weights, DataType dtype) { if (weights) { return *weights; } - return Weights{DataType::kFLOAT, nullptr, 0}; + return Weights{dtype, nullptr, 0}; } static const auto get_dynamic_range = [] (ITensor const& self) -> py::object { @@ -148,7 +148,7 @@ namespace tensorrt static const auto add_convolution_nd = [](INetworkDefinition& self, ITensor& input, int32_t numOutputMaps, Dims kernelSize, Weights kernel, Weights* bias) { - return self.addConvolutionNd(input, numOutputMaps, kernelSize, kernel, optionalWeights(bias)); + return self.addConvolutionNd(input, numOutputMaps, kernelSize, kernel, optionalWeights(bias, input.getType())); }; IGridSampleLayer* add_grid_sample(INetworkDefinition& self, ITensor& input, ITensor& grid) @@ -158,17 +158,17 @@ namespace tensorrt static const auto add_scale = [](INetworkDefinition& self, ITensor& input, ScaleMode mode, Weights* shift, Weights* scale, Weights* power) { - return self.addScale(input, mode, optionalWeights(shift), optionalWeights(scale), optionalWeights(power)); + return self.addScale(input, mode, optionalWeights(shift, input.getType()), optionalWeights(scale, input.getType()), optionalWeights(power, input.getType())); }; static const auto add_scale_nd = [](INetworkDefinition& self, ITensor& input, ScaleMode mode, Weights* shift, Weights* scale, Weights* power, int32_t channelAxis) { - return self.addScaleNd(input, mode, optionalWeights(shift), optionalWeights(scale), optionalWeights(power), channelAxis); + return self.addScaleNd(input, mode, optionalWeights(shift, input.getType()), optionalWeights(scale, input.getType()), optionalWeights(power, input.getType()), channelAxis); }; static const auto add_deconvolution_nd = [](INetworkDefinition& self, ITensor& input, int32_t numOutputMaps, Dims kernelSize, Weights kernel, Weights* bias) { - return self.addDeconvolutionNd(input, numOutputMaps, kernelSize, kernel, optionalWeights(bias)); + return self.addDeconvolutionNd(input, numOutputMaps, kernelSize, kernel, optionalWeights(bias, input.getType())); }; static const auto add_einsum = [] (INetworkDefinition& self, const std::vector& inputs, const char* equation) { @@ -271,6 +271,12 @@ namespace tensorrt PY_ASSERT_RUNTIME_ERROR(status, "Failed to set CumulativeLayer's CumulativeOperation"); } + static void attention_set_operation(IAttention& self, AttentionNormalizationOp op) + { + bool const status = self.setNormalizationOperation(op); + PY_ASSERT_RUNTIME_ERROR(status, "Failed to set Attention's AttentionNormalizationOp"); + } + void bindGraph(py::module& m) { // Bind to a Python enum called LayerType. @@ -326,6 +332,8 @@ namespace tensorrt .value("UNSQUEEZE", LayerType::kUNSQUEEZE, LayerTypeDoc::UNSQUEEZE) .value("CUMULATIVE", LayerType::kCUMULATIVE, LayerTypeDoc::CUMULATIVE) .value("DYNAMIC_QUANTIZE", LayerType::kDYNAMIC_QUANTIZE, LayerTypeDoc::DYNAMIC_QUANTIZE) + .value("ATTENTION_INPUT", LayerType::kATTENTION_INPUT, LayerTypeDoc::ATTENTION_INPUT) + .value("ATTENTION_OUTPUT", LayerType::kATTENTION_OUTPUT, LayerTypeDoc::ATTENTION_OUTPUT) ; // LayerType py::enum_(m, "TensorFormat", TensorFormatDoc::descr, py::arithmetic{}, py::module_local()) @@ -670,6 +678,7 @@ namespace tensorrt .def_property("k", &ITopKLayer::getK, &ITopKLayer::setK) .def_property("axes", &ITopKLayer::getReduceAxes, &ITopKLayer::setReduceAxes) .def("set_input", &ITopKLayer::setInput, "index"_a, "tensor"_a, ITopKLayerDoc::set_input) + .def_property("indices_type", &ITopKLayer::getIndicesType, &ITopKLayer::setIndicesType) ; py::enum_(m, "MatrixOperation", MatrixOperationDoc::descr, py::module_local()) @@ -793,6 +802,7 @@ namespace tensorrt .def_property("bounding_box_format", &INMSLayer::getBoundingBoxFormat, &INMSLayer::setBoundingBoxFormat) .def_property("topk_box_limit", &INMSLayer::getTopKBoxLimit, &INMSLayer::setTopKBoxLimit) .def("set_input", &INMSLayer::setInput, "index"_a, "tensor"_a, INMSLayerDoc::set_input) + .def_property("indices_type", &INMSLayer::getIndicesType, &INMSLayer::setIndicesType) ; py::enum_(m, "FillOperation", FillOperationDoc::descr, py::module_local()) @@ -840,6 +850,7 @@ namespace tensorrt ; py::class_>(m, "INonZeroLayer", INonZeroLayerDoc::descr, py::module_local()) + .def_property("indices_type", &INonZeroLayer::getIndicesType, &INonZeroLayer::setIndicesType) ; py::class_>(m, "IReverseSequenceLayer", IReverseSequenceLayerDoc::descr, py::module_local()) @@ -871,6 +882,36 @@ namespace tensorrt .def_property("reverse", &ICumulativeLayer::getReverse, &ICumulativeLayer::setReverse) ; + py::enum_(m, "AttentionNormalizationOp", AttentionNormalizationOpDoc::descr, py::module_local()) + .value("NONE", AttentionNormalizationOp::kNONE, AttentionNormalizationOpDoc::NONE) + .value("SOFTMAX", AttentionNormalizationOp::kSOFTMAX, AttentionNormalizationOpDoc::SOFTMAX) + ; + + py::class_>(m, "IAttention", IAttentionDoc::descr, py::module_local()) + .def_property("mask", &IAttention::getMask, &IAttention::setMask) + .def_property("norm_op", &IAttention::getNormalizationOperation, &attention_set_operation) + .def_property("decomposable", &IAttention::getDecomposable, &IAttention::setDecomposable) + .def_property("causal", &IAttention::getCausal, &IAttention::setCausal) + .def_property("name", &IAttention::getName, &IAttention::setName) + .def_property("normalization_quantize_scale", &IAttention::getNormalizationQuantizeScale, &IAttention::setNormalizationQuantizeScale) + .def_property("normalization_quantize_to_type", &IAttention::getNormalizationQuantizeToType, &IAttention::setNormalizationQuantizeToType) + .def_property_readonly("num_inputs", &IAttention::getNbInputs) + .def_property_readonly("num_outputs", &IAttention::getNbOutputs) + .def("set_input", &IAttention::setInput, "index"_a, "tensor"_a, IAttentionDoc::set_input) + .def("get_input", &IAttention::getInput, "index"_a, IAttentionDoc::get_input) + .def("get_output", &IAttention::getOutput, "index"_a, IAttentionDoc::get_output) + ; + + py::class_>(m, "IAttentionBoundaryLayer", IAttentionBoundaryLayerDoc::descr, py::module_local()) + .def_property_readonly("attention", &IAttentionBoundaryLayer::getAttention, py::return_value_policy::reference_internal) + ; + + py::class_>(m, "IAttentionInputLayer", IAttentionInputLayerDoc::descr, py::module_local()) + ; + + py::class_>(m, "IAttentionOutputLayer", IAttentionOutputLayerDoc::descr, py::module_local()) + ; + // Weights must be kept alive for the duration of the network. py::keep_alive is critical here! // Additionally, we use reference_internal so that pybind11 does not free layers when they go out of scope. py::class_(m, "INetworkDefinition", INetworkDefinitionDoc::descr, py::module_local()) @@ -921,7 +962,9 @@ namespace tensorrt INetworkDefinitionDoc::add_slice, py::return_value_policy::reference_internal) .def("add_reduce", &INetworkDefinition::addReduce, "input"_a, "op"_a, "axes"_a, "keep_dims"_a, INetworkDefinitionDoc::add_reduce, py::return_value_policy::reference_internal) - .def("add_topk", &INetworkDefinition::addTopK, "input"_a, "op"_a, "k"_a, "axes"_a, + .def("add_topk", static_cast(&INetworkDefinition::addTopK), "input"_a, "op"_a, "k"_a, "axes"_a, + INetworkDefinitionDoc::add_topk, py::return_value_policy::reference_internal) + .def("add_topk", static_cast(&INetworkDefinition::addTopK), "input"_a, "op"_a, "k"_a, "axes"_a, "indices_type"_a, INetworkDefinitionDoc::add_topk, py::return_value_policy::reference_internal) .def("add_gather", &INetworkDefinition::addGather, "input"_a, "indices"_a, "axis"_a, INetworkDefinitionDoc::add_gather, py::return_value_policy::reference_internal) @@ -965,8 +1008,10 @@ namespace tensorrt py::return_value_policy::reference_internal) .def("add_grid_sample", &INetworkDefinition::addGridSample, "input"_a, "grid"_a, INetworkDefinitionDoc::add_grid_sample, py::return_value_policy::reference_internal) - .def("add_nms", &INetworkDefinition::addNMS, "boxes"_a, - "scores"_a, "max_output_boxes_per_class"_a, INetworkDefinitionDoc::add_nms, py::return_value_policy::reference_internal) + .def("add_nms", static_cast(&INetworkDefinition::addNMS), "boxes"_a, "scores"_a, "max_output_boxes_per_class"_a, + INetworkDefinitionDoc::add_nms, py::return_value_policy::reference_internal) + .def("add_nms", static_cast(&INetworkDefinition::addNMS), "boxes"_a, "scores"_a, "max_output_boxes_per_class"_a, "indices_type"_a, + INetworkDefinitionDoc::add_nms, py::return_value_policy::reference_internal) .def("add_fill", static_cast(&INetworkDefinition::addFill), "shape"_a, "op"_a, "output_type"_a, INetworkDefinitionDoc::add_fill) .def("add_fill", static_cast(&INetworkDefinition::addFill), "shape"_a, "op"_a, INetworkDefinitionDoc::add_fill) .def("add_quantize", static_cast(&INetworkDefinition::addQuantize), "input"_a, "scale"_a, @@ -985,14 +1030,17 @@ namespace tensorrt py::return_value_policy::reference_internal) .def("add_one_hot", &INetworkDefinition::addOneHot, "indices"_a, "values"_a, "depth"_a, "axis"_a, INetworkDefinitionDoc::add_one_hot, py::return_value_policy::reference_internal) - .def("add_non_zero", &INetworkDefinition::addNonZero, "input"_a, INetworkDefinitionDoc::add_non_zero, - py::return_value_policy::reference_internal) + .def("add_non_zero", static_cast(&INetworkDefinition::addNonZero), "input"_a, + INetworkDefinitionDoc::add_non_zero, py::return_value_policy::reference_internal) + .def("add_non_zero", static_cast(&INetworkDefinition::addNonZero), "input"_a, "indices_type"_a, + INetworkDefinitionDoc::add_non_zero, py::return_value_policy::reference_internal) .def("add_reverse_sequence", &INetworkDefinition::addReverseSequence, "input"_a, "sequence_lens"_a, INetworkDefinitionDoc::add_reverse_sequence, py::return_value_policy::reference_internal) .def("add_normalization", &INetworkDefinition::addNormalization, "input"_a, "scale"_a, "bias"_a, "axesMask"_a, INetworkDefinitionDoc::add_normalization, py::return_value_policy::reference_internal) .def("add_cumulative", &INetworkDefinition::addCumulative, "input"_a, "axis"_a, "op"_a, "exclusive"_a, "reverse"_a, INetworkDefinitionDoc::add_cumulative, py::return_value_policy::reference_internal) + .def("add_attention", &INetworkDefinition::addAttention, "query"_a, "key"_a, "value"_a, "norm_op"_a, "causal"_a, INetworkDefinitionDoc::add_attention, py::return_value_policy::reference_internal) .def("remove_tensor", &INetworkDefinition::removeTensor, "tensor"_a, INetworkDefinitionDoc::remove_tensor) .def("unmark_output", &INetworkDefinition::unmarkOutput, "tensor"_a, INetworkDefinitionDoc::unmark_output) .def("mark_output_for_shapes", &INetworkDefinition::markOutputForShapes, "tensor"_a, INetworkDefinitionDoc::mark_output_for_shapes) diff --git a/python/src/infer/pyPlugin.cpp b/python/src/infer/pyPlugin.cpp index 1a181d0f..184d4d76 100644 --- a/python/src/infer/pyPlugin.cpp +++ b/python/src/infer/pyPlugin.cpp @@ -2000,7 +2000,8 @@ public: { auto result = pyResult.cast>(); - *kernelName = std::get<0>(result).c_str(); + mKernelName = std::get<0>(result); + *kernelName = mKernelName.c_str(); mCompiledKernel = std::get<1>(result); py::buffer_info buffer(py::buffer(mCompiledKernel).request()); *compiledKernel = static_cast(buffer.ptr); @@ -2081,6 +2082,7 @@ public: private: py::bytes mCompiledKernel; + std::string mKernelName; }; class PyIPluginV3QuickRuntimeImpl : public IPluginV3QuickRuntime diff --git a/quickstart/IntroNotebooks/2. Using PyTorch through ONNX.ipynb b/quickstart/IntroNotebooks/2. Using PyTorch through ONNX.ipynb index 895366ba..c27a9d34 100644 --- a/quickstart/IntroNotebooks/2. Using PyTorch through ONNX.ipynb +++ b/quickstart/IntroNotebooks/2. Using PyTorch through ONNX.ipynb @@ -233,7 +233,7 @@ "\n", "with torch.no_grad():\n", " preds = np.array(resnet50_gpu_half(input_half).cpu()) # Warm Up\n", - " \n", + "\n", "preds.shape" ] }, @@ -426,7 +426,7 @@ "if USE_FP16:\n", " !trtexec --onnx=resnet50_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --fp16\n", "else:\n", - " !trtexec --onnx=resnet50_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt " + " !trtexec --onnx=resnet50_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt" ] }, { @@ -454,14 +454,14 @@ "%%time\n", "\n", "import tensorrt as trt\n", - "from cuda import cudart\n", + "from cuda.bindings import runtime as cudart\n", "import numpy as np\n", "\n", "err, = cudart.cudaSetDevice(0)\n", "assert err == cudart.cudaError_t.cudaSuccess\n", "\n", "f = open(\"resnet_engine_pytorch.trt\", \"rb\")\n", - "runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING)) \n", + "runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))\n", "\n", "engine = runtime.deserialize_cuda_engine(f.read())\n", "context = engine.create_execution_context()" @@ -483,7 +483,7 @@ "import numpy as np\n", "\n", "# need to set input and output precisions to FP16 to fully enable it\n", - "output = np.empty([BATCH_SIZE, 1000], dtype = target_dtype) \n", + "output = np.empty([BATCH_SIZE, 1000], dtype = target_dtype)\n", "\n", "# allocate device memory\n", "err, d_input = cudart.cudaMalloc(input_batch.nbytes)\n", @@ -518,22 +518,22 @@ "source": [ "def predict(batch): # result gets copied into output\n", " # transfer input data to device\n", - " err, = cudart.cudaMemcpyAsync(d_input, batch.ctypes.data, batch.nbytes, \n", + " err, = cudart.cudaMemcpyAsync(d_input, batch.ctypes.data, batch.nbytes,\n", " cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)\n", " assert err == cudart.cudaError_t.cudaSuccess\n", - " \n", + "\n", " # execute model\n", " context.execute_async_v3(stream)\n", - " \n", + "\n", " # transfer predictions back\n", " err, = cudart.cudaMemcpyAsync(output.ctypes.data, d_output, output.nbytes,\n", " cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)\n", " assert err == cudart.cudaError_t.cudaSuccess\n", - " \n", + "\n", " # synchronize stream\n", " err, = cudart.cudaStreamSynchronize(stream)\n", " assert err == cudart.cudaError_t.cudaSuccess\n", - " \n", + "\n", " return output" ] }, @@ -655,5 +655,7 @@ "nbformat": 4, "nbformat_minor": 4 } - } + }, + "nbformat": 4, + "nbformat_minor": 5 } diff --git a/quickstart/IntroNotebooks/onnx_helper.py b/quickstart/IntroNotebooks/onnx_helper.py index afa950cc..4b1875be 100644 --- a/quickstart/IntroNotebooks/onnx_helper.py +++ b/quickstart/IntroNotebooks/onnx_helper.py @@ -19,7 +19,7 @@ import numpy as np import tensorrt as trt import weakref -from cuda import cudart +from cuda.bindings import runtime as cudart def _cleanup_cuda_resources(d_input, d_output, stream): diff --git a/quickstart/SemanticSegmentation/tutorial-runtime.ipynb b/quickstart/SemanticSegmentation/tutorial-runtime.ipynb index 1f5e4df4..5ccb1102 100644 --- a/quickstart/SemanticSegmentation/tutorial-runtime.ipynb +++ b/quickstart/SemanticSegmentation/tutorial-runtime.ipynb @@ -87,7 +87,7 @@ "import numpy as np\n", "import os\n", "import ctypes\n", - "from cuda import cudart\n", + "from cuda.bindings import runtime as cudart\n", "import tensorrt as trt\n", "\n", "import matplotlib.pyplot as plt\n", @@ -207,13 +207,13 @@ " with engine.create_execution_context() as context:\n", " input_buffers = {}\n", " input_memories = {}\n", - " \n", + "\n", " # Allocate host and device buffers\n", " tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]\n", " for tensor in tensor_names:\n", " size = trt.volume(context.get_tensor_shape(tensor))\n", " dtype = trt.nptype(engine.get_tensor_dtype(tensor))\n", - " \n", + "\n", " if engine.get_tensor_mode(tensor) == trt.TensorIOMode.INPUT:\n", " context.set_input_shape(tensor, (1, 3, image_height, image_width))\n", " input_buffers[tensor] = np.ascontiguousarray(input_image)\n", @@ -225,24 +225,24 @@ " assert err == cudart.cudaError_t.cudaSuccess\n", " pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype))\n", " output_buffer = np.ctypeslib.as_array(ctypes.cast(output_buffer_ptr, pointer_type), (size,))\n", - " \n", + "\n", " err, output_memory = cudart.cudaMalloc(output_buffer.nbytes)\n", " assert err == cudart.cudaError_t.cudaSuccess\n", " context.set_tensor_address(tensor, output_memory)\n", "\n", " err, stream = cudart.cudaStreamCreate()\n", " assert err == cudart.cudaError_t.cudaSuccess\n", - " \n", + "\n", " # Transfer input data to the GPU for all input tensors\n", " for tensor_name, input_buffer in input_buffers.items():\n", " input_memory = input_memories[tensor_name]\n", " err, = cudart.cudaMemcpyAsync(input_memory, input_buffer.ctypes.data, input_buffer.nbytes,\n", " cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream)\n", " assert err == cudart.cudaError_t.cudaSuccess\n", - " \n", + "\n", " # Run inference\n", " context.execute_async_v3(stream)\n", - " \n", + "\n", " # Transfer prediction output from the GPU.\n", " err, = cudart.cudaMemcpyAsync(output_buffer.ctypes.data, output_memory, output_buffer.nbytes,\n", " cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream)\n", @@ -263,8 +263,7 @@ " cudart.cudaFree(input_memory)\n", " cudart.cudaFree(output_memory)\n", " cudart.cudaFreeHost(output_buffer_ptr)\n", - " cudart.cudaStreamDestroy(stream)\n", - " " + " cudart.cudaStreamDestroy(stream)\n" ] }, { diff --git a/samples/CMakeLists.txt b/samples/CMakeLists.txt index 45871b32..d2b90432 100644 --- a/samples/CMakeLists.txt +++ b/samples/CMakeLists.txt @@ -22,6 +22,21 @@ endif() if(${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) +# Setup aliases for ease of swapping between static/dynamic TRT. +if(${TRT_BUILD_SAMPLES_LINK_STATIC_TRT}) + add_library(TRT_SAMPLES::tensorrt ALIAS tensorrt_static) + add_library(TRT_SAMPLES::onnxparser ALIAS nvonnxparser_static) + + # Limit linking jobs to reduce memory usage. + # libnvinfer_static.a is multiple GB, so linking it into everything in parallel is a recipe for OOM kills. + # Observation is peak ~6GB per thread on x86 Linux. + set_property(GLOBAL PROPERTY JOB_POOLS three_jobs=3) + set(CMAKE_JOB_POOL_LINK three_jobs) +else() + add_library(TRT_SAMPLES::tensorrt ALIAS tensorrt) + add_library(TRT_SAMPLES::onnxparser ALIAS nvonnxparser) +endif() + # OSS samples need the ONNX parser path to be included when each sample is built add_subdirectory(common) @@ -48,9 +63,7 @@ if(${TRT_BUILD_SAMPLES}) # Public (OSS) Samples add_sample( sampleCharRNN - sampleDynamicReshape sampleEditableTimingCache - sampleINT8API sampleIOFormats sampleNamedDimensions sampleOnnxMNIST @@ -58,7 +71,11 @@ if(${TRT_BUILD_SAMPLES}) ) if (NOT ${TRT_BUILD_WINML}) - list(APPEND SAMPLE_FOLDERS sampleNonZeroPlugin) + add_sample( + sampleDynamicReshape + sampleINT8API + sampleNonZeroPlugin + ) endif() # This sample needs to link against nvinfer_plugin. @@ -94,6 +111,3 @@ foreach(SAMPLE_ITER ${OPENSOURCE_SAMPLES_LIST}) add_subdirectory(${SAMPLE_ITER}) endforeach(SAMPLE_ITER) endif() - - - diff --git a/samples/README.md b/samples/README.md index 4e969355..83c95245 100644 --- a/samples/README.md +++ b/samples/README.md @@ -33,7 +33,29 @@ | Sample | Language | Format | Description | |---|---|---|---| | [detectron2](python/detectron2) | Python | ONNX | Support for Detectron 2 Mask R-CNN R50-FPN 3x model in TensorRT | -| [efficientdet](python/efficientdet) | Python | ONNX | EfficientDet Object Detection with TensorRT | -| [efficientnet](python/efficientnet) | Python | ONNX | EfficientNet V1 and V2 Classification with TensorRT | -| [tensorflow_object_detection_api](python/tensorflow_object_detection_api) | Python | ONNX | TensorFlow Object Detection API Models in TensorRT | +| [[DEPRECATED] efficientdet](python/efficientdet) | Python | ONNX | EfficientDet Object Detection with TensorRT | +| [[DEPRECATED] tensorflow_object_detection_api](python/tensorflow_object_detection_api) | Python | ONNX | TensorFlow Object Detection API Models in TensorRT | | [[DEPRECATED] yolov3_onnx](python/yolov3_onnx) | Python | ONNX | Object Detection Using YOLOv3 With TensorRT ONNX Backend | + +## Preparing sample data + +Many samples require the TensorRT sample data package. If not already mounted under `/usr/src/tensorrt/data` (NVIDIA NGC containers), download and extract it: + +1. Download the sample data from [TensorRT GitHub Releases](https://github.com/NVIDIA/TensorRT/releases). + +2. Extract and set up the data: + ```bash + unzip tensorrt_sample_data_xxx.zip + mkdir -p /usr/src/tensorrt/data + cp -r tensorrt_sample_data_*/* /usr/src/tensorrt/data/ + export TRT_DATADIR=/usr/src/tensorrt/data + ``` + +After extraction, the data directory structure should be: +``` +$TRT_DATADIR/ +├── char-rnn/ +├── int8_api/ +├── mnist/ +└── resnet50/ +``` diff --git a/samples/common/CMakeLists.txt b/samples/common/CMakeLists.txt index efdd537d..6f1b0ed0 100644 --- a/samples/common/CMakeLists.txt +++ b/samples/common/CMakeLists.txt @@ -10,9 +10,6 @@ add_library(trt_samples_common STATIC) - -set(SAFE_EXECUTOR_INCLUDE_PATH "${TRT_BUILD_CONAN_DEP_PATH}/myelin_include/lwe") - target_sources(trt_samples_common PRIVATE argsParser.h BatchStream.h @@ -56,35 +53,28 @@ if (MSVC) ) endif() -if(${TRT_BUILD_SAMPLES_LINK_STATIC_TRT}) - target_link_libraries(trt_samples_common PUBLIC - tensorrt_static - nvonnxparser_static - ) -else() - target_link_libraries(trt_samples_common PUBLIC - tensorrt - nvonnxparser - ) -endif() +target_include_directories(trt_samples_common PUBLIC + ${CMAKE_CURRENT_LIST_DIR} +) + target_link_libraries(trt_samples_common PUBLIC trt_global_definitions trt_shared + Threads::Threads dl + TRT::cudart + $ # Each sample individually must determine its linkage to TRT. + TRT_SAMPLES::onnxparser ) - -if(NOT MSVC) +# For statically-linked samples, we need to upgrade the link to always link TRT rather than letting the samples decide. +if(${TRT_BUILD_SAMPLES_LINK_STATIC_TRT}) target_link_libraries(trt_samples_common PUBLIC - Threads::Threads + $ # Has to be whole archive so we keep the builder resources correctly. ) endif() -target_include_directories(trt_samples_common PUBLIC - ${CMAKE_CURRENT_LIST_DIR} - ${SAFE_EXECUTOR_INCLUDE_PATH} -) if(${TRT_BUILD_ENABLE_DLA}) target_link_libraries(trt_samples_common PUBLIC NVDLA::compiler) diff --git a/samples/common/common.h b/samples/common/common.h index 2f9f57e4..7ff2ed01 100644 --- a/samples/common/common.h +++ b/samples/common/common.h @@ -450,7 +450,7 @@ inline float getMaxValue(const float* buffer, int64_t size) // All tensors in a network must have a dynamic range specified if a calibrator is not used. // This function is just a utility to globally fill in missing scales and zero-points for the entire network. // -// If a tensor does not have a dyanamic range set, it is assigned inRange or outRange as follows: +// If a tensor does not have a dynamic range set, it is assigned inRange or outRange as follows: // // * If the tensor is the input to a layer or output of a pooling node, its dynamic range is derived from inRange. // * Otherwise its dynamic range is derived from outRange. @@ -478,7 +478,7 @@ inline void setAllDynamicRanges(nvinfer1::INetworkDefinition* network, float inR } // Ensure that all layer outputs have a scale. - // Tensors that are also inputs to layers are ingored here + // Tensors that are also inputs to layers are ignored here // since the previous loop nest assigned scales to them. for (int i = 0; i < network->getNbLayers(); i++) { @@ -918,11 +918,11 @@ inline int getW(const nvinfer1::Dims& d) class DynamicLibrary { public: - explicit DynamicLibrary(std::string const& name) - : mLibName{name} + explicit DynamicLibrary(std::string name) + : mLibName{std::move(name)} { #if defined(_WIN32) - mHandle = LoadLibraryA(name.c_str()); + mHandle = LoadLibraryA(mLibName.c_str()); #else // defined(_WIN32) int32_t flags{RTLD_LAZY}; #if ENABLE_ASAN @@ -933,7 +933,7 @@ public: flags |= RTLD_NODELETE; #endif // ENABLE_ASAN - mHandle = dlopen(name.c_str(), flags); + mHandle = dlopen(mLibName.c_str(), flags); #endif // defined(_WIN32) if (mHandle == nullptr) @@ -942,7 +942,7 @@ public: #if !defined(_WIN32) errorStr = std::string{" due to "} + std::string{dlerror()}; #endif - throw std::runtime_error("Unable to open library: " + name + errorStr); + throw std::runtime_error("Unable to open library: " + mLibName + errorStr); } } @@ -997,10 +997,9 @@ private: void* mHandle{}; //!< Handle to the DynamicLibrary }; -inline std::unique_ptr loadLibrary(std::string const& path) +[[nodiscard]] inline std::unique_ptr loadLibrary(std::string name) { - // make_unique not available until C++14 - we still need to support C++11 builds. - return std::unique_ptr(new DynamicLibrary{path}); + return std::make_unique(std::move(name)); } //! Represents the compute capability of a device. @@ -1029,7 +1028,7 @@ struct ComputeCapability } }; -inline int32_t getSMVersion() +inline int32_t getSmVersion() { int32_t deviceIndex = 0; CHECK(cudaGetDevice(&deviceIndex)); @@ -1038,9 +1037,9 @@ inline int32_t getSMVersion() return ((cc.major << 8) | cc.minor); } -inline bool isSMSafe() +inline bool isSmSafe() { - const int32_t smVersion = getSMVersion(); + const int32_t smVersion = getSmVersion(); return smVersion == 0x0705 || smVersion == 0x0800 || smVersion == 0x0806 || smVersion == 0x0807; } diff --git a/samples/common/parserOnnxConfig.h b/samples/common/parserOnnxConfig.h index 67ee6c71..53da93b3 100644 --- a/samples/common/parserOnnxConfig.h +++ b/samples/common/parserOnnxConfig.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,6 +18,7 @@ #ifndef PARSER_ONNX_CONFIG_H #define PARSER_ONNX_CONFIG_H + #include #include #include @@ -142,4 +143,4 @@ public: } }; // class ParserOnnxConfig -#endif +#endif // PARSER_ONNX_CONFIG_H diff --git a/samples/common/sampleDevice.h b/samples/common/sampleDevice.h index 6a5000bd..323379f5 100644 --- a/samples/common/sampleDevice.h +++ b/samples/common/sampleDevice.h @@ -124,7 +124,7 @@ public: CHECK(cudaEventRecord(mEvent, stream.get())); } - void synchronize() + void synchronize() const { CHECK(cudaEventSynchronize(mEvent)); } @@ -132,6 +132,10 @@ public: // Returns time elapsed time in milliseconds float operator-(const TrtCudaEvent& e) const { + // Synchronize both events to ensure they have completed before calculating elapsed time + synchronize(); + e.synchronize(); + float time{0}; CHECK(cudaEventElapsedTime(&time, e.get(), get())); return time; diff --git a/samples/common/sampleEngines.cpp b/samples/common/sampleEngines.cpp index c26f2b71..e54c0fa5 100644 --- a/samples/common/sampleEngines.cpp +++ b/samples/common/sampleEngines.cpp @@ -219,6 +219,7 @@ void setTensorScalesFromCalibration(nvinfer1::INetworkDefinition& network, std:: } } + //! //! \brief Generate a network definition for a given model //! @@ -605,6 +606,26 @@ void setLayerDeviceTypes( } } +void setDecomposables(INetworkDefinition& network, DecomposableAttentions const& decomposableAttentions) +{ + for (int32_t layerIdx = 0; layerIdx < network.getNbLayers(); ++layerIdx) + { + auto* layer = network.getLayer(layerIdx); + if (layer->getType() == LayerType::kATTENTION_INPUT) + { + auto* attention = static_cast(layer)->getAttention(); + auto const attentionName = attention->getName(); + auto match = findPlausible(decomposableAttentions, attentionName); + if (match != decomposableAttentions.end()) + { + attention->setDecomposable(match->second); + sample::gLogInfo << "Set attention " << attentionName + << " to decomposable = " << ((match->second) ? "true" : "false") << std::endl; + } + } + } +} + void markDebugTensors(INetworkDefinition& network, StringSet const& debugTensors) { for (int64_t inputIndex = 0; inputIndex < network.getNbInputs(); ++inputIndex) @@ -1141,6 +1162,11 @@ bool setupNetworkAndConfig(BuildOptions const& build, SystemOptions const& sys, setLayerDeviceTypes(network, config, build.layerDeviceTypes); } + if (!build.decomposableAttentions.empty()) + { + setDecomposables(network, build.decomposableAttentions); + } + if (!build.debugTensors.empty()) { markDebugTensors(network, build.debugTensors); @@ -1223,7 +1249,7 @@ bool setupNetworkAndConfig(BuildOptions const& build, SystemOptions const& sys, if (!build.remoteAutoTuningConfig.empty()) { SMP_RETVAL_IF_FALSE(config.setRemoteAutoTuningConfig(build.remoteAutoTuningConfig.c_str()), - "Failed to set remote auto tuning configuration", false, err); + "Failed to set remote auto tuning config", false, err); } return true; @@ -1304,7 +1330,13 @@ bool networkToSerializedEngine( if (build.safe && build.consistency) { - if (!checkSafeEngine(serializedEngine->data(), serializedEngine->size())) + std::vector pluginBuildLibPaths; +#if ENABLE_UNIFIED_BUILDER + pluginBuildLibPaths.reserve(sys.safetyPlugins.size()); + std::transform(sys.safetyPlugins.begin(), sys.safetyPlugins.end(), std::back_inserter(pluginBuildLibPaths), + [](auto const& sp) { return sp.libraryName; }); +#endif + if (!checkSafeEngine(serializedEngine->data(), serializedEngine->size(), pluginBuildLibPaths)) { return false; } @@ -1326,6 +1358,7 @@ bool networkToSerializedEngine( return true; } + //! //! \brief Parse a given model, create a network and an engine. //! @@ -1440,8 +1473,8 @@ bool loadAsyncStreamingEngineToBuildEnv(std::string const& filepath, BuildEnviro } -bool loadEngineToBuildEnv( - std::string const& filepath, BuildEnvironment& env, std::ostream& err, bool const enableConsistency) +bool loadEngineToBuildEnv(std::string const& filepath, BuildEnvironment& env, std::ostream& err, + SystemOptions const& sys, bool const enableConsistency) { auto const tBegin = std::chrono::high_resolution_clock::now(); std::ifstream engineFile(filepath, std::ios::binary); @@ -1460,7 +1493,13 @@ bool loadEngineToBuildEnv( if (enableConsistency) { - if (!checkSafeEngine(engineBlob.data(), fsize)) + std::vector pluginBuildLibPaths; +#if ENABLE_UNIFIED_BUILDER + pluginBuildLibPaths.reserve(sys.safetyPlugins.size()); + std::transform(sys.safetyPlugins.begin(), sys.safetyPlugins.end(), std::back_inserter(pluginBuildLibPaths), + [](auto const& sp) { return sp.libraryName; }); +#endif + if (!checkSafeEngine(engineBlob.data(), fsize, pluginBuildLibPaths)) { sample::gLogError << "Consistency validation is not enabled." << std::endl; return false; @@ -1537,7 +1576,8 @@ void dumpRefittable(nvinfer1::ICudaEngine& engine) ICudaEngine* loadEngine(std::string const& engine, int32_t DLACore, std::ostream& err) { BuildEnvironment env(/* isSafe */ false, /* versionCompatible */ false, DLACore, "", getTempfileControlDefaults()); - return loadEngineToBuildEnv(engine, env, err, false) ? env.engine.release() : nullptr; + SystemOptions sys; + return loadEngineToBuildEnv(engine, env, err, sys, false) ? env.engine.release() : nullptr; } bool saveEngine(ICudaEngine const& engine, std::string const& fileName, std::ostream& err) @@ -1569,7 +1609,7 @@ bool getEngineBuildEnv( { if (build.safe) { - createEngineSuccess = loadEngineToBuildEnv(build.engine, env, err, build.safe && build.consistency); + createEngineSuccess = loadEngineToBuildEnv(build.engine, env, err, sys, build.safe && build.consistency); } else { @@ -1606,6 +1646,7 @@ bool getEngineBuildEnv( engineFile.close(); if (!build.safe) { + env.engine.releaseBlob(); if (build.asyncFileReader) { SMP_RETVAL_IF_FALSE(loadAsyncStreamingEngineToBuildEnv(build.engine, env, err), @@ -1628,7 +1669,6 @@ bool getEngineBuildEnv( engineTextFile.close(); } } - env.engine.releaseBlob(); } return true; @@ -1681,6 +1721,8 @@ std::vector> getAllRefitWeightsForLayer(ILayer c std::make_pair(WeightsRole::kSHIFT, layer.getShift())}; } case LayerType::kACTIVATION: + case LayerType::kATTENTION_INPUT: + case LayerType::kATTENTION_OUTPUT: case LayerType::kASSERTION: case LayerType::kCAST: case LayerType::kCONCATENATION: @@ -1731,6 +1773,32 @@ std::vector> getAllRefitWeightsForLayer(ILayer c return {}; } +bool refitFromOnnx(nvinfer1::ICudaEngine& engine, std::string onnxModelFile, bool multiThreading) +{ + sample::gLogInfo << "Refitting engine from ONNX model " << onnxModelFile << std::endl; + std::unique_ptr refitter{createRefitter(engine)}; + if (multiThreading && !refitter->setMaxThreads(10)) + { + sample::gLogError << "Failed to set max threads to refitter." << std::endl; + return false; + } + std::unique_ptr parserRefitter{createONNXRefitter(*refitter)}; + + if (!parserRefitter->refitFromFile(onnxModelFile.c_str())) + { + return false; + } + TrtCudaStream stream; + if (!refitter->refitCudaEngineAsync(stream.get())) + { + return false; + } + stream.synchronize(); + + sample::gLogInfo << "Engine successfully refitted from ONNX model " << onnxModelFile << std::endl; + return true; +} + bool timeRefit(INetworkDefinition const& network, nvinfer1::ICudaEngine& engine, bool multiThreading) { using time_point = std::chrono::time_point; @@ -1891,8 +1959,8 @@ bool hasConsistencyChecker() #if ENABLE_UNIFIED_BUILDER -nvinfer2::safe::consistency::IConsistencyChecker* createConsistencyChecker( - sample::SampleSafeRecorder& recorder, void const* serializedEngine, int32_t const engineSize) noexcept +nvinfer2::safe::consistency::IConsistencyChecker* createConsistencyChecker(sample::SampleSafeRecorder& recorder, + void const* serializedEngine, int32_t const engineSize, std::vector const& pluginBuildLibPath) noexcept { nvinfer2::safe::consistency::IConsistencyChecker* checker{nullptr}; @@ -1904,13 +1972,14 @@ nvinfer2::safe::consistency::IConsistencyChecker* createConsistencyChecker( #if !defined(_WIN32) constexpr char symbolName[] = "createConsistencyChecker"; typedef ErrorCode (*CreateCheckerFn)(nvinfer2::safe::consistency::IConsistencyChecker * &checker, - sample::SampleSafeRecorder & recorder, void const* data, size_t size); + sample::SampleSafeRecorder & recorder, void const* data, size_t size, + std::vector const& pluginBuildLibPath); if (hasSafeRuntime()) { auto createFn = reinterpret_cast(dlsym(consistencyCheckerLibrary.get(), symbolName)); if (createFn != nullptr) { - ErrorCode errorCode = createFn(checker, recorder, serializedEngine, engineSize); + ErrorCode errorCode = createFn(checker, recorder, serializedEngine, engineSize, pluginBuildLibPath); if (errorCode != ErrorCode::kSUCCESS) { return nullptr; @@ -1922,7 +1991,8 @@ nvinfer2::safe::consistency::IConsistencyChecker* createConsistencyChecker( } #endif -bool checkSafeEngine(void const* serializedEngine, int64_t const engineSize) +bool checkSafeEngine( + void const* serializedEngine, int64_t const engineSize, std::vector const& pluginBuildLibPath) { if (!hasConsistencyChecker()) { @@ -1932,9 +2002,8 @@ bool checkSafeEngine(void const* serializedEngine, int64_t const engineSize) #if ENABLE_UNIFIED_BUILDER sample::SampleSafeRecorder recorder{nvinfer2::safe::Severity::kINFO}; - auto checker = std::unique_ptr( - createConsistencyChecker(recorder, serializedEngine, engineSize)); + createConsistencyChecker(recorder, serializedEngine, engineSize, pluginBuildLibPath)); if (checker.get() == nullptr) { sample::gLogError << "Failed to create consistency checker." << std::endl; diff --git a/samples/common/sampleEngines.h b/samples/common/sampleEngines.h index e7ecf8c8..cc02a7c0 100644 --- a/samples/common/sampleEngines.h +++ b/samples/common/sampleEngines.h @@ -331,6 +331,18 @@ nvinfer1::IHostMemory* modelToSerialized( bool serializeAndSave( const ModelOptions& model, const BuildOptions& build, const SystemOptions& sys, std::ostream& err); +//! +//! \brief Refit an engine using the weights from the specified ONNX model. +//! +//! \return boolean Return true if the engine was successfully refit from the model. +//! +bool refitFromOnnx(nvinfer1::ICudaEngine& engine, std::string onnxModelFile, bool multiThreading); + +//! +//! \brief Refit an engine using the weights from the INetworkDefintiion and report the amount of time it took. +//! +//! \return boolean Return true if the engine was successfully refit from the INetworkDefinition. +//! bool timeRefit(const nvinfer1::INetworkDefinition& network, nvinfer1::ICudaEngine& engine, bool multiThreading); //! @@ -347,12 +359,13 @@ bool hasSafeRuntime(); //! //! \brief Run consistency check on serialized engine. //! -bool checkSafeEngine(void const* serializedEngine, int64_t const engineSize); +bool checkSafeEngine( + void const* serializedEngine, int64_t const engineSize, std::vector const& pluginBuildLibPath); bool loadStreamingEngineToBuildEnv(std::string const& engine, BuildEnvironment& env, std::ostream& err); -bool loadEngineToBuildEnv( - std::string const& engine, BuildEnvironment& env, std::ostream& err, bool const enableConsistency); +bool loadEngineToBuildEnv(std::string const& engine, BuildEnvironment& env, std::ostream& err, SystemOptions const& sys, + bool const enableConsistency); } // namespace sample #endif // TRT_SAMPLE_ENGINES_H diff --git a/samples/common/sampleEntrypoints.h b/samples/common/sampleEntrypoints.h index e77cd443..306f1b96 100644 --- a/samples/common/sampleEntrypoints.h +++ b/samples/common/sampleEntrypoints.h @@ -35,8 +35,8 @@ extern nvinfer1::IBuilder* createBuilder(); extern nvinfer1::IRuntime* createRuntime(); extern nvinfer1::IRefitter* createRefitter(nvinfer1::ICudaEngine& engine); - extern nvonnxparser::IParser* createONNXParser(nvinfer1::INetworkDefinition& network); +extern nvonnxparser::IParserRefitter* createONNXRefitter(nvinfer1::IRefitter& refitter); #if !defined(DEFINE_TRT_ENTRYPOINTS) #define DEFINE_TRT_ENTRYPOINTS 0 @@ -55,9 +55,6 @@ extern nvonnxparser::IParser* createONNXParser(nvinfer1::INetworkDefinition& net #if !defined(DEFINE_TRT_ONNX_PARSER_ENTRYPOINT) #define DEFINE_TRT_ONNX_PARSER_ENTRYPOINT 1 #endif -#if !defined(DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT) -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 1 -#endif #if DEFINE_TRT_ENTRYPOINTS nvinfer1::IBuilder* createBuilder() @@ -96,6 +93,15 @@ nvonnxparser::IParser* createONNXParser(nvinfer1::INetworkDefinition& network) #endif } +nvonnxparser::IParserRefitter* createONNXRefitter(nvinfer1::IRefitter& refitter) +{ +#if DEFINE_TRT_ONNX_PARSER_ENTRYPOINT + return nvonnxparser::createParserRefitter(refitter, sample::gLogger.getTRTLogger()); +#else + return {}; +#endif +} + #endif // DEFINE_TRT_ENTRYPOINTS #endif // TRT_SAMPLE_ENTRYPOINTS_H diff --git a/samples/common/sampleInference.cpp b/samples/common/sampleInference.cpp index c9be06f7..d9f3d388 100644 --- a/samples/common/sampleInference.cpp +++ b/samples/common/sampleInference.cpp @@ -193,6 +193,7 @@ nvinfer2::safe::TypedArray createTypedArray(void* const ptr, DataType const type case DataType::kINT32: return nvinfer2::safe::TypedArray(static_cast(ptr), bufferSize); case DataType::kINT8: return nvinfer2::safe::TypedArray(static_cast(ptr), bufferSize); case DataType::kINT64: return nvinfer2::safe::TypedArray(static_cast(ptr), bufferSize); + case DataType::kBOOL: return nvinfer2::safe::TypedArray(static_cast(ptr), bufferSize); default: { sample::gLogError << "Invalid tensor DataType encountered." << std::endl; @@ -350,8 +351,9 @@ bool allocateContextMemory(InferenceEnvironmentStd& iEnv, InferenceOptions const auto const& ec = iEnv.contexts.at(i); if (inference.memoryAllocationStrategy == MemoryAllocationStrategy::kSTATIC) { - sample::gLogInfo << "Created execution context with device memory size: " - << (engine->getDeviceMemorySize() / 1.0_MiB) << " MiB" << std::endl; + sample::gLogInfo << "Created execution context with device memory size: " << + (engine->getDeviceMemorySize() / 1.0_MiB) + << " MiB" << std::endl; } else { @@ -816,13 +818,14 @@ TaskInferenceEnvironment::TaskInferenceEnvironment(std::string engineFile, Infer , batch(bs) { BuildEnvironment bEnv(/* isSafe */ false, /* versionCompatible */ false, DLACore, "", getTempfileControlDefaults()); - loadEngineToBuildEnv(engineFile, bEnv, sample::gLogError, false); - iEnv = std::make_unique(bEnv); - - CHECK(cudaSetDevice(device)); SystemOptions system{}; system.device = device; system.DLACore = DLACore; + loadEngineToBuildEnv(engineFile, bEnv, sample::gLogError, system, false); + iEnv = std::make_unique(bEnv); + + CHECK(cudaSetDevice(device)); + if (!setUpStdInference(*iEnv, iOptions, system)) { sample::gLogError << "Inference set up failed" << std::endl; @@ -1142,6 +1145,8 @@ public: void wait(TrtCudaEvent& gpuStart) { getStream(StreamType::kINPUT).wait(gpuStart); + getStream(StreamType::kCOMPUTE).wait(gpuStart); + getStream(StreamType::kOUTPUT).wait(gpuStart); } void setInputData(bool sync) diff --git a/samples/common/sampleOptions.cpp b/samples/common/sampleOptions.cpp index 061511a0..2ea4df9a 100644 --- a/samples/common/sampleOptions.cpp +++ b/samples/common/sampleOptions.cpp @@ -552,6 +552,23 @@ void getLayerDeviceTypes(Arguments& arguments, char const* argument, LayerDevice } } +void getDecomposableAttentions( + Arguments& arguments, char const* argument, DecomposableAttentions& decomposableAttentions) +{ + std::string list; + if (!getAndDelOption(arguments, argument, list)) + { + return; + } + + std::vector attentionList{splitToStringVec(list, ',')}; + for (auto& s : attentionList) + { + auto const attentionName = removeSingleQuotationMarks(s); + decomposableAttentions[attentionName] = true; + } +} + void getAndDelStringsSet(Arguments& arguments, char const* argument, StringSet& stringSet) { std::string list; @@ -1208,11 +1225,11 @@ void BuildOptions::parse(Arguments& arguments) getAndDelOption(arguments, "--best", best); if (best) { - int8 = (samplesCommon::getSMVersion() != 0x0a03); + int8 = (samplesCommon::getSmVersion() != 0x0a03); fp16 = true; // BF16 only supported on Ampere+ - if (samplesCommon::getSMVersion() >= 0x0800) + if (samplesCommon::getSmVersion() >= 0x0800) { bf16 = true; } @@ -1283,9 +1300,17 @@ void BuildOptions::parse(Arguments& arguments) disableAndLog(fp8, "fp8", "kFP8"); disableAndLog(int4, "int4", "kINT4"); } - + // Print a message to tell users that weakly-typed networks have been deprecated in TensorRT. + if (fp16 || bf16 || int8 || fp8 || int4 || best) + { + sample::gLogWarning << "Weakly-typed networks have been deprecated in TensorRT. " + "You can use the AutoCast tool " + "(https://nvidia.github.io/TensorRT-Model-Optimizer/guides/8_autocast.html) to convert " + "the network to be strongly typed." + << std::endl; + } // Print a message to tell users that --noTF32 can be added to improve accuracy with performance cost. - if (samplesCommon::getSMVersion() >= 0x0800) + if (samplesCommon::getSmVersion() >= 0x0800) { if (!(stronglyTyped || fp16 || bf16 || int8 || fp8 || int4)) { @@ -1336,6 +1361,7 @@ void BuildOptions::parse(Arguments& arguments) getLayerPrecisions(arguments, "--layerPrecisions", layerPrecisions); getLayerOutputTypes(arguments, "--layerOutputTypes", layerOutputTypes); getLayerDeviceTypes(arguments, "--layerDeviceTypes", layerDeviceTypes); + getDecomposableAttentions(arguments, "--decomposableAttentions", decomposableAttentions); if (layerPrecisions.empty() && layerOutputTypes.empty() && precisionConstraints != PrecisionConstraints::kNONE) { @@ -1451,22 +1477,22 @@ void BuildOptions::parse(Arguments& arguments) source = nvinfer1::TacticSource::kCUBLAS_LT; } else - if (t == "CUDNN") - { - source = nvinfer1::TacticSource::kCUDNN; - } - else if (t == "EDGE_MASK_CONVOLUTIONS") - { - source = nvinfer1::TacticSource::kEDGE_MASK_CONVOLUTIONS; - } - else if (t == "JIT_CONVOLUTIONS") - { - source = nvinfer1::TacticSource::kJIT_CONVOLUTIONS; - } - else - { - throw std::invalid_argument(std::string("Unknown tactic source: ") + t); - } + if (t == "CUDNN") + { + source = nvinfer1::TacticSource::kCUDNN; + } + else if (t == "EDGE_MASK_CONVOLUTIONS") + { + source = nvinfer1::TacticSource::kEDGE_MASK_CONVOLUTIONS; + } + else if (t == "JIT_CONVOLUTIONS") + { + source = nvinfer1::TacticSource::kJIT_CONVOLUTIONS; + } + else + { + throw std::invalid_argument(std::string("Unknown tactic source: ") + t); + } uint32_t sourceBit = 1U << static_cast(source); @@ -1517,8 +1543,9 @@ void BuildOptions::parse(Arguments& arguments) } else { - throw std::invalid_argument(std::string("Unknown runtime platform: ") + runtimePlatformArgs - + ". Valid options: SameAsBuild, WindowsAMD64."); + std::string validOptions = "SameAsBuild, WindowsAMD64"; + throw std::invalid_argument( + std::string("Unknown runtime platform: ") + runtimePlatformArgs + ". Valid options: " + validOptions); } std::string hardwareCompatibleArgs; @@ -1748,6 +1775,7 @@ void InferenceOptions::parse(Arguments& arguments) std::string debugFormats; getAndDelOption(arguments, "--saveAllDebugTensors", debugFormats); dumpAlldebugTensorFormats = splitToStringVec(debugFormats, ','); + getAndDelOption(arguments, "--refitFromOnnx", refitOnnxModel); } void ReportingOptions::parse(Arguments& arguments) @@ -1938,7 +1966,6 @@ void SafeBuilderOptions::parse(Arguments& arguments) getFormats(outputFormats, "--outputIOFormats"); getAndDelOption(arguments, "--int8", int8); getAndDelOption(arguments, "--calib", calibFile); - getAndDelOption(arguments, "--consistency", consistency); getAndDelOption(arguments, "--std", standard); std::string pluginName; while (getAndDelOption(arguments, "--plugins", pluginName)) @@ -2233,6 +2260,17 @@ std::ostream& operator<<(std::ostream& os, LayerDeviceTypes const& layerDeviceTy return os; } +std::ostream& operator<<(std::ostream& os, DecomposableAttentions const& decomposableAttentions) +{ + char const* sep = ""; + for (auto const& attentionDecomposablePair : decomposableAttentions) + { + os << sep << attentionDecomposablePair.first << ":" << attentionDecomposablePair.second; + sep = ", "; + } + return os; +} + std::ostream& operator<<(std::ostream& os, StringSet const& stringSet) { int64_t i = 0; @@ -2259,6 +2297,7 @@ std::ostream& operator<<(std::ostream& os, const BuildOptions& options) "Precision: "; printPrecision(os, options) << std::endl << "LayerPrecisions: " << options.layerPrecisions << std::endl << "Layer Device Types: " << options.layerDeviceTypes << std::endl << + "Decomposable Attentions: " << options.decomposableAttentions << std::endl << "Calibration: " << (options.int8 && options.calibration.empty() ? "Dynamic" : options.calibration.c_str()) << std::endl << "Refit: " << boolToEnabled(options.refittable) << std::endl << "Strip weights: " << boolToEnabled(options.stripWeights) << std::endl << @@ -2586,7 +2625,7 @@ void BuildOptions::help(std::ostream& os) R"( --useRuntime=runtime TensorRT runtime to execute engine. "lean" and "dispatch" require loading VC engine and do)" "\n" " not support building an engine." "\n" R"( runtime::= "full"|"lean"|"dispatch")" "\n" - " --leanDLLPath= External lean runtime DLL to use in version compatiable mode." "\n" + " --leanDLLPath= External lean runtime DLL to use in version compatible mode." "\n" " --excludeLeanRuntime When --versionCompatible is enabled, this flag indicates that the generated engine should" "\n" " not include an embedded lean runtime. If this is set, the user must explicitly specify a" "\n" " valid lean runtime to use when loading the engine." "\n" @@ -2608,6 +2647,9 @@ void BuildOptions::help(std::ostream& os) " --fp8 Enable fp8 precision, in addition to fp32 (default = disabled)" "\n" " --int4 Enable int4 precision, in addition to fp32 (default = disabled)" "\n" " --best Enable all precisions to achieve the best performance (default = disabled)" "\n" + " Note: --fp16, --bf16, --int8, --fp8, --int4, --best are deprecated and superseded by strong typing.""\n" + " The AutoCast tool (https://nvidia.github.io/TensorRT-Model-Optimizer/guides/8_autocast.html)" "\n" + " can be used to convert the network to be strongly typed." "\n" " --stronglyTyped Create a strongly typed network. (default = disabled)" "\n" " --directIO [Deprecated] Avoid reformatting at network boundaries. (default = disabled)" "\n" " --precisionConstraints=spec Control precision constraint setting. (default = none)" "\n" @@ -2637,9 +2679,12 @@ void BuildOptions::help(std::ostream& os) R"( Per-layer device type spec ::= layerDeviceTypePair[","spec])" "\n" R"( layerDeviceTypePair ::= layerName":"deviceType)" "\n" R"( deviceType ::= "GPU"|"DLA")" "\n" + " --decomposableAttentions=spec Specify decomposable attentions by comma-separated names." "\n" + R"( The specs are read left-to-right, and later ones override earlier ones. Each layer name can)" "\n" + " contain at most one wildcard ('*') character." "\n" " --calib= Read INT8 calibration cache file" "\n" - " --safe Enable build safety certified engine, if DLA is enable, --buildDLAStandalone will be specified" "\n" - " automatically (default = disabled)" "\n" + " --safe Enable build safety certified engine, --stronglyTyped will be enabled by default with this option." "\n" + " If DLA is enabled, --buildDLAStandalone will be specified" "\n" " --dumpKernelText Dump the kernel text to a file, only available when --safe is enabled" "\n" " --buildDLAStandalone Enable build DLA standalone loadable which can be loaded by cuDLA, when this option is enabled, " "\n" " --allowGPUFallback is disallowed and --skipInference is enabled by default. Additionally, " "\n" @@ -2720,6 +2765,8 @@ void BuildOptions::help(std::ostream& os) " --remoteAutoTuningConfig Set the remote auto tuning config. Must be specified with --safe." "\n" " Format: protocol://username[:password]@hostname[:port]?param1=value1¶m2=value2" "\n" " Example: ssh://user:pass@192.0.2.100:22?remote_exec_path=/opt/tensorrt/bin&remote_lib_path=/opt/tensorrt/lib" "\n" + " --refitFromOnnx Refit the loaded engine with the weights from the provided ONNX model." "\n" + " The model should be identical to the one used to generate the engine." "\n" ; // clang-format on os << std::flush; @@ -2911,7 +2958,6 @@ void SafeBuilderOptions::printHelp(std::ostream& os) R"( fmt ::= ("chw"|"chw2"|"hwc8"|"chw4"|"chw16"|"chw32"|"dhwc8"|)" << std::endl << R"( "cdhw32"|"hwc"|"dla_linear"|"dla_hwc4"|"hwc16"|"dhwc")["+"fmt])" << std::endl << " --int8 Enable int8 precision, in addition to fp16 (default = disabled)" << std::endl << - " --consistency Perform consistency checking on safety certified engine" << std::endl << " --std Build standard serialized engine, (default = disabled)" << std::endl << " --calib= Read INT8 calibration cache file" << std::endl << " --serialized= Save the serialized network" << std::endl << diff --git a/samples/common/sampleOptions.h b/samples/common/sampleOptions.h index 9b18e429..321b3b8f 100644 --- a/samples/common/sampleOptions.h +++ b/samples/common/sampleOptions.h @@ -154,6 +154,7 @@ using ShapeRange = std::array, nvinfer1::EnumMax; using LayerOutputTypes = std::unordered_map>; using LayerDeviceTypes = std::unordered_map; +using DecomposableAttentions = std::unordered_map; using StringSet = std::unordered_set; @@ -235,6 +236,7 @@ public: LayerPrecisions layerPrecisions; LayerOutputTypes layerOutputTypes; LayerDeviceTypes layerDeviceTypes; + DecomposableAttentions decomposableAttentions; StringSet debugTensors; bool markUnfusedTensorsAsDebugTensors{false}; StringSet debugTensorStates; @@ -344,6 +346,7 @@ public: std::unordered_map debugTensorFileNames; std::vector dumpAlldebugTensorFormats; WeightStreamingBudget weightStreamingBudget; + std::string refitOnnxModel; void parse(Arguments& arguments) override; diff --git a/samples/common/streamReader.h b/samples/common/streamReader.h index cd17a2d7..57a1c45a 100644 --- a/samples/common/streamReader.h +++ b/samples/common/streamReader.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/samples/python/aliased_io_plugin/requirements.txt b/samples/python/aliased_io_plugin/requirements.txt index 4e3febb5..783d4c79 100644 --- a/samples/python/aliased_io_plugin/requirements.txt +++ b/samples/python/aliased_io_plugin/requirements.txt @@ -6,5 +6,5 @@ polygraphy colored numpy==1.26.4 pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 diff --git a/samples/python/common.py b/samples/python/common.py index 10b2c323..31a12887 100644 --- a/samples/python/common.py +++ b/samples/python/common.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,11 +21,7 @@ import os import tensorrt as trt from common_runtime import * -try: - # Sometimes python does not understand FileNotFoundError - FileNotFoundError -except NameError: - FileNotFoundError = IOError +# FileNotFoundError is available in Python 3.3+ def GiB(val): diff --git a/samples/python/common_runtime.py b/samples/python/common_runtime.py index 34ac6b5d..64c3db67 100644 --- a/samples/python/common_runtime.py +++ b/samples/python/common_runtime.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,71 +20,269 @@ from typing import Optional, List, Union import numpy as np import tensorrt as trt -from cuda.bindings import driver as cuda, runtime as cudart -def check_cuda_err(err): - if isinstance(err, cuda.CUresult): - if err != cuda.CUresult.CUDA_SUCCESS: - raise RuntimeError("Cuda Error: {}".format(err)) - if isinstance(err, cudart.cudaError_t): - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError("Cuda Runtime Error: {}".format(err)) - else: - raise RuntimeError("Unknown error type: {}".format(err)) +from cuda.bindings import driver as cuda, runtime as cudart, nvrtc + + +class ArrayWithOwner(np.ndarray): + """Numpy array that holds a reference to its owner object""" + def __new__(cls, input_array, owner): + obj = np.asarray(input_array).view(cls) + obj._owner = owner + return obj + + def __array_finalize__(self, obj): + if obj is None: + return + self._owner = getattr(obj, '_owner', None) + + def cuda_call(call): + """Helper function to make CUDA calls and check for errors""" + def _cudaGetErrorEnum(error): + if isinstance(error, cuda.CUresult): + err, name = cuda.cuGetErrorName(error) + return name if err == cuda.CUresult.CUDA_SUCCESS else "" + elif isinstance(error, cudart.cudaError_t): + return cudart.cudaGetErrorName(error)[1] + elif isinstance(error, nvrtc.nvrtcResult): + return nvrtc.nvrtcGetErrorString(error)[1] + else: + raise RuntimeError("Unknown error type: {}".format(error)) + err, res = call[0], call[1:] - check_cuda_err(err) + if err.value: + raise RuntimeError( + "CUDA error code={}({})".format( + err.value, _cudaGetErrorEnum(err) + ) + ) if len(res) == 1: - res = res[0] - return res + return res[0] + elif len(res) == 0: + return None + else: + return res + + +def create_cuda_context(device): + """ + Create CUDA context with version-aware API handling. + + Handles different CUDA API versions based on actual documented signatures: + - CUDA 11.8-12.9: cuCtxCreate(flags, device) - 2 arguments + - CUDA 13.0+: cuCtxCreate(ctxCreateParams, flags, device) - 3 arguments + + Args: + device: CUDA device handle from cuDeviceGet + + Returns: + CUDA context handle + """ + # Try different API versions + try: + # Try CUDA 13.0+ API first (3 arguments with ctxCreateParams) + # cuCtxCreate(ctxCreateParams, flags, device) + return cuda_call(cuda.cuCtxCreate(None, 0, device)) + except TypeError: + # CUDA 11.8-12.9 API: cuCtxCreate(flags, device) + return cuda_call(cuda.cuCtxCreate(0, device)) + class HostDeviceMem: - """Pair of host and device memory, where the host memory is wrapped in a numpy array""" + """Pair of host and device memory using RAII composition""" def __init__(self, size: int, dtype: Optional[np.dtype] = None): - dtype = dtype or np.dtype(np.uint8) - nbytes = size * dtype.itemsize - host_mem = cuda_call(cudart.cudaMallocHost(nbytes)) - pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype)) + if dtype is None: + dtype = np.dtype(np.uint8) + else: + dtype = np.dtype(dtype) + self._size = size + self._dtype = dtype - self._host = np.ctypeslib.as_array(ctypes.cast(host_mem, pointer_type), (size,)) - self._device = cuda_call(cudart.cudaMalloc(nbytes)) - self._nbytes = nbytes + # Use RAII classes for memory management + self._host_mem = PinnedHostMem(size, dtype) + self._device_mem = DeviceMem(size * dtype.itemsize) @property def host(self) -> np.ndarray: - return self._host + # Return the array directly - ArrayWithOwner ensures proper lifetime management + return self._host_mem.array @host.setter def host(self, data: Union[np.ndarray, bytes]): - if isinstance(data, np.ndarray): - if data.size > self.host.size: - raise ValueError( - f"Tried to fit an array of size {data.size} into host memory of size {self.host.size}" - ) - np.copyto(self.host[:data.size], data.flat, casting='safe') - else: - assert self.host.dtype == np.uint8 - self.host[:self.nbytes] = np.frombuffer(data, dtype=np.uint8) + # Delegate to PinnedHostMem for proper data handling + self._host_mem.array = data @property - def device(self) -> int: - return self._device + def device_ptr(self) -> int: + """Device memory pointer""" + return self._device_mem.device_ptr + + @property + def nbytes(self) -> int: + return self._host_mem.nbytes + + + def __str__(self): + return f"Host:\n{self.host}\nDevice:\n{self.device_ptr}\nSize:\n{self.nbytes}\n" + + def __repr__(self): + return self.__str__() + + +class DeviceMem: + """Device-only memory allocation for cases where host memory is not needed""" + def __init__(self, size: int): + self._device_ptr = cuda_call(cudart.cudaMalloc(size)) + self._nbytes = size + + @property + def device_ptr(self) -> int: + """Device memory pointer""" + return self._device_ptr @property def nbytes(self) -> int: return self._nbytes + def free(self): + """Explicitly free device memory""" + if self._device_ptr is not None: + try: + cuda_call(cudart.cudaFree(self._device_ptr)) + self._device_ptr = None + except Exception: + # Log but don't raise - cleanup should be best effort + pass + def __str__(self): - return f"Host:\n{self.host}\nDevice:\n{self.device}\nSize:\n{self.nbytes}\n" + return f"Device:\n{self.device_ptr}\nSize:\n{self.nbytes}\n" def __repr__(self): return self.__str__() + def __del__(self): + # Fallback cleanup - not guaranteed to be called + self.free() + + +class PinnedHostMem: + """Pinned host memory allocation for faster GPU transfers""" + def __init__(self, size: int, dtype: Optional[np.dtype] = None): + if dtype is None: + dtype = np.dtype(np.uint8) + else: + dtype = np.dtype(dtype) + nbytes = size * dtype.itemsize + host_mem = cuda_call(cudart.cudaMallocHost(nbytes)) + + self._host_ptr = host_mem + self._host_size = size + self._nbytes = nbytes + self._dtype = dtype + + @property + def array(self) -> np.ndarray: + # Create view with proper memory ownership + pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(self._dtype)) + host_array = np.ctypeslib.as_array(ctypes.cast(self._host_ptr, pointer_type), (self._host_size,)) + return ArrayWithOwner(host_array, self) + + @array.setter + def array(self, data: Union[np.ndarray, bytes]): + """Set the array data with proper bounds checking""" + host_array = self.array # Get the numpy array view + if isinstance(data, np.ndarray): + if data.size > self._host_size: + raise ValueError( + f"Tried to fit an array of size {data.size} into host memory of size {self._host_size}" + ) + np.copyto(host_array[:data.size], data.flat, casting='safe') + else: + assert self._dtype == np.uint8 + host_array[:self.nbytes] = np.frombuffer(data, dtype=np.uint8) + + + @property + def nbytes(self) -> int: + return self._nbytes + def free(self): - cuda_call(cudart.cudaFree(self.device)) - cuda_call(cudart.cudaFreeHost(self.host.ctypes.data)) + """Explicitly free pinned host memory""" + if self._host_ptr is not None: + try: + cuda_call(cudart.cudaFreeHost(self._host_ptr)) + self._host_ptr = None + except Exception: + # Log but don't raise - cleanup should be best effort + pass + + def __str__(self): + return f"PinnedHost:\n{self.array}\nSize:\n{self.nbytes}\n" + + def __repr__(self): + return self.__str__() + + def __del__(self): + # Fallback cleanup - not guaranteed to be called + self.free() + +class CudaStreamContext: + """CUDA stream lifecycle management with context manager support""" + def __init__(self): + """Initialize CUDA stream""" + self._stream = cuda_call(cudart.cudaStreamCreate()) + + def __enter__(self): + """Create CUDA stream when entering context (if not already created)""" + if self._stream is None: + self._stream = cuda_call(cudart.cudaStreamCreate()) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Destroy CUDA stream when exiting context""" + if self._stream is not None: + try: + cuda_call(cudart.cudaStreamDestroy(self._stream)) + except Exception: + # Silently handle cleanup failures + pass + self._stream = None + + @property + def stream(self) -> cudart.cudaStream_t: + if self._stream is None: + raise RuntimeError("Stream not created. Use 'with' statement.") + return self._stream + + def synchronize(self): + """Synchronize the stream""" + if self._stream is None: + raise RuntimeError("Stream not created. Use 'with' statement.") + cuda_call(cudart.cudaStreamSynchronize(self._stream)) + + def free(self): + """Explicitly free the CUDA stream""" + if self._stream is not None: + try: + cuda_call(cudart.cudaStreamDestroy(self._stream)) + self._stream = None + except Exception: + # Log but don't raise - cleanup should be best effort + pass + + def __del__(self): + """Cleanup stream on destruction""" + if hasattr(self, '_stream') and self._stream is not None: + self.free() + + def __str__(self): + return f"CudaStreamContext: {self._stream}" + + def __repr__(self): + return self.__str__() # Allocates all buffers required for an engine, i.e. host/device inputs/outputs. @@ -93,7 +291,6 @@ def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None) inputs = [] outputs = [] bindings = [] - stream = cuda_call(cudart.cudaStreamCreate()) tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)] for binding in tensor_names: # get_tensor_profile_shape returns (min_shape, optimal_shape, max_shape) @@ -115,56 +312,100 @@ def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None) bindingMemory = HostDeviceMem(size) # Append the device buffer to device bindings. - bindings.append(int(bindingMemory.device)) + bindings.append(int(bindingMemory.device_ptr)) # Append to the appropriate list. if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT: inputs.append(bindingMemory) else: outputs.append(bindingMemory) - return inputs, outputs, bindings, stream + return inputs, outputs, bindings # Frees the resources allocated in allocate_buffers -def free_buffers(inputs: List[HostDeviceMem], outputs: List[HostDeviceMem], stream: cudart.cudaStream_t): - for mem in inputs + outputs: - mem.free() - cuda_call(cudart.cudaStreamDestroy(stream)) +def free_buffers(inputs: List[HostDeviceMem], outputs: List[HostDeviceMem]): + """ + Explicitly free CUDA memory resources. + + While __del__ methods provide automatic cleanup, they are not guaranteed to be called. + This function provides explicit resource management for critical applications. + """ + for inp in inputs: + if hasattr(inp, '_device_mem') and hasattr(inp._device_mem, 'free'): + inp._device_mem.free() + if hasattr(inp, '_host_mem') and hasattr(inp._host_mem, 'free'): + inp._host_mem.free() + + for out in outputs: + if hasattr(out, '_device_mem') and hasattr(out._device_mem, 'free'): + out._device_mem.free() + if hasattr(out, '_host_mem') and hasattr(out._host_mem, 'free'): + out._host_mem.free() # Wrapper for cudaMemcpy which infers copy size and does error checking def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray): - nbytes = host_arr.size * host_arr.itemsize - cuda_call(cudart.cudaMemcpy(device_ptr, host_arr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)) + cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)) # Wrapper for cudaMemcpy which infers copy size and does error checking def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int): - nbytes = host_arr.size * host_arr.itemsize - cuda_call(cudart.cudaMemcpy(host_arr, device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) + cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) + + +# Additional CUDA wrapper functions for common operations + + +def cuda_init(): + """Initialize CUDA driver API with error checking.""" + cuda_call(cuda.cuInit(0)) + + +def cuda_get_device(device_id: int = 0): + """Get CUDA device handle with error checking.""" + return cuda_call(cuda.cuDeviceGet(device_id)) + + +# CUDA Runtime API functions (preferred over driver API when available) + + +def cuda_memcpy_htod(device_ptr: int, host_data: np.ndarray): + """Copy data from host to device using CUDA runtime API with error checking. + + Note: Consider using HostDeviceMem.host setter for integrated memory management. + """ + cuda_call(cudart.cudaMemcpy(device_ptr, host_data, host_data.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)) def _do_inference_base(inputs, outputs, stream, execute_async_func): # Transfer input data to the GPU. kind = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice - [cuda_call(cudart.cudaMemcpyAsync(inp.device, inp.host, inp.nbytes, kind, stream)) for inp in inputs] + [cuda_call(cudart.cudaMemcpyAsync(inp.device_ptr, inp.host.ctypes.data, inp.nbytes, kind, stream)) for inp in inputs] # Run inference. execute_async_func() # Transfer predictions back from the GPU. kind = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost - [cuda_call(cudart.cudaMemcpyAsync(out.host, out.device, out.nbytes, kind, stream)) for out in outputs] + [cuda_call(cudart.cudaMemcpyAsync(out.host.ctypes.data, out.device_ptr, out.nbytes, kind, stream)) for out in outputs] # Synchronize the stream cuda_call(cudart.cudaStreamSynchronize(stream)) # Return only the host outputs. - return [out.host for out in outputs] + return [out.host.copy() for out in outputs] # This function is generalized for multiple inputs/outputs. # inputs and outputs are expected to be lists of HostDeviceMem objects. def do_inference(context, engine, bindings, inputs, outputs, stream): + """ + Perform inference using the provided context and stream. + + Usage with context manager: + with stream: # Ensures proper stream lifecycle + outputs = do_inference(context, engine, bindings, inputs, outputs, stream) + """ + stream_handle = stream.stream def execute_async_func(): - context.execute_async_v3(stream_handle=stream) + context.execute_async_v3(stream_handle=stream_handle) # Setup context tensor address. num_io = engine.num_io_tensors for i in range(num_io): context.set_tensor_address(engine.get_tensor_name(i), bindings[i]) - return _do_inference_base(inputs, outputs, stream, execute_async_func) + return _do_inference_base(inputs, outputs, stream_handle, execute_async_func) diff --git a/samples/python/dds_faster_rcnn/infer.py b/samples/python/dds_faster_rcnn/infer.py index ae539330..5b5fc0f2 100644 --- a/samples/python/dds_faster_rcnn/infer.py +++ b/samples/python/dds_faster_rcnn/infer.py @@ -119,13 +119,17 @@ class MyOutputAllocator(trt.IOutputAllocator): print(f"Updated shape for tensor '{tensor_name}': {dims}") def __del__(self): - with self.lock: - for tensor_name, item in self.states.items(): - if item.ptr is not None: - cuda_call(cudart.cudaFree(item.ptr)) - if self.verbose: - print(f"Freed memory for tensor '{tensor_name}'") - self.states.clear() + try: + with self.lock: + for tensor_name, item in self.states.items(): + if item.ptr is not None: + cuda_call(cudart.cudaFree(item.ptr)) + if self.verbose: + print(f"Freed memory for tensor '{tensor_name}'") + self.states.clear() + except Exception: + # Silently handle cleanup failures to prevent exceptions during object deletion + pass class PoolAllocator(trt.IGpuAsyncAllocator): @@ -183,8 +187,12 @@ class PoolAllocator(trt.IGpuAsyncAllocator): return True def __del__(self): - if self.pool: - cuda_call(cudart.cudaMemPoolDestroy(self.pool)) + try: + if self.pool: + cuda_call(cudart.cudaMemPoolDestroy(self.pool)) + except Exception: + # Silently handle cleanup failures to prevent exceptions during object deletion + pass class TensorRTInfer: diff --git a/samples/python/dds_faster_rcnn/requirements.txt b/samples/python/dds_faster_rcnn/requirements.txt index 4e522444..c5c4d496 100644 --- a/samples/python/dds_faster_rcnn/requirements.txt +++ b/samples/python/dds_faster_rcnn/requirements.txt @@ -1,5 +1,5 @@ Pillow==11.3.0 cuda-python==12.9.0 -onnx +onnx==1.18.0 onnx-graphsurgeon --index-url https://pypi.ngc.nvidia.com numpy==1.26.4 diff --git a/samples/python/detectron2/build_engine.py b/samples/python/detectron2/build_engine.py index f4c50fb6..32d776cf 100644 --- a/samples/python/detectron2/build_engine.py +++ b/samples/python/detectron2/build_engine.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -46,7 +46,7 @@ class EngineCalibrator(trt.IInt8MinMaxCalibrator): super().__init__() self.cache_file = cache_file self.image_batcher = None - self.batch_allocation = None + self.batch_memory = None self.batch_generator = None def set_image_batcher(self, image_batcher: ImageBatcher): @@ -60,9 +60,11 @@ class EngineCalibrator(trt.IInt8MinMaxCalibrator): np.dtype(self.image_batcher.dtype).itemsize * np.prod(self.image_batcher.shape) ) - self.batch_allocation = common.cuda_call(cudart.cudaMalloc(self.size)) + self.batch_memory = common.DeviceMem(self.size) self.batch_generator = self.image_batcher.get_batch() + + def get_batch_size(self): """ Overrides from trt.IInt8MinMaxCalibrator. @@ -90,10 +92,10 @@ class EngineCalibrator(trt.IInt8MinMaxCalibrator): ) ) common.memcpy_host_to_device( - self.batch_allocation, np.ascontiguousarray(batch) + self.batch_memory.device_ptr, np.ascontiguousarray(batch) ) - return [int(self.batch_allocation)] + return [int(self.batch_memory.device_ptr)] except StopIteration: log.info("Finished calibration batches") return None diff --git a/samples/python/detectron2/infer.py b/samples/python/detectron2/infer.py index da99f469..45edaf3c 100644 --- a/samples/python/detectron2/infer.py +++ b/samples/python/detectron2/infer.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -51,7 +51,7 @@ class TensorRTInfer: # Setup I/O bindings self.inputs = [] self.outputs = [] - self.allocations = [] + self.device_memories = [] for i in range(self.engine.num_io_tensors): name = self.engine.get_tensor_name(i) is_input = False @@ -64,16 +64,16 @@ class TensorRTInfer: size = np.dtype(trt.nptype(dtype)).itemsize for s in shape: size *= s - allocation = common.cuda_call(cudart.cudaMalloc(size)) + device_mem = common.DeviceMem(size) binding = { "index": i, "name": name, "dtype": np.dtype(trt.nptype(dtype)), "shape": list(shape), - "allocation": allocation, + "allocation": device_mem.device_ptr, "size": size, } - self.allocations.append(allocation) + self.device_memories.append(device_mem) if is_input: self.inputs.append(binding) else: @@ -82,7 +82,7 @@ class TensorRTInfer: assert self.batch_size > 0 assert len(self.inputs) > 0 assert len(self.outputs) > 0 - assert len(self.allocations) > 0 + assert len(self.device_memories) > 0 def input_spec(self): """ diff --git a/samples/python/detectron2/requirements.txt b/samples/python/detectron2/requirements.txt index 5a552559..2c34b757 100644 --- a/samples/python/detectron2/requirements.txt +++ b/samples/python/detectron2/requirements.txt @@ -1,4 +1,4 @@ -onnx==1.16.0 +onnx==1.18.0 onnxruntime==1.18.1 Pillow==11.3.0 git+https://github.com/facebookresearch/detectron2.git @@ -6,6 +6,6 @@ git+https://github.com/NVIDIA/TensorRT#subdirectory=tools/onnx-graphsurgeon cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/efficientdet/README.md b/samples/python/efficientdet/README.md index d3a36a8f..82bb4826 100644 --- a/samples/python/efficientdet/README.md +++ b/samples/python/efficientdet/README.md @@ -1,4 +1,6 @@ -# EfficientDet Object Detection in TensorRT +# [DEPRECATED] EfficientDet Object Detection in TensorRT + +> **Notice:** This sample has been deprecated as of TensorRT 10.14 due to compatibility issues with outdated dependencies in the [tf2onnx](https://github.com/onnx/tensorflow-onnx) conversion pipeline. Users are advised to use earlier TensorRT versions if this sample is required for legacy workflows. ![efficientdet](https://drive.google.com/uc?export=view&id=1Le98wETvmKKj0fUKoCFLsld7o8QPJq9C) diff --git a/samples/python/efficientdet/build_engine.py b/samples/python/efficientdet/build_engine.py index ccd264cc..8b9139fa 100644 --- a/samples/python/efficientdet/build_engine.py +++ b/samples/python/efficientdet/build_engine.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -46,7 +46,7 @@ class EngineCalibrator(trt.IInt8EntropyCalibrator2): super().__init__() self.cache_file = cache_file self.image_batcher = None - self.batch_allocation = None + self.batch_memory = None self.batch_generator = None def set_image_batcher(self, image_batcher: ImageBatcher): @@ -60,9 +60,11 @@ class EngineCalibrator(trt.IInt8EntropyCalibrator2): np.dtype(self.image_batcher.dtype).itemsize * np.prod(self.image_batcher.shape) ) - self.batch_allocation = common.cuda_call(cudart.cudaMalloc(size)) + self.batch_memory = common.DeviceMem(size) self.batch_generator = self.image_batcher.get_batch() + + def get_batch_size(self): """ Overrides from trt.IInt8EntropyCalibrator2. @@ -90,9 +92,9 @@ class EngineCalibrator(trt.IInt8EntropyCalibrator2): ) ) common.memcpy_host_to_device( - self.batch_allocation, np.ascontiguousarray(batch) + self.batch_memory.device_ptr, np.ascontiguousarray(batch) ) - return [int(self.batch_allocation)] + return [int(self.batch_memory.device_ptr)] except StopIteration: log.info("Finished calibration batches") return None diff --git a/samples/python/efficientdet/requirements.txt b/samples/python/efficientdet/requirements.txt index 3e44409c..c129eaf8 100644 --- a/samples/python/efficientdet/requirements.txt +++ b/samples/python/efficientdet/requirements.txt @@ -1,10 +1,10 @@ Pillow==11.3.0 -onnx==1.16.1 +onnx==1.18.0 onnxruntime==1.18.1 tf2onnx==1.16.0 cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/efficientnet/README.md b/samples/python/efficientnet/README.md deleted file mode 100644 index 3d8ad6ff..00000000 --- a/samples/python/efficientnet/README.md +++ /dev/null @@ -1,297 +0,0 @@ -# EfficientNet V1 and V2 in TensorRT - -> NOTE: This sample will be deprecated after TensorRT 10.13.3. - -These scripts help with conversion and execution of Google [EfficientNet V1](https://arxiv.org/abs/1905.11946) and [EfficientNet V2](https://arxiv.org/abs/2104.00298) models with [NVIDIA TensorRT](https://developer.nvidia.com/tensorrt). - -## Contents -- [Changelog](#changelog) -- [Setup](#setup) -- [Model Conversion](#model-conversion) - * [TensorFlow Saved Model](#tensorflow-saved-model) - * [Create ONNX Graph](#create-onnx-graph) - * [Build TensorRT Engine](#build-tensorrt-engine) - * [Benchmark TensorRT Engine](#benchmark-tensorrt-engine) -- [Inference](#inference) - * [Input Preprocessing](#input-preprocessing) - * [Inference in Python](#inference-in-python) - * [Validate against Ground Truth](#validate-against-ground-truth) - * [Compare against TensorFlow](#compare-against-tensorflow) - -# Changelog - -August 2025: - - Removed support for Python versions < 3.10. - - Added deprecation notice - -August 2023: - - Update ONNX version support to 1.14.0 - - Removed support for Python versions < 3.8. - -## Setup - -Note: The sample is not compatible with Python-3.12 because tensorflow-addons does not support Python-3.12. - -For best results, we recommend running these scripts on an environment with TensorRT >= 8.0.1 and TensorFlow 2.12.0. - -Install TensorRT as per the [TensorRT Install Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html). You will need to make sure the Python bindings for TensorRT are also installed correctly, these are available by installing the `python3-libnvinfer` and `python3-libnvinfer-dev` packages on your TensorRT download. - -Make sure all other packages listed in `requirements.txt`: - -```bash -pip3 install -r requirements.txt -``` - -You will also need the latest `onnx_graphsurgeon` python module. If not already installed by TensorRT, you can install it manually by running: - -```bash -pip3 install onnx-graphsurgeon --index-url https://pypi.ngc.nvidia.com -``` - -## Model Conversion - -The workflow to convert an EfficientNet model is basically TensorFlow → ONNX → TensorRT, and so parts of this process require TensorFlow to be installed. If you are performing this conversion to run inference on the edge, such as for NVIDIA Jetson devices, it might be easier to do the ONNX conversion on a PC first. - -### TensorFlow Saved Model - -The starting point of conversion is a TensorFlow saved model. This can be exported from your own trained models, or you can download a pre-trained model. This conversion script is compatible with two types of models: - -1. EfficientNet V1 models trained with the [TensorFlow TPU Models](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet) framework. -2. EfficientNet V2 models trained with the [AutoML](https://github.com/google/automl/tree/master/efficientnetv2) framework. - -#### 1. EfficientNet V1 - -You can download one of the pre-trained saved models from the [EfficientNet TFHub](https://tfhub.dev/google/collections/efficientnet), such as: - -```bash -wget https://storage.googleapis.com/tfhub-modules/tensorflow/efficientnet/b0/classification/1.tar.gz -``` - -The contents of this package, when extracted, will hold a saved model ready for conversion. - -Alternatively, if you are training your own model, or if you need to re-export the saved model manually, you will need the training checkpoint (or a pre-trained "ckpt" from the [EfficientNet Repository](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet) such as [this](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b0.tar.gz)). - -To export a saved model from the checkpoint, clone and install the [TensorFlow TPU Models](https://github.com/tensorflow/tpu) repository, and run: - -```bash -cd /path/to/tpu/models/official/efficientnet -python3 export_model.py \ - --ckpt_dir /path/to/efficientnet-b0 \ - --image_size 224 \ - --model_name efficientnet-b0 \ - --output_tflite /dev/null \ - --noquantize \ - --output_saved_model_dir /path/to/saved_model -``` - -Adapt `--image_size` and `--model_name` according to the checkpoint model being used. The `--ckpt_dir` argument points to the directory holding the checkpoint as described above. The TF saved model will be exported to the path given by `--output_saved_model_dir`. - -#### 2. EfficientNet V2 - -At the time of this writing, there exist no EfficientNet V2 saved models in TFHub yet. So you will need to download a pre-trained checkpoint, or use your own trained model of course. - -To do so, you will need your training checkpoint (or a pre-trained "ckpt" from the [EfficientNet V2 Repository](https://github.com/google/automl/tree/master/efficientnetv2) such as [this](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-s.tgz)): - -```bash -wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/v2/efficientnetv2-s.tgz -``` - -To export a saved model from here, clone and install the [AutoML](https://github.com/google/automl) repository, and run: - -```bash -cd /path/to/automl/efficientnetv2 -python3 infer.py \ - --mode tf2bm \ - --model_name efficientnetv2-s \ - --model_dir ../../efficientnetv2-s/ \ - --export_dir ../../efficientnetv2-s/saved_model -``` - -Where you should adapt `--model_name` to the corresponding model for the checkpoint used. The `--model_dir` argument should point to the downloaded or trained checkpoint as described above. The exported saved model will then be available in the directory pointed by the `--export_dir` argument. - -### Create ONNX Graph - -To generate an ONNX model file, find the saved model as described above, select a batch size and input size, and run: - -```bash -python3 create_onnx.py \ - --saved_model /path/to/saved_model \ - --onnx /path/to/model.onnx \ - --batch_size 1 \ - --input_size 384 -``` - -You may need to adapt the argument `--input_size` to explicitly define the exact input image dimensions to use in the graph. Consult the model definitions in the corresponding training system, to find the expected input size for the model you are working with. - -This will create the file `model.onnx` which is ready to convert to TensorRT. - -Optionally, you may wish to visualize the resulting ONNX graph with a tool such as [Netron](https://netron.app/). - -### Build TensorRT Engine - -It is possible to build the TensorRT engine directly with `trtexec` using the ONNX graph generated in the previous step. However, the script `build_engine.py` is provided for convenience, as it has been tailored to EfficientNet engine building and calibration. Run `python3 build_engine.py --help` for details on available settings. - -#### FP16 Precision - -To build the TensorRT engine file with FP16 precision, run: - -```bash -python3 build_engine.py \ - --onnx /path/to/model.onnx \ - --engine /path/to/engine.trt \ - --precision fp16 -``` - -The file `engine.trt` will be created, which can now be used to infer with TensorRT. - -For best results, make sure no other processes are using the GPU during engine build, as it may affect the optimal tactic selection process. - -#### INT8 Precision - -To build and calibrate an engine for INT8 precision, run: - -```bash -python3 build_engine.py \ - --onnx /path/to/model.onnx \ - --engine /path/to/engine.trt \ - --precision int8 \ - --calib_input /path/to/calibration/images \ - --calib_cache /path/to/calibration.cache \ - --calib_preprocessor V2 -``` - -Where `--calib_input` points to a directory with several thousands of images. For example, this could be a subset of the training or validation datasets that were used for the model. It's important that this data represents the runtime data distribution relatively well, therefore, the more images that are used for calibration, the better accuracy that will be achieved in INT8 precision. For ImageNet networks, we have found that 25,000 images gives a good result. - -The `--calib_cache` argument controls where the calibration cache file will be written to. This is useful to keep a cached copy of the calibration results. Next time you need to build the engine for the same network, if this file exists, it will skip the calibration step and use the cached values instead. - -Finally, the `--calib_preprocessor` option sets the preprocessing algorithm to apply on calibration images. Please refer to the [Input Preprocessing](#input-preprocessing) section below for more details. - -Run `python3 build_engine.py --help` for additional build options. - -### Benchmark TensorRT Engine - -Optionally, you can obtain execution timing information for the built engine by using the `trtexec` utility, as: - -```bash -trtexec \ - --loadEngine=/path/to/engine.trt \ - --useCudaGraph --noDataTransfers \ - --iterations=100 --avgRuns=100 -``` - -If it's not already in your `$PATH`, the `trtexec` binary is usually found in `/usr/src/tensorrt/bin/trtexec`, depending on your TensorRT installation method. - -An inference benchmark will run, with GPU Compute latency times printed out to the console. Depending on the version of TensorRT, you should see something similar to: - -``` -GPU Compute Time: min = 1.79895 ms, max = 1.9209 ms, mean = 1.80589 ms, median = 1.80493 ms, percentile(99%) = 1.81396 ms -``` - -## Inference - -For optimal performance, inference should be done in a C++ application that takes advantage of CUDA Graphs to launch the inference request. Alternatively, the TensorRT engine built with this process can also be executed through either [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server) or [DeepStream SDK](https://developer.nvidia.com/deepstream-sdk). - -However, for convenience, a python inference script is provided here for quick testing of the built TensorRT engine. - -### Input Preprocessing - -An important concept for computer vision models is the preprocessing applied to an image before feeding it to the classifier network. The various EfficientNet models supported by this converter use different preprocessing algorithms. - -We have implemented three different preprocessor algorithms, as defined in `image_batcher.py`. They are: - -| **Preprocessing** | **Resizing** | **Normalization** | **Mean Subtract** | -| ----------------- | ------------------------ | ----------------- | ----------------- | -| **V2** | Bilinear Resize | [-1 to +1] Range | No | -| **V1** | Bicubic Resize + PadCrop | [0 to +1] Range | No | -| **V1MS** | Bicubic Resize + PadCrop | [0 to +1] Range | Yes | - -**V2:** This is the preprocessor to be used with all EfficientNet V2 models. EfficientNet V2 does not require mean subtraction, so it is never performed for these models. - -**V1:** This is the default preprocessor to be used with most EfficientNet V1 models. EfficientNet V1 normally expects mean subtraction to be applied. However, some TensorFlow saved models, such as those downloaded from TFHub, already perform this operation within the graph itself, so it is not required to do it during preprocessing. - -**V1MS:** Depending on the saved model exporter, some EfficientNet V1 models may not have the integrated mean subtraction. This is often the case with models exported from the pre-trained *checkpoints*. For those cases, this preprocessor will apply mean subtraction during preprocessing. - -These are the supported values for `--preprocessor` and `--calib_preprocessor` arguments used throughout these scripts. Note that choosing an incorrect preprocessor for a model will considerably impact its accuracy. Please take a moment to choose the correct preprocessor to use before performing inference or validation of a model. - -### Inference in Python - -To classify a set of images with TensorRT, run: - -```bash -python3 infer.py \ - --engine /path/to/engine.trt \ - --input /path/to/images \ - --preprocessor V2 -``` - -Where the input path can be either a single image file, or a directory of jpg/png/bmp images. The classification results will be printed out to the console, one image per line, as: - -``` - -``` - -You can also redirect these results to a file, and optionally set a separator character (such as for CSV file creation): - -```bash -python3 infer.py \ - --engine /path/to/engine.trt \ - --input /path/to/ILSVRC2012_img_val \ - --preprocessor V2 \ - --separator ',' > results.csv -``` - -### Validate against Ground Truth - -To validate the TensorRT inference results accuracy against ground truth labels, run: - -```bash -python3 eval_gt.py \ - --engine /path/to/engine.trt \ - --annotations /path/to/val.txt \ - --input /path/to/images \ - --preprocessor V2 -``` - -The annotations file is expected to have one line per image, where the first column is the image filename, and the second column is the ground truth class label. For example: - -``` -ILSVRC2012_val_00000001.JPEG 65 -ILSVRC2012_val_00000002.JPEG 970 -ILSVRC2012_val_00000003.JPEG 230 -ILSVRC2012_val_00000004.JPEG 809 -[...] -``` - -> **NOTE:** The ImageNet pre-trained models follow the label mapping introduced by [Caffe](https://github.com/BVLC/caffe/blob/master/data/ilsvrc12/get_ilsvrc_aux.sh), which indexes labels according to their synset number. The validation file for this format can be downloaded from Caffe's ILSVRC2012 auxiliary package at: -> -> http://dl.caffe.berkeleyvision.org/caffe_ilsvrc12.tar.gz -> -> You can use the `val.txt` file bundled in this package for ImageNet evaluation purposes. - -Upon a successful run of `EfficientNet V2-S` on the `ILSVRC2012_img_val` [ImageNet](https://www.image-net.org/download.php) dataset, for example, you should see something like: - -``` -Top-1 Accuracy: 83.710% -Top-5 Accuracy: 96.615% -``` - -### Compare against TensorFlow - -Another method to validate the engine is to compare the TensorRT inference results with what TensorFlow produces, to make sure both frameworks give similar results. For this, run: - -```bash -python3 compare_tf.py \ - --engine /path/to/engine.trt \ - --saved_model /path/to/saved_model \ - --input /path/to/images \ - --preprocessor V2 -``` - -This can be performed on any set of images, no ground truth is required. The script executes both the TensorFlow saved model and the TensorRT engine simultaneously on the given input images. It then computes the class prediction similarity and RMSE in confidence scores between both outputs. - -Upon a successful run, you should see something like: - -``` -Matching Top-1 class predictions for 4999 out of 5000 images: 99.98% -RMSE between TensorFlow and TensorRT confidence scores: 0.006 -``` diff --git a/samples/python/efficientnet/build_engine.py b/samples/python/efficientnet/build_engine.py deleted file mode 100644 index 38b1cc12..00000000 --- a/samples/python/efficientnet/build_engine.py +++ /dev/null @@ -1,321 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -import sys -import logging -import argparse - -import numpy as np -import tensorrt as trt -from cuda.bindings import runtime as cudart - -sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)) -import common - -from image_batcher import ImageBatcher - -logging.basicConfig(level=logging.INFO) -logging.getLogger("EngineBuilder").setLevel(logging.INFO) -log = logging.getLogger("EngineBuilder") - - -class EngineCalibrator(trt.IInt8EntropyCalibrator2): - """ - Implements the INT8 Entropy Calibrator 2. - """ - - def __init__(self, cache_file): - """ - :param cache_file: The location of the cache file. - """ - super().__init__() - self.cache_file = cache_file - self.image_batcher = None - self.batch_allocation = None - self.batch_generator = None - - def set_image_batcher(self, image_batcher: ImageBatcher): - """ - Define the image batcher to use, if any. If using only the cache file, an image batcher doesn't need - to be defined. - :param image_batcher: The ImageBatcher object - """ - self.image_batcher = image_batcher - size = int( - np.dtype(self.image_batcher.dtype).itemsize - * np.prod(self.image_batcher.shape) - ) - self.batch_allocation = common.cuda_call(cudart.cudaMalloc(size)) - self.batch_generator = self.image_batcher.get_batch() - - def get_batch_size(self): - """ - Overrides from trt.IInt8EntropyCalibrator2. - Get the batch size to use for calibration. - :return: Batch size. - """ - if self.image_batcher: - return self.image_batcher.batch_size - return 1 - - def get_batch(self, names): - """ - Overrides from trt.IInt8EntropyCalibrator2. - Get the next batch to use for calibration, as a list of device memory pointers. - :param names: The names of the inputs, if useful to define the order of inputs. - :return: A list of int-casted memory pointers. - """ - if not self.image_batcher: - return None - try: - batch, _ = next(self.batch_generator) - log.info( - "Calibrating image {} / {}".format( - self.image_batcher.image_index, self.image_batcher.num_images - ) - ) - common.memcpy_host_to_device( - self.batch_allocation, np.ascontiguousarray(batch) - ) - return [int(self.batch_allocation)] - except StopIteration: - log.info("Finished calibration batches") - return None - - def read_calibration_cache(self): - """ - Overrides from trt.IInt8EntropyCalibrator2. - Read the calibration cache file stored on disk, if it exists. - :return: The contents of the cache file, if any. - """ - if os.path.exists(self.cache_file): - with open(self.cache_file, "rb") as f: - log.info("Using calibration cache file: {}".format(self.cache_file)) - return f.read() - - def write_calibration_cache(self, cache): - """ - Overrides from trt.IInt8EntropyCalibrator2. - Store the calibration cache to a file on disk. - :param cache: The contents of the calibration cache to store. - """ - with open(self.cache_file, "wb") as f: - log.info("Writing calibration cache data to: {}".format(self.cache_file)) - f.write(cache) - - -class EngineBuilder: - """ - Parses an ONNX graph and builds a TensorRT engine from it. - """ - - def __init__(self, verbose=False): - """ - :param verbose: If enabled, a higher verbosity level will be set on the TensorRT logger. - """ - self.trt_logger = trt.Logger(trt.Logger.INFO) - if verbose: - self.trt_logger.min_severity = trt.Logger.Severity.VERBOSE - - trt.init_libnvinfer_plugins(self.trt_logger, namespace="") - - self.builder = trt.Builder(self.trt_logger) - self.config = self.builder.create_builder_config() - self.config.set_memory_pool_limit( - trt.MemoryPoolType.WORKSPACE, 8 * (2**30) - ) # 8 GB - - self.batch_size = None - self.network = None - self.parser = None - - def create_network(self, onnx_path): - """ - Parse the ONNX graph and create the corresponding TensorRT network definition. - :param onnx_path: The path to the ONNX graph to load. - """ - - self.network = self.builder.create_network(0) - self.parser = trt.OnnxParser(self.network, self.trt_logger) - - onnx_path = os.path.realpath(onnx_path) - with open(onnx_path, "rb") as f: - if not self.parser.parse(f.read()): - log.error("Failed to load ONNX file: {}".format(onnx_path)) - for error in range(self.parser.num_errors): - log.error(self.parser.get_error(error)) - sys.exit(1) - - inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)] - outputs = [self.network.get_output(i) for i in range(self.network.num_outputs)] - - log.info("Network Description") - for input in inputs: - self.batch_size = input.shape[0] - log.info( - "Input '{}' with shape {} and dtype {}".format( - input.name, input.shape, input.dtype - ) - ) - for output in outputs: - log.info( - "Output '{}' with shape {} and dtype {}".format( - output.name, output.shape, output.dtype - ) - ) - assert self.batch_size > 0 - - def create_engine( - self, - engine_path, - precision, - calib_input=None, - calib_cache=None, - calib_num_images=25000, - calib_batch_size=8, - calib_preprocessor=None, - ): - """ - Build the TensorRT engine and serialize it to disk. - :param engine_path: The path where to serialize the engine to. - :param precision: The datatype to use for the engine, either 'fp32', 'fp16' or 'int8'. - :param calib_input: The path to a directory holding the calibration images. - :param calib_cache: The path where to write the calibration cache to, or if it already exists, load it from. - :param calib_num_images: The maximum number of images to use for calibration. - :param calib_batch_size: The batch size to use for the calibration process. - :param calib_preprocessor: The ImageBatcher preprocessor algorithm to use. - """ - engine_path = os.path.realpath(engine_path) - engine_dir = os.path.dirname(engine_path) - os.makedirs(engine_dir, exist_ok=True) - log.info("Building {} Engine in {}".format(precision, engine_path)) - - inputs = [self.network.get_input(i) for i in range(self.network.num_inputs)] - - log.info("Reading timing cache from file: {:}".format(args.timing_cache)) - common.setup_timing_cache(self.config, args.timing_cache) - - if precision == "fp16": - if not self.builder.platform_has_fast_fp16: - log.warning("FP16 is not supported natively on this platform/device") - else: - self.config.set_flag(trt.BuilderFlag.FP16) - elif precision == "int8": - if not self.builder.platform_has_fast_int8: - log.warning("INT8 is not supported natively on this platform/device") - else: - self.config.set_flag(trt.BuilderFlag.INT8) - self.config.int8_calibrator = EngineCalibrator(calib_cache) - if not os.path.exists(calib_cache): - calib_shape = [calib_batch_size] + list(inputs[0].shape[1:]) - calib_dtype = trt.nptype(inputs[0].dtype) - self.config.int8_calibrator.set_image_batcher( - ImageBatcher( - calib_input, - calib_shape, - calib_dtype, - max_num_images=calib_num_images, - exact_batches=True, - preprocessor=calib_preprocessor, - ) - ) - - engine_bytes = self.builder.build_serialized_network(self.network, self.config) - if engine_bytes is None: - log.error("Failed to create engine") - sys.exit(1) - - log.info("Serializing timing cache to file: {:}".format(args.timing_cache)) - common.save_timing_cache(self.config, args.timing_cache) - - with open(engine_path, "wb") as f: - log.info("Serializing engine to file: {:}".format(engine_path)) - f.write(engine_bytes) - - -def main(args): - builder = EngineBuilder(args.verbose) - builder.create_network(args.onnx) - builder.create_engine( - args.engine, - args.precision, - args.calib_input, - args.calib_cache, - args.calib_num_images, - args.calib_batch_size, - args.calib_preprocessor, - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-o", "--onnx", help="The input ONNX model file to load") - parser.add_argument("-e", "--engine", help="The output path for the TRT engine") - parser.add_argument( - "-p", - "--precision", - default="fp16", - choices=["fp32", "fp16", "int8"], - help="The precision mode to build in, either 'fp32', 'fp16' or 'int8', default: 'fp16'", - ) - parser.add_argument( - "-v", "--verbose", action="store_true", help="Enable more verbose log output" - ) - parser.add_argument( - "--calib_input", help="The directory holding images to use for calibration" - ) - parser.add_argument( - "--calib_cache", - default="./calibration.cache", - help="The file path for INT8 calibration cache to use, default: ./calibration.cache", - ) - parser.add_argument( - "--calib_num_images", - default=25000, - type=int, - help="The maximum number of images to use for calibration, default: 25000", - ) - parser.add_argument( - "--calib_batch_size", - default=8, - type=int, - help="The batch size for the calibration process, default: 1", - ) - parser.add_argument( - "--calib_preprocessor", - default="V2", - choices=["V1", "V1MS", "V2"], - help="Set the calibration image preprocessor to use, either 'V2', 'V1' or 'V1MS', default: V2", - ) - parser.add_argument( - "--timing_cache", - default="./timing.cache", - help="The file path for timing cache, default: ./timing.cache", - ) - args = parser.parse_args() - if not all([args.onnx, args.engine]): - parser.print_help() - log.error("These arguments are required: --onnx and --engine") - sys.exit(1) - if args.precision == "int8" and not any([args.calib_input, args.calib_cache]): - parser.print_help() - log.error( - "When building in int8 precision, either --calib_input or --calib_cache are required" - ) - sys.exit(1) - main(args) diff --git a/samples/python/efficientnet/compare_tf.py b/samples/python/efficientnet/compare_tf.py deleted file mode 100644 index 2671572e..00000000 --- a/samples/python/efficientnet/compare_tf.py +++ /dev/null @@ -1,196 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -import sys -import argparse - -import numpy as np -import tensorflow as tf - -from infer import TensorRTInfer -from image_batcher import ImageBatcher - - -class TensorFlowInfer: - """ - Implements TensorFlow inference of a saved model, following the same API as the TensorRTInfer class. - """ - - def __init__(self, saved_model_path): - gpus = tf.config.experimental.list_physical_devices("GPU") - for gpu in gpus: - tf.config.experimental.set_memory_growth(gpu, True) - - self.model = tf.saved_model.load(saved_model_path) - self.pred_fn = self.model.signatures["serving_default"] - - # Setup I/O bindings - self.inputs = [] - fn_inputs = self.pred_fn.structured_input_signature[1] - for i, input in enumerate(list(fn_inputs.values())): - self.inputs.append( - { - "index": i, - "name": input.name, - "dtype": np.dtype(input.dtype.as_numpy_dtype()), - "shape": input.shape.as_list(), - } - ) - self.outputs = [] - fn_outputs = self.pred_fn.structured_outputs - for i, output in enumerate(list(fn_outputs.values())): - self.outputs.append( - { - "index": i, - "name": output.name, - "dtype": np.dtype(output.dtype.as_numpy_dtype()), - "shape": output.shape.as_list(), - } - ) - - def input_spec(self): - return self.inputs[0]["shape"], self.inputs[0]["dtype"] - - def output_spec(self): - return self.outputs[0]["shape"], self.outputs[0]["dtype"] - - def infer(self, batch, top=1): - # Process I/O and execute the network - input = {self.inputs[0]["name"]: tf.convert_to_tensor(batch)} - output = self.pred_fn(**input) - output = output[self.outputs[0]["name"]].numpy() - - # Read and process the results - classes = np.argmax(output, axis=1) - scores = np.max(output, axis=1) - top = max(top, output.shape[1]) - top_classes = np.flip(np.argsort(output, axis=1), axis=1)[:, 0:top] - top_scores = np.flip(np.sort(output, axis=1), axis=1)[:, 0:top] - - return classes, scores, [top_classes, top_scores] - - -def main(args): - # Initialize TRT and TF infer objects. - tf_infer = TensorFlowInfer(args.saved_model) - trt_infer = TensorRTInfer(args.engine) - - batcher = ImageBatcher( - args.input, - *trt_infer.input_spec(), - max_num_images=args.num_images, - preprocessor=args.preprocessor - ) - - # Make sure both systems use the same input spec, so we can use the exact same image batches with both - tf_shape, tf_dtype = tf_infer.input_spec() - trt_shape, trt_dtype = trt_infer.input_spec() - if trt_dtype != tf_dtype: - print("Input datatype does not match") - print("TRT Engine Input Dtype: {} {}".format(trt_dtype)) - print("TF Saved Model Input Dtype: {} {}".format(tf_dtype)) - print( - "Please use the same TensorFlow saved model that the TensorRT engine was built with" - ) - sys.exit(1) - - if (tf_shape[1] and trt_shape[1] != tf_shape[1]) or ( - tf_shape[2] and trt_shape[2] != tf_shape[2] - ): - print("Input shapes do not match") - print("TRT Engine Input Shape: {} {}".format(trt_shape[1:])) - print("TF Saved Model Input Shape: {} {}".format(tf_shape[1:])) - print( - "Please use the same TensorFlow saved model that the TensorRT engine was built with" - ) - sys.exit(1) - - match = 0 - error = 0 - for batch, images in batcher.get_batch(): - # Run inference on the same batch with both inference systems - tf_classes, tf_scores, _ = tf_infer.infer(batch) - trt_classes, trt_scores, _ = trt_infer.infer(batch) - - # The last batch may not have all image slots filled, so limit the results to only the amount of actual images - tf_classes = tf_classes[0 : len(images)] - tf_scores = tf_scores[0 : len(images)] - trt_classes = trt_classes[0 : len(images)] - trt_scores = trt_scores[0 : len(images)] - - # Track how many images match on top-1 class id predictions - match += np.sum(trt_classes == tf_classes) - # Track the mean square error in confidence score - error += np.sum((trt_scores - tf_scores) * (trt_scores - tf_scores)) - - print( - "Processing {} / {} images: {:.2f}% match ".format( - batcher.image_index, - batcher.num_images, - (100 * (match / batcher.image_index)), - ), - end="\r", - ) - - print() - pc = 100 * (match / batcher.num_images) - print( - "Matching Top-1 class predictions for {} out of {} images: {:.2f}%".format( - match, batcher.num_images, pc - ) - ) - avgerror = np.sqrt(error / batcher.num_images) - print( - "RMSE between TensorFlow and TensorRT confidence scores: {:.3f}".format( - avgerror - ) - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-e", "--engine", help="The TensorRT engine to infer with") - parser.add_argument( - "-m", - "--saved_model", - help="The TensorFlow saved model path to validate against", - ) - parser.add_argument( - "-i", - "--input", - help="The input to infer, either a single image path, or a directory of images", - ) - parser.add_argument( - "-n", - "--num_images", - default=5000, - type=int, - help="The maximum number of images to use for validation, default: 5000", - ) - parser.add_argument( - "-p", - "--preprocessor", - default="V2", - choices=["V1", "V1MS", "V2"], - help="Select the image preprocessor to use, either 'V2', 'V1' or 'V1MS', default: V2", - ) - args = parser.parse_args() - if not all([args.engine, args.saved_model, args.input]): - parser.print_help() - sys.exit(1) - main(args) diff --git a/samples/python/efficientnet/create_onnx.py b/samples/python/efficientnet/create_onnx.py deleted file mode 100644 index c0e7d109..00000000 --- a/samples/python/efficientnet/create_onnx.py +++ /dev/null @@ -1,122 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -import sys -import argparse - -import onnx -import onnx_graphsurgeon as gs -from onnx import shape_inference - -import numpy as np -import tensorflow as tf -from tf2onnx import tfonnx, optimizer, tf_loader - - -def main(args): - # Load saved model - saved_model_path = os.path.realpath(args.saved_model) - assert os.path.isdir(saved_model_path) - graph_def, inputs, outputs = tf_loader.from_saved_model( - saved_model_path, None, None, "serve", ["serving_default"] - ) - with tf.Graph().as_default() as tf_graph: - tf.import_graph_def(graph_def, name="") - with tf_loader.tf_session(graph=tf_graph): - onnx_graph = tfonnx.process_tf_graph( - tf_graph, input_names=inputs, output_names=outputs, opset=11 - ) - onnx_model = optimizer.optimize_graph(onnx_graph).make_model( - "Converted from {}".format(saved_model_path) - ) - graph = gs.import_onnx(onnx_model) - assert graph - print() - print("ONNX graph created successfully") - - # Set the I/O tensor shapes - graph.inputs[0].shape[0] = args.batch_size - graph.outputs[0].shape[0] = args.batch_size - if args.input_size and args.input_size > 0: - if graph.inputs[0].shape[3] == 3: - # Format NHWC - graph.inputs[0].shape[1] = args.input_size - graph.inputs[0].shape[2] = args.input_size - elif graph.inputs[0].shape[1] == 3: - # Format NCHW - graph.inputs[0].shape[2] = args.input_size - graph.inputs[0].shape[3] = args.input_size - print( - "ONNX input named '{}' with shape {}".format( - graph.inputs[0].name, graph.inputs[0].shape - ) - ) - print( - "ONNX output named '{}' with shape {}".format( - graph.outputs[0].name, graph.outputs[0].shape - ) - ) - for i in range(4): - if type(graph.inputs[0].shape[i]) != int or graph.inputs[0].shape[i] <= 0: - print( - "The input shape of the graph is invalid, try overriding it by giving a fixed size with --input_size" - ) - sys.exit(1) - - # Fix Clip Nodes (ReLU6) - for node in [n for n in graph.nodes if n.op == "Clip"]: - for input in node.inputs[1:]: - # In TensorRT, the min/max inputs on a Clip op *must* have fp32 datatype - input.values = np.float32(input.values) - - # Run tensor shape inference - graph.cleanup().toposort() - model = shape_inference.infer_shapes(gs.export_onnx(graph)) - graph = gs.import_onnx(model) - - # Save updated model - graph.cleanup().toposort() - model = gs.export_onnx(graph) - onnx_path = os.path.realpath(args.onnx) - os.makedirs(os.path.dirname(onnx_path), exist_ok=True) - onnx.save(model, onnx_path) - engine_path = os.path.join(os.path.dirname(onnx_path), "engine.trt") - print("ONNX model saved to {}".format(onnx_path)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "-m", "--saved_model", help="The TensorFlow saved model directory to load" - ) - parser.add_argument("-o", "--onnx", help="The output ONNX model file to write") - parser.add_argument( - "-b", "--batch_size", type=int, default=1, help="Set the batch size, default: 1" - ) - parser.add_argument( - "-i", - "--input_size", - type=int, - help="Override the input height and width, e.g. '380', default: keep original size", - ) - args = parser.parse_args() - if not all([args.saved_model, args.onnx]): - parser.print_help() - print("\nThese arguments are required: --saved_model and --onnx") - sys.exit(1) - main(args) diff --git a/samples/python/efficientnet/eval_gt.py b/samples/python/efficientnet/eval_gt.py deleted file mode 100644 index 9f57aaa5..00000000 --- a/samples/python/efficientnet/eval_gt.py +++ /dev/null @@ -1,116 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -import sys -import argparse - -import numpy as np - -from infer import TensorRTInfer -from image_batcher import ImageBatcher - - -def main(args): - annotations = {} - for line in open(args.annotations, "r"): - line = line.strip().split(args.separator) - if len(line) < 2 or not line[1].isnumeric(): - print( - "Could not parse the annotations file correctly, make sure the correct separator is used" - ) - sys.exit(1) - annotations[os.path.basename(line[0])] = int(line[1]) - - trt_infer = TensorRTInfer(args.engine) - batcher = ImageBatcher( - args.input, - *trt_infer.input_spec(), - max_num_images=args.num_images, - preprocessor=args.preprocessor - ) - top1 = 0 - top5 = 0 - total = 0 - for batch, images in batcher.get_batch(): - classes, scores, top = trt_infer.infer(batch, top=5) - for i in range(len(images)): - image = os.path.basename(images[i]) - if image not in annotations.keys(): - print( - "Image '{}' does not appear in the annotations file, please make sure all evaluated " - "images have a corresponding ground truth label".format(image) - ) - sys.exit(1) - if annotations[image] == classes[i]: - top1 += 1 - if annotations[image] in top[0][i]: - top5 += 1 - total += 1 - top1_acc = 100 * (top1 / total) - top5_acc = 100 * (top5 / total) - print( - "Processing {} / {} : Top-1 {:0.1f}% , Top-5: {:0.1f}% ".format( - total, batcher.num_images, top1_acc, top5_acc - ), - end="\r", - ) - print() - print("Top-1 Accuracy: {:0.3f}%".format(top1_acc)) - print("Top-5 Accuracy: {:0.3f}%".format(top5_acc)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-e", "--engine", help="The TensorRT engine to infer with") - parser.add_argument( - "-i", - "--input", - help="The input to infer, either a single image path, or a directory of images", - ) - parser.add_argument( - "-a", - "--annotations", - help="Set the file to use for classification ground truth annotations", - ) - parser.add_argument( - "-s", - "--separator", - default=" ", - help="Separator to use between columns when parsing the annotations file, default: ' ' (space)", - ) - parser.add_argument( - "-p", - "--preprocessor", - default="V2", - choices=["V1", "V1MS", "V2"], - help="Select the image preprocessor to use, either 'V2', 'V1' or 'V1MS', default: V2", - ) - parser.add_argument( - "-n", - "--num_images", - default=5000, - type=int, - help="The maximum number of images to use for validation, default: 5000", - ) - args = parser.parse_args() - if not all([args.engine, args.input, args.annotations]): - parser.print_help() - print("\nThese arguments are required: --engine --input and --annotations") - sys.exit(1) - - main(args) diff --git a/samples/python/efficientnet/image_batcher.py b/samples/python/efficientnet/image_batcher.py deleted file mode 100644 index 63d37784..00000000 --- a/samples/python/efficientnet/image_batcher.py +++ /dev/null @@ -1,191 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -import sys - -import numpy as np -from PIL import Image - - -class ImageBatcher: - """ - Creates batches of pre-processed images. - """ - - def __init__( - self, - input, - shape, - dtype, - max_num_images=None, - exact_batches=False, - preprocessor="V2", - ): - """ - :param input: The input directory to read images from. - :param shape: The tensor shape of the batch to prepare, either in NCHW or NHWC format. - :param dtype: The (numpy) datatype to cast the batched data to. - :param max_num_images: The maximum number of images to read from the directory. - :param exact_batches: This defines how to handle a number of images that is not an exact multiple of the batch - size. If false, it will pad the final batch with zeros to reach the batch size. If true, it will *remove* the - last few images in excess of a batch size multiple, to guarantee batches are exact (useful for calibration). - :param preprocessor: Set the preprocessor to use, V1 or V2, depending on which network is being used. - """ - # Find images in the given input path - input = os.path.realpath(input) - self.images = [] - - extensions = [".jpg", ".jpeg", ".png", ".bmp"] - - def is_image(path): - return ( - os.path.isfile(path) and os.path.splitext(path)[1].lower() in extensions - ) - - if os.path.isdir(input): - self.images = [ - os.path.join(input, f) - for f in os.listdir(input) - if is_image(os.path.join(input, f)) - ] - self.images.sort() - elif os.path.isfile(input): - if is_image(input): - self.images.append(input) - self.num_images = len(self.images) - if self.num_images < 1: - print("No valid {} images found in {}".format("/".join(extensions), input)) - sys.exit(1) - - # Handle Tensor Shape - self.dtype = dtype - self.shape = shape - assert len(self.shape) == 4 - self.batch_size = shape[0] - assert self.batch_size > 0 - self.format = None - self.width = -1 - self.height = -1 - if self.shape[1] == 3: - self.format = "NCHW" - self.height = self.shape[2] - self.width = self.shape[3] - elif self.shape[3] == 3: - self.format = "NHWC" - self.height = self.shape[1] - self.width = self.shape[2] - assert all([self.format, self.width > 0, self.height > 0]) - - # Adapt the number of images as needed - if max_num_images and 0 < max_num_images < len(self.images): - self.num_images = max_num_images - if exact_batches: - self.num_images = self.batch_size * (self.num_images // self.batch_size) - if self.num_images < 1: - print("Not enough images to create batches") - sys.exit(1) - self.images = self.images[0 : self.num_images] - - # Subdivide the list of images into batches - self.num_batches = 1 + int((self.num_images - 1) / self.batch_size) - self.batches = [] - for i in range(self.num_batches): - start = i * self.batch_size - end = min(start + self.batch_size, self.num_images) - self.batches.append(self.images[start:end]) - - # Indices - self.image_index = 0 - self.batch_index = 0 - - self.preprocessor = preprocessor - - def preprocess_image(self, image_path): - """ - The image preprocessor loads an image from disk and prepares it as needed for batching. This includes cropping, - resizing, normalization, data type casting, and transposing. - This Image Batcher implements two algorithms: - * V2: The algorithm for EfficientNet V2, as defined in automl/efficientnetv2/preprocessing.py. - * V1: The algorithm for EfficientNet V1, aka "Legacy", as defined in automl/efficientnetv2/preprocess_legacy.py. - :param image_path: The path to the image on disk to load. - :return: A numpy array holding the image sample, ready to be contacatenated into the rest of the batch. - """ - - def pad_crop(image): - """ - A subroutine to implement padded cropping. This will create a center crop of the image, padded by 32 pixels. - :param image: The PIL image object - :return: The PIL image object already padded and cropped. - """ - # Assume square images - assert self.height == self.width - width, height = image.size - ratio = self.height / (self.height + 32) - crop_size = int(ratio * min(height, width)) - y = (height - crop_size) // 2 - x = (width - crop_size) // 2 - return image.crop((x, y, x + crop_size, y + crop_size)) - - image = Image.open(image_path) - image = image.convert(mode="RGB") - if self.preprocessor == "V2": - # For EfficientNet V2: Bilinear Resize and [-1,+1] Normalization - if self.height < 320: - # Padded crop only on smaller sizes - image = pad_crop(image) - image = image.resize((self.width, self.height), resample=Image.BILINEAR) - image = np.asarray(image, dtype=self.dtype) - image = (image - 128.0) / 128.0 - elif self.preprocessor == "V1": - # For EfficientNet V1: Padded Crop, Bicubic Resize, and [0,1] Normalization - # (Mean subtraction and Std Dev scaling will be part of the graph, so not done here) - image = pad_crop(image) - image = image.resize((self.width, self.height), resample=Image.BICUBIC) - image = np.asarray(image, dtype=self.dtype) - image = image / 255.0 - elif self.preprocessor == "V1MS": - # For EfficientNet V1: Padded Crop, Bicubic Resize, and [0,1] Normalization - # Mean subtraction and Std dev scaling are applied as a pre-processing step outside the graph. - image = pad_crop(image) - image = image.resize((self.width, self.height), resample=Image.BICUBIC) - image = np.asarray(image, dtype=self.dtype) - image = image - np.asarray([123.68, 116.28, 103.53]) - image = image / np.asarray([58.395, 57.120, 57.375]) - else: - print("Preprocessing method {} not supported".format(self.preprocessor)) - sys.exit(1) - if self.format == "NCHW": - image = np.transpose(image, (2, 0, 1)) - return image - - def get_batch(self): - """ - Retrieve the batches. This is a generator object, so you can use it within a loop as: - for batch, images in batcher.get_batch(): - ... - Or outside of a batch with the next() function. - :return: A generator yielding two items per iteration: a numpy array holding a batch of images, and the list of - paths to the images loaded within this batch. - """ - for i, batch_images in enumerate(self.batches): - batch_data = np.zeros(self.shape, dtype=self.dtype) - for i, image in enumerate(batch_images): - self.image_index += 1 - batch_data[i] = self.preprocess_image(image) - self.batch_index += 1 - yield batch_data, batch_images diff --git a/samples/python/efficientnet/infer.py b/samples/python/efficientnet/infer.py deleted file mode 100644 index 5d469223..00000000 --- a/samples/python/efficientnet/infer.py +++ /dev/null @@ -1,182 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -import sys -import argparse - -import numpy as np -import tensorrt as trt -from cuda.bindings import runtime as cudart - -sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)) -import common - -from image_batcher import ImageBatcher - - -class TensorRTInfer: - """ - Implements inference for the EfficientNet TensorRT engine. - """ - - def __init__(self, engine_path): - """ - :param engine_path: The path to the serialized engine to load from disk. - """ - # Load TRT engine - self.logger = trt.Logger(trt.Logger.ERROR) - with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime: - assert runtime - self.engine = runtime.deserialize_cuda_engine(f.read()) - assert self.engine - self.context = self.engine.create_execution_context() - assert self.context - - # Setup I/O bindings - self.inputs = [] - self.outputs = [] - self.allocations = [] - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - is_input = False - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: - is_input = True - dtype = self.engine.get_tensor_dtype(name) - shape = self.engine.get_tensor_shape(name) - if is_input: - self.batch_size = shape[0] - size = np.dtype(trt.nptype(dtype)).itemsize - for s in shape: - size *= s - allocation = common.cuda_call(cudart.cudaMalloc(size)) - binding = { - "index": i, - "name": name, - "dtype": np.dtype(trt.nptype(dtype)), - "shape": list(shape), - "allocation": allocation, - } - self.allocations.append(allocation) - if is_input: - self.inputs.append(binding) - else: - self.outputs.append(binding) - - assert self.batch_size > 0 - assert len(self.inputs) > 0 - assert len(self.outputs) > 0 - assert len(self.allocations) > 0 - - def input_spec(self): - """ - Get the specs for the input tensor of the network. Useful to prepare memory allocations. - :return: Two items, the shape of the input tensor and its (numpy) datatype. - """ - return self.inputs[0]["shape"], self.inputs[0]["dtype"] - - def output_spec(self): - """ - Get the specs for the output tensor of the network. Useful to prepare memory allocations. - :return: Two items, the shape of the output tensor and its (numpy) datatype. - """ - return self.outputs[0]["shape"], self.outputs[0]["dtype"] - - def infer(self, batch, top=1): - """ - Execute inference on a batch of images. The images should already be batched and preprocessed, as prepared by - the ImageBatcher class. Memory copying to and from the GPU device will be performed here. - :param batch: A numpy array holding the image batch. - :param top: The number of classes to return as top_predicitons, in descending order by their score. By default, - setting to one will return the same as the maximum score class. Useful for Top-5 accuracy metrics in validation. - :return: Three items, as numpy arrays for each batch image: The maximum score class, the corresponding maximum - score, and a list of the top N classes and scores. - """ - # Prepare the output data - output = np.zeros(*self.output_spec()) - - # Process I/O and execute the network - common.memcpy_host_to_device( - self.inputs[0]["allocation"], np.ascontiguousarray(batch) - ) - self.context.execute_v2(self.allocations) - common.memcpy_device_to_host(output, self.outputs[0]["allocation"]) - - # Process the results - classes = np.argmax(output, axis=1) - scores = np.max(output, axis=1) - top = min(top, output.shape[1]) - top_classes = np.flip(np.argsort(output, axis=1), axis=1)[:, 0:top] - top_scores = np.flip(np.sort(output, axis=1), axis=1)[:, 0:top] - - return classes, scores, [top_classes, top_scores] - - -def main(args): - trt_infer = TensorRTInfer(args.engine) - batcher = ImageBatcher( - args.input, *trt_infer.input_spec(), preprocessor=args.preprocessor - ) - for batch, images in batcher.get_batch(): - classes, scores, top = trt_infer.infer(batch) - for i in range(len(images)): - if args.top == 1: - print(images[i], classes[i], scores[i], sep=args.separator) - else: - line = [images[i]] - assert args.top <= top[0].shape[1] - for t in range(args.top): - line.append(str(top[0][i][t])) - for t in range(args.top): - line.append(str(top[1][i][t])) - print(args.separator.join(line)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-e", "--engine", help="The TensorRT engine to infer with") - parser.add_argument( - "-i", - "--input", - help="The input to infer, either a single image path, or a directory of images", - ) - parser.add_argument( - "-t", - "--top", - default=1, - type=int, - help="The amount of top classes and scores to output per image, default: 1", - ) - parser.add_argument( - "-s", - "--separator", - default="\t", - help="Separator to use between columns when printing the results, default: \\t", - ) - parser.add_argument( - "-p", - "--preprocessor", - default="V2", - choices=["V1", "V1MS", "V2"], - help="Select the image preprocessor to use, either 'V2', 'V1' or 'V1MS', default: V2", - ) - args = parser.parse_args() - if not all([args.engine, args.input]): - parser.print_help() - print("\nThese arguments are required: --engine and --input") - sys.exit(1) - main(args) diff --git a/samples/python/efficientnet/requirements.txt b/samples/python/efficientnet/requirements.txt deleted file mode 100644 index b355362e..00000000 --- a/samples/python/efficientnet/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ -Pillow==11.3.0 -onnx==1.16.1 -tensorrt>=7.1.0.0 -tf2onnx==1.16.0 -cuda-python==12.9.0 -pywin32; platform_system == "Windows" -pyyaml==6.0.1 -requests==2.32.2 -tqdm==4.66.4 -numpy==1.26.4 diff --git a/samples/python/engine_refit_onnx_bidaf/build_and_refit_engine.py b/samples/python/engine_refit_onnx_bidaf/build_and_refit_engine.py index fe0b0551..3de02557 100644 --- a/samples/python/engine_refit_onnx_bidaf/build_and_refit_engine.py +++ b/samples/python/engine_refit_onnx_bidaf/build_and_refit_engine.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,7 +23,7 @@ import numpy as np import argparse import tensorrt as trt -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) from cuda.bindings import runtime as cudart TRT_LOGGER = trt.Logger() @@ -201,7 +201,7 @@ def main(): if args.weights_location == "GPU": for name, weights in refit_weights_dict.items(): nbytes = weights.size * weights.itemsize - device_mem_dict[name] = common.cuda_call(cudart.cudaMalloc(nbytes)) + device_mem_dict[name] = common.DeviceMem(nbytes) execution_context = engine.create_execution_context() refitter = trt.Refitter(engine, TRT_LOGGER) @@ -213,7 +213,7 @@ def main(): if args.weights_location == "GPU": for name, device_mem in device_mem_dict.items(): device_weights = trt.Weights( - trt.DataType.FLOAT, device_mem, refit_weights_dict[name].size + trt.DataType.FLOAT, device_mem.device_ptr, refit_weights_dict[name].size ) weights_prototype = refitter.get_weights_prototype(name) assert device_weights.dtype == weights_prototype.dtype @@ -236,7 +236,7 @@ def main(): location = trt.TensorLocation.HOST refitter.set_named_weights(name, weights, location) else: - common.memcpy_host_to_device(device_mem_dict[name], host_weights) + common.memcpy_host_to_device(device_mem_dict[name].device_ptr, host_weights) # Get missing weights names. This should return empty lists in this case. missing_weights = refitter.get_missing_weights() @@ -254,40 +254,42 @@ def main(): for profile_idx in range(engine.num_optimization_profiles): print("Doing inference...") # Do inference - inputs, outputs, bindings, stream = common.allocate_buffers( + inputs, outputs, bindings = common.allocate_buffers( engine, profile_idx ) padding_bindings = [0] * (len(bindings) * profile_idx) new_bindings = padding_bindings + bindings - # Set host input. The common.do_inference function will copy the input to the GPU before executing. - inputs[0].host = cw - inputs[1].host = cc - inputs[2].host = qw - inputs[3].host = qc - execution_context.set_optimization_profile_async(profile_idx, stream) - execution_context.set_input_shape("CategoryMapper_4", (10, 1)) - execution_context.set_input_shape("CategoryMapper_5", (10, 1, 1, 16)) - execution_context.set_input_shape("CategoryMapper_6", (6, 1)) - execution_context.set_input_shape("CategoryMapper_7", (6, 1, 1, 16)) + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + # Set host input. The common.do_inference function will copy the input to the GPU before executing. + inputs[0].host = cw + inputs[1].host = cc + inputs[2].host = qw + inputs[3].host = qc + execution_context.set_optimization_profile_async(profile_idx, stream.stream) + execution_context.set_input_shape("CategoryMapper_4", (10, 1)) + execution_context.set_input_shape("CategoryMapper_5", (10, 1, 1, 16)) + execution_context.set_input_shape("CategoryMapper_6", (6, 1)) + execution_context.set_input_shape("CategoryMapper_7", (6, 1, 1, 16)) - trt_outputs = common.do_inference( - execution_context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) + trt_outputs = common.do_inference( + execution_context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) start = trt_outputs[0].item() end = trt_outputs[1].item() answer = [w.encode() for w in cw_str[start : end + 1].reshape(-1)] assert answer_correct == (answer == [b"brown"]), answer - common.free_buffers(inputs, outputs, stream) + common.free_buffers(inputs, outputs) for _, device_mem in device_mem_dict.items(): - common.cuda_call(cudart.cudaFree(device_mem)) + device_mem.free() print("Passed") diff --git a/samples/python/engine_refit_onnx_bidaf/prepare_model.py b/samples/python/engine_refit_onnx_bidaf/prepare_model.py index eb45226e..8f2ed178 100644 --- a/samples/python/engine_refit_onnx_bidaf/prepare_model.py +++ b/samples/python/engine_refit_onnx_bidaf/prepare_model.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,7 +22,7 @@ import json import sys, os -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) from downloader import getFilePath diff --git a/samples/python/engine_refit_onnx_bidaf/requirements.txt b/samples/python/engine_refit_onnx_bidaf/requirements.txt index 17da94c1..b25548e7 100644 --- a/samples/python/engine_refit_onnx_bidaf/requirements.txt +++ b/samples/python/engine_refit_onnx_bidaf/requirements.txt @@ -1,9 +1,9 @@ -onnx==1.16.0 +onnx==1.18.0 nltk==3.9.1 wget==3.2 cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/introductory_parser_samples/README.md b/samples/python/introductory_parser_samples/README.md index 34a47318..51da07e7 100644 --- a/samples/python/introductory_parser_samples/README.md +++ b/samples/python/introductory_parser_samples/README.md @@ -31,12 +31,16 @@ This sample demonstrates how to build an engine from an ONNX model file using th pip3 install -r requirements.txt ``` +2. Preparing sample data + +See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README. + ## Running the sample 1. Run the sample to create a TensorRT inference engine and run inference: `python3 onnx_resnet50.py` - **Note:** If the TensorRT sample data is not installed in the default location, for example `/usr/src/tensorrt/data/`, the `data` directory must be specified. For example: `python3 onnx_resnet50.py -d /path/to/my/data/` + **Note:** If the TensorRT sample data is not installed in the default location, the `data` directory must be specified. For example: `python3 onnx_resnet50.py -d $TRT_DATADIR` 2. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following: `Correctly recognized data/samples/resnet50/reflex_camera.jpeg as reflex camera` diff --git a/samples/python/introductory_parser_samples/onnx_resnet50.py b/samples/python/introductory_parser_samples/onnx_resnet50.py index fd69cc48..ed60eaed 100644 --- a/samples/python/introductory_parser_samples/onnx_resnet50.py +++ b/samples/python/introductory_parser_samples/onnx_resnet50.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,7 +26,7 @@ import numpy as np import tensorrt as trt from PIL import Image -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) import common @@ -102,27 +102,29 @@ def main(): # Build a TensorRT engine. engine = build_engine_onnx(onnx_model_file) # Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same. - # Allocate buffers and create a CUDA stream. - inputs, outputs, bindings, stream = common.allocate_buffers(engine) + # Allocate buffers + inputs, outputs, bindings = common.allocate_buffers(engine) # Contexts are used to perform inference. context = engine.create_execution_context() - # Load a normalized test case into the host input page-locked buffer. - test_image = random.choice(test_images) - test_case = load_normalized_test_case(test_image, inputs[0].host) - # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the - # probability that the image corresponds to that label - trt_outputs = common.do_inference( - context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) - # We use the highest probability as our prediction. Its index corresponds to the predicted label. - pred = labels[np.argmax(trt_outputs[0])] - common.free_buffers(inputs, outputs, stream) + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + # Load a normalized test case into the host input page-locked buffer. + test_image = random.choice(test_images) + test_case = load_normalized_test_case(test_image, inputs[0].host) + # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the + # probability that the image corresponds to that label + trt_outputs = common.do_inference( + context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) + # We use the highest probability as our prediction. Its index corresponds to the predicted label. + pred = labels[np.argmax(trt_outputs[0])] + common.free_buffers(inputs, outputs) if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]: print("Correctly recognized " + test_case + " as " + pred) else: diff --git a/samples/python/introductory_parser_samples/requirements.txt b/samples/python/introductory_parser_samples/requirements.txt index 04e24992..240683e4 100644 --- a/samples/python/introductory_parser_samples/requirements.txt +++ b/samples/python/introductory_parser_samples/requirements.txt @@ -2,6 +2,6 @@ Pillow==11.3.0 cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/network_api_pytorch_mnist/README.md b/samples/python/network_api_pytorch_mnist/README.md index 862ef795..ac11126a 100644 --- a/samples/python/network_api_pytorch_mnist/README.md +++ b/samples/python/network_api_pytorch_mnist/README.md @@ -54,7 +54,11 @@ To run this sample you must be using Python 3.6 or newer. On PowerPC systems, you will need to manually install PyTorch using IBM's [PowerAI](https://www.ibm.com/support/knowledgecenter/SS5SF7_1.6.0/navigation/pai_install.htm). -2. The MNIST dataset can be found under the data directory (usually `/usr/src/tensorrt/data/mnist`) if using the TensorRT containers. It is also bundled along with the [TensorRT tarball](https://developer.nvidia.com/nvidia-tensorrt-download). +2. Preparing sample data + +See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README. + +The MNIST dataset can be found under `$TRT_DATADIR/mnist`. ## Running the sample diff --git a/samples/python/network_api_pytorch_mnist/requirements.txt b/samples/python/network_api_pytorch_mnist/requirements.txt index 38146f56..0f1fc412 100644 --- a/samples/python/network_api_pytorch_mnist/requirements.txt +++ b/samples/python/network_api_pytorch_mnist/requirements.txt @@ -4,6 +4,6 @@ torchvision cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/network_api_pytorch_mnist/sample.py b/samples/python/network_api_pytorch_mnist/sample.py index a695ee9a..d56b04a4 100644 --- a/samples/python/network_api_pytorch_mnist/sample.py +++ b/samples/python/network_api_pytorch_mnist/sample.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -24,7 +24,7 @@ import numpy as np import tensorrt as trt -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) import common # You can set the logger severity higher to suppress messages (or lower to display more messages). @@ -153,22 +153,24 @@ def main(): # Build an engine, allocate buffers and create a stream. # For more information on buffer allocation, refer to the introductory samples. - inputs, outputs, bindings, stream = common.allocate_buffers(engine) + inputs, outputs, bindings = common.allocate_buffers(engine) context = engine.create_execution_context() - case_num = load_random_test_case(mnist_model, pagelocked_buffer=inputs[0].host) - # For more information on performing inference, refer to the introductory samples. - # The common.do_inference function will return a list of outputs - we only have one in this case. - [output] = common.do_inference( - context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) - pred = np.argmax(output) - common.free_buffers(inputs, outputs, stream) + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + case_num = load_random_test_case(mnist_model, pagelocked_buffer=inputs[0].host) + # For more information on performing inference, refer to the introductory samples. + # The common.do_inference function will return a list of outputs - we only have one in this case. + [output] = common.do_inference( + context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) + pred = np.argmax(output) + common.free_buffers(inputs, outputs) print("Test Case: " + str(case_num)) print("Prediction: " + str(pred)) diff --git a/samples/python/non_zero_plugin/non_zero_plugin.py b/samples/python/non_zero_plugin/non_zero_plugin.py index 89ef3826..cdbdfcbc 100644 --- a/samples/python/non_zero_plugin/non_zero_plugin.py +++ b/samples/python/non_zero_plugin/non_zero_plugin.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -36,11 +36,12 @@ import argparse from polygraphy import mod sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)) -from plugin_utils import checkCudaErrors, KernelHelper, UnownedMemory, volume +from plugin_utils import cuda_call, KernelHelper, UnownedMemory, volume -cuda = mod.lazy_import("cuda.cuda") -cudart = mod.lazy_import("cuda.cudart") -nvrtc = mod.lazy_import("cuda.nvrtc") + +cuda = mod.lazy_import("cuda.bindings.driver") +cudart = mod.lazy_import("cuda.bindings.runtime") +nvrtc = mod.lazy_import("cuda.bindings.nvrtc") torch = mod.lazy_import("torch") cp = mod.lazy_import("cupy") @@ -145,11 +146,11 @@ class NonZeroPlugin(trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, def configure_plugin(self, inp, out): if self.backend == "cuda_python": - err, self.cuDevice = cuda.cuDeviceGet(0) + self.cuDevice = cuda_call(cuda.cuDeviceGet(0)) def on_shape_change(self, inp, out): if self.backend == "cuda_python": - err, self.cuDevice = cuda.cuDeviceGet(0) + self.cuDevice = cuda_call(cuda.cuDeviceGet(0)) def supports_format_combination(self, pos, in_out, num_inputs): assert num_inputs == 1 @@ -190,7 +191,7 @@ class NonZeroPlugin(trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, if inp_dtype == np.float32: kernelHelper = KernelHelper(non_zero_float_kernel, int(self.cuDevice)) _non_zero_float_kernel = kernelHelper.getFunction(b'find_non_zero_indices_float') - checkCudaErrors(cuda.cuLaunchKernel(_non_zero_float_kernel, + cuda_call(cuda.cuLaunchKernel(_non_zero_float_kernel, numBlocks, 1, 1, blockSize, 1, 1, 0, @@ -199,7 +200,7 @@ class NonZeroPlugin(trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, elif inp_dtype == np.float16: kernelHelper = KernelHelper(non_zero_half_kernel, int(self.cuDevice)) _non_zero_half_kernel = kernelHelper.getFunction(b'find_non_zero_indices_half') - checkCudaErrors(cuda.cuLaunchKernel(_non_zero_half_kernel, + cuda_call(cuda.cuLaunchKernel(_non_zero_half_kernel, numBlocks, 1, 1, blockSize, 1, 1, 0, @@ -277,12 +278,12 @@ if __name__ == "__main__": args = parser.parse_args() if args.backend == "cuda_python": - # Initialize CUDA Driver API - err, = cuda.cuInit(0) - # Retrieve handle for device 0 - err, cuDevice = cuda.cuDeviceGet(0) - # Create context - _, cudaCtx = cuda.cuCtxCreate(0, cuDevice) + # Initialize CUDA and create default context + cuda_call(cudart.cudaFree(0)) + + elif args.backend == "torch": + # Initialize CUDA and create default context + torch.cuda.init() precision = np.float32 if args.precision == "fp32" else np.float16 @@ -348,5 +349,4 @@ if __name__ == "__main__": else: print("Inference result incorrect!") - if args.backend == "cuda_python": - checkCudaErrors(cuda.cuCtxDestroy(cudaCtx)) + diff --git a/samples/python/non_zero_plugin/requirements.txt b/samples/python/non_zero_plugin/requirements.txt index dcdfcb98..350fd80f 100644 --- a/samples/python/non_zero_plugin/requirements.txt +++ b/samples/python/non_zero_plugin/requirements.txt @@ -9,5 +9,5 @@ numpy==1.26.4 onnx-graphsurgeon pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 diff --git a/samples/python/onnx_custom_plugin/CMakeLists.txt b/samples/python/onnx_custom_plugin/CMakeLists.txt index 58c7dce5..21bc09f1 100644 --- a/samples/python/onnx_custom_plugin/CMakeLists.txt +++ b/samples/python/onnx_custom_plugin/CMakeLists.txt @@ -15,8 +15,8 @@ # limitations under the License. # -# We need cmake >= 3.8, since 3.8 introduced CUDA as a first class language -cmake_minimum_required(VERSION 3.8 FATAL_ERROR) +# We need cmake >= 3.19 for REAL_PATH support +cmake_minimum_required(VERSION 3.19 FATAL_ERROR) project(CustomHardMax LANGUAGES CXX CUDA) if(NOT MSVC) @@ -76,16 +76,19 @@ set_ifndef(CUDA_LIB ${_CUDA_LIB}) add_definitions(-DTENSORRT_BUILD_LIB) # Add include directories -get_filename_component(SAMPLES_COMMON_DIR ${CMAKE_SOURCE_DIR}/../../common/ ABSOLUTE) -get_filename_component(SAMPLES_DIR ${CMAKE_SOURCE_DIR}/../../ ABSOLUTE) -include_directories(${CUDA_INC_DIR} ${TRT_INCLUDE} ${CMAKE_SOURCE_DIR}/plugin/ - ${SAMPLES_COMMON_DIR} ${SAMPLES_DIR}) +file(REAL_PATH ${CMAKE_SOURCE_DIR} CMAKE_SOURCE_DIR_REALPATH) + +get_filename_component(SAMPLES_COMMON_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../common/ ABSOLUTE) +get_filename_component(SHARED_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../../shared ABSOLUTE) +get_filename_component(SAMPLES_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../ ABSOLUTE) +include_directories(${CUDA_INC_DIR} ${TRT_INCLUDE} ${CMAKE_SOURCE_DIR_REALPATH}/plugin/ + ${SAMPLES_COMMON_DIR} ${SAMPLES_DIR} ${SHARED_DIR}) # Define Hardmax plugin library target add_library( customHardmaxPlugin MODULE - ${SAMPLES_COMMON_DIR}/logger.cpp ${SAMPLES_DIR}/utils/fileLock.cpp - ${CMAKE_SOURCE_DIR}/plugin/customHardmaxPlugin.cpp ${CMAKE_SOURCE_DIR}/plugin/customHardmaxPlugin.h) + ${SAMPLES_COMMON_DIR}/logger.cpp ${SHARED_DIR}/utils/fileLock.cpp + ${CMAKE_SOURCE_DIR_REALPATH}/plugin/customHardmaxPlugin.cpp ${CMAKE_SOURCE_DIR_REALPATH}/plugin/customHardmaxPlugin.h) # Use C++11 target_compile_features(customHardmaxPlugin PUBLIC cxx_std_17) diff --git a/samples/python/onnx_custom_plugin/requirements.txt b/samples/python/onnx_custom_plugin/requirements.txt index 40d186ff..37221dab 100644 --- a/samples/python/onnx_custom_plugin/requirements.txt +++ b/samples/python/onnx_custom_plugin/requirements.txt @@ -1,11 +1,11 @@ nltk==3.9.1 -onnx==1.16.0 +onnx==1.18.0 --extra-index-url https://pypi.ngc.nvidia.com onnx-graphsurgeon>=0.3.20 wget>=3.2 cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/onnx_custom_plugin/sample.py b/samples/python/onnx_custom_plugin/sample.py index 25d4ca36..4442910c 100644 --- a/samples/python/onnx_custom_plugin/sample.py +++ b/samples/python/onnx_custom_plugin/sample.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -143,7 +143,7 @@ def main(): print("Engine plan not saved. Building new engine...") engine = build_engine(TRT_MODEL_PATH) - inputs, outputs, bindings, stream = common.allocate_buffers(engine, profile_idx=0) + inputs, outputs, bindings = common.allocate_buffers(engine, profile_idx=0) testcases = [ ( @@ -162,28 +162,33 @@ def main(): testcases = [(context_text, query_text)] trt_context = engine.create_execution_context() - for context_text, query_text in testcases: + + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + for context_text, query_text in testcases: - context_words, _ = preprocess(context_text) + context_words, _ = preprocess(context_text) - load_test_case(inputs, context_text, query_text, trt_context) - if not interactive: - print(f"Input context: {context_text}") - print(f"Input query: {query_text}") - trt_outputs = common.do_inference( - trt_context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) - start = trt_outputs[1].item() - end = trt_outputs[0].item() - answer = context_words[start : end + 1].flatten() - print(f"Model prediction: ", " ".join(answer)) - print() - common.free_buffers(inputs, outputs, stream) + load_test_case(inputs, context_text, query_text, trt_context) + if not interactive: + print(f"Input context: {context_text}") + print(f"Input query: {query_text}") + trt_outputs = common.do_inference( + trt_context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) + start = trt_outputs[1].item() + end = trt_outputs[0].item() + answer = context_words[start : end + 1].flatten() + print(f"Model prediction: ", " ".join(answer)) + print() + + # Note: free_buffers no longer needs stream parameter + common.free_buffers(inputs, outputs) print("Passed") diff --git a/samples/python/onnx_custom_plugin/test_custom_hardmax_plugin.py b/samples/python/onnx_custom_plugin/test_custom_hardmax_plugin.py index 59b08b06..88a622c2 100644 --- a/samples/python/onnx_custom_plugin/test_custom_hardmax_plugin.py +++ b/samples/python/onnx_custom_plugin/test_custom_hardmax_plugin.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -68,19 +68,22 @@ def make_trt_network_and_engine(input_shape, axis): def custom_plugin_impl(input_arr, engine): - inputs, outputs, bindings, stream = common.allocate_buffers(engine) + inputs, outputs, bindings = common.allocate_buffers(engine) context = engine.create_execution_context() inputs[0].host = input_arr.astype(trt.nptype(trt.float32)) - trt_outputs = common.do_inference( - context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) - output = trt_outputs[0].copy() - common.free_buffers(inputs, outputs, stream) + + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + trt_outputs = common.do_inference( + context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) + output = trt_outputs[0].copy() + common.free_buffers(inputs, outputs) return output diff --git a/samples/python/onnx_packnet/requirements.txt b/samples/python/onnx_packnet/requirements.txt index 566aa942..df7596ce 100644 --- a/samples/python/onnx_packnet/requirements.txt +++ b/samples/python/onnx_packnet/requirements.txt @@ -1,9 +1,9 @@ -onnx==1.16.0 +onnx==1.18.0 --extra-index-url https://pypi.ngc.nvidia.com onnx-graphsurgeon>=0.3.20 torch torchvision pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/plugin_utils.py b/samples/python/plugin_utils.py index 77f14e69..88c8fc5e 100644 --- a/samples/python/plugin_utils.py +++ b/samples/python/plugin_utils.py @@ -16,8 +16,10 @@ # from cuda.bindings import driver as cuda, runtime as cudart, nvrtc + import numpy as np import os +from common_runtime import cuda_call, create_cuda_context, cuda_init, cuda_get_device, cuda_memcpy_htod import argparse import threading @@ -44,42 +46,18 @@ def volume(d): return np.prod(d) -# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/helper_cuda.py -def checkCudaErrors(result): - def _cudaGetErrorEnum(error): - if isinstance(error, cuda.CUresult): - err, name = cuda.cuGetErrorName(error) - return name if err == cuda.CUresult.CUDA_SUCCESS else "" - elif isinstance(error, cudart.cudaError_t): - return cudart.cudaGetErrorName(error)[1] - elif isinstance(error, nvrtc.nvrtcResult): - return nvrtc.nvrtcGetErrorString(error)[1] - else: - raise RuntimeError("Unknown error type: {}".format(error)) - if result[0].value: - raise RuntimeError( - "CUDA error code={}({})".format( - result[0].value, _cudaGetErrorEnum(result[0]) - ) - ) - if len(result) == 1: - return None - elif len(result) == 2: - return result[1] - else: - return result[1:] def getComputeCapacity(devID): - major = checkCudaErrors(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID)) - minor = checkCudaErrors(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID)) + major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID)) + minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID)) return (major, minor) # Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/common.py class KernelHelper: def __init__(self, code, devID): - prog = checkCudaErrors( + prog = cuda_call( nvrtc.nvrtcCreateProgram(str.encode(code), b"sourceCode.cu", 0, [], []) ) CUDA_HOME = os.getenv("CUDA_HOME") @@ -90,10 +68,10 @@ class KernelHelper: include_dirs = os.path.join(CUDA_HOME, "include") # Initialize CUDA - checkCudaErrors(cudart.cudaFree(0)) + cuda_call(cudart.cudaFree(0)) major, minor = getComputeCapacity(devID) - _, nvrtc_minor = checkCudaErrors(nvrtc.nvrtcVersion()) + _, nvrtc_minor = cuda_call(nvrtc.nvrtcVersion()) use_cubin = nvrtc_minor >= 1 prefix = "sm" if use_cubin else "compute" arch_arg = bytes(f"--gpu-architecture={prefix}_{major}{minor}", "ascii") @@ -106,28 +84,28 @@ class KernelHelper: b"--std=c++11", b"-default-device", ] - checkCudaErrors(nvrtc.nvrtcCompileProgram(prog, len(opts), opts)) + cuda_call(nvrtc.nvrtcCompileProgram(prog, len(opts), opts)) except RuntimeError as err: - logSize = checkCudaErrors(nvrtc.nvrtcGetProgramLogSize(prog)) + logSize = cuda_call(nvrtc.nvrtcGetProgramLogSize(prog)) log = b" " * logSize - checkCudaErrors(nvrtc.nvrtcGetProgramLog(prog, log)) + cuda_call(nvrtc.nvrtcGetProgramLog(prog, log)) print(log.decode()) print(err) exit(-1) if use_cubin: - dataSize = checkCudaErrors(nvrtc.nvrtcGetCUBINSize(prog)) + dataSize = cuda_call(nvrtc.nvrtcGetCUBINSize(prog)) data = b" " * dataSize - checkCudaErrors(nvrtc.nvrtcGetCUBIN(prog, data)) + cuda_call(nvrtc.nvrtcGetCUBIN(prog, data)) else: - dataSize = checkCudaErrors(nvrtc.nvrtcGetPTXSize(prog)) + dataSize = cuda_call(nvrtc.nvrtcGetPTXSize(prog)) data = b" " * dataSize - checkCudaErrors(nvrtc.nvrtcGetPTX(prog, data)) + cuda_call(nvrtc.nvrtcGetPTX(prog, data)) - self.module = checkCudaErrors(cuda.cuModuleLoadData(np.char.array(data))) + self.module = cuda_call(cuda.cuModuleLoadData(np.char.array(data))) def getFunction(self, name): - return checkCudaErrors(cuda.cuModuleGetFunction(self.module, name)) + return cuda_call(cuda.cuModuleGetFunction(self.module, name)) class CudaCtxManager(trt.IPluginResource): @@ -141,11 +119,11 @@ class CudaCtxManager(trt.IPluginResource): cloned.__dict__.update(self.__dict__) # Delay the CUDA ctx creation until clone() # since only a cloned resource is registered by TRT - _, cloned.cuda_ctx = cuda.cuCtxCreate(0, self.device) + cloned.cuda_ctx = create_cuda_context(self.device) return cloned def release(self): - checkCudaErrors(cuda.cuCtxDestroy(self.cuda_ctx)) + cuda_call(cuda.cuCtxDestroy(self.cuda_ctx)) class UnownedMemory: def __init__(self, ptr, shape, dtype): diff --git a/samples/python/python_plugin/CMakeLists.txt b/samples/python/python_plugin/CMakeLists.txt index 5120998f..7156d859 100644 --- a/samples/python/python_plugin/CMakeLists.txt +++ b/samples/python/python_plugin/CMakeLists.txt @@ -88,8 +88,14 @@ if(NOT MSVC) cuda_ge(12 ${minor} ${result_var}) endforeach() + # Add checks for CUDA 13.x versions + foreach(minor RANGE 0 9) + set(result_var "CUDA_GE_13_${minor}") + cuda_ge(13 ${minor} ${result_var}) + endforeach() + # Define THOR_SM based on CUDA version - if(CUDA_GE_13_0) + if(${CUDA_GE_13_0}) set(THOR_SM 110) else() set(THOR_SM 101) @@ -97,24 +103,24 @@ if(NOT MSVC) set(SAMPLE_SMS "75") - if(CUDA_GE_11_0) + if(${CUDA_GE_11_0}) list(APPEND SAMPLE_SMS "80") endif() - if(CUDA_GE_11_1) + if(${CUDA_GE_11_1}) list(APPEND SAMPLE_SMS "86") endif() - if(CUDA_GE_11_4) + if(${CUDA_GE_11_4}) list(APPEND SAMPLE_SMS "87") endif() - if(CUDA_GE_11_8) + if(${CUDA_GE_11_8}) list(APPEND SAMPLE_SMS "89" "90") endif() # Blackwell support - if(CUDA_GE_12_8) + if(${CUDA_GE_12_8}) list(APPEND SAMPLE_SMS "100" ${THOR_SM} "120") endif() diff --git a/samples/python/python_plugin/circ_pad_plugin_cuda_python.py b/samples/python/python_plugin/circ_pad_plugin_cuda_python.py index c3bb73a1..60b06d58 100755 --- a/samples/python/python_plugin/circ_pad_plugin_cuda_python.py +++ b/samples/python/python_plugin/circ_pad_plugin_cuda_python.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -31,7 +31,8 @@ from polygraphy.backend.trt import ( from polygraphy.json import to_json, from_json sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)) -from plugin_utils import checkCudaErrors, KernelHelper, parseArgs, CudaCtxManager +from plugin_utils import cuda_call, KernelHelper, parseArgs, CudaCtxManager, cuda_init, cuda_get_device, cuda_memcpy_htod +import common_runtime as common from cuda.bindings import driver as cuda circ_pad_half_kernel = r""" @@ -121,17 +122,13 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): self.N = int(f.data) def initialize(self): - err, self.cuDevice = cuda.cuDeviceGet(0) + self.cuDevice = cuda_get_device(0) trt.get_plugin_registry().acquire_plugin_resource( "cuda_ctx", CudaCtxManager(self.cuDevice) ) - self.all_pads_d = checkCudaErrors( - cuda.cuMemAlloc(np.int32().itemsize * self.N * 2) - ) - self.orig_dims_d = checkCudaErrors( - cuda.cuMemAlloc(np.int32().itemsize * self.N) - ) - self.Y_shape_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N)) + self.all_pads_d = common.DeviceMem(np.int32().itemsize * self.N * 2) + self.orig_dims_d = common.DeviceMem(np.int32().itemsize * self.N) + self.Y_shape_d = common.DeviceMem(np.int32().itemsize * self.N) def get_output_datatype(self, index, input_types): return input_types[0] @@ -169,17 +166,11 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): # Copy vectors from host memory to device memory if self.all_pads_d: - checkCudaErrors( - cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes) - ) + cuda_memcpy_htod(self.all_pads_d.device_ptr, all_pads) if self.orig_dims_d: - checkCudaErrors( - cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes) - ) + cuda_memcpy_htod(self.orig_dims_d.device_ptr, orig_dims) if self.Y_shape_d: - checkCudaErrors( - cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes) - ) + cuda_memcpy_htod(self.Y_shape_d.device_ptr, out_dims) self.Y_len_d = np.prod(out_dims) @@ -211,9 +202,9 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): da = np.array([inputs[0]], dtype=np.uint64) dc = np.array([outputs[0]], dtype=np.uint64) - d_all_pads = np.array([int(self.all_pads_d)], dtype=np.uint64) - d_orig_dims = np.array([int(self.orig_dims_d)], dtype=np.uint64) - d_Y_shape = np.array([int(self.Y_shape_d)], dtype=np.uint64) + d_all_pads = np.array([int(self.all_pads_d.device_ptr)], dtype=np.uint64) + d_orig_dims = np.array([int(self.orig_dims_d.device_ptr)], dtype=np.uint64) + d_Y_shape = np.array([int(self.Y_shape_d.device_ptr)], dtype=np.uint64) Y_len = np.array(self.Y_len_d, dtype=np.uint32) args = [da, d_all_pads, d_orig_dims, dc, d_Y_shape, Y_len] @@ -224,7 +215,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): if inp_dtype == np.float32: kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice)) _circ_pad_float_kernel = kernelHelper.getFunction(b"circ_pad_float") - checkCudaErrors( + cuda_call( cuda.cuLaunchKernel( _circ_pad_float_kernel, numBlocks, @@ -242,7 +233,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): elif inp_dtype == np.float16: kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice)) _circ_pad_half_kernel = kernelHelper.getFunction(b"circ_pad_half") - checkCudaErrors( + cuda_call( cuda.cuLaunchKernel( _circ_pad_half_kernel, numBlocks, @@ -266,12 +257,10 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): return cloned_plugin def terminate(self): - if self.all_pads_d: - checkCudaErrors(cuda.cuMemFree(self.all_pads_d)) - if self.orig_dims_d: - checkCudaErrors(cuda.cuMemFree(self.orig_dims_d)) - if self.Y_shape_d: - checkCudaErrors(cuda.cuMemFree(self.Y_shape_d)) + # Release DeviceMem objects - automatic cleanup via __del__ when reference count reaches 0 + self.all_pads_d = None + self.orig_dims_d = None + self.Y_shape_d = None trt.get_plugin_registry().release_plugin_resource("cuda_ctx") @@ -317,10 +306,10 @@ if __name__ == "__main__": args = parseArgs() # Initialize CUDA Driver API - (err,) = cuda.cuInit(0) + cuda_init() # Retrieve handle for device 0 - err, cuDevice = cuda.cuDeviceGet(0) + cuDevice = cuda_get_device(0) plg_registry = trt.get_plugin_registry() diff --git a/samples/python/python_plugin/circ_pad_plugin_inetdef_cuda_python.py b/samples/python/python_plugin/circ_pad_plugin_inetdef_cuda_python.py index 7be9108b..1efc2a7c 100644 --- a/samples/python/python_plugin/circ_pad_plugin_inetdef_cuda_python.py +++ b/samples/python/python_plugin/circ_pad_plugin_inetdef_cuda_python.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -31,7 +31,7 @@ from polygraphy.backend.trt import ( from polygraphy.json import to_json, from_json sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)) -from plugin_utils import checkCudaErrors, KernelHelper, parseArgs, CudaCtxManager +from plugin_utils import cuda_call, KernelHelper, parseArgs, CudaCtxManager from cuda.bindings import driver as cuda circ_pad_half_kernel = r""" @@ -121,17 +121,17 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): self.N = int(f.data) def initialize(self): - err, self.cuDevice = cuda.cuDeviceGet(0) + self.cuDevice = cuda_call(cuda.cuDeviceGet(0)) trt.get_plugin_registry().acquire_plugin_resource( "cuda_ctx", CudaCtxManager(self.cuDevice) ) - self.all_pads_d = checkCudaErrors( + self.all_pads_d = cuda_call( cuda.cuMemAlloc(np.int32().itemsize * self.N * 2) ) - self.orig_dims_d = checkCudaErrors( + self.orig_dims_d = cuda_call( cuda.cuMemAlloc(np.int32().itemsize * self.N) ) - self.Y_shape_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N)) + self.Y_shape_d = cuda_call(cuda.cuMemAlloc(np.int32().itemsize * self.N)) def get_output_datatype(self, index, input_types): return input_types[0] @@ -169,15 +169,15 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): # Copy vectors from host memory to device memory if self.all_pads_d: - checkCudaErrors( + cuda_call( cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes) ) if self.orig_dims_d: - checkCudaErrors( + cuda_call( cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes) ) if self.Y_shape_d: - checkCudaErrors( + cuda_call( cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes) ) @@ -224,7 +224,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): if inp_dtype == np.float32: kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice)) _circ_pad_float_kernel = kernelHelper.getFunction(b"circ_pad_float") - checkCudaErrors( + cuda_call( cuda.cuLaunchKernel( _circ_pad_float_kernel, numBlocks, @@ -242,7 +242,7 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): elif inp_dtype == np.float16: kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice)) _circ_pad_half_kernel = kernelHelper.getFunction(b"circ_pad_half") - checkCudaErrors( + cuda_call( cuda.cuLaunchKernel( _circ_pad_half_kernel, numBlocks, @@ -267,11 +267,11 @@ class CircPadPlugin(trt.IPluginV2DynamicExt): def terminate(self): if self.all_pads_d: - checkCudaErrors(cuda.cuMemFree(self.all_pads_d)) + cuda_call(cuda.cuMemFree(self.all_pads_d)) if self.orig_dims_d: - checkCudaErrors(cuda.cuMemFree(self.orig_dims_d)) + cuda_call(cuda.cuMemFree(self.orig_dims_d)) if self.Y_shape_d: - checkCudaErrors(cuda.cuMemFree(self.Y_shape_d)) + cuda_call(cuda.cuMemFree(self.Y_shape_d)) plg_registry.release_plugin_resource("cuda_ctx") @@ -318,10 +318,10 @@ if __name__ == "__main__": precision = np.float32 if args.precision == "fp32" else np.float16 # Initialize CUDA Driver API - (err,) = cuda.cuInit(0) + cuda_call(cuda.cuInit(0)) # Retrieve handle for device 0 - err, cuDevice = cuda.cuDeviceGet(0) + cuDevice = cuda_call(cuda.cuDeviceGet(0)) plg_registry = trt.get_plugin_registry() diff --git a/samples/python/python_plugin/requirements.txt b/samples/python/python_plugin/requirements.txt index fc79419a..9cf6aec0 100644 --- a/samples/python/python_plugin/requirements.txt +++ b/samples/python/python_plugin/requirements.txt @@ -1,16 +1,17 @@ cuda-python==12.9.0 cupy-cuda12x numba +numba-cuda[cu12] triton; platform_system != "Windows" torch --extra-index-url https://pypi.ngc.nvidia.com polygraphy colored numpy==1.26.4 -onnx==1.16.0; platform_system == "Windows" +onnx==1.18.0; platform_system == "Windows" --extra-index-url https://pypi.ngc.nvidia.com onnx-graphsurgeon pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 diff --git a/samples/python/quickly_deployable_plugins/requirements.txt b/samples/python/quickly_deployable_plugins/requirements.txt index 0932ba1c..ceecc880 100644 --- a/samples/python/quickly_deployable_plugins/requirements.txt +++ b/samples/python/quickly_deployable_plugins/requirements.txt @@ -4,9 +4,9 @@ torch polygraphy colored numpy==1.26.4 -onnx==1.16.0; platform_system == "Windows" +onnx==1.18.0; platform_system == "Windows" --extra-index-url https://pypi.ngc.nvidia.com onnx-graphsurgeon pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 diff --git a/samples/python/quickly_deployable_plugins/requirements.yml b/samples/python/quickly_deployable_plugins/requirements.yml index 4ee24890..6b3b3035 100644 --- a/samples/python/quickly_deployable_plugins/requirements.yml +++ b/samples/python/quickly_deployable_plugins/requirements.yml @@ -16,7 +16,7 @@ conditions: onnx-graphsurgeon: - onnx-graphsurgeon onnx: - - onnx==1.16.0; platform_system == "Windows" + - onnx==1.18.0; platform_system == "Windows" triton: - 'triton==3.2.0; (platform_system != "Windows")' numpy: diff --git a/samples/python/requirements.txt b/samples/python/requirements.txt index d83e67ff..3798df75 100644 --- a/samples/python/requirements.txt +++ b/samples/python/requirements.txt @@ -1,4 +1,4 @@ pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/sample_weight_stripping/README.md b/samples/python/sample_weight_stripping/README.md index 5b47c96d..b1966d2b 100644 --- a/samples/python/sample_weight_stripping/README.md +++ b/samples/python/sample_weight_stripping/README.md @@ -28,6 +28,9 @@ This sample demonstrates how to build a weight-stripped engine from an ONNX mode pip3 install -r requirements.txt ``` +2. Preparing sample data +See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README. + ## Running the sample 1. Build and save both normal engine and weight-stripped engine: diff --git a/samples/python/sample_weight_stripping/build_engines.py b/samples/python/sample_weight_stripping/build_engines.py index 6f1e3936..8d52f9f9 100644 --- a/samples/python/sample_weight_stripping/build_engines.py +++ b/samples/python/sample_weight_stripping/build_engines.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -24,7 +24,7 @@ import datetime import tensorrt as trt -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) import common # You can set the logger severity higher to suppress messages (or lower to display more messages). diff --git a/samples/python/sample_weight_stripping/notebooks/weight_stripping.ipynb b/samples/python/sample_weight_stripping/notebooks/weight_stripping.ipynb index 5c6fcb36..997cb4b0 100644 --- a/samples/python/sample_weight_stripping/notebooks/weight_stripping.ipynb +++ b/samples/python/sample_weight_stripping/notebooks/weight_stripping.ipynb @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,6 +17,29 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "license-cell", + "metadata": {}, + "source": [ + "```\n", + "SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n", + "SPDX-License-Identifier: Apache-2.0\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\"); \n", + "you may not use this file except in compliance with the License. \n", + "You may obtain a copy of the License at\n", + "\n", + "http://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License.\n", + "```" + ] + }, { "cell_type": "markdown", "id": "bbb82290-07cd-4b09-afa4-cf4ad3002f06", diff --git a/samples/python/sample_weight_stripping/refit_engine_and_infer.py b/samples/python/sample_weight_stripping/refit_engine_and_infer.py index 19d4df61..b315e597 100644 --- a/samples/python/sample_weight_stripping/refit_engine_and_infer.py +++ b/samples/python/sample_weight_stripping/refit_engine_and_infer.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,7 +26,7 @@ import numpy as np import tensorrt as trt from PIL import Image -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) import common @@ -100,9 +100,9 @@ def main(args): engine = load_normal_engine(args.normal_engine) refitted_engine = load_stripped_engine_and_refit(args.stripped_engine, onnx_model_file) - # Allocate buffers and create a CUDA stream. - inputs, outputs, bindings, stream = common.allocate_buffers(engine) - inputs_1, outputs_1, bindings_1, stream_1 = common.allocate_buffers(refitted_engine) + # Allocate buffers + inputs, outputs, bindings = common.allocate_buffers(engine) + inputs_1, outputs_1, bindings_1 = common.allocate_buffers(refitted_engine) # Contexts are used to perform inference. context = engine.create_execution_context() @@ -115,17 +115,22 @@ def main(args): # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the # probability that the image corresponds to that label - start_time = time.time() - for i in range(100): # count time for 100 times of inference - trt_outputs = common.do_inference(context, engine=engine, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) - total_time = time.time() - start_time - print("Normal engine inference time on 100 cases: {:.4f} seconds".format(total_time)) + + # Use context manager for proper stream lifecycle management - Normal engine + with common.CudaStreamContext() as stream: + start_time = time.time() + for i in range(100): # count time for 100 times of inference + trt_outputs = common.do_inference(context, engine=engine, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) + total_time = time.time() - start_time + print("Normal engine inference time on 100 cases: {:.4f} seconds".format(total_time)) - start_time = time.time() - for i in range(100): - trt_outputs_refitted = common.do_inference(context_1, engine=refitted_engine, bindings=bindings_1, inputs=inputs_1, outputs=outputs_1, stream=stream_1) - total_time = time.time() - start_time - print("Refitted stripped engine inference time on 100 cases: {:.4f} seconds".format(total_time)) + # Use context manager for proper stream lifecycle management - Refitted engine + with common.CudaStreamContext() as stream_1: + start_time = time.time() + for i in range(100): + trt_outputs_refitted = common.do_inference(context_1, engine=refitted_engine, bindings=bindings_1, inputs=inputs_1, outputs=outputs_1, stream=stream_1) + total_time = time.time() - start_time + print("Refitted stripped engine inference time on 100 cases: {:.4f} seconds".format(total_time)) # We use the highest probability as our prediction. Its index corresponds to the predicted label. pred = labels[np.argmax(trt_outputs[0])] diff --git a/samples/python/sample_weight_stripping/requirements.txt b/samples/python/sample_weight_stripping/requirements.txt index 04e24992..240683e4 100644 --- a/samples/python/sample_weight_stripping/requirements.txt +++ b/samples/python/sample_weight_stripping/requirements.txt @@ -2,6 +2,6 @@ Pillow==11.3.0 cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/simple_progress_monitor/README.md b/samples/python/simple_progress_monitor/README.md index bbab175c..ad5871ec 100644 --- a/samples/python/simple_progress_monitor/README.md +++ b/samples/python/simple_progress_monitor/README.md @@ -30,13 +30,15 @@ This sample demonstrates how to build an engine from an ONNX model file using th ```bash pip3 install -r requirements.txt ``` +2. Preparing sample data +See [Preparing sample data](../../README.md#preparing-sample-data) in the main samples README. ## Running the sample 1. Run the sample from a terminal to create a TensorRT inference engine and run inference: `python3 simple_progress_monitor.py` - **Note:** If the TensorRT sample data is not installed in the default location, for example `/usr/src/tensorrt/data/`, the `data` directory must be specified. For example: `python3 simple_progress_monitor.py -d /path/to/my/data/` + **Note:** If the TensorRT sample data is not installed in the default location, the `data` directory must be specified. For example: `python3 simple_progress_monitor.py -d $TRT_DATADIR` **Note:** Do not redirect the output of this script to a file or pipe. @@ -84,6 +86,9 @@ For terms and conditions for use, reproduction, and distribution, see the [Tenso # Changelog +October 2025 +Migrate to strongly typed APIs. + August 2025 Removed support for Python versions < 3.10. diff --git a/samples/python/simple_progress_monitor/requirements.txt b/samples/python/simple_progress_monitor/requirements.txt index 04e24992..240683e4 100644 --- a/samples/python/simple_progress_monitor/requirements.txt +++ b/samples/python/simple_progress_monitor/requirements.txt @@ -2,6 +2,6 @@ Pillow==11.3.0 cuda-python==12.9.0 pywin32; platform_system == "Windows" pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/simple_progress_monitor/simple_progress_monitor.py b/samples/python/simple_progress_monitor/simple_progress_monitor.py index fe54f720..ac6f7453 100644 --- a/samples/python/simple_progress_monitor/simple_progress_monitor.py +++ b/samples/python/simple_progress_monitor/simple_progress_monitor.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,7 +26,7 @@ import numpy as np import tensorrt as trt from PIL import Image -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) import common @@ -189,27 +189,29 @@ def main(): # Build a TensorRT engine. engine = build_engine_onnx(onnx_model_file) # Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same. - # Allocate buffers and create a CUDA stream. - inputs, outputs, bindings, stream = common.allocate_buffers(engine) + # Allocate buffers + inputs, outputs, bindings = common.allocate_buffers(engine) # Contexts are used to perform inference. context = engine.create_execution_context() - # Load a normalized test case into the host input page-locked buffer. - test_image = random.choice(test_images) - test_case = load_normalized_test_case(test_image, inputs[0].host) - # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the - # probability that the image corresponds to that label - trt_outputs = common.do_inference( - context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) - # We use the highest probability as our prediction. Its index corresponds to the predicted label. - pred = labels[np.argmax(trt_outputs[0])] - common.free_buffers(inputs, outputs, stream) + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + # Load a normalized test case into the host input page-locked buffer. + test_image = random.choice(test_images) + test_case = load_normalized_test_case(test_image, inputs[0].host) + # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the + # probability that the image corresponds to that label + trt_outputs = common.do_inference( + context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) + # We use the highest probability as our prediction. Its index corresponds to the predicted label. + pred = labels[np.argmax(trt_outputs[0])] + common.free_buffers(inputs, outputs) if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]: print("Correctly recognized " + test_case + " as " + pred) else: diff --git a/samples/python/tensorflow_object_detection_api/README.md b/samples/python/tensorflow_object_detection_api/README.md index 9832987c..0e95233f 100644 --- a/samples/python/tensorflow_object_detection_api/README.md +++ b/samples/python/tensorflow_object_detection_api/README.md @@ -1,5 +1,6 @@ -# TensorFlow Object Detection API Models in TensorRT +# [DEPRECATED] TensorFlow Object Detection API Models in TensorRT +> **Notice:** This sample has been deprecated as of TensorRT 10.14 due to compatibility issues with outdated dependencies in the [tf2onnx](https://github.com/onnx/tensorflow-onnx) conversion pipeline. Users are advised to use earlier TensorRT versions if this sample is required for legacy workflows. Support for [TensorFlow Object Detection (TFOD) API](https://github.com/tensorflow/models/tree/master/research/object_detection) models in TensorRT, including Single Shot Detector, Faster R-CNN and Mask R-CNN models. This script helps with converting, running and validating these models with TensorRT. diff --git a/samples/python/tensorflow_object_detection_api/build_engine.py b/samples/python/tensorflow_object_detection_api/build_engine.py index 811e878f..2c9cb647 100644 --- a/samples/python/tensorflow_object_detection_api/build_engine.py +++ b/samples/python/tensorflow_object_detection_api/build_engine.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -46,7 +46,7 @@ class EngineCalibrator(trt.IInt8EntropyCalibrator2): super().__init__() self.cache_file = cache_file self.image_batcher = None - self.batch_allocation = None + self.batch_memory = None self.batch_generator = None def set_image_batcher(self, image_batcher: ImageBatcher): @@ -60,9 +60,11 @@ class EngineCalibrator(trt.IInt8EntropyCalibrator2): np.dtype(self.image_batcher.dtype).itemsize * np.prod(self.image_batcher.shape) ) - self.batch_allocation = common.cuda_call(cudart.cudaMalloc(size)) + self.batch_memory = common.DeviceMem(size) self.batch_generator = self.image_batcher.get_batch() + + def get_batch_size(self): """ Overrides from trt.IInt8EntropyCalibrator2. @@ -90,9 +92,9 @@ class EngineCalibrator(trt.IInt8EntropyCalibrator2): ) ) common.memcpy_host_to_device( - self.batch_allocation, np.ascontiguousarray(batch) + self.batch_memory.device_ptr, np.ascontiguousarray(batch) ) - return [int(self.batch_allocation)] + return [int(self.batch_memory.device_ptr)] except StopIteration: log.info("Finished calibration batches") return None diff --git a/samples/python/tensorflow_object_detection_api/infer.py b/samples/python/tensorflow_object_detection_api/infer.py index c7dc7bfd..0f903eb0 100644 --- a/samples/python/tensorflow_object_detection_api/infer.py +++ b/samples/python/tensorflow_object_detection_api/infer.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -54,7 +54,7 @@ class TensorRTInfer: # Setup I/O bindings self.inputs = [] self.outputs = [] - self.allocations = [] + self.device_memories = [] for i in range(self.engine.num_io_tensors): name = self.engine.get_tensor_name(i) is_input = False @@ -67,15 +67,15 @@ class TensorRTInfer: size = np.dtype(trt.nptype(dtype)).itemsize for s in shape: size *= s - allocation = common.cuda_call(cudart.cudaMalloc(size)) + device_mem = common.DeviceMem(size) binding = { "index": i, "name": name, "dtype": np.dtype(trt.nptype(dtype)), "shape": list(shape), - "allocation": allocation, + "allocation": device_mem.device_ptr, } - self.allocations.append(allocation) + self.device_memories.append(device_mem) if is_input: self.inputs.append(binding) else: @@ -84,7 +84,9 @@ class TensorRTInfer: assert self.batch_size > 0 assert len(self.inputs) > 0 assert len(self.outputs) > 0 - assert len(self.allocations) > 0 + assert len(self.device_memories) > 0 + + def input_spec(self): """ diff --git a/samples/python/tensorflow_object_detection_api/requirements.txt b/samples/python/tensorflow_object_detection_api/requirements.txt index 4313c96d..aaa00e25 100644 --- a/samples/python/tensorflow_object_detection_api/requirements.txt +++ b/samples/python/tensorflow_object_detection_api/requirements.txt @@ -1,4 +1,4 @@ -onnx==1.16.1 +onnx==1.18.0 onnxruntime==1.18.1 Pillow==11.3.0 tf2onnx==1.16.0 @@ -7,6 +7,6 @@ cuda-python==12.9.0 pywin32; platform_system == "Windows" Cython<3.0 pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 numpy==1.26.4 diff --git a/samples/python/yolov3_onnx/onnx_to_tensorrt.py b/samples/python/yolov3_onnx/onnx_to_tensorrt.py index 2ba322bc..b1b18004 100644 --- a/samples/python/yolov3_onnx/onnx_to_tensorrt.py +++ b/samples/python/yolov3_onnx/onnx_to_tensorrt.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,7 +26,7 @@ import tensorrt as trt from data_processing import ALL_CATEGORIES, PostprocessYOLO, PreprocessYOLO from PIL import ImageDraw -sys.path.insert(1, os.path.join(sys.path[0], "..")) +sys.path.insert(1, os.path.join(sys.path[0], os.path.pardir)) from downloader import getFilePath import common @@ -148,19 +148,25 @@ def main(): with get_engine( onnx_file_path, engine_file_path ) as engine, engine.create_execution_context() as context: - inputs, outputs, bindings, stream = common.allocate_buffers(engine) - # Do inference - print("Running inference on image {}...".format(input_image_path)) - # Set host input to the image. The common.do_inference function will copy the input to the GPU before executing. - inputs[0].host = image - trt_outputs = common.do_inference( - context, - engine=engine, - bindings=bindings, - inputs=inputs, - outputs=outputs, - stream=stream, - ) + inputs, outputs, bindings = common.allocate_buffers(engine) + + # Use context manager for proper stream lifecycle management + with common.CudaStreamContext() as stream: + # Do inference + print("Running inference on image {}...".format(input_image_path)) + # Set host input to the image. The common.do_inference function will copy the input to the GPU before executing. + inputs[0].host = image + trt_outputs = common.do_inference( + context, + engine=engine, + bindings=bindings, + inputs=inputs, + outputs=outputs, + stream=stream, + ) + + # Free host and device memory used for inputs and outputs + common.free_buffers(inputs, outputs) # Before doing post-processing, we need to reshape the outputs as the common.do_inference will give us flat arrays. trt_outputs = [ @@ -203,8 +209,7 @@ def main(): ) ) - # Free host and device memory used for inputs and outputs - common.free_buffers(inputs, outputs, stream) + if __name__ == "__main__": diff --git a/samples/python/yolov3_onnx/requirements.txt b/samples/python/yolov3_onnx/requirements.txt index fbca5658..dada880d 100644 --- a/samples/python/yolov3_onnx/requirements.txt +++ b/samples/python/yolov3_onnx/requirements.txt @@ -1,9 +1,9 @@ cuda-python==12.9.0 pywin32; platform_system == "Windows" numpy==1.26.4 -onnx==1.16.0 +onnx==1.18.0 Pillow==11.3.0 protobuf==3.20.3 pyyaml==6.0.1 -requests==2.32.2 +requests==2.32.4 tqdm==4.66.4 diff --git a/samples/sampleCharRNN/CMakeLists.txt b/samples/sampleCharRNN/CMakeLists.txt index 7fa61f6b..89feaf88 100644 --- a/samples/sampleCharRNN/CMakeLists.txt +++ b/samples/sampleCharRNN/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_char_rnn sampleCharRNN.cpp) -target_link_libraries(sample_char_rnn PRIVATE trt_samples_common) +target_link_libraries(sample_char_rnn PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_char_rnn) installLibraries( TARGETS sample_char_rnn OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleCharRNN/README.md b/samples/sampleCharRNN/README.md index 40c81489..4bd3edc9 100644 --- a/samples/sampleCharRNN/README.md +++ b/samples/sampleCharRNN/README.md @@ -25,7 +25,7 @@ There are also many GitHub repositories that contain CharRNN implementations tha The CharRNN network is a fairly simple RNN network. The input into the network is a single character that is embedded into a vector of size 512. This embedded input is then supplied to a RNN layer containing two stacked LSTM cells. The output from the RNN layer is then supplied to a fully connected layer, which can be represented in TensorRT by a Matrix Multiply layer followed by an ElementWise sum layer. Constant layers are used to supply the weights and biases to the Matrix Multiply and ElementWise Layers, respectively. A TopK operation is then performed on the output of the ElementWise sum layer where `K = 1` to find the next predicted character in the sequence. For more information about these layers, see the [TensorRT API](http://docs.nvidia.com/deeplearning/sdk/tensorrt-api/index.html) documentation. -This sample provides a pre-trained model called `model-20080.data-00000-of-00001` located in the `/usr/src/tensorrt/data/samples/char-rnn/model` directory, therefore, training is not required for this sample. The model used by this sample was trained using [tensorflow-char-rnn](https://github.com/crazydonkey200/tensorflow-char-rnn). This GitHub repository includes instructions on how to train and produce checkpoint that can be used by TensorRT. +This sample provides a pre-trained model called `model-20080.data-00000-of-00001` located in the `$TRT_DATADIR/char-rnn/model` directory, therefore, training is not required for this sample. The model used by this sample was trained using [tensorflow-char-rnn](https://github.com/crazydonkey200/tensorflow-char-rnn). This GitHub repository includes instructions on how to train and produce checkpoint that can be used by TensorRT. **Note:** If you wanted to train your own model and then perform inference with TensorRT, you will simply need to do a char to char comparison between TensorFlow and TensorRT. @@ -43,12 +43,11 @@ The MatrixMultiply layer implements matrix multiplication for a collection of ma [TopK](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#topk-layer) The TopK layer is used to identify the character that has the maximum probability of appearing next. The TopK layer finds the top K maximum (or minimum) elements along a dimension, returning a reduced tensor and a tensor of index positions. -## Preparing sample data +## Prerequisites -1. Download the sample data from [TensorRT release tarball](https://developer.nvidia.com/nvidia-tensorrt-download#), if not already mounted under `/usr/src/tensorrt/data` (NVIDIA NGC containers) and set it to `$TRT_DATADIR`. - ```bash - export TRT_DATADIR=/usr/src/tensorrt/data - ``` +1. Preparing sample data + +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. ## Converting TensorFlow weights @@ -122,6 +121,9 @@ documentation. # Changelog +September 2025 +* Use new addTopK API with indices type. + January 2024 * Removed RNNv2Layer based addLSTMLayer implementation. addLSTMLayer is now implemented with ILoop only. * Default to use ILoop in paramaters. diff --git a/samples/sampleCharRNN/sampleCharRNN.cpp b/samples/sampleCharRNN/sampleCharRNN.cpp index 680e8943..2fdbb9f6 100644 --- a/samples/sampleCharRNN/sampleCharRNN.cpp +++ b/samples/sampleCharRNN/sampleCharRNN.cpp @@ -718,7 +718,8 @@ void SampleCharRNNBase::constructNetwork(SampleUniquePtr& bu // Add TopK layer to determine which character has highest probability. int reduceAxis = 0x1; // reduce across vocab axis - auto pred = network->addTopK(*addBiasLayer->getOutput(0), nvinfer1::TopKOperation::kMAX, 1, reduceAxis); + auto pred = network->addTopK( + *addBiasLayer->getOutput(0), nvinfer1::TopKOperation::kMAX, 1, reduceAxis, nvinfer1::DataType::kINT32); ASSERT(pred != nullptr); pred->getOutput(1)->setName(mParams.bindingNames.OUTPUT_BLOB_NAME); diff --git a/samples/sampleDynamicReshape/CMakeLists.txt b/samples/sampleDynamicReshape/CMakeLists.txt index d56bb803..45f7c9ae 100644 --- a/samples/sampleDynamicReshape/CMakeLists.txt +++ b/samples/sampleDynamicReshape/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_dynamic_reshape sampleDynamicReshape.cpp) -target_link_libraries(sample_dynamic_reshape PRIVATE trt_samples_common) +target_link_libraries(sample_dynamic_reshape PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_dynamic_reshape) installLibraries( TARGETS sample_dynamic_reshape OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleDynamicReshape/README.md b/samples/sampleDynamicReshape/README.md index 4216dbc6..302848d6 100644 --- a/samples/sampleDynamicReshape/README.md +++ b/samples/sampleDynamicReshape/README.md @@ -181,15 +181,8 @@ In this sample, the following layers are used. For more information about these [Resize layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#resize-layer) The IResizeLayer implements the resize operation on an input tensor. -## Preparing sample data - -1. Download the sample data from [TensorRT release tarball](https://developer.nvidia.com/nvidia-tensorrt-download#), if not already mounted under `/usr/src/tensorrt/data` (NVIDIA NGC containers) and set it to `$TRT_DATADIR`. - ```bash - export TRT_DATADIR=/usr/src/tensorrt/data - pushd $TRT_DATADIR/mnist - pip3 install Pillow - popd - ``` +## Prerequisites +1. See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. ## Running the sample diff --git a/samples/sampleDynamicReshape/sampleDynamicReshape.cpp b/samples/sampleDynamicReshape/sampleDynamicReshape.cpp index b66ca5b8..0fc9a92d 100644 --- a/samples/sampleDynamicReshape/sampleDynamicReshape.cpp +++ b/samples/sampleDynamicReshape/sampleDynamicReshape.cpp @@ -26,7 +26,6 @@ // Define TRT entrypoints used in common code #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "BatchStream.h" #include "EntropyCalibrator.h" diff --git a/samples/sampleEditableTimingCache/CMakeLists.txt b/samples/sampleEditableTimingCache/CMakeLists.txt index e95c46ef..106e472c 100644 --- a/samples/sampleEditableTimingCache/CMakeLists.txt +++ b/samples/sampleEditableTimingCache/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_editable_timing_cache sampleEditableTimingCache.cpp) -target_link_libraries(sample_editable_timing_cache PRIVATE trt_samples_common) +target_link_libraries(sample_editable_timing_cache PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_editable_timing_cache) installLibraries( TARGETS sample_editable_timing_cache OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleEditableTimingCache/sampleEditableTimingCache.cpp b/samples/sampleEditableTimingCache/sampleEditableTimingCache.cpp index ef8744d7..d0dce55e 100644 --- a/samples/sampleEditableTimingCache/sampleEditableTimingCache.cpp +++ b/samples/sampleEditableTimingCache/sampleEditableTimingCache.cpp @@ -42,7 +42,6 @@ #include // for strtoull #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "NvInfer.h" #include "common.h" #include "logger.h" diff --git a/samples/sampleINT8API/CMakeLists.txt b/samples/sampleINT8API/CMakeLists.txt index 6acc3dad..e62dcee8 100644 --- a/samples/sampleINT8API/CMakeLists.txt +++ b/samples/sampleINT8API/CMakeLists.txt @@ -18,13 +18,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_int8_api sampleINT8API.cpp) -target_link_libraries(sample_int8_api PRIVATE trt_samples_common) +target_link_libraries(sample_int8_api PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_int8_api) installLibraries( TARGETS sample_int8_api OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleINT8API/README.md b/samples/sampleINT8API/README.md index 4c2264a8..189bb709 100644 --- a/samples/sampleINT8API/README.md +++ b/samples/sampleINT8API/README.md @@ -145,13 +145,15 @@ Set the computational precision of this layer. Setting the precision forces Tens [ILayer::SetOutputType](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/classnvinfer1_1_1_i_layer.html#a85aded4e3ff0867e392602551d5b5dc7) Set the output type of this layer. Setting the output type forces TensorRT to choose the implementations which generate output data with the given type. If the output type is not set, TensorRT will select the implementation based on performance considerations and the flags specified to the builder. -## Preparing sample data +## Prerequisites +1. Preparing sample data +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. -`ResNet50.onnx` is located in the `data/resnet50` directory. +`ResNet50.onnx` is located in the `$TRT_DATADIR/resnet50` directory. In addition to the model file and input image, you will need per-tensor dynamic range stored in a text file along with the ImageNet label reference file. -The following required files are included in the package and are located in the `data/int8_api` directory. +The following required files are included in the package and are located in the `$TRT_DATADIR/int8_api` directory. `reference_labels.txt` The ImageNet reference label file. diff --git a/samples/sampleINT8API/sampleINT8API.cpp b/samples/sampleINT8API/sampleINT8API.cpp index 352d66a0..b6b7e02a 100644 --- a/samples/sampleINT8API/sampleINT8API.cpp +++ b/samples/sampleINT8API/sampleINT8API.cpp @@ -26,7 +26,6 @@ // Define TRT entrypoints used in common code #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "argsParser.h" #include "buffers.h" diff --git a/samples/sampleIOFormats/CMakeLists.txt b/samples/sampleIOFormats/CMakeLists.txt index cec34a1c..7f70a9bc 100644 --- a/samples/sampleIOFormats/CMakeLists.txt +++ b/samples/sampleIOFormats/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_io_formats sampleIOFormats.cpp) -target_link_libraries(sample_io_formats PRIVATE trt_samples_common) +target_link_libraries(sample_io_formats PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_io_formats) installLibraries( TARGETS sample_io_formats OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleIOFormats/README.md b/samples/sampleIOFormats/README.md index 1d5fd80e..ee8e0264 100644 --- a/samples/sampleIOFormats/README.md +++ b/samples/sampleIOFormats/README.md @@ -30,15 +30,9 @@ This sample, sampleIOFormats, uses a Onnx model that was trained on the [MNIST d } ``` -## Preparing sample data - -1. Download the sample data from [TensorRT release tarball](https://developer.nvidia.com/nvidia-tensorrt-download#), if not already mounted under `/usr/src/tensorrt/data` (NVIDIA NGC containers) and set it to `$TRT_DATADIR`. - ```bash - export TRT_DATADIR=/usr/src/tensorrt/data - pushd $TRT_DATADIR/mnist - pip3 install Pillow - popd - ``` +## Prerequisites +1. Preparing sample data +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. ## Running the sample diff --git a/samples/sampleNamedDimensions/CMakeLists.txt b/samples/sampleNamedDimensions/CMakeLists.txt index 882f88d1..645815b2 100644 --- a/samples/sampleNamedDimensions/CMakeLists.txt +++ b/samples/sampleNamedDimensions/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_named_dimensions sampleNamedDimensions.cpp) -target_link_libraries(sample_named_dimensions PRIVATE trt_samples_common) +target_link_libraries(sample_named_dimensions PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_named_dimensions) installLibraries( TARGETS sample_named_dimensions OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleNamedDimensions/sampleNamedDimensions.cpp b/samples/sampleNamedDimensions/sampleNamedDimensions.cpp index 07b83eff..9c3c4026 100644 --- a/samples/sampleNamedDimensions/sampleNamedDimensions.cpp +++ b/samples/sampleNamedDimensions/sampleNamedDimensions.cpp @@ -25,7 +25,6 @@ // Define TRT entrypoints used in common code #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "argsParser.h" #include "buffers.h" diff --git a/samples/sampleNonZeroPlugin/CMakeLists.txt b/samples/sampleNonZeroPlugin/CMakeLists.txt index 313f785a..355fc706 100644 --- a/samples/sampleNonZeroPlugin/CMakeLists.txt +++ b/samples/sampleNonZeroPlugin/CMakeLists.txt @@ -20,13 +20,13 @@ add_executable(sample_non_zero_plugin sampleNonZeroPlugin.cpp nonZeroKernel.cu ) -target_link_libraries(sample_non_zero_plugin PRIVATE trt_samples_common) +target_link_libraries(sample_non_zero_plugin PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_non_zero_plugin) installLibraries( TARGETS sample_non_zero_plugin OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleNonZeroPlugin/README.md b/samples/sampleNonZeroPlugin/README.md index 411dbb76..efe1beaf 100644 --- a/samples/sampleNonZeroPlugin/README.md +++ b/samples/sampleNonZeroPlugin/README.md @@ -85,9 +85,10 @@ interface. As sample inputs, random images from MNIST dataset are selected and scaled to between `[0,1]`. The network will output both the non-zero indices, as well as the non-zero count. -## Preparing sample data +## Prerequisites +1. Preparing sample data -Download the sample data from the [TensorRT release tarball](https://developer.nvidia.com/nvidia-tensorrt-download#). +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. ## Running the sample diff --git a/samples/sampleNonZeroPlugin/sampleNonZeroPlugin.cpp b/samples/sampleNonZeroPlugin/sampleNonZeroPlugin.cpp index 4de0b5b5..ee8c68a6 100644 --- a/samples/sampleNonZeroPlugin/sampleNonZeroPlugin.cpp +++ b/samples/sampleNonZeroPlugin/sampleNonZeroPlugin.cpp @@ -24,7 +24,6 @@ // Define TRT entrypoints used in common code #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "argsParser.h" #include "buffers.h" diff --git a/samples/sampleOnnxMNIST/CMakeLists.txt b/samples/sampleOnnxMNIST/CMakeLists.txt index 8288f05d..deab5ad9 100644 --- a/samples/sampleOnnxMNIST/CMakeLists.txt +++ b/samples/sampleOnnxMNIST/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_onnx_mnist sampleOnnxMNIST.cpp) -target_link_libraries(sample_onnx_mnist PRIVATE trt_samples_common) +target_link_libraries(sample_onnx_mnist PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_onnx_mnist) installLibraries( TARGETS sample_onnx_mnist OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleOnnxMNIST/README.md b/samples/sampleOnnxMNIST/README.md index 3196e806..e0da026b 100644 --- a/samples/sampleOnnxMNIST/README.md +++ b/samples/sampleOnnxMNIST/README.md @@ -91,9 +91,10 @@ The Scale layer implements a per-tensor, per-channel, or per-element affine tran [Shuffle layer](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#shuffle-layer) The Shuffle layer implements a reshape and transpose operator for tensors. -## Preparing sample data +## Prerequisites +1. Preparing sample data -Download the sample data from the [TensorRT release tarball](https://developer.nvidia.com/nvidia-tensorrt-download#). +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. ## Running the sample diff --git a/samples/sampleOnnxMNIST/sampleOnnxMNIST.cpp b/samples/sampleOnnxMNIST/sampleOnnxMNIST.cpp index 8da7215e..5c55e68d 100644 --- a/samples/sampleOnnxMNIST/sampleOnnxMNIST.cpp +++ b/samples/sampleOnnxMNIST/sampleOnnxMNIST.cpp @@ -26,7 +26,6 @@ // Define TRT entrypoints used in common code #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "argsParser.h" #include "buffers.h" diff --git a/samples/sampleOnnxMnistCoordConvAC/CMakeLists.txt b/samples/sampleOnnxMnistCoordConvAC/CMakeLists.txt index 2665c6dd..ce0d8608 100644 --- a/samples/sampleOnnxMnistCoordConvAC/CMakeLists.txt +++ b/samples/sampleOnnxMnistCoordConvAC/CMakeLists.txt @@ -17,7 +17,7 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_onnx_mnist_coord_conv_ac sampleOnnxMnistCoordConvAC.cpp) -target_link_libraries(sample_onnx_mnist_coord_conv_ac PRIVATE trt_samples_common) +target_link_libraries(sample_onnx_mnist_coord_conv_ac PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_onnx_mnist_coord_conv_ac) if(${TRT_BUILD_SAMPLES_LINK_STATIC_TRT}) @@ -33,7 +33,7 @@ endif() installLibraries( TARGETS sample_onnx_mnist_coord_conv_ac OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleOnnxMnistCoordConvAC/README.md b/samples/sampleOnnxMnistCoordConvAC/README.md index d3022496..3dd2b4ee 100644 --- a/samples/sampleOnnxMnistCoordConvAC/README.md +++ b/samples/sampleOnnxMnistCoordConvAC/README.md @@ -118,6 +118,10 @@ The Shuffle layer implements a reshape and transpose operator for tensors. +## Preparing sample data + +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. + ## Running the sample 1. The sample gets compiled when building the TensorRT OSS following the [instructions](https://github.com/NVIDIA/TensorRT). The binary named sample_onnx_mnist_coord_conv_ac will be created in the output directory. diff --git a/samples/sampleOnnxMnistCoordConvAC/sampleOnnxMnistCoordConvAC.cpp b/samples/sampleOnnxMnistCoordConvAC/sampleOnnxMnistCoordConvAC.cpp index 54b29424..8223cc6a 100644 --- a/samples/sampleOnnxMnistCoordConvAC/sampleOnnxMnistCoordConvAC.cpp +++ b/samples/sampleOnnxMnistCoordConvAC/sampleOnnxMnistCoordConvAC.cpp @@ -26,7 +26,6 @@ // Define TRT entrypoints used in common code #define DEFINE_TRT_ENTRYPOINTS 1 -#define DEFINE_TRT_LEGACY_PARSER_ENTRYPOINT 0 #include "argsParser.h" #include "buffers.h" diff --git a/samples/sampleProgressMonitor/CMakeLists.txt b/samples/sampleProgressMonitor/CMakeLists.txt index 6b5759b0..6a53bfe0 100644 --- a/samples/sampleProgressMonitor/CMakeLists.txt +++ b/samples/sampleProgressMonitor/CMakeLists.txt @@ -17,13 +17,13 @@ if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) add_executable(sample_progress_monitor sampleProgressMonitor.cpp) -target_link_libraries(sample_progress_monitor PRIVATE trt_samples_common) +target_link_libraries(sample_progress_monitor PRIVATE trt_samples_common TRT_SAMPLES::tensorrt) add_dependencies(tensorrt_samples sample_progress_monitor) installLibraries( TARGETS sample_progress_monitor OPTIONAL - COMPONENT full + COMPONENT internal ) else() diff --git a/samples/sampleProgressMonitor/README.md b/samples/sampleProgressMonitor/README.md index 35477677..51fa61b6 100644 --- a/samples/sampleProgressMonitor/README.md +++ b/samples/sampleProgressMonitor/README.md @@ -42,15 +42,9 @@ This sample implements an `IProgressMonitor` to display progress while building The progress bars are drawn using virtual terminal escape sequences to manipulate the terminal's cursor and clear lines. -## Preparing sample data - -1. Download the sample data from [TensorRT release tarball](https://developer.nvidia.com/nvidia-tensorrt-download#), if not already mounted under `/usr/src/tensorrt/data` (NVIDIA NGC containers) and set it to `$TRT_DATADIR`. - ```bash - export TRT_DATADIR=/usr/src/tensorrt/data - pushd $TRT_DATADIR/mnist - pip3 install Pillow - popd - ``` +## Prerequisites +1. Preparing sample data +See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README. ## Running the sample diff --git a/samples/trtexec/CMakeLists.txt b/samples/trtexec/CMakeLists.txt index 6628a1c1..0807dc38 100644 --- a/samples/trtexec/CMakeLists.txt +++ b/samples/trtexec/CMakeLists.txt @@ -15,19 +15,67 @@ # limitations under the License. # if (${TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW}) + add_executable(trtexec trtexec.cpp) target_link_libraries(trtexec PRIVATE trt_samples_common) if (TRT_BUILD_SAMPLES) add_dependencies(tensorrt_samples trtexec) endif() +# Windows-specific configuration for long path support +if(MSVC) + # For long path support to work on Windows, users must + # additionally enable a system-level setting. + # Adding manifest as a source file works on CMake 3.4+. + target_sources(trtexec PRIVATE trtexec.manifest) +endif() + installLibraries( TARGETS trtexec OPTIONAL - COMPONENT release + COMPONENT external ) -else() +# When statically linked, trtexec requires the plugins library. +# Otherwise it will resolve it via dlopen and doesn't need a link-time dependency. +if(${TRT_BUILD_SAMPLES_LINK_STATIC_TRT}) + target_link_libraries(trtexec PRIVATE + $ + ) +endif() + +# Change the file name if TRT_WINML variable is set +set(sample_name "trtexec") +if (${TRT_BUILD_WINML}) + set(sample_name "tensorrt_rtx") +endif() + +set_target_properties(trtexec + PROPERTIES + OUTPUT_NAME ${sample_name} +) + +# In this mode, we build an additional binary trtexec_static that always links tensorrt_static. +if(${TRT_BUILD_TRTEXEC_STATIC}) + add_executable(trtexec_static trtexec.cpp) + target_link_libraries(trtexec_static PRIVATE trt_samples_common) + if (TRT_BUILD_SAMPLES) + add_dependencies(tensorrt_samples trtexec_static) + endif() + + installLibraries( + TARGETS trtexec_static + OPTIONAL + COMPONENT internal + ) + + target_link_libraries(trtexec_static PRIVATE + $ + $ + ) +endif() + +else() # TRT_BUILD_ENABLE_NEW_SAMPLES_FLOW - old flow below. set(SAMPLE_SOURCES ../common/sampleDevice.cpp @@ -44,10 +92,4 @@ set(SAMPLE_PARSERS "onnx") include(../CMakeSamplesTemplate.txt) -# Change the file name if TRT_WINML variable is set -if (${TRT_BUILD_WINML}) - set_target_properties(trtexec PROPERTIES - OUTPUT_NAME tensorrt_rtx) -endif() - endif() diff --git a/samples/trtexec/trtexec.cpp b/samples/trtexec/trtexec.cpp index 9456562b..9ec08e20 100644 --- a/samples/trtexec/trtexec.cpp +++ b/samples/trtexec/trtexec.cpp @@ -55,6 +55,7 @@ std::function pCreateInferRuntimeInternal{}; std::function pCreateInferRefitterInternal{}; std::function pCreateInferBuilderInternal{}; std::function pCreateNvOnnxParserInternal{}; +std::function pCreateNvOnnxRefitterInternal{}; //! Track runtime used for the execution of trtexec. //! Must be tracked as a global variable due to how library init functions APIs are organized. @@ -96,10 +97,13 @@ bool initNvonnxparser() static LibraryPtr libnvonnxparserPtr{}; auto fetchPtrs = [](DynamicLibrary* l) { pCreateNvOnnxParserInternal = l->symbolAddress("createNvOnnxParser_INTERNAL"); + pCreateNvOnnxRefitterInternal + = l->symbolAddress("createNvOnnxParserRefitter_INTERNAL"); }; return initLibrary(libnvonnxparserPtr, kNVONNXPARSER_LIBNAME, fetchPtrs); #else pCreateNvOnnxParserInternal = createNvOnnxParser_INTERNAL; + pCreateNvOnnxRefitterInternal = createNvOnnxParserRefitter_INTERNAL; return true; #endif // !TRT_STATIC } @@ -147,6 +151,17 @@ nvonnxparser::IParser* createONNXParser(INetworkDefinition& network) pCreateNvOnnxParserInternal(&network, &gLogger.getTRTLogger(), NV_ONNX_PARSER_VERSION)); } +nvonnxparser::IParserRefitter* createONNXRefitter(nvinfer1::IRefitter& refitter) +{ + if (!initNvonnxparser()) + { + return {}; + } + ASSERT(pCreateNvOnnxRefitterInternal != nullptr); + return static_cast( + pCreateNvOnnxRefitterInternal(&refitter, &gLogger.getTRTLogger(), NV_ONNX_PARSER_VERSION)); +} + #if ENABLE_UNIFIED_BUILDER bool processSafetyPluginLibrary(nvinfer2::safe::ISafePluginRegistry* safetyPluginRegistry, DynamicLibrary* libPtr, @@ -265,7 +280,7 @@ int main(int argc, char** argv) LibraryPtr nvinferPluginLib{}; #endif /* TRT_STATIC */ std::vector pluginLibs; - if (gUseRuntime == RuntimeMode::kFULL) + if (gUseRuntime == RuntimeMode::kFULL && !options.build.safe) { sample::gLogInfo << "Loading standard plugins" << std::endl; #if !TRT_STATIC @@ -283,6 +298,10 @@ int main(int argc, char** argv) pluginLibs.emplace_back(loadLibrary(pluginPath)); } } + else if (gUseRuntime == RuntimeMode::kFULL && options.build.safe) + { + sample::gLogInfo << "Skipping standard plugin loading due to --safe flag" << std::endl; + } else if (!options.system.plugins.empty()) { throw std::runtime_error("TRT-18412: Plugins require --useRuntime=full."); @@ -315,6 +334,12 @@ int main(int argc, char** argv) options.build.consistency = false; } + if (options.build.safe) + { + sample::gLogInfo << "StronglyTyped is enabled by default on safety mode." << std::endl; + options.build.stronglyTyped = true; + } + // Start engine building phase. std::unique_ptr bEnv(new BuildEnvironment(options.build.safe, options.build.versionCompatible, options.system.DLACore, options.build.tempdir, options.build.tempfileControls, options.build.leanDLLPath, @@ -353,6 +378,16 @@ int main(int argc, char** argv) { dumpRefittable(*engine); } + // Refit from ONNX model + if (!options.inference.refitOnnxModel.empty()) + { + bool const success = refitFromOnnx(*engine, options.inference.refitOnnxModel, options.inference.threads); + if (!success) + { + sample::gLogError << "Engine refit from ONNX model failed." << std::endl; + return sample::gLogger.reportFail(sampleTest); + } + } if (options.inference.timeRefit) { if (bEnv->network.operator bool()) diff --git a/samples/trtexec/trtexec.manifest b/samples/trtexec/trtexec.manifest new file mode 100644 index 00000000..feb2ae4a --- /dev/null +++ b/samples/trtexec/trtexec.manifest @@ -0,0 +1,21 @@ + + + + + trtexec with long path support + + + true + + + diff --git a/shared/CMakeLists.txt b/shared/CMakeLists.txt index 093da544..1e015c63 100644 --- a/shared/CMakeLists.txt +++ b/shared/CMakeLists.txt @@ -22,7 +22,6 @@ endfunction() target_link_libraries(trt_shared PRIVATE $ trt_global_definitions - TRT::cudart ) add_subdirectory(utils) diff --git a/shared/utils/fileLock.cpp b/shared/utils/fileLock.cpp index 177aaa53..24fe640b 100644 --- a/shared/utils/fileLock.cpp +++ b/shared/utils/fileLock.cpp @@ -23,9 +23,9 @@ namespace nvinfer1::utils { -FileLock::FileLock(ILogger& logger, std::string const& fileName) +FileLock::FileLock(ILogger& logger, std::string fileName) : mLogger(logger) - , mFileName(fileName) + , mFileName(std::move(fileName)) { std::string lockFileName = mFileName + ".lock"; #ifdef _MSC_VER @@ -42,6 +42,7 @@ FileLock::FileLock(ILogger& logger, std::string const& fileName) } #elif defined(__QNX__) // Calling lockf(F_TLOCK) on QNX returns -1; the reported error is 89 (function not implemented). + mLogger.log(ILogger::Severity::kVERBOSE, "FileLock is not supported on QNX or GOS."); #else mHandle = fopen(lockFileName.c_str(), "wb+"); if (mHandle == nullptr) @@ -74,6 +75,7 @@ FileLock::~FileLock() } #elif defined(__QNX__) // Calling lockf(F_TLOCK) on QNX returns -1; the reported error is 89 (function not implemented). + mLogger.log(ILogger::Severity::kVERBOSE, "FileLock is not supported on QNX or GOS."); #else if (mDescriptor != -1) { diff --git a/shared/utils/fileLock.h b/shared/utils/fileLock.h index 6a595473..774cb5a0 100644 --- a/shared/utils/fileLock.h +++ b/shared/utils/fileLock.h @@ -41,7 +41,7 @@ namespace nvinfer1::utils class FileLock { public: - FileLock(nvinfer1::ILogger& logger, std::string const& fileName); + explicit FileLock(nvinfer1::ILogger& logger, std::string fileName); ~FileLock(); FileLock() = delete; // no default ctor FileLock(FileLock const&) = delete; // no copy ctor @@ -66,7 +66,7 @@ private: //! The file handle on windows for the file lock. //! HANDLE mHandle{}; -#else +#elif !defined(__QNX__) //! //! The file handle on linux for the file lock. //! diff --git a/tools/Polygraphy/CHANGELOG.md b/tools/Polygraphy/CHANGELOG.md index 730f576e..35c33eaf 100644 --- a/tools/Polygraphy/CHANGELOG.md +++ b/tools/Polygraphy/CHANGELOG.md @@ -2,6 +2,37 @@ Dates are in YYYY-MM-DD format. +## v0.49.27 +### Added +- Added `polygraphy template shard-hints` to generate hints file for `polygraphy multi-device shard`. +- Added support for inserting transposes to `polygraphy multi-device shard` when the sequence length dimension of sharded tensors is not 0. + +### Changed +- Changed hints file format for `polygraphy multi-device shard` to support new additions. +- Changed `polygraphy multi-device shard` example and README to reflect new additions. + +### Fixed +- Fixed issue when `polygraphy multi-device shard` would exceed python recursive depth limit on large models. + + +## v0.49.26 (2025-07-16) +### Added +- Added support for dumping unfused tensors when running TensorRT using `--mark-unfused-tensors-as-debug-tensors`. + int4, fp4, fp8, and bfloat16 tensors are not supported and will be skipped. +- Added support for combining tensor statistics from the tensor JSON file with the output of the engine inspector. + +### Changed +- Relaxed `onnxconverter_common` version requirements to allow newer versions to be installed. + +### Fixed +- Fixed a bug where `--save-heatmaps` and `--save-error-metrics-plot` would not handle slashes in output names well. +- Fixed a bug where `--save-error-metrics-plot` would not work if the minimum error was 0. + + +## v0.49.25 (2025-06-11) +### Added +- Added `multi-device shard` tool for converting SD models to MD. + ## v0.49.24 (2025-05-27) ### Fixed diff --git a/tools/Polygraphy/examples/cli/inspect/02_inspecting_a_tensorrt_engine/README.md b/tools/Polygraphy/examples/cli/inspect/02_inspecting_a_tensorrt_engine/README.md index c9e14069..37593a80 100644 --- a/tools/Polygraphy/examples/cli/inspect/02_inspecting_a_tensorrt_engine/README.md +++ b/tools/Polygraphy/examples/cli/inspect/02_inspecting_a_tensorrt_engine/README.md @@ -18,6 +18,9 @@ about TensorRT engines, i.e. plan files: --save-engine dynamic_identity.engine ``` + You can also dump unfused intermediate tensors by adding `--mark-unfused-tensors-as-debug-tensors` and + `--save-outputs output.json` options. Later, this tensor information can be combined with the inspector output. + 2. Inspect the engine: ```bash @@ -65,3 +68,15 @@ about TensorRT engines, i.e. plan files: ``` It is also possible to show more detailed layer information using `--show layers attrs`. + + You can also combine tensor value statistics using `--combine-tensor-info output.json` where the JSON file is got from + `--mark-unfused-tensors-as-debug-tensors` and `--save-outputs output.json`. + + The statistics will be added to the input and output tensors of each layer: + + + ``` + {X [dtype=float32, shape=(1, 2, -1, -1), Format: Float, min=0.42, max=0.72, avg=0.57]} + -> {Y [dtype=float32, shape=(1, 2, -1, -1), Format: Float, min=0.42, max=0.72, avg=0.57]} + ``` + diff --git a/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/README.md b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/README.md new file mode 100644 index 00000000..eec5c773 --- /dev/null +++ b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/README.md @@ -0,0 +1,36 @@ +# Using Shard To Convert a SD Model to MD + + +## Introduction + +The `shard` tool can be used to convert single-device (SD) models containing attention layers into multi-device (MD) models intended to be run on multiple GPUs using a hints file. + +In this example, we'll show how to shard a simple model containing an attention layer + +![./model.png](./model.png) + +## Hint Configuration + +For this example we'll be using [this hints file](./hint.json). + +See the [Shard README](../../../../polygraphy/tools/multi_device/README.md#sharding-hints-file-format) for an explanation of the hints file format. + + +## Running the Example + +```bash +polygraphy multi-device shard \ + ../attention.onnx \ + -s hint.json \ + -o attention_md.onnx +``` + +Looking at the result, we can now see the model is ready to be run on multiple GPUs through TensorRT + +![./model_md.png](./model_md.png) + + +### A Note On Gathering Q + +If we changed `gather_q` in the hints to `true` the model effectively becomes SD, and a final all-gather will not be inserted. All attention layers must have Q consistently sharded, as it affects whether or not to place an all-gather at the output of the model + diff --git a/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/hint.json b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/hint.json new file mode 100644 index 00000000..af2566ac --- /dev/null +++ b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/hint.json @@ -0,0 +1,35 @@ +{ + "parallelism": "CP", + "group_size": 0, + "root": 0, + "groups": [], + "attention_layers": [ + { + "q": "q", + "gather_kv": true, + "gather_q": false, + "polygraphy_class": "AttentionLayerHint" + } + ], + "inputs": [ + { + "name": "input", + "seq_len_idx": 0, + "rank": 3, + "polygraphy_class": "ShardTensor" + } + ], + "outputs": [ + { + "name": "output", + "seq_len_idx": 0, + "rank": 3, + "polygraphy_class": "ShardTensor" + } + ], + "k_seq_len_idx": 0, + "v_seq_len_idx": 0, + "kv_rank": null, + "reduce_scatter_reduce_op": "max", + "polygraphy_class": "ShardHints" +} \ No newline at end of file diff --git a/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/model.png b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/model.png new file mode 100644 index 00000000..6ec565f9 Binary files /dev/null and b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/model.png differ diff --git a/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/model_md.png b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/model_md.png new file mode 100644 index 00000000..32f44f04 Binary files /dev/null and b/tools/Polygraphy/examples/cli/multi_device/01_sharding_attention_layers/model_md.png differ diff --git a/tools/Polygraphy/examples/cli/multi_device/attention_head.onnx b/tools/Polygraphy/examples/cli/multi_device/attention_head.onnx new file mode 100644 index 00000000..03df332b Binary files /dev/null and b/tools/Polygraphy/examples/cli/multi_device/attention_head.onnx differ diff --git a/tools/Polygraphy/polygraphy/__init__.py b/tools/Polygraphy/polygraphy/__init__.py index 7d0b861c..80c084ae 100644 --- a/tools/Polygraphy/polygraphy/__init__.py +++ b/tools/Polygraphy/polygraphy/__init__.py @@ -1,3 +1,3 @@ import polygraphy.config -__version__ = "0.49.24" +__version__ = "0.49.27" diff --git a/tools/Polygraphy/polygraphy/backend/onnx/loader.py b/tools/Polygraphy/polygraphy/backend/onnx/loader.py index ea1179a9..01accbf5 100644 --- a/tools/Polygraphy/polygraphy/backend/onnx/loader.py +++ b/tools/Polygraphy/polygraphy/backend/onnx/loader.py @@ -29,7 +29,7 @@ np = mod.lazy_import("numpy") onnx = mod.lazy_import("onnx>=1.8.1") onnxrt = mod.lazy_import("onnxruntime>=1.10.0") onnxmltools = mod.lazy_import( - "onnxmltools==1.11.1", requires=["onnxconverter_common==1.12.2"] + "onnxmltools==1.11.1", requires=["onnxconverter_common>=1.12.2"] ) tf = mod.lazy_import("tensorflow<2.0") tf2onnx = mod.lazy_import("tf2onnx") diff --git a/tools/Polygraphy/polygraphy/backend/tensorrt_rtx/__init__.py b/tools/Polygraphy/polygraphy/backend/tensorrt_rtx/__init__.py new file mode 100644 index 00000000..b4144534 --- /dev/null +++ b/tools/Polygraphy/polygraphy/backend/tensorrt_rtx/__init__.py @@ -0,0 +1 @@ +from polygraphy.backend.tensorrt_rtx.config import * diff --git a/tools/Polygraphy/polygraphy/backend/tensorrt_rtx/config.py b/tools/Polygraphy/polygraphy/backend/tensorrt_rtx/config.py new file mode 100644 index 00000000..79666ef1 --- /dev/null +++ b/tools/Polygraphy/polygraphy/backend/tensorrt_rtx/config.py @@ -0,0 +1,117 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from polygraphy import config as polygraphy_config, mod, util +from polygraphy.backend.trt.config import _CreateConfigCommon +from polygraphy.backend.trt.util import inherit_and_extend_docstring +from polygraphy.logger import G_LOGGER +from polygraphy.mod.trt_importer import lazy_import_trt + +trt = lazy_import_trt() + + +@mod.export(funcify=True, func_name="create_config_rtx") +class CreateConfigRTX(_CreateConfigCommon): + """ + Functor that creates an IBuilderConfig with TensorRT-RTX specific features. + """ + + @inherit_and_extend_docstring(_CreateConfigCommon.__init__) + def __init__( + self, + use_gpu=None, + compute_capabilities=None, + **kwargs + ): + """ + Creates an IBuilderConfig with TensorRT-RTX specific features. + + Args: + use_gpu (bool): + Whether to use the current GPU device as target for engine compilation. + Equivalent to setting ComputeCapability.CURRENT. This is mutually exclusive with compute_capabilities. + Defaults to False. + compute_capabilities (List[Tuple[int, int]]): + List of (major, minor) compute capability tuples to target for engine compilation. + This is mutually exclusive with use_gpu. When specified, the engine can only run on devices + with the specified compute capabilities. + Defaults to None. + """ + super().__init__(**kwargs) + self.use_gpu = util.default(use_gpu, False) + self.compute_capabilities = compute_capabilities + + if self.use_gpu and self.compute_capabilities: + G_LOGGER.critical("use_gpu and compute_capabilities are mutually exclusive.") + + self._validator() + + def _validator(self): + """ + Validates initialization parameters for TensorRT-RTX specific features. + """ + if self.use_gpu or self.compute_capabilities is not None: + if not polygraphy_config.USE_TENSORRT_RTX: + G_LOGGER.critical("--compute-capabilities and --use-gpu settings are only supported with USE_TENSORRT_RTX=1.") + + # Validate compute capabilities format and availability + if self.compute_capabilities: + for major, minor in self.compute_capabilities: + cap_name = f"SM{major}{minor}" + if not hasattr(trt.ComputeCapability, cap_name): + G_LOGGER.critical(f"Compute capability {major}.{minor} ({cap_name})" + " not supported by this TensorRT-RTX version.") + + def _configure_flags(self, builder, network, config): + """ + Validates and configures TensorRT-RTX-specific features. + + Args: + builder (trt.Builder): The TensorRT builder + network (trt.INetworkDefinition): The TensorRT network + config (trt.IBuilderConfig): The TensorRT builder config to modify + """ + # Set compute capabilities if specified + if self.use_gpu or self.compute_capabilities is not None: + try: + if self.use_gpu: + # Use current GPU device + config.num_compute_capabilities = 1 + config.set_compute_capability(trt.ComputeCapability.CURRENT, 0) + G_LOGGER.info("Using current GPU device for engine compilation (ComputeCapability.CURRENT)") + elif self.compute_capabilities: + # Set specific compute capabilities + config.num_compute_capabilities = len(self.compute_capabilities) + G_LOGGER.info(f"Setting {len(self.compute_capabilities)} target compute capabilities: {self.compute_capabilities}") + for i, (major, minor) in enumerate(self.compute_capabilities): + cap_name = f"SM{major}{minor}" + compute_cap = getattr(trt.ComputeCapability, cap_name) + config.set_compute_capability(compute_cap, i) + except Exception as e: + G_LOGGER.critical(f"Failed to set compute capabilities: {e}. You are likely not using a TensorRT-RTX build.") + + @util.check_called_by("__call__") + def call_impl(self, builder, network): + """ + Callable implementation that creates and configures the IBuilderConfig with TensorRT-RTX features. + """ + # Enable all common config options + config = super().call_impl(builder, network) + + self._configure_flags(builder, network, config) + + return config diff --git a/tools/Polygraphy/polygraphy/backend/trt/config.py b/tools/Polygraphy/polygraphy/backend/trt/config.py index 8be0079a..b80c56c9 100644 --- a/tools/Polygraphy/polygraphy/backend/trt/config.py +++ b/tools/Polygraphy/polygraphy/backend/trt/config.py @@ -16,38 +16,33 @@ # import contextlib import copy +import re -from polygraphy import mod, util +from polygraphy import config as polygraphy_config, mod, util from polygraphy.backend.base import BaseLoader from polygraphy.backend.trt import util as trt_util from polygraphy.backend.trt.profile import Profile +from polygraphy.backend.trt.util import inherit_and_extend_docstring from polygraphy.mod.trt_importer import lazy_import_trt from polygraphy.logger import G_LOGGER trt = lazy_import_trt() -@mod.export(funcify=True) -class CreateConfig(BaseLoader): +class _CreateConfigCommon(BaseLoader): """ - Functor that creates a TensorRT IBuilderConfig. + Generic TensorRT IBuilderConfig. """ def __init__( self, - tf32=None, - fp16=None, - int8=None, profiles=None, - calibrator=None, precision_constraints=None, load_timing_cache=None, algorithm_selector=None, sparse_weights=None, tactic_sources=None, restricted=None, - use_dla=None, - allow_gpu_fallback=None, profiling_verbosity=None, memory_pool_limits=None, refittable=None, @@ -56,14 +51,12 @@ class CreateConfig(BaseLoader): engine_capability=None, direct_io=None, builder_optimization_level=None, - fp8=None, hardware_compatibility_level=None, max_aux_streams=None, version_compatible=None, exclude_lean_runtime=None, quantization_flags=None, error_on_timing_cache_miss=None, - bf16=None, disable_compilation_cache=None, progress_monitor=None, weight_streaming=None, @@ -71,18 +64,9 @@ class CreateConfig(BaseLoader): tiling_optimization_level=None, ): """ - Creates a TensorRT IBuilderConfig that can be used by EngineFromNetwork. + Creates an IBuilderConfig that can be used by EngineFromNetwork. Args: - tf32 (bool): - Whether to build the engine with TF32 precision enabled. - Defaults to False. - fp16 (bool): - Whether to build the engine with FP16 precision enabled. - Defaults to False. - int8 (bool): - Whether to build the engine with INT8 precision enabled. - Defaults to False. profiles (List[Profile]): A list of optimization profiles to add to the configuration. Only needed for networks with dynamic input shapes. If this is omitted for a network with @@ -90,11 +74,6 @@ class CreateConfig(BaseLoader): replaced with Polygraphy's DEFAULT_SHAPE_VALUE (defined in constants.py). A partially populated profile will be automatically filled using values from ``Profile.fill_defaults()`` See ``Profile`` for details. - calibrator (trt.IInt8Calibrator): - An int8 calibrator. Only required in int8 mode when - the network does not have explicit precision. For networks with - dynamic shapes, the last profile provided (or default profile if - no profiles are provided) is used during calibration. precision_constraints (Optional[str]): If set to "obey", require that layers execute in specified precisions. If set to "prefer", prefer that layers execute in specified precisions but allow TRT to fall back to @@ -124,13 +103,6 @@ class CreateConfig(BaseLoader): Whether to enable safety scope checking in the builder. This will check if the network and builder configuration are compatible with safety scope. Defaults to False. - use_dla (bool): - [EXPERIMENTAL] Whether to enable DLA as the default device type. - Defaults to False. - allow_gpu_fallback (bool): - [EXPERIMENTAL] When DLA is enabled, whether to allow layers to fall back to GPU if they cannot be run on DLA. - Has no effect if DLA is not enabled. - Defaults to False. profiling_verbosity (trt.ProfilingVerbosity): The verbosity of NVTX annotations in the generated engine. Higher verbosity allows you to determine more information about the engine. @@ -161,9 +133,6 @@ class CreateConfig(BaseLoader): to an engine built with a lower optimization level. Refer to the TensorRT API documentation for details. Defaults to TensorRT's default optimization level. - fp8 (bool): - Whether to build the engine with FP8 precision enabled. - Defaults to False. hardware_compatibility_level (trt.HardwareCompatibilityLevel): The hardware compatibility level. This allows engines built on one GPU architecture to work on GPUs of other architectures. @@ -187,9 +156,6 @@ class CreateConfig(BaseLoader): Emit error when a tactic being timed is not present in the timing cache. This flag has an effect only when IBuilderConfig has an associated ITimingCache. Defaults to False. - bf16 (bool): - Whether to build the engine with BF16 precision enabled. - Defaults to False. disable_compilation_cache (bool): Whether to disable caching JIT-compiled code. Defaults to False. @@ -205,13 +171,7 @@ class CreateConfig(BaseLoader): The tiling optimization level. Setting a higher optimization level allows TensorRT to spend more building time for more tiling strategies. Defaults to TensorRT's default tiling optimization level. Refer to the TensorRT API documentation for details. """ - self.tf32 = util.default(tf32, False) - self.fp16 = util.default(fp16, False) - self.bf16 = util.default(bf16, False) - self.int8 = util.default(int8, False) - self.fp8 = util.default(fp8, False) self.profiles = util.default(profiles, [Profile()]) - self.calibrator = calibrator self.precision_constraints = precision_constraints self.restricted = util.default(restricted, False) self.refittable = util.default(refittable, False) @@ -220,8 +180,6 @@ class CreateConfig(BaseLoader): self.algorithm_selector = algorithm_selector self.sparse_weights = util.default(sparse_weights, False) self.tactic_sources = tactic_sources - self.use_dla = util.default(use_dla, False) - self.allow_gpu_fallback = util.default(allow_gpu_fallback, False) self.profiling_verbosity = profiling_verbosity self.memory_pool_limits = memory_pool_limits self.preview_features = preview_features @@ -242,18 +200,6 @@ class CreateConfig(BaseLoader): self.runtime_platform = runtime_platform self.tiling_optimization_level = tiling_optimization_level - if self.calibrator is not None and not self.int8: - G_LOGGER.warning( - "A calibrator was provided to `CreateConfig`, but int8 mode was not enabled. " - "Did you mean to set `int8=True` to enable building with int8 precision?" - ) - - # Print a message to tell users that TF32 can be enabled to improve perf with minor accuracy differences. - if not self.tf32: - G_LOGGER.info( - "TF32 is disabled by default. Turn on TF32 for better performance with minor accuracy differences." - ) - @util.check_called_by("__call__") def call_impl(self, builder, network): """ @@ -273,7 +219,7 @@ class CreateConfig(BaseLoader): try: return func() except AttributeError: - trt_util.fail_unavailable(f"{name} in CreateConfig") + trt_util.fail_unavailable(f"{name} in {self.__class__.__name__}") def try_set_flag(flag_name): return try_run( @@ -330,61 +276,9 @@ class CreateConfig(BaseLoader): if self.direct_io: try_set_flag("DIRECT_IO") - if self.tf32: - try_set_flag("TF32") - else: # TF32 is on by default - with contextlib.suppress(AttributeError): - config.clear_flag(trt.BuilderFlag.TF32) - - if self.fp16: - try_set_flag("FP16") - - if self.bf16: - try_set_flag("BF16") - - if self.fp8: - try_set_flag("FP8") - - if self.int8: - try_set_flag("INT8") - # No Q/DQ layers means that we will need to calibrate. - if not any( - layer.type in [trt.LayerType.QUANTIZE, trt.LayerType.DEQUANTIZE] - for layer in network - ): - if self.calibrator is not None: - config.int8_calibrator = self.calibrator - try: - config.set_calibration_profile( - calib_profile.to_trt(builder, network) - ) - G_LOGGER.info(f"Using calibration profile: {calib_profile}") - except AttributeError: - G_LOGGER.extra_verbose( - "Cannot set calibration profile on TensorRT 7.0 and older." - ) - - trt_util.try_setup_polygraphy_calibrator( - config, - network, - calib_profile=calib_profile.to_trt(builder, network), - ) - else: - G_LOGGER.warning( - "Network does not have explicit precision and no calibrator was provided. Please ensure " - "that tensors in the network have dynamic ranges set, or provide a calibrator in order to use int8 mode." - ) - if self.sparse_weights: try_set_flag("SPARSE_WEIGHTS") - if self.use_dla: - config.default_device_type = trt.DeviceType.DLA - config.DLA_core = 0 - - if self.allow_gpu_fallback: - try_set_flag("GPU_FALLBACK") - if self.profiling_verbosity is not None: def set_profiling_verbosity(): @@ -434,7 +328,7 @@ class CreateConfig(BaseLoader): cache = config.create_timing_cache(b"") except AttributeError: if self.timing_cache_path: - trt_util.fail_unavailable("load_timing_cache in CreateConfig") + trt_util.fail_unavailable(f"load_timing_cache in {self.__class__.__name__}") else: config.set_timing_cache(cache, ignore_mismatch=False) @@ -541,6 +435,180 @@ class CreateConfig(BaseLoader): return config +@mod.export(funcify=True) +class CreateConfig(_CreateConfigCommon): + """ + Functor that creates an IBuilderConfig with TensorRT features. + """ + + @inherit_and_extend_docstring(_CreateConfigCommon.__init__) + def __init__( + self, + tf32=None, + fp16=None, + int8=None, + fp8=None, + bf16=None, + calibrator=None, + use_dla=None, + allow_gpu_fallback=None, + **kwargs + ): + """ + Creates an IBuilderConfig with TensorRT-specific features. + + Args: + tf32 (bool): + Whether to enable TF32 precision. Defaults to False. + fp16 (bool): + Whether to enable FP16 precision. Defaults to False. + int8 (bool): + Whether to enable INT8 precision. Defaults to False. + fp8 (bool): + Whether to enable FP8 precision. Defaults to False. + bf16 (bool): + Whether to enable BF16 precision. Defaults to False. + calibrator (trt.IInt8Calibrator): + An int8 calibrator. Only required in int8 mode when + the network does not have explicit precision. For networks with + dynamic shapes, the last profile provided (or default profile if + no profiles are provided) is used during calibration. + use_dla (bool): + [EXPERIMENTAL] Whether to enable DLA as the default device type. + Defaults to False. + allow_gpu_fallback (bool): + [EXPERIMENTAL] When DLA is enabled, whether to allow layers to fall back to GPU if they cannot be run on DLA. + Has no effect if DLA is not enabled. + Defaults to False. + **kwargs: All other arguments from _CreateConfigCommon. + """ + super().__init__(**kwargs) + self.tf32 = util.default(tf32, False) + self.fp16 = util.default(fp16, False) + self.bf16 = util.default(bf16, False) + self.int8 = util.default(int8, False) + self.fp8 = util.default(fp8, False) + self.calibrator = calibrator + self.use_dla = util.default(use_dla, False) + self.allow_gpu_fallback = util.default(allow_gpu_fallback, False) + + if self.calibrator is not None and not self.int8: + G_LOGGER.warning( + "A calibrator was provided to `CreateConfig`, but int8 mode was not enabled. " + "Did you mean to set `int8=True` to enable building with int8 precision?" + ) + + # Print a message to tell users that TF32 can be enabled to improve perf with minor accuracy differences. + if not self.tf32: + G_LOGGER.info( + "TF32 is disabled by default. Turn on TF32 for better performance with minor accuracy differences." + ) + + self._validator() + + def _validator(self): + """ + Validates initialization parameters for TensorRT-specific features. + """ + # Validate that TensorRT-RTX specific flags are not used in regular TensorRT mode + if polygraphy_config.USE_TENSORRT_RTX: + if self.fp16 or self.int8 or self.bf16 or self.fp8: + G_LOGGER.critical("Precision flags (fp16, int8, bf16, fp8) are not supported with USE_TENSORRT_RTX=1.") + if self.use_dla: + G_LOGGER.critical("DLA is not supported with USE_TENSORRT_RTX=1.") + if self.calibrator is not None: + G_LOGGER.critical("Custom calibrator is not supported with USE_TENSORRT_RTX=1.") + + def _configure_flags(self, builder, network, config): + """ + Validates and configures TensorRT-specific features. + + Args: + builder (trt.Builder): The TensorRT builder + network (trt.INetworkDefinition): The TensorRT network + config (trt.IBuilderConfig): The TensorRT builder config to modify + """ + def try_run(func, name): + try: + return func() + except AttributeError: + trt_util.fail_unavailable(f"{name} in CreateConfig") + + def try_set_flag(flag_name): + return try_run( + lambda: config.set_flag(getattr(trt.BuilderFlag, flag_name)), + flag_name.lower(), + ) + + # Add precision-related logic + if self.tf32: + try_set_flag("TF32") + else: # TF32 is on by default + with contextlib.suppress(AttributeError): + config.clear_flag(trt.BuilderFlag.TF32) + + if self.fp16: + try_set_flag("FP16") + + if self.bf16: + try_set_flag("BF16") + + if self.fp8: + try_set_flag("FP8") + + if self.int8: + try_set_flag("INT8") + + if self.int8: + # No Q/DQ layers means that we will need to calibrate. + if not any( + layer.type in [trt.LayerType.QUANTIZE, trt.LayerType.DEQUANTIZE] + for layer in network + ): + if self.calibrator is not None: + config.int8_calibrator = self.calibrator + try: + profiles = copy.deepcopy(self.profiles) + calib_profile = profiles[-1].fill_defaults(network) + config.set_calibration_profile( + calib_profile.to_trt(builder, network) + ) + G_LOGGER.info(f"Using calibration profile: {calib_profile}") + except AttributeError: + G_LOGGER.extra_verbose( + "Cannot set calibration profile on TensorRT 7.0 and older." + ) + + trt_util.try_setup_polygraphy_calibrator( + config, + network, + calib_profile=calib_profile.to_trt(builder, network), + ) + else: + G_LOGGER.warning( + "Network does not have explicit precision and no calibrator was provided. Please ensure " + "that tensors in the network have dynamic ranges set, or provide a calibrator in order to use int8 mode." + ) + + if self.use_dla: + config.default_device_type = trt.DeviceType.DLA + config.DLA_core = 0 + + if self.allow_gpu_fallback: + try_set_flag("GPU_FALLBACK") + + @util.check_called_by("__call__") + def call_impl(self, builder, network): + """ + Callable implementation that creates and configures the IBuilderConfig with TensorRT features. + """ + config = super().call_impl(builder, network) + + self._configure_flags(builder, network, config) + + return config + + @mod.export(funcify=True) class PostprocessConfig(BaseLoader): """ diff --git a/tools/Polygraphy/polygraphy/backend/trt/file_reader.py b/tools/Polygraphy/polygraphy/backend/trt/file_reader.py index 10bedf6b..857a5cf5 100644 --- a/tools/Polygraphy/polygraphy/backend/trt/file_reader.py +++ b/tools/Polygraphy/polygraphy/backend/trt/file_reader.py @@ -61,6 +61,19 @@ def FileReader( def read(self, size: int) -> bytes: return self.file.read(size) + def seek(self, offset: int, whence: int = 0) -> int: + """ + Seek to a position in the stream. Required for IStreamReaderV2. + + Args: + offset: The offset to seek to + whence: How to interpret the offset (0=absolute, 1=relative to current, 2=relative to end) + + Returns: + The new absolute position + """ + return self.file.seek(offset, whence) + def free(self): if self.file: self.file.close() diff --git a/tools/Polygraphy/polygraphy/backend/trt/loader.py b/tools/Polygraphy/polygraphy/backend/trt/loader.py index 79bff841..215d6703 100644 --- a/tools/Polygraphy/polygraphy/backend/trt/loader.py +++ b/tools/Polygraphy/polygraphy/backend/trt/loader.py @@ -84,7 +84,7 @@ class CreateNetwork(BaseLoader): Functor that creates an empty TensorRT network. """ - def __init__(self, explicit_batch=None, strongly_typed=None): + def __init__(self, explicit_batch=None, strongly_typed=None, mark_unfused_tensors_as_debug_tensors=None): """ Creates an empty TensorRT network. @@ -95,12 +95,16 @@ class CreateNetwork(BaseLoader): strongly_typed (bool): Whether to mark the network as being strongly typed. Defaults to False. + mark_unfused_tensors_as_debug_tensors (bool): + Whether to mark unfused tensors as debug tensors. + Defaults to False. """ self.explicit_batch = util.default( explicit_batch, True if mod.version(trt.__version__) < mod.version("10.0") else None, ) self.strongly_typed = util.default(strongly_typed, False) + self.mark_unfused_tensors_as_debug_tensors = util.default(mark_unfused_tensors_as_debug_tensors, False) @util.check_called_by("__call__") def call_impl(self): @@ -130,11 +134,15 @@ class CreateNetwork(BaseLoader): network = builder.create_network(flags=network_flags) if network is None: G_LOGGER.critical("Invalid network. See logging output above for details.") + + if self.mark_unfused_tensors_as_debug_tensors: + network.mark_unfused_tensors_as_debug_tensors() + return builder, network class BaseNetworkFromOnnx(BaseLoader): - def __init__(self, flags=None, plugin_instancenorm=None, strongly_typed=None): + def __init__(self, flags=None, plugin_instancenorm=None, strongly_typed=None, mark_unfused_tensors_as_debug_tensors=None): """ Args: flags (List[trt.OnnxParserFlag]): @@ -152,10 +160,11 @@ class BaseNetworkFromOnnx(BaseLoader): self.flags = flags self.plugin_instancenorm = util.default(plugin_instancenorm, False) self.strongly_typed = util.default(strongly_typed, False) + self.mark_unfused_tensors_as_debug_tensors = util.default(mark_unfused_tensors_as_debug_tensors, False) @util.check_called_by("__call__") def call_impl(self): - builder, network = create_network(strongly_typed=self.strongly_typed) + builder, network = create_network(strongly_typed=self.strongly_typed, mark_unfused_tensors_as_debug_tensors=self.mark_unfused_tensors_as_debug_tensors) # Initialize plugin library for the parser. trt.init_libnvinfer_plugins(trt_util.get_trt_logger(), "") parser = trt.OnnxParser(network, trt_util.get_trt_logger()) @@ -178,7 +187,7 @@ class NetworkFromOnnxBytes(BaseNetworkFromOnnx): """ def __init__( - self, model_bytes, flags=None, plugin_instancenorm=None, strongly_typed=None + self, model_bytes, flags=None, plugin_instancenorm=None, strongly_typed=None, mark_unfused_tensors_as_debug_tensors=None ): """ Parses an ONNX model. @@ -203,6 +212,7 @@ class NetworkFromOnnxBytes(BaseNetworkFromOnnx): flags=flags, plugin_instancenorm=plugin_instancenorm, strongly_typed=strongly_typed, + mark_unfused_tensors_as_debug_tensors=mark_unfused_tensors_as_debug_tensors ) self._model_bytes = model_bytes @@ -227,7 +237,7 @@ class NetworkFromOnnxPath(BaseNetworkFromOnnx): This loader supports models with weights stored in an external location. """ - def __init__(self, path, flags=None, plugin_instancenorm=None, strongly_typed=None): + def __init__(self, path, flags=None, plugin_instancenorm=None, strongly_typed=None, mark_unfused_tensors_as_debug_tensors=None): """ Parses an ONNX model from a file. @@ -250,6 +260,7 @@ class NetworkFromOnnxPath(BaseNetworkFromOnnx): flags=flags, plugin_instancenorm=plugin_instancenorm, strongly_typed=strongly_typed, + mark_unfused_tensors_as_debug_tensors=mark_unfused_tensors_as_debug_tensors ) self.path = path @@ -753,8 +764,14 @@ class EngineFromPath(BaseLoader): except AttributeError: pass - file_reader = FileReader(path) - engine = runtime.deserialize_cuda_engine(file_reader) + if config.USE_TENSORRT_RTX: + # Read the entire file into memory for buffer-based deserialization + with open(path, 'rb') as f: + buffer_data = f.read() + engine = runtime.deserialize_cuda_engine(buffer_data) + else: + file_reader = FileReader(path) + engine = runtime.deserialize_cuda_engine(file_reader) if not engine: G_LOGGER.critical("Could not deserialize engine. See log for details.") return engine @@ -1003,3 +1020,4 @@ class MarkDebug(PostprocessNetwork): """ func = lambda network: MarkDebug._apply(network, mark_debug) super().__init__(network, func, "MarkDebug") + diff --git a/tools/Polygraphy/polygraphy/backend/trt/runner.py b/tools/Polygraphy/polygraphy/backend/trt/runner.py index 8b660e0c..54b12e0d 100644 --- a/tools/Polygraphy/polygraphy/backend/trt/runner.py +++ b/tools/Polygraphy/polygraphy/backend/trt/runner.py @@ -16,6 +16,7 @@ # import math import time +import ctypes from collections import OrderedDict from polygraphy import config, cuda, mod, util @@ -38,19 +39,26 @@ def _make_debug_listener(): self.debug_tensor_outputs = {} def process_debug_tensor(self, addr, location, type, shape, name, stream): + if type in [util.try_getattr(trt, "fp8"), util.try_getattr(trt, "int4"), util.try_getattr(trt, "fp4"), util.try_getattr(trt, "bfloat16")]: + G_LOGGER.warning(f"Not supported datatype for debug tensor in polygraphy: {type}") + return + cuda.wrapper().stream_synchronize(stream) datatype = DataType.from_dtype(type) size = util.volume(shape) buffer = np.zeros(shape, dtype=DataType.to_dtype(datatype, "numpy")) buffer = util.array.resize_or_reallocate(buffer, size) - cuda.wrapper().memcpy( - dst=util.array.data_ptr(buffer), - src=addr, - nbytes=size * datatype.itemsize, - kind=cuda.MemcpyKind.DeviceToHost, - stream_ptr=stream, - ) - cuda.wrapper().stream_synchronize(stream) + if location == trt.TensorLocation.HOST: + ctypes.memmove(util.array.data_ptr(buffer), addr, size * datatype.itemsize) + else: + cuda.wrapper().memcpy( + dst=util.array.data_ptr(buffer), + src=addr, + nbytes=size * datatype.itemsize, + kind=cuda.MemcpyKind.DeviceToHost, + stream_ptr=stream, + ) + cuda.wrapper().stream_synchronize(stream) self.debug_tensor_outputs[name] = util.array.resize_or_reallocate(buffer, shape) return DebugTensorWriter() diff --git a/tools/Polygraphy/polygraphy/backend/trt/util.py b/tools/Polygraphy/polygraphy/backend/trt/util.py index 21c4d6e0..ed1af7b1 100644 --- a/tools/Polygraphy/polygraphy/backend/trt/util.py +++ b/tools/Polygraphy/polygraphy/backend/trt/util.py @@ -17,6 +17,7 @@ import contextlib import json import os +import re import signal from polygraphy import config, mod, util, cuda @@ -25,6 +26,8 @@ from polygraphy.common import TensorMetadata from polygraphy.datatype import DataType from polygraphy.exception import PolygraphyException from polygraphy.logger import G_LOGGER, LogMode +from polygraphy.json import load_json +from polygraphy.comparator import RunResults trt = lazy_import_trt() np = mod.lazy_import("numpy") @@ -159,6 +162,8 @@ def get_layer_class_mapping(): try_add("UNSQUEEZE", "IUnsqueezeLayer") try_add("CUMULATIVE", "ICumulativeLayer") try_add("DYNAMIC_QUANTIZE", "IDynamicQuantizeLayer") + try_add("ATTENTION_INPUT", "IAttentionInputLayer") + try_add("ATTENTION_OUTPUT", "IAttentionOutputLayer") return layer_class_mapping @@ -672,7 +677,50 @@ def get_metadata_from_engine(engine, context, mode): return meta -def str_from_engine(engine, context, show_layers=None, show_attrs=None): +class TensorInfo: + def __init__(self, json_path: str = None): + self.tensors = {} + if json_path: + self.load_json(json_path) + + def load_json(self, json_path: str) -> None: + data = load_json(json_path) + if isinstance(data, RunResults): + # Handle RunResults format + for runner_name, iterations in data.items(): + if not iterations: + G_LOGGER.warning(f"No iterations found for runner: {runner_name}") + continue + + if len(iterations) > 1: + G_LOGGER.warning( + f"Found {len(iterations)} iterations in tensor info file, only using the first one" + ) + + iter_data = iterations[0] + for name, tensor in iter_data.items(): + if not isinstance(tensor, np.ndarray): + tensor = np.array(tensor) + + self.tensors[name] = { + "min": float(np.min(tensor)), + "max": float(np.max(tensor)), + "avg": float(np.mean(tensor)), + } + break # Only use first runner + else: + G_LOGGER.warning(f"Unsupported tensor info format: {json_path}") + + def get_tensor_statistics(self, tensor_name: str) -> str: + tensor = self.tensors.get(tensor_name) + if not tensor: + return "" + return f", min={tensor['min']:.2f}, max={tensor['max']:.2f}, avg={tensor['avg']:.2f}" + + +def str_from_engine( + engine, context, show_layers=None, show_attrs=None, combine_tensor_info=None +): show_layers = util.default(show_layers, False) show_attrs = util.default(show_attrs, False) @@ -759,6 +807,7 @@ def str_from_engine(engine, context, show_layers=None, show_attrs=None): if num_profiles_to_print > 1: indent_level = 1 engine_str += f"- Profile: {profile_idx}\n" + tensor_info = TensorInfo(combine_tensor_info) offset = profile_idx * layers_per_profile for index in range(layers_per_profile): @@ -807,6 +856,9 @@ def str_from_engine(engine, context, show_layers=None, show_attrs=None): return meta for elem in info: names.append(elem["Name"]) + tensor_statistics = tensor_info.get_tensor_statistics( + elem["Name"] + ) meta.add( name=elem["Name"], dtype=dtype_from_fmt_dtype(elem["Format/Datatype"]), @@ -814,8 +866,9 @@ def str_from_engine(engine, context, show_layers=None, show_attrs=None): docstring=( f"Format: {elem['Format/Datatype']}" if "N/A" not in elem["Format/Datatype"] - else None - ), + else "" + ) + + tensor_statistics, ) return names, meta @@ -901,3 +954,59 @@ def _get_array_on_gpu(arr, name, device_buffers, stream=None): device_buffers[name].resize(shape) device_buffers[name].copy_from(util.array.view(arr, DataType.UINT8, shape), stream) return device_buffers[name].ptr + + +def inherit_and_extend_docstring(parent_method): + """ + Decorator to inherit and extend docstrings from parent class methods. + + Combines the parent method's description and Args with the child method's + description and Args, preserving proper formatting for Sphinx documentation. + + Args: + parent_method: The parent method to inherit docstring from + + Returns: + Decorator function that combines parent and child docstrings + """ + + def decorator(child_method): + parent_doc = parent_method.__doc__ or "" + child_doc = child_method.__doc__ or "" + + if not parent_doc: + return child_method + if not child_doc: + child_method.__doc__ = parent_doc + return child_method + + def extract_description_and_args(docstring): + """Extract description and Args section from a docstring.""" + desc = re.split(r"\n\s*Args:", docstring, 1)[0].strip() + args_match = re.search( + r"\n\s*Args:\s*\n(.*?)(?=\n\s*[A-Z][a-z]*:|\Z)", docstring, re.DOTALL + ) + args = args_match.group(1).rstrip() if args_match else "" + return desc, args + + # Extract components from both docstrings + parent_desc, parent_args = extract_description_and_args(parent_doc) + child_desc, child_args = extract_description_and_args(child_doc) + + # Combine descriptions + combined_desc = f"{parent_desc}\n\n{child_desc}" if child_desc else parent_desc + + # Combine Args sections + args_parts = [ + args for args in [parent_args, child_args] if args + ] # Filter for non-empty argument strings + combined_doc = ( + f"{combined_desc}\n\nArgs:\n" + "\n".join(args_parts) + if args_parts + else combined_desc + ) + + child_method.__doc__ = combined_doc + return child_method + + return decorator diff --git a/tools/Polygraphy/polygraphy/common/struct.py b/tools/Polygraphy/polygraphy/common/struct.py index cd60d43c..666d631b 100644 --- a/tools/Polygraphy/polygraphy/common/struct.py +++ b/tools/Polygraphy/polygraphy/common/struct.py @@ -62,7 +62,7 @@ class MetadataTuple: meta_items.append(f"dtype={self.dtype}") if self.shape is not None: meta_items.append(f"shape={tuple(self.shape)}") - if self.docstring is not None: + if self.docstring: meta_items.append(self.docstring) if meta_items: ret += "[" + ", ".join(meta_items) + "]" diff --git a/tools/Polygraphy/polygraphy/comparator/util.py b/tools/Polygraphy/polygraphy/comparator/util.py index ae24f4cd..7da7a0b4 100644 --- a/tools/Polygraphy/polygraphy/comparator/util.py +++ b/tools/Polygraphy/polygraphy/comparator/util.py @@ -289,7 +289,9 @@ def build_heatmaps( fig.colorbar(images[0], ax=axs, shrink=0.7) if save_dir is not None: - path = os.path.join(save_dir, f"{fig_title}.svg") + path = os.path.join( + save_dir, f"{util.sanitize_filename(fig_title)}.svg" + ) util.makedirs(path) G_LOGGER.info(f"Saving '{prefix}' heatmap to: '{path}'") fig.savefig(path) @@ -334,7 +336,7 @@ def scatter_plot_error_magnitude( ) with G_LOGGER.indent(): title = f"Error metrics between output0 and output1\noutput0: {runner0_name:35} | {out0_name}\noutput1: {runner1_name:35} | {out1_name}" - fname = f"error_metrics_{out0_name}.png" + fname = util.sanitize_filename(f"error_metrics_{out0_name}.png") TICK_FONT_SIZE = 6 TITLE_FONT_SIZE = 7 NUM_X_TICKS = 20 @@ -358,7 +360,9 @@ def scatter_plot_error_magnitude( ax.set_yscale("log") xticks = ax.get_xticks() - yrange = np.log10(np.array([min_diff, max_diff])) + # Add a very small epsilon to prevent division by 0: + eps = 1e-15 + yrange = np.log10(np.array([min_diff + eps, max_diff + eps])) yrange[0] = math.floor(yrange[0]) yrange[1] = math.ceil(yrange[1]) diff --git a/tools/Polygraphy/polygraphy/tools/args/backend/trt/config.py b/tools/Polygraphy/polygraphy/tools/args/backend/trt/config.py index 2b6001cd..04000491 100644 --- a/tools/Polygraphy/polygraphy/tools/args/backend/trt/config.py +++ b/tools/Polygraphy/polygraphy/tools/args/backend/trt/config.py @@ -19,6 +19,7 @@ import os from polygraphy import constants, mod, util from polygraphy.common import TensorMetadata +from polygraphy import config as polygraphy_config from polygraphy.logger import G_LOGGER, LogMode from polygraphy.mod.trt_importer import tensorrt_module_and_version_string from polygraphy.tools.args import util as args_util @@ -116,6 +117,7 @@ class TrtConfigArgs(BaseArgs): allow_custom_input_shapes: bool = None, allow_engine_capability: bool = None, allow_tensor_formats: bool = None, + allow_compute_capabilities: bool = None, ): """ Args: @@ -134,6 +136,9 @@ class TrtConfigArgs(BaseArgs): allow_tensor_formats (bool): Whether to allow tensor formats and related options to be set. Defaults to False. + allow_compute_capabilities (bool): + Whether to allow compute capabilities options to be set. + Defaults to False. """ super().__init__() self._precision_constraints_default = util.default( @@ -145,6 +150,7 @@ class TrtConfigArgs(BaseArgs): self._allow_custom_input_shapes = util.default(allow_custom_input_shapes, True) self._allow_engine_capability = util.default(allow_engine_capability, False) self._allow_tensor_formats = util.default(allow_tensor_formats, False) + self._allow_compute_capabilities = util.default(allow_compute_capabilities, False) def add_parser_args_impl(self): self.group.add_argument( @@ -442,7 +448,7 @@ class TrtConfigArgs(BaseArgs): help="The verbosity of NVTX annotations in the generated engine." "Values come from the names of values in the `trt.ProfilingVerbosity` enum and are case-insensitive. " "For example, `--profiling-verbosity detailed`. " - "Defaults to 'verbose'.", + "Defaults to 'detailed'.", default=None, ) @@ -486,6 +492,25 @@ class TrtConfigArgs(BaseArgs): default=None, ) + if polygraphy_config.USE_TENSORRT_RTX and self._allow_compute_capabilities: + compute_capabilities_group = self.group.add_mutually_exclusive_group() + + compute_capabilities_group.add_argument( + "--use-gpu", + help="Use the current GPU device as target for engine compilation. " + "Equivalent to setting ComputeCapability.CURRENT.", + action="store_true", + default=None, + ) + + compute_capabilities_group.add_argument( + "--compute-capabilities", + help="Specify target compute capabilities for engine compilation. " + "Values should be major.minor versions (e.g., '7.5 8.0 8.6').", + nargs="+", + default=None, + ) + def parse_impl(self, args): """ Parses command-line arguments and populates the following attributes: @@ -532,6 +557,8 @@ class TrtConfigArgs(BaseArgs): weight_streaming (bool): Whether to enable weight streaming for the TensorRT Engine. runtime_platform (str): A string representing the target runtime platform enum value. tiling_optimization_level (str): The tiling optimization level. + use_gpu (bool): Whether to use the current GPU device as target for engine compilation. + compute_capabilities (List[Tuple[int, int]]): List of (major, minor) compute capability tuples to target for engine compilation. """ trt_min_shapes = args_util.get(args, "trt_min_shapes", default=[]) @@ -709,6 +736,26 @@ class TrtConfigArgs(BaseArgs): "TilingOptimizationLevel", tiling_optimization_level ) + # Parse compute capabilities arguments if enabled and TensorRT-RTX is available + self.use_gpu = False + self.compute_capabilities = None + + if self._allow_compute_capabilities and polygraphy_config.USE_TENSORRT_RTX: + self.use_gpu = args_util.get(args, "use_gpu", default=False) + compute_capabilities_list = args_util.get(args, "compute_capabilities") + + if compute_capabilities_list: + # Parse compute capabilities from list of strings + try: + capabilities = [] + for cap_str in compute_capabilities_list: + major, minor = map(int, cap_str.split('.')) + capabilities.append((major, minor)) + self.compute_capabilities = capabilities + except ValueError: + G_LOGGER.critical(f"Invalid compute capabilities format: {compute_capabilities_list}. " + "Expected format: space-separated 'major.minor' versions (e.g., '7.5 8.0').") + def add_to_script_impl(self, script): profiles = [] for profile_dict in self.profile_dicts: @@ -798,6 +845,8 @@ class TrtConfigArgs(BaseArgs): self.runtime_platform, self.quantization_flags, self.tiling_optimization_level, + self.use_gpu, + self.compute_capabilities, ] ): script.add_import(imports=tensorrt_module_and_version_string(), imp_as="trt") @@ -812,23 +861,37 @@ class TrtConfigArgs(BaseArgs): name=self.trt_config_func_name, ) else: + # Use CreateConfigRTX if TensorRT-RTX is enabled, otherwise use CreateConfig + if polygraphy_config.USE_TENSORRT_RTX: + config_class = "CreateConfigRTX" + config_alias = "CreateTrtConfigRTX" + extra_args = { + "use_gpu": self.use_gpu, + "compute_capabilities": self.compute_capabilities, + } + else: + config_class = "CreateConfig" + config_alias = "CreateTrtConfig" + extra_args = { + "tf32": self.tf32, + "fp16": self.fp16, + "bf16": self.bf16, + "int8": self.int8, + "fp8": self.fp8, + "calibrator": calibrator, + "use_dla": self.use_dla, + "allow_gpu_fallback": self.allow_gpu_fallback, + } + config_loader_str = make_invocable_if_nondefault( - "CreateTrtConfig", - tf32=self.tf32, - fp16=self.fp16, - bf16=self.bf16, - int8=self.int8, - fp8=self.fp8, + config_alias, precision_constraints=self.precision_constraints, restricted=self.restricted, profiles=profile_name, - calibrator=calibrator, load_timing_cache=self.load_timing_cache, algorithm_selector=algo_selector, sparse_weights=self.sparse_weights, tactic_sources=self.tactic_sources, - use_dla=self.use_dla, - allow_gpu_fallback=self.allow_gpu_fallback, memory_pool_limits=self.memory_pool_limits, refittable=self.refittable, strip_plan=self.strip_plan, @@ -847,13 +910,25 @@ class TrtConfigArgs(BaseArgs): weight_streaming=self.weight_streaming, runtime_platform=self.runtime_platform, tiling_optimization_level=self.tiling_optimization_level, + **extra_args ) + + if config_loader_str is None and polygraphy_config.USE_TENSORRT_RTX: + config_loader_str = make_invocable(config_alias) + if config_loader_str is not None: - script.add_import( - imports="CreateConfig", - frm="polygraphy.backend.trt", - imp_as="CreateTrtConfig", - ) + if polygraphy_config.USE_TENSORRT_RTX: + script.add_import( + imports=config_class, + frm="polygraphy.backend.tensorrt_rtx", + imp_as=config_alias, + ) + else: + script.add_import( + imports=config_class, + frm="polygraphy.backend.trt", + imp_as=config_alias, + ) if config_loader_str is not None: config_loader_name = script.add_loader( @@ -908,7 +983,13 @@ class TrtConfigArgs(BaseArgs): Returns: trt.IBuilderConfig: The TensorRT builder configuration. """ - from polygraphy.backend.trt import CreateConfig + # Use CreateConfigRTX if TensorRT-RTX is enabled, otherwise use CreateConfig + if polygraphy_config.USE_TENSORRT_RTX: + from polygraphy.backend.tensorrt_rtx import CreateConfigRTX + default_loader = CreateConfigRTX() + else: + from polygraphy.backend.trt import CreateConfig + default_loader = CreateConfig() - loader = util.default(args_util.run_script(self.add_to_script), CreateConfig()) + loader = util.default(args_util.run_script(self.add_to_script), default_loader) return loader(builder, network) diff --git a/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py b/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py index 5cabe180..e78b8fa3 100644 --- a/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py +++ b/tools/Polygraphy/polygraphy/tools/args/backend/trt/loader.py @@ -260,6 +260,13 @@ class TrtLoadNetworkArgs(BaseArgs): default=None, ) + self.group.add_argument( + "--mark-unfused-tensors-as-debug-tensors", + help="Mark unfused tensors as debug tensors.", + action="store_true", + default=None, + ) + def parse_impl(self, args): """ Parses command-line arguments and populates the following attributes: @@ -275,6 +282,7 @@ class TrtLoadNetworkArgs(BaseArgs): A list of tuples specifying a path to a network postprocessing script and the name of the postprocessing function. strongly_typed (bool): Whether to mark the network as being strongly typed. mark_debug (List[str]): Names of tensors which should be marked as debug tensors. + mark_unfused_tensors_as_debug_tensors (bool): Whether to mark unfused tensors as debug tensors. """ self.outputs = args_util.get_outputs(args, "trt_outputs") @@ -335,6 +343,9 @@ class TrtLoadNetworkArgs(BaseArgs): self.strongly_typed = args_util.get(args, "strongly_typed") self.mark_debug = args_util.get(args, "mark_debug") + self.mark_unfused_tensors_as_debug_tensors = args_util.get( + args, "mark_unfused_tensors_as_debug_tensors" + ) def add_to_script_impl(self, script): network_func_name = self.arg_groups[ModelArgs].extra_model_info @@ -394,6 +405,7 @@ class TrtLoadNetworkArgs(BaseArgs): flags=parser_flags, plugin_instancenorm=plugin_instancenorm, strongly_typed=self.strongly_typed, + mark_unfused_tensors_as_debug_tensors=self.mark_unfused_tensors_as_debug_tensors, ) loader_name = script.add_loader(loader_str, "parse_network_from_onnx") else: @@ -408,6 +420,7 @@ class TrtLoadNetworkArgs(BaseArgs): flags=parser_flags, plugin_instancenorm=plugin_instancenorm, strongly_typed=self.strongly_typed, + mark_unfused_tensors_as_debug_tensors=self.mark_unfused_tensors_as_debug_tensors, ) loader_name = script.add_loader(loader_str, "parse_network_from_onnx") else: diff --git a/tools/Polygraphy/polygraphy/tools/convert/convert.py b/tools/Polygraphy/polygraphy/tools/convert/convert.py index 219fcc8a..50ba59bb 100644 --- a/tools/Polygraphy/polygraphy/tools/convert/convert.py +++ b/tools/Polygraphy/polygraphy/tools/convert/convert.py @@ -56,7 +56,7 @@ class Convert(Tool): OnnxLoadArgs(allow_from_tf=True), OnnxSaveArgs(output_opt=False), DataLoaderArgs(), # For int8 calibration - TrtConfigArgs(allow_engine_capability=True, allow_tensor_formats=True), + TrtConfigArgs(allow_engine_capability=True, allow_tensor_formats=True, allow_compute_capabilities=True), TrtLoadPluginsArgs(), TrtLoadNetworkArgs(allow_tensor_formats=True), TrtLoadEngineBytesArgs(), diff --git a/tools/Polygraphy/polygraphy/tools/inspect/subtool/model.py b/tools/Polygraphy/polygraphy/tools/inspect/subtool/model.py index 9fe01241..e7d98367 100644 --- a/tools/Polygraphy/polygraphy/tools/inspect/subtool/model.py +++ b/tools/Polygraphy/polygraphy/tools/inspect/subtool/model.py @@ -91,6 +91,18 @@ class Model(Tool): dest="show_unbounded_dds", ) + parser.add_argument( + "--combine-tensor-info", + help=""" + Set the path to the tensor JSON file to combine information from the file into layers' input and output information. + This is only supported when --model-type is "engine" and --show includes "layers". + To get the tensor JSON file, use '--mark-unfused-tensors-as-debug-tensors' and '--save-outputs' when running model with TensorRT. + """, + type=str, + default=None, + dest="combine_tensor_info", + ) + def run_impl(self, args): def show(aspect): return aspect in args.show @@ -104,6 +116,7 @@ class Model(Tool): context, show_layers=show("layers"), show_attrs=show("attrs"), + combine_tensor_info=args.combine_tensor_info ) G_LOGGER.info(f"==== TensorRT Engine ====\n{engine_str}") else: diff --git a/tools/Polygraphy/polygraphy/tools/multi_device/README.md b/tools/Polygraphy/polygraphy/tools/multi_device/README.md new file mode 100644 index 00000000..81860170 --- /dev/null +++ b/tools/Polygraphy/polygraphy/tools/multi_device/README.md @@ -0,0 +1,87 @@ +# Multi-Device + +## Table of Contents + +- [Introduction](#introduction) +- [Subtools](#subtools) +- [Usage](#usage) +- [Examples](#examples) + + +## Subtools + +### Shard + +The `shard` tool in Polygraphy is used to convert a single-device (SD) model to a multi-device (MD) model using a sharding hints file. + +#### Sharding Hints File Format + +The hints file is a JSON file that describes how to shard the model. Example: + +```json +{ + "parallelism": "CP", + "group_size": 4, + "root": 0, + "groups": [], + "attention_layers": [ + { + "q": "q", + "gather_kv": true, + "gather_q": false, + "polygraphy_class": "AttentionLayerHint" + } + ], + "inputs": [ + { + "name": "input", + "seq_len_idx": 0, + "rank": 3, + "polygraphy_class": "ShardTensor" + } + ], + "outputs": [ + { + "name": "output", + "seq_len_idx": 0, + "rank": 3, + "polygraphy_class": "ShardTensor" + } + ], + "k_seq_len_idx": 0, + "v_seq_len_idx": 0, + "kv_rank": null, + "reduce_scatter_reduce_op": "max", + "polygraphy_class": "ShardHints" +} +``` + +- `parallelism`: Type of parallelism (e.g., CP/DP/PP). Currently, only CP is supported +- `group_size`: Number of GPUs model will be run on. 0 indiciates all available GPUs will run. +- `root`: Root rank for collectives +- `groups`: Indices of NCCL groups in which collective operations will run. A value of [] indicates collective operations will run on all ranks with no grouping. +- `attention_layers`: List of attention layer configs: + - `q`: Name of the Q tensor + - `gather_kv`: Whether to all-gather K/V + - `gather_q`: Whether to all-gather Q +- `inputs`: List of tensors that should be reduce-scattered. + - `name`: Name of tensor. + - `seq_len_idx`: Index of dimension that represents sequence length for this input tensor. A non-zero index will cause transpose tensors to be inserted before and after the DistCollective node to transpose the tensor to have sequence length be the first dimension. + - `rank`: Rank of this tensor. Used as a fallback if no dimension can be obtained from the model and `seq_len_idx` != 0. +- `outputs`: List of tensors that should be all-gathered. + - `name`: Name of tensor. + - `seq_len_idx`: Index of dimension that represents sequence length for this output tensor. A non-zero index will cause transpose tensors to be inserted before and after the DistCollective node to transpose the tensor to have sequence length be the first dimension. + - `rank`: Rank of this tensor. Used as a fallback if no dimension can be obtained from the model and `seq_len_idx` != 0. +- `k_seq_len_idx`: Index of dimension that represents sequence length for all K tensor(s). A non-zero index will cause transpose tensors to be inserted before and after the DistCollective node to transpose the tensor to have sequence length be the first dimension. +- `v_seq_len_idx`: Index of dimension that represents sequence length for all V tensor(s). A non-zero index will cause transpose tensors to be inserted before and after the DistCollective node to transpose the tensor to have sequence length be the first dimension. +- `kv_rank`: Rank all K and V tensor(s). Used as a fallback if no dimension can be obtained from the model and `k_seq_len_idx` or `v_seq_len_idx` != 0. +- `reduce_scatter_reduce_op`: Reduction operator to be used on reduce-scatter nodes. + +## Usage + +See `polygraphy multi-device -h` for usage information. + + +## Examples + +For examples, see [this directory](../../../examples/cli/multi_device/) diff --git a/tools/Polygraphy/polygraphy/tools/multi_device/__init__.py b/tools/Polygraphy/polygraphy/tools/multi_device/__init__.py new file mode 100644 index 00000000..c20b0313 --- /dev/null +++ b/tools/Polygraphy/polygraphy/tools/multi_device/__init__.py @@ -0,0 +1 @@ +from polygraphy.tools.multi_device.multi_device import MultiDevice, ShardHints, ShardTensor, AttentionLayerHint, get_attention_pattern diff --git a/tools/Polygraphy/polygraphy/tools/multi_device/multi_device.py b/tools/Polygraphy/polygraphy/tools/multi_device/multi_device.py new file mode 100644 index 00000000..2320cca7 --- /dev/null +++ b/tools/Polygraphy/polygraphy/tools/multi_device/multi_device.py @@ -0,0 +1,160 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from polygraphy import mod +from polygraphy.json import Decoder, Encoder, add_json_methods +from polygraphy.tools.base import Tool + +gs = mod.lazy_import("onnx_graphsurgeon") + +@add_json_methods("shard tensor") +class ShardTensor: + def __init__(self, name, seq_len_idx, rank = None): + self.name = name + self.seq_len_idx = seq_len_idx + self.rank = rank + +@Decoder.register(ShardTensor) +def decode(dct): + return ShardTensor( + name = dct["name"], + seq_len_idx=dct["seq_len_idx"], + rank=dct.get("rank") + ) + +@Encoder.register(ShardTensor) +def encode(shard_tensor): + return { + "name" : shard_tensor.name, + "seq_len_idx" : shard_tensor.seq_len_idx, + "rank" : shard_tensor.rank + } + +@add_json_methods("attention layer hint") +class AttentionLayerHint: + def __init__(self, q, gather_kv, gather_q): + self.q = q + self.gather_kv = gather_kv + self.gather_q = gather_q + +@Decoder.register(AttentionLayerHint) +def decode(dct): + return AttentionLayerHint( + q=dct["q"], + gather_kv=dct["gather_kv"], + gather_q=dct["gather_q"], + ) + +@Encoder.register(AttentionLayerHint) +def encode(attention_layer_hint): + return { + "q" : attention_layer_hint.q, + "gather_kv" : attention_layer_hint.gather_kv, + "gather_q" : attention_layer_hint.gather_q, + } + +@add_json_methods("shard hints") +class ShardHints: + def __init__(self, parallelism, group_size, root, groups, attention_layers, inputs, outputs, k_seq_len_idx, v_seq_len_idx, kv_rank, scatter_op): + self.parallelism = parallelism + self.group_size = group_size + self.root = root + self.groups = groups + self.attention_layers = attention_layers + self.inputs = inputs + self.outputs = outputs + self.k_seq_len_idx = k_seq_len_idx + self.v_seq_len_idx = v_seq_len_idx + self.kv_rank = kv_rank + self.scatter_op = scatter_op + +@Decoder.register(ShardHints) +def decode(dct): + return ShardHints( + parallelism=dct["parallelism"], + group_size=dct["group_size"], + root=dct["root"], + groups=dct["groups"], + attention_layers=dct["attention_layers"], + inputs=dct["inputs"], + outputs=dct["outputs"], + k_seq_len_idx=dct["k_seq_len_idx"], + v_seq_len_idx=dct["v_seq_len_idx"], + kv_rank=dct["kv_rank"], + scatter_op=dct["reduce_scatter_reduce_op"], + ) + +@Encoder.register(ShardHints) +def encode(shard_hints): + return { + "parallelism" : shard_hints.parallelism, + "group_size" : shard_hints.group_size, + "root" : shard_hints.root, + "groups" : shard_hints.groups, + "attention_layers" : shard_hints.attention_layers, + "inputs" : shard_hints.inputs, + "outputs" : shard_hints.outputs, + "k_seq_len_idx" : shard_hints.k_seq_len_idx, + "v_seq_len_idx" : shard_hints.v_seq_len_idx, + "kv_rank" : shard_hints.kv_rank, + "reduce_scatter_reduce_op" : shard_hints.scatter_op, + } + +def get_attention_pattern(): + """ + Returns the pattern for canonical attention layers. + + Attention layers follow the pattern: + + Q K + | | + MatMul + | + SoftMax + | + | V + | | + MatMul + | + Output + """ + + pattern = gs.GraphPattern() + q = pattern.variable() + k = pattern.variable() + v = pattern.variable() + + matmul_1 = pattern.add("MatMul1", "MatMul", inputs=[q, k]) + softmax = pattern.add("Softmax", "Softmax", inputs=[matmul_1]) + matmul_2 = pattern.add("MatMul2", "MatMul", inputs=[softmax, v]) + pattern.set_output_tensors([matmul_2]) + return pattern + +class MultiDevice(Tool): + """ + Multi-Device related operations on an onnx model. + """ + + def __init__(self): + super().__init__("multi-device") + + def get_subtools_impl(self): + # Avoid circular dependency + from polygraphy.tools.multi_device.subtool.shard import Shard + + return "Multi-Device Subtools", [ + Shard() + ] diff --git a/tools/Polygraphy/polygraphy/tools/multi_device/subtool/__init__.py b/tools/Polygraphy/polygraphy/tools/multi_device/subtool/__init__.py new file mode 100644 index 00000000..14d9addc --- /dev/null +++ b/tools/Polygraphy/polygraphy/tools/multi_device/subtool/__init__.py @@ -0,0 +1 @@ +from polygraphy.tools.multi_device.subtool.shard import Shard diff --git a/tools/Polygraphy/polygraphy/tools/multi_device/subtool/shard.py b/tools/Polygraphy/polygraphy/tools/multi_device/subtool/shard.py new file mode 100644 index 00000000..9b373832 --- /dev/null +++ b/tools/Polygraphy/polygraphy/tools/multi_device/subtool/shard.py @@ -0,0 +1,225 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import argparse +from polygraphy import mod +from polygraphy.logger.logger import G_LOGGER +from polygraphy.tools.args.backend.onnx.loader import OnnxInferShapesArgs, OnnxLoadArgs +from polygraphy.tools.base import Tool +from polygraphy.tools.args import ModelArgs +from polygraphy.tools.args import OnnxSaveArgs +from polygraphy.tools.multi_device import get_attention_pattern, ShardHints + +onnx = mod.lazy_import("onnx>=1.17") +gs = mod.lazy_import("onnx_graphsurgeon") +onnx_backend = mod.lazy_import("polygraphy.backend.onnx") + + +class Shard(Tool): + """ + Convert a SD model to a MD model using a sharding hints file. + """ + + def __init__(self): + super().__init__("shard") + + def _make_shuffle(self, graph, dist_node, n, rank): + input = dist_node.inputs[0] + output = dist_node.outputs[0] + dist_name = dist_node.name + dtype = input.dtype + + # Perm is static at runtime, we need a rank to actually do the shuffle permutation + # DistCollective won't change rank, so either get it from input/output or use rank + # as a fallback + size = len(input.shape) if input.shape else (len(output.shape) if output.shape else rank) + if size: + if n >= size: + G_LOGGER.critical(f"Specified seq_len_idx {n} is out of actual range {size} for DistCollective node {dist_name}") + + perm = [i for i in range(0, size)] + perm[0], perm[n] = perm[n], perm[0] + attrs = {"perm" : perm} + + tensor_pre = gs.Variable(name = "TensorPre_" + dist_name, shape = None, dtype = dtype) + tensor_post = gs.Variable(name = "TensorPost_" + dist_name, shape = None, dtype = dtype) + + transpose_pre = gs.Node(name = "TransposePre_" + dist_name, op= "Transpose", inputs = [input], outputs = [tensor_pre], attrs = attrs) + transpose_post = gs.Node(name = "TransposePost_" + dist_name, op = "Transpose", inputs = [tensor_post], outputs = [output], attrs = attrs) + + dist_node.inputs = [tensor_pre] + dist_node.outputs = [tensor_post] + graph.nodes.extend([transpose_pre, transpose_post]) + else: + G_LOGGER.critical("Shape inference needs to be run with --shape-inference if any seq_len_idx is not 0 and rank is not specified") + + def _make_dist_node(self, graph, tensor, attrs, inputs, outputs, tensors, seq_len_idx, fallback): + name = tensor.name + "_md" + + # Make new tensor and node needed + tensor_md = gs.Variable(name = name, shape = tensor.shape, dtype = tensor.dtype) + + if inputs is None: + inputs = [tensor_md] + if outputs is None: + outputs = [tensor_md] + + dist_name = "DistCollective_" + str(self.dist_count) + + node_md = gs.Node(op = "DistCollective", name = dist_name, inputs = inputs, outputs = outputs, attrs = attrs) + + # Update nodes affected by tensor + for node in [n for n in graph.nodes if tensor in tensors(n)]: + for i, t in enumerate(tensors(node)): + if tensor == t: + tensors(node)[i] = tensor_md + + # Change layers that have scattered input as tensor + for layer in self.hints.attention_layers: + if layer.q == tensor.name: + layer.q = tensor_md.name + + if seq_len_idx: + self._make_shuffle(graph, node_md, seq_len_idx, fallback) + + self.dist_count += 1 + graph.nodes.insert(0, node_md) + + return tensor_md + + def _make_attrs(self, collective_operation, reduce_op): + return { + "collective_operation": collective_operation, + "reduce_op": reduce_op, + "root": self.hints.root, + "group_size": self.hints.group_size, + } + + def _make_all_gather(self, graph, tensor, seq_len_idx = None, rank = None): + """ + Insert an all gather operation to tensor + T -> T'--AG--T + """ + + G_LOGGER.info(f"Inserting all-gather for tensor: {tensor.name}") + attrs = self._make_attrs("all_gather", "sum") + return self._make_dist_node(graph, tensor, attrs, None, [tensor], lambda n : n.outputs, seq_len_idx, rank) + + def _make_reduce_scatter(self, graph, tensor, seq_len_idx = None, rank = None): + """ + Insert a reduce scatter operation to tensor + T -> T--RS--T' + """ + + G_LOGGER.info(f"Inserting reduce-scatter for tensor: {tensor.name}") + attrs = self._make_attrs("reduce_scatter", self.hints.scatter_op) + return self._make_dist_node(graph, tensor, attrs, [tensor], None, lambda n : n.inputs, seq_len_idx, rank) + + def get_subscriptions_impl(self): + return [ + ModelArgs( + model_opt_required=True, + input_shapes_opt_name=False, + required_model_type="onnx", + ), + OnnxInferShapesArgs(), + OnnxLoadArgs(outputs_opt_prefix=False, allow_shape_inference=True), + OnnxSaveArgs(allow_shape_inference=False, output_opt_required=True), + ] + + def add_parser_args_impl(self, parser): + parser.add_argument( + "-s", + "--hint", + help = "Hints file to describe shardable layers.", + type = argparse.FileType("r"), + dest = "hint_file", + required = True + ) + + def run_impl(self, args): + # Reset state for each run to avoid interference between tests + self.dist_count = 0 + self.gather_output = False + gathered_q = False + + graph = onnx_backend.gs_from_onnx(self.arg_groups[OnnxLoadArgs].load_onnx()) + + G_LOGGER.info(f"Loading sharding hints from: {args.hint_file.name if hasattr(args.hint_file, 'name') else args.hint_file}") + self.hints = ShardHints.load(args.hint_file) + G_LOGGER.info(f"Loaded hints: parallelism={self.hints.parallelism}, group_size={self.hints.group_size}, root={self.hints.root}, groups={self.hints.groups}") + + + # Shard dependent inputs/outputs + tensors = graph.tensors() + for input in self.hints.inputs: + tensor_md = self._make_reduce_scatter(graph, tensors[input.name], input.seq_len_idx, input.rank) + + # Update attention layers if new input will be scattered tensor + for a_l in [a_l for a_l in self.hints.attention_layers if a_l.q == tensor_md.name]: + a_l.q = tensor_md.name + + # Get all attention layers that match supported pattern(s) + pattern = get_attention_pattern() + matches = pattern.match_all(graph) + G_LOGGER.info(f"Found {len(matches)} attention pattern matches in the graph.") + + # Perform sharding + for layer in self.hints.attention_layers: + G_LOGGER.info(f"Processing attention layer: q={layer.q}, gather_q={layer.gather_q}, gather_kv={layer.gather_kv}") + # Find configuration for matching attenion layer (inputs and outputs match) + match = next((match for match in matches if match.inputs[0].name == layer.q), None) + if match is not None: + sharded = set() + gather_q = layer.gather_q + gather_kv = layer.gather_kv + + # Get tensors directly + q = match["MatMul1"].onnx_node.inputs[0] + k = match["MatMul1"].onnx_node.inputs[1] + v = match["MatMul2"].onnx_node.inputs[1] + + # Insert collective ops as specified + for i, (tensor, seq_len_idx, rank) in enumerate([(q, None, None), (k, self.hints.k_seq_len_idx, self.hints.kv_rank), (v, self.hints.v_seq_len_idx, self.hints.kv_rank)]): + if [gather_q, gather_kv, gather_kv][i] and tensor.name not in sharded: + + # If any q is gathered (on purpose or k == q or v == q), prevent final output from being all gathered + if tensor.name == layer.q and not gathered_q: + G_LOGGER.info(f"Q {layer.q} was gathered") + gathered_q = True + + self._make_all_gather(graph, tensor, seq_len_idx, rank) + sharded.add(tensor.name) + else: + G_LOGGER.warning(f"No matching attention pattern found for layer with q={layer.q}") + + if not gathered_q: + for output in self.hints.outputs: + self._make_all_gather(graph, tensors[output.name], output.seq_len_idx, output.rank) + + # Cleanup and save + graph.cleanup() + graph.toposort() + + # Manually add in groups attribute since graph surgeon doesn't support + # type inference for an empty list, which is necessary for a group configuration of '[]' + model = gs.export_onnx(graph) + for node in model.graph.node: + if node.op_type == "DistCollective": + node.attribute.append(onnx.helper.make_attribute("groups", self.hints.groups, attr_type = onnx.AttributeProto.INTS)) + + self.arg_groups[OnnxSaveArgs].save_onnx(model) diff --git a/tools/Polygraphy/polygraphy/tools/registry.py b/tools/Polygraphy/polygraphy/tools/registry.py index efb655e5..b96f0953 100644 --- a/tools/Polygraphy/polygraphy/tools/registry.py +++ b/tools/Polygraphy/polygraphy/tools/registry.py @@ -57,6 +57,7 @@ try_register_tool("polygraphy.tools.template", "Template") try_register_tool("polygraphy.tools.debug", "Debug") try_register_tool("polygraphy.tools.data", "Data") try_register_tool("polygraphy.tools.plugin", "Plugin") +try_register_tool("polygraphy.tools.multi_device", "MultiDevice") # Check that tool names are unique tool_names = [tool.name for tool in TOOL_REGISTRY] diff --git a/tools/Polygraphy/polygraphy/tools/template/subtool/__init__.py b/tools/Polygraphy/polygraphy/tools/template/subtool/__init__.py index c3642f99..f32ca18c 100644 --- a/tools/Polygraphy/polygraphy/tools/template/subtool/__init__.py +++ b/tools/Polygraphy/polygraphy/tools/template/subtool/__init__.py @@ -1,3 +1,4 @@ from polygraphy.tools.template.subtool.trt_network import TrtNetwork from polygraphy.tools.template.subtool.trt_config import TrtConfig from polygraphy.tools.template.subtool.onnx_gs import OnnxGs +from polygraphy.tools.template.subtool.shard_hint import ShardHint diff --git a/tools/Polygraphy/polygraphy/tools/template/subtool/shard_hint.py b/tools/Polygraphy/polygraphy/tools/template/subtool/shard_hint.py new file mode 100644 index 00000000..51a41534 --- /dev/null +++ b/tools/Polygraphy/polygraphy/tools/template/subtool/shard_hint.py @@ -0,0 +1,270 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from polygraphy import mod +from polygraphy.logger.logger import G_LOGGER +from polygraphy.tools.template.subtool.base import BaseTemplateTool +from polygraphy.tools.args.backend.onnx.loader import OnnxLoadArgs +from polygraphy.tools.args import ModelArgs +from polygraphy.tools.multi_device import get_attention_pattern, ShardHints, AttentionLayerHint, ShardTensor + + +onnx_backend = mod.lazy_import("polygraphy.backend.onnx") + + +class GraphTraverser(): + def __init__(self, graph): + self.inputs = {input.name : input for input in graph.inputs} + self.outputs = {output.name : output for output in graph.outputs} + self.nodes = {node.name : node for node in graph.nodes} + self.visited_inputs = set() + self.visited_outputs = set() + self.reached_terminal = set() + + if len(self.nodes) != len(graph.nodes): + G_LOGGER.critical(f"All nodes in graph need to have unique names") + + # Tensors are the output of only one node + self.node_outputs = {output.name: {node.name} for node in graph.nodes for output in node.outputs} + + # Have to handle this slightly differently + self.node_inputs = {input.name : set() for node in graph.nodes for input in node.inputs} + for node in graph.nodes: + for input in node.inputs: + self.node_inputs[input.name].add(node.name) + + + def _traverse(self, node, edges, terminal, relative, dependent, visited): + queue = [node] + while queue: + cur = queue.pop() + if cur.name in visited: + continue + + visited.add(cur.name) + + for tensor in edges(cur): + if tensor in terminal and tensor not in self.reached_terminal: + G_LOGGER.info(f"Found dependent tensor {tensor}") + dependent.add(tensor) + self.reached_terminal.add(tensor) + elif tensor in relative: + for r in relative[tensor]: + if r not in visited: + queue.append(self.nodes[r]) + + def get_dep_inputs(self, node): + inputs = set() + self._traverse( + node, + lambda n: [input.name for input in n.inputs], + self.inputs, + self.node_outputs, + inputs, + self.visited_inputs + ) + return [self.inputs[input] for input in inputs] + + def get_dep_outputs(self, node): + outputs = set() + self._traverse( + node, + lambda n: [output.name for output in n.outputs], + self.outputs, + self.node_inputs, + outputs, + self.visited_outputs + ) + return [self.outputs[output] for output in outputs] + +class ShardHint(BaseTemplateTool): + """ + Generate a sharding hints file + """ + def __init__(self): + super().__init__("shard-hints") + + def get_subscriptions_impl(self): + return [ + ModelArgs( + model_opt_required=True, + input_shapes_opt_name=False, + required_model_type="onnx", + ), + OnnxLoadArgs(outputs_opt_prefix=False, allow_shape_inference=False), + ] + + def add_parser_args_impl(self, parser): + reduce_ops = ["sum", "prod", "min", "max", "avg"] + + super().add_parser_args_impl(parser) + + parser.add_argument( + "--parallelism", + help = "Type of parallelism to use", + type = str, + choices = ["CP"], + default = "CP" + ) + + parser.add_argument( + "--root", + help = "Rank of root process", + type=int, + default=0 + ) + + parser.add_argument( + "--gpus", + help = "Number of participating gpus (0 is all gpus)", + type=int, + default=0 + ) + + parser.add_argument( + "--groups", + help="Space-separated list of NCCL group indices (omit for all groups)", + nargs="*", + type=int, + default=[], + ) + + parser.add_argument( + "--cp-type", + help = "Sharding strategy for attention layers", + type=str, + choices=["native", "ring_attention", "fused"], + default = "native" + ) + + parser.add_argument( + "--no-suggest-io", + help = "Disable suggestions of which input/output tensors need to be sharded based on attention layer dependencies", + action='store_true', + default=False + ) + + parser.add_argument( + "--i-idx", + help = "Default index of sequence length on input tensor(s)", + type=int, + default=0 + ) + + parser.add_argument( + "--o-idx", + help = "Default index of sequence length on output tensor(s)", + type=int, + default=0 + ) + + parser.add_argument( + "--k-idx", + help = "Default index of sequence length on K tensor(s)", + type=int, + default=0 + ) + + parser.add_argument( + "--v-idx", + help = "Default index of sequence length on V tensor(s)", + type=int, + default=0 + ) + + parser.add_argument( + "--i-rank", + help = "Fallback rank of input shapes if sequence length index is > 0 and shape inference is not run", + type=int, + ) + + parser.add_argument( + "--o-rank", + help = "Fallback rank of output shapes if sequence length index is > 0 and shape inference is not run", + type=int, + ) + + parser.add_argument( + "--kv-rank", + help = "Fallback rank of KV shapes if sequence length index is > 0 and shape inference is not run", + type=int, + ) + + parser.add_argument( + "--scatter-op", + help = "reduce_op for reduce_scatter operations", + type=str, + choices = reduce_ops, + default = "max" + ) + + @staticmethod + def guess_seq_len_idx(tensor, default): + if (shape := tensor.shape) is not None: + for i, dim in enumerate(shape): + if dim == "sequence_length" or dim == "seq_len": + G_LOGGER.info(f"Found sequence_length index at {i} for tensor {tensor.name}") + return i + + return default + + def run_impl(self, args): + graph = onnx_backend.gs_from_onnx(self.arg_groups[OnnxLoadArgs].load_onnx()) + + if not args.output.name.endswith(".json"): + G_LOGGER.critical("Output file must be a json") + + traverser = GraphTraverser(graph) + gather_kv = args.cp_type == "native" + attention_layers = [] + inputs = [] + outputs = [] + kv_rank = None + get_rank = lambda t, default: len(t.shape) if t.shape else default + + for match in get_attention_pattern().match_all(graph): + q = match.inputs[0].name + G_LOGGER.info(f"Found attention layer with Q tensor {q}") + + k = match.inputs[1] + v = match.inputs[2] + + # Find rank of kv (all assumed to be same) + if not kv_rank: + # K or V will be same rank, use whichever has a shape (if any) + tensor = k if k.shape else v + G_LOGGER.info(f"Trying to find KV rank...") + kv_rank = get_rank(tensor, args.kv_rank) + + if kv_rank: + G_LOGGER.info(f"KV has rank of {kv_rank}") + + # Find dependent inputs/outputs (suggestion) + if not args.no_suggest_io: + inputs.extend(traverser.get_dep_inputs(match["MatMul2"].onnx_node)) + outputs.extend(traverser.get_dep_outputs(match["MatMul2"].onnx_node)) + + # TODO TRT-26378 https://jirasw.nvidia.com/browse/TRT-26378 add support for plugin, later myelin + attention_layers.append(AttentionLayerHint(q, gather_kv, False)) + + + convert_io = lambda tensors, default, rank=None: list(map(lambda t : ShardTensor(t.name, ShardHint.guess_seq_len_idx(t, default), get_rank(t, rank)), tensors)) + inputs = convert_io(inputs, args.i_idx, args.i_rank) + outputs = convert_io(outputs, args.o_idx, args.o_rank) + hints = ShardHints(args.parallelism, args.gpus, args.root, args.groups, attention_layers, inputs, outputs, args.k_idx, args.v_idx, kv_rank if kv_rank else args.kv_rank, args.scatter_op) + + hints.save(args.output) diff --git a/tools/Polygraphy/polygraphy/tools/template/template.py b/tools/Polygraphy/polygraphy/tools/template/template.py index 5c899493..7abd943a 100644 --- a/tools/Polygraphy/polygraphy/tools/template/template.py +++ b/tools/Polygraphy/polygraphy/tools/template/template.py @@ -15,7 +15,7 @@ # limitations under the License. # from polygraphy.tools.base import Tool -from polygraphy.tools.template.subtool import TrtNetwork, TrtConfig, OnnxGs +from polygraphy.tools.template.subtool import TrtNetwork, TrtConfig, OnnxGs, ShardHint class Template(Tool): @@ -31,4 +31,5 @@ class Template(Tool): TrtNetwork(), TrtConfig(), OnnxGs(), + ShardHint() ] diff --git a/tools/Polygraphy/polygraphy/util/util.py b/tools/Polygraphy/polygraphy/util/util.py index 64fd72b6..0821447f 100644 --- a/tools/Polygraphy/polygraphy/util/util.py +++ b/tools/Polygraphy/polygraphy/util/util.py @@ -290,6 +290,14 @@ def unpack_args(args, num): ## +@mod.export() +def sanitize_filename(path): + """ + Sanitizes a path so it can be used as a filename + """ + return path.replace(os.path.sep, "_") + + @mod.export() class NamedTemporaryFile: """ @@ -560,7 +568,9 @@ def _get_num_bytes(contents: Union[str, bytes, trt.IHostMemory]) -> int: try: memory_view = memoryview(contents) except Exception: - raise TypeError(f"`contents` is {contents}, which is not bytes-like. Cannot get number of bytes.") + raise TypeError( + f"`contents` is {contents}, which is not bytes-like. Cannot get number of bytes." + ) return len(memory_view) @@ -1089,6 +1099,7 @@ def try_getattr(obj, attr, default=None): return getattr(obj, attr) return default + @mod.export() def contains_wildcard(target): """ @@ -1101,16 +1112,17 @@ def contains_wildcard(target): """ return any(ch in target for ch in "*?[]!") + @mod.export() def match_keys(keys, targets): """ - Matching targets to keys, all matched targets will be return as a dict of the corresponding keys. The keys + Matching targets to keys, all matched targets will be return as a dict of the corresponding keys. The keys are allowed to contain wildcards Args: keys (iterable): Contains normal string names and wildcards targets (iterable): Targets list for matching - + Returns: Tuple[dict, list]: A tuple including matched target to key dict and unmatched keys list @@ -1123,6 +1135,4 @@ def match_keys(keys, targets): matched_keys.append(key) target_to_key[target] = key - return target_to_key, [name for name in keys if name not in matched_keys] - diff --git a/tools/Polygraphy/tests/backend/onnx/test_loader.py b/tools/Polygraphy/tests/backend/onnx/test_loader.py index 0b8d8a00..57e1eea5 100644 --- a/tools/Polygraphy/tests/backend/onnx/test_loader.py +++ b/tools/Polygraphy/tests/backend/onnx/test_loader.py @@ -22,7 +22,7 @@ import onnx import onnx_graphsurgeon as gs import pytest -from polygraphy import constants, mod, util +from polygraphy import constants from polygraphy.backend.onnx import ( ConvertToFp16, FoldConstants, @@ -178,17 +178,24 @@ class TestInferShapes: class TestConvertToFp16: @pytest.mark.parametrize("copy", [True, False]) def test_basic(self, copy): + # Precondition. original_model = onnx_from_path(ONNX_MODELS["identity_identity"].path) + assert original_model.graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT or not copy + + # Under test. loader = ConvertToFp16(original_model, copy=copy) model = loader() - assert original_model.graph.input[0].type.tensor_type.elem_type == 1 or not copy - assert model.graph.input[0].type.tensor_type.elem_type == 1 - assert model.graph.node[2].op_type == "Cast" - assert model.graph.node[0].op_type == "Identity" - assert model.graph.node[1].op_type == "Identity" - assert model.graph.node[3].op_type == "Cast" - assert model.graph.output[0].type.tensor_type.elem_type == 1 + # Postcondition. + graph = gs_from_onnx(model) + graph.toposort() + + assert graph.inputs[0].dtype == "float32" + assert graph.nodes[0].op == "Cast" + assert graph.nodes[1].op == "Identity" + assert graph.nodes[2].op == "Identity" + assert graph.nodes[3].op == "Cast" + assert graph.outputs[0].dtype == "float32" class TestFoldConstants: diff --git a/tools/Polygraphy/tests/backend/trt/test_algorithm_selector.py b/tools/Polygraphy/tests/backend/trt/test_algorithm_selector.py index 0e6703a0..5cadf67c 100644 --- a/tools/Polygraphy/tests/backend/trt/test_algorithm_selector.py +++ b/tools/Polygraphy/tests/backend/trt/test_algorithm_selector.py @@ -19,7 +19,7 @@ from collections import namedtuple import pytest import tensorrt as trt -from polygraphy import mod, util +from polygraphy import config, util from polygraphy.backend.trt import ( Algorithm, TacticRecorder, @@ -29,6 +29,10 @@ from polygraphy.backend.trt import ( ) from polygraphy.exception import PolygraphyException +# Skip all tests in this file if TensorRT-RTX is enabled +if config.USE_TENSORRT_RTX: + pytest.skip("Algorithm selector tests are not compatible with TensorRT-RTX", allow_module_level=True) + FakeAlgorithmContext = namedtuple( "FakeAlgorithmContext", ["name", "num_inputs", "num_outputs"] diff --git a/tools/Polygraphy/tests/backend/trt/test_calibrator.py b/tools/Polygraphy/tests/backend/trt/test_calibrator.py index 94f514a8..4ce763be 100644 --- a/tools/Polygraphy/tests/backend/trt/test_calibrator.py +++ b/tools/Polygraphy/tests/backend/trt/test_calibrator.py @@ -19,7 +19,7 @@ import pytest import tensorrt as trt import torch -from polygraphy import cuda, util +from polygraphy import config, cuda, util from polygraphy.backend.trt import ( Calibrator, CreateConfig, @@ -35,6 +35,10 @@ from polygraphy.exception import PolygraphyException from tests.helper import get_file_size, is_file_non_empty from tests.models.meta import ONNX_MODELS +# Skip all tests in this file if TensorRT-RTX is enabled +if config.USE_TENSORRT_RTX: + pytest.skip("Calibrator tests are not compatible with TensorRT-RTX", allow_module_level=True) + @pytest.fixture(scope="session") def identity_builder_network(): diff --git a/tools/Polygraphy/tests/backend/trt/test_config.py b/tools/Polygraphy/tests/backend/trt/test_config.py index b78e093b..2850c7d5 100644 --- a/tools/Polygraphy/tests/backend/trt/test_config.py +++ b/tools/Polygraphy/tests/backend/trt/test_config.py @@ -3,12 +3,10 @@ import os import tempfile import pytest -import tensorrt as trt from polygraphy import mod, util from polygraphy.backend.trt import ( Calibrator, - CreateConfig, Profile, network_from_onnx_bytes, postprocess_config, @@ -16,9 +14,17 @@ from polygraphy.backend.trt import ( from polygraphy.common.struct import BoundedShape from polygraphy.comparator import DataLoader from polygraphy.datatype import DataType +from polygraphy import config as polygraphy_config from tests.helper import has_dla from tests.models.meta import ONNX_MODELS +# Import CreateConfigRTX conditionally for TensorRT-RTX builds +if polygraphy_config.USE_TENSORRT_RTX: + import tensorrt_rtx as trt + from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig +else: + import tensorrt as trt + from polygraphy.backend.trt import CreateConfig @pytest.fixture(scope="session") def identity_builder_network(): @@ -36,7 +42,10 @@ class TestCreateConfig: with loader(builder, network) as config: assert not config.get_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE) with contextlib.suppress(AttributeError): - assert not config.get_flag(trt.BuilderFlag.TF32) + if polygraphy_config.USE_TENSORRT_RTX: + assert config.get_flag(trt.BuilderFlag.TF32) + else: + assert not config.get_flag(trt.BuilderFlag.TF32) with contextlib.suppress(AttributeError): assert not config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS) assert not config.get_flag(trt.BuilderFlag.FP16) @@ -47,19 +56,21 @@ class TestCreateConfig: assert not config.get_flag(trt.BuilderFlag.FP8) assert not config.get_flag(trt.BuilderFlag.VERSION_COMPATIBLE) assert not config.get_flag(trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME) - assert ( - config.hardware_compatibility_level - == trt.HardwareCompatibilityLevel.NONE - ) - if mod.version(trt.__version__) >= mod.version("10.2"): + if not polygraphy_config.USE_TENSORRT_RTX: + assert ( + config.hardware_compatibility_level + == trt.HardwareCompatibilityLevel.NONE + ) + if mod.version(trt.__version__) >= mod.version("10.2") and not polygraphy_config.USE_TENSORRT_RTX: assert ( config.runtime_platform == trt.RuntimePlatform.SAME_AS_BUILD ) assert config.num_optimization_profiles == 1 - assert config.int8_calibrator is None + if not polygraphy_config.USE_TENSORRT_RTX: + assert config.int8_calibrator is None with contextlib.suppress(AttributeError): - if mod.version(trt.__version__) >= mod.version("10.0"): + if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX: assert config.get_tactic_sources() == 24 elif mod.version(trt.__version__) >= mod.version("8.7"): assert config.get_tactic_sources() == 29 @@ -72,7 +83,10 @@ class TestCreateConfig: with contextlib.suppress(AttributeError): assert not config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) with contextlib.suppress(AttributeError): - assert config.engine_capability == trt.EngineCapability.STANDARD + if polygraphy_config.USE_TENSORRT_RTX: + assert config.engine_capability == trt.EngineCapability.STANDARD + else: + assert config.engine_capability == trt.EngineCapability.DEFAULT with contextlib.suppress(AttributeError): assert not config.get_flag(trt.BuilderFlag.DIRECT_IO) @@ -105,7 +119,7 @@ class TestCreateConfig: assert not obey_set and not prefer_set @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("8.6"), + mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX, reason="Unsupported before TRT 8.6", ) @pytest.mark.parametrize( @@ -142,26 +156,8 @@ class TestCreateConfig: @pytest.mark.parametrize( "arg_name, flag_type", [ - ("fp16", trt.BuilderFlag.FP16), - ("int8", trt.BuilderFlag.INT8), - ("allow_gpu_fallback", trt.BuilderFlag.GPU_FALLBACK), ("refittable", trt.BuilderFlag.REFIT), - ("tf32", trt.BuilderFlag.TF32), ] - + ( - [ - ("bf16", trt.BuilderFlag.BF16), - ] - if mod.version(trt.__version__) >= mod.version("8.7") - else [] - ) - + ( - [ - ("fp8", trt.BuilderFlag.FP8), - ] - if mod.version(trt.__version__) >= mod.version("8.6") - else [] - ) + ( [ ( @@ -178,6 +174,30 @@ class TestCreateConfig: ] if mod.version(trt.__version__) >= mod.version("10.0") else [] + ) + + ( + [ + ("fp16", trt.BuilderFlag.FP16), + ("int8", trt.BuilderFlag.INT8), + ("allow_gpu_fallback", trt.BuilderFlag.GPU_FALLBACK), + ("tf32", trt.BuilderFlag.TF32), + ] + + ( + [ + ("bf16", trt.BuilderFlag.BF16), + ] + if mod.version(trt.__version__) >= mod.version("8.7") + else [] + ) + + ( + [ + ("fp8", trt.BuilderFlag.FP8), + ] + if mod.version(trt.__version__) >= mod.version("8.6") + else [] + ) + if not polygraphy_config.USE_TENSORRT_RTX + else [] ), ) @pytest.mark.parametrize("value", [True, False]) @@ -194,6 +214,10 @@ class TestCreateConfig: with loader(builder, network) as config: assert config.get_flag(trt.BuilderFlag.SPARSE_WEIGHTS) == flag + @pytest.mark.skipif( + polygraphy_config.USE_TENSORRT_RTX, + reason="TensorRT-RTX does not support DLA" + ) def test_use_dla(self, identity_builder_network): builder, network = identity_builder_network loader = CreateConfig(use_dla=True) @@ -241,7 +265,7 @@ class TestCreateConfig: ), ] - if mod.version(trt.__version__) >= mod.version("10.0"): + if mod.version(trt.__version__) >= mod.version("10.0") or polygraphy_config.USE_TENSORRT_RTX: TACTIC_SOURCES_CASES[0] = (None, 24) elif mod.version(trt.__version__) >= mod.version("8.7"): TACTIC_SOURCES_CASES[0] = (None, 29) @@ -254,7 +278,7 @@ class TestCreateConfig: assert config.get_tactic_sources() == expected @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("8.7"), + mod.version(trt.__version__) < mod.version("8.7") and not polygraphy_config.USE_TENSORRT_RTX, reason="API was added in TRT 8.7", ) @pytest.mark.parametrize("flag", [True, False]) @@ -264,6 +288,10 @@ class TestCreateConfig: with loader(builder, network) as config: assert config.get_flag(trt.BuilderFlag.ERROR_ON_TIMING_CACHE_MISS) == flag + @pytest.mark.skipif( + polygraphy_config.USE_TENSORRT_RTX, + reason="TensorRT-RTX does not support calibrators" + ) def test_calibrator_metadata_set(self, identity_builder_network): builder, network = identity_builder_network calibrator = Calibrator(DataLoader()) @@ -336,6 +364,10 @@ class TestCreateConfig: }, ] + # @pytest.mark.skipif( + # config.USE_TENSORRT_RTX, + # reason="TensorRT-RTX does not support DLA memory pools" + # ) @pytest.mark.parametrize("pool_limits", POOL_LIMITS) def test_memory_pool_limits(self, pool_limits, identity_builder_network): if any("dla" in key.name.lower() for key in pool_limits) and not has_dla(): @@ -352,7 +384,11 @@ class TestCreateConfig: [ [trt.PreviewFeature.PROFILE_SHARING_0806] if mod.version(trt.__version__) >= mod.version("10.0") - else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805], + else ( + [trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03] + if polygraphy_config.USE_TENSORRT_RTX + else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805] + ), ], ) def test_preview_features(self, identity_builder_network, preview_features): @@ -361,8 +397,16 @@ class TestCreateConfig: with loader(builder, network) as config: # Check that only the enabled preview features are on. for pf in trt.PreviewFeature.__members__.values(): - assert config.get_preview_feature(pf) == (pf in preview_features) + expected = pf in preview_features + # TensorRT-RTX enables PROFILE_SHARING_0806 by default and can't be disabled + if polygraphy_config.USE_TENSORRT_RTX and pf == trt.PreviewFeature.PROFILE_SHARING_0806: + expected = True + assert config.get_preview_feature(pf) == expected + @pytest.mark.skipif( + polygraphy_config.USE_TENSORRT_RTX, + reason="TensorRT-RTX does not support quantization_flag API" + ) @pytest.mark.parametrize( "quantization_flags", [ @@ -379,7 +423,7 @@ class TestCreateConfig: assert config.get_quantization_flag(qf) == (qf in quantization_flags) @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("8.6"), + mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX, reason="Unsupported for TRT versions prior to 8.6", ) @pytest.mark.parametrize("level", range(6)) @@ -420,7 +464,7 @@ class TestCreateConfig: assert config.runtime_platform == platform @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("8.6"), + mod.version(trt.__version__) < mod.version("8.6") and not polygraphy_config.USE_TENSORRT_RTX, reason="Unsupported for TRT versions prior to 8.6", ) @pytest.mark.parametrize("num_streams", range(3)) @@ -431,7 +475,7 @@ class TestCreateConfig: assert config.max_aux_streams == num_streams @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("9.0"), + mod.version(trt.__version__) < mod.version("9.0") and not polygraphy_config.USE_TENSORRT_RTX, reason="API was added in TRT 9.0", ) def test_progress_monitor(self, identity_builder_network): @@ -454,7 +498,7 @@ class TestCreateConfig: with loader(builder, network) as config: assert config.progress_monitor == progress_monitor - if mod.version(trt.__version__) >= mod.version("10.8"): + if mod.version(trt.__version__) >= mod.version("10.8") and not polygraphy_config.USE_TENSORRT_RTX: @pytest.mark.parametrize( "level", [ diff --git a/tools/Polygraphy/tests/backend/trt/test_loader.py b/tools/Polygraphy/tests/backend/trt/test_loader.py index 4dcf8075..a617b6dc 100644 --- a/tools/Polygraphy/tests/backend/trt/test_loader.py +++ b/tools/Polygraphy/tests/backend/trt/test_loader.py @@ -17,14 +17,11 @@ from __future__ import annotations import sys - import pytest -import tensorrt as trt -from polygraphy import constants, mod, util +from polygraphy import config, constants, mod, util from polygraphy.backend.trt import ( Calibrator, - CreateConfig, EngineBytesFromNetwork, EngineFromBytes, EngineFromNetwork, @@ -57,6 +54,15 @@ from polygraphy.exception import PolygraphyException from tests.helper import get_file_size, is_file_non_empty from tests.models.meta import ONNX_MODELS +# Import CreateConfigRTX conditionally for TensorRT-RTX builds +if config.USE_TENSORRT_RTX: + import tensorrt_rtx as trt + from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig +else: + import tensorrt as trt + from polygraphy.backend.trt import CreateConfig + + ## ## Fixtures ## @@ -124,6 +130,10 @@ def modifiable_reshape_network(): class TestLoadPlugins: + @pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="Plugin tests are not compatible with TensorRT-RTX" + ) def test_can_load_libnvinfer_plugins(self): def get_plugin_names(): return [pc.name for pc in trt.get_plugin_registry().plugin_creator_list] @@ -165,7 +175,7 @@ class TestSerializedEngineLoader: @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("10.0"), reason="API was added in TRT 10.0" + mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX, reason="API was added in TRT 10.0" ) class TestSerializedEngineLoaderFromDisk: def test_serialized_engine_loader_from_lambda(self, identity_engine): @@ -190,6 +200,10 @@ class TestSerializedEngineLoaderFromDisk: @pytest.mark.skipif( mod.version(trt.__version__) < mod.version("8.6"), reason="API was added in TRT 8.6" ) +@pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="TensorRT-RTX does not have lean runtime shared objects" +) class TestLoadRuntime: def test_load_lean_runtime(self, nvinfer_lean_path): loader = LoadRuntime(nvinfer_lean_path) @@ -198,7 +212,11 @@ class TestLoadRuntime: @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("8.6"), reason="API was added in TRT 8.6" + mod.version(trt.__version__) < mod.version("8.6") and not config.USE_TENSORRT_RTX, reason="API was added in TRT 8.6" +) +@pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="TensorRT-RTX does not have libnvinfer_lean.so.1" ) class TestSerializedVCEngineLoader: def test_serialized_vc_engine_loader_from_lambda(self, identity_vc_engine_bytes): @@ -221,7 +239,8 @@ class TestNetworkFromOnnxBytes: builder, network, parser = network_from_onnx_bytes( ONNX_MODELS["identity"].loader ) - assert not network.has_implicit_batch_dimension + if not config.USE_TENSORRT_RTX: + assert not network.has_implicit_batch_dimension @pytest.mark.parametrize( "kwargs, flag", @@ -232,7 +251,7 @@ class TestNetworkFromOnnxBytes: trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED, ) ] - if mod.version(trt.__version__) >= mod.version("8.7") + if mod.version(trt.__version__) >= mod.version("8.7") and not config.USE_TENSORRT_RTX else [] ), ) @@ -246,7 +265,8 @@ class TestNetworkFromOnnxBytes: class TestNetworkFromOnnxPath: def test_loader(self): builder, network, parser = network_from_onnx_path(ONNX_MODELS["identity"].path) - assert not network.has_implicit_batch_dimension + if not config.USE_TENSORRT_RTX: + assert not network.has_implicit_batch_dimension @pytest.mark.parametrize( "kwargs, flag", @@ -257,7 +277,7 @@ class TestNetworkFromOnnxPath: trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED, ) ] - if mod.version(trt.__version__) >= mod.version("8.7") + if mod.version(trt.__version__) >= mod.version("8.7") and not config.USE_TENSORRT_RTX else [] ), ) @@ -360,6 +380,10 @@ class TestPostprocessNetwork: builder, network, parser = postprocess_network(modifiable_network, func) assert func_called + @pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="TensorRT-RTX uses strongly typed networks where layer precision cannot be set" + ) def test_modify_network(self, modifiable_network): """Tests that the network passed in is properly modified by the callback.""" @@ -385,6 +409,10 @@ class TestPostprocessNetwork: class TestSetLayerPrecisions: + @pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="TensorRT-RTX uses strongly typed networks where layer precision cannot be set" + ) def test_basic(self, modifiable_network): builder, network, parser = set_layer_precisions( modifiable_network, @@ -399,6 +427,10 @@ class TestSetLayerPrecisions: class TestSetTensorDatatypes: + @pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="TensorRT-RTX uses strongly typed networks where tensor datatypes cannot be set" + ) def test_basic(self, modifiable_network): builder, network, parser = set_tensor_datatypes( modifiable_network, @@ -459,6 +491,10 @@ class TestEngineFromNetwork: with loader() as engine: assert isinstance(engine, trt.ICudaEngine) + @pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="TensorRT-RTX does not support calibrators" + ) @pytest.mark.parametrize( "use_config_loader, set_calib_profile", [(True, None), (False, False), (False, True)], @@ -617,6 +653,10 @@ class TestOnnxLikeFromNetwork: class TestDefaultPlugins: + @pytest.mark.skipif( + config.USE_TENSORRT_RTX, + reason="Plugin tests are not compatible with TensorRT-RTX" + ) def test_default_plugins(self): network_loader = NetworkFromOnnxBytes(ONNX_MODELS["roialign"].loader) engine_loader = EngineFromNetwork(network_loader) diff --git a/tools/Polygraphy/tests/backend/trt/test_profile.py b/tools/Polygraphy/tests/backend/trt/test_profile.py index 7419edf6..1abe4e67 100644 --- a/tools/Polygraphy/tests/backend/trt/test_profile.py +++ b/tools/Polygraphy/tests/backend/trt/test_profile.py @@ -17,7 +17,12 @@ import numpy as np import pytest -import tensorrt as trt + +from polygraphy import config +if config.USE_TENSORRT_RTX: + import tensorrt_rtx as trt +else: + import tensorrt as trt from polygraphy.backend.trt import Profile, create_network, network_from_onnx_bytes from tests.models.meta import ONNX_MODELS @@ -55,7 +60,10 @@ class TestProfile: fill_shape = network.add_input("fill_shape", shape=tuple(), dtype=trt.int32) # Need to add some other operations so TensorRT treats `fill_shape` as a shape tensor. - fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE) + if config.USE_TENSORRT_RTX: + fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE, trt.int32) + else: + fill = network.add_fill(tuple(), trt.FillOperation.LINSPACE) fill.set_input(0, fill_shape) fill.set_input( 1, @@ -109,4 +117,3 @@ class TestProfile: trt_prof = profile.to_trt(builder, network) res = [trt_prof.get_shape(case) == [(2, 2, 3), (2, 2, 3), (2, 2, 3)] for case in match_case] assert res == should_match - diff --git a/tools/Polygraphy/tests/backend/trt/test_runner.py b/tools/Polygraphy/tests/backend/trt/test_runner.py index dc3fe0e0..8e8a9844 100644 --- a/tools/Polygraphy/tests/backend/trt/test_runner.py +++ b/tools/Polygraphy/tests/backend/trt/test_runner.py @@ -18,12 +18,10 @@ import threading import numpy as np import pytest -import tensorrt as trt import torch -from polygraphy import cuda, mod +from polygraphy import config, cuda, mod from polygraphy.backend.trt import ( - CreateConfig, EngineFromNetwork, NetworkFromOnnxBytes, Profile, @@ -36,6 +34,14 @@ from polygraphy.exception import PolygraphyException from polygraphy.logger import G_LOGGER from tests.models.meta import ONNX_MODELS +# Import CreateConfigRTX conditionally for TensorRT-RTX builds +if config.USE_TENSORRT_RTX: + import tensorrt_rtx as trt + from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig +else: + import tensorrt as trt + from polygraphy.backend.trt import CreateConfig + class TestLoggerCallbacks: @pytest.mark.parametrize("sev", G_LOGGER.SEVERITY_LETTER_MAPPING.keys()) @@ -79,9 +85,8 @@ class TestTrtRunner: assert not runner.is_active @pytest.mark.serial - def test_warn_if_impl_methods_called( - self, check_warnings_on_runner_impl_methods, identity_engine - ): + @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX has different warning output behavior") + def test_warn_if_impl_methods_called(self, check_warnings_on_runner_impl_methods, identity_engine): runner = TrtRunner(identity_engine) check_warnings_on_runner_impl_methods(runner) @@ -93,39 +98,30 @@ class TestTrtRunner: ([0, 0, 0, 1], [[3]]), ], ) + @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support data dependent shapes") def test_data_dependent_shapes(self, nonzero_engine, inp, expected): with TrtRunner(nonzero_engine) as runner: outputs = runner.infer( { "input": np.array( inp, - dtype=( - np.int32 - if mod.version(trt.__version__) < mod.version("9.0") - else np.int64 - ), + dtype=(np.int32 if mod.version(trt.__version__) < mod.version("9.0") else np.int64), ) } ) - assert np.array_equal( - outputs["nonzero_out_0"], np.array(expected, dtype=np.int32) - ) + assert np.array_equal(outputs["nonzero_out_0"], np.array(expected, dtype=np.int32)) @pytest.mark.parametrize("copy_outputs_to_host", [True, False]) @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_torch_tensors(self, copy_outputs_to_host, identity_engine, device): with TrtRunner(identity_engine) as runner: arr = torch.ones([1, 1, 2, 2], dtype=torch.float32, device=device) - outputs = runner.infer( - {"x": arr}, copy_outputs_to_host=copy_outputs_to_host - ) + outputs = runner.infer({"x": arr}, copy_outputs_to_host=copy_outputs_to_host) assert all(isinstance(t, torch.Tensor) for t in outputs.values()) assert torch.equal(outputs["y"].to("cpu"), arr.to("cpu")) - assert outputs["y"].device.type == ( - "cpu" if copy_outputs_to_host else "cuda" - ) + assert outputs["y"].device.type == ("cpu" if copy_outputs_to_host else "cuda") def test_context(self, identity_engine): with TrtRunner(identity_engine.create_execution_context) as runner: @@ -144,15 +140,9 @@ class TestTrtRunner: model.check_runner(runner) def test_multithreaded_runners_from_engine(self, identity_engine): - with TrtRunner(identity_engine) as runner0, TrtRunner( - identity_engine - ) as runner1: - t1 = threading.Thread( - target=ONNX_MODELS["identity"].check_runner, args=(runner0,) - ) - t2 = threading.Thread( - target=ONNX_MODELS["identity"].check_runner, args=(runner1,) - ) + with TrtRunner(identity_engine) as runner0, TrtRunner(identity_engine) as runner1: + t1 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner0,)) + t2 = threading.Thread(target=ONNX_MODELS["identity"].check_runner, args=(runner1,)) t1.start() t2.start() t1.join() @@ -160,7 +150,8 @@ class TestTrtRunner: @pytest.mark.parametrize("use_optimization_profile", [True, False]) @pytest.mark.skipif( - mod.version(trt.__version__) >= mod.version("8.6") + not config.USE_TENSORRT_RTX + and mod.version(trt.__version__) >= mod.version("8.6") and mod.version(trt.__version__) < mod.version("8.7"), reason="Bug in TRT 8.6", ) @@ -183,9 +174,7 @@ class TestTrtRunner: engine = engine_from_network(network_loader, config_loader) context = engine.create_execution_context() - for index, shapes in enumerate( - [profile0_shapes, profile1_shapes, profile2_shapes] - ): + for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]): with TrtRunner( context, optimization_profile=index if use_optimization_profile else None, @@ -198,13 +187,14 @@ class TestTrtRunner: model.check_runner(runner, {"X": shape}) @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("10.0"), + not config.USE_TENSORRT_RTX and mod.version(trt.__version__) < mod.version("10.0"), reason="Feature not present before 10.0", ) - @pytest.mark.parametrize( - "allocation_strategy", [None, "static", "profile", "runtime"] - ) + @pytest.mark.parametrize("allocation_strategy", [None, "static", "profile", "runtime"]) def test_allocation_strategies(self, allocation_strategy): + if config.USE_TENSORRT_RTX and allocation_strategy == "runtime": + pytest.skip("TensorRT-RTX issues with runtime allocation strategy") + model = ONNX_MODELS["residual_block"] profile0_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (1, 3, 224, 224)] profile1_shapes = [(1, 3, 224, 224), (1, 3, 224, 224), (2, 3, 224, 224)] @@ -218,9 +208,7 @@ class TestTrtRunner: config_loader = CreateConfig(profiles=profiles) engine = engine_from_network(network_loader, config_loader) - for index, shapes in enumerate( - [profile0_shapes, profile1_shapes, profile2_shapes] - ): + for index, shapes in enumerate([profile0_shapes, profile1_shapes, profile2_shapes]): with TrtRunner( engine, optimization_profile=index, @@ -252,12 +240,7 @@ class TestTrtRunner: def test_error_on_wrong_name_feed_dict(self, names, err, identity_engine, module): with TrtRunner(identity_engine) as runner: with pytest.raises(PolygraphyException, match=err): - runner.infer( - { - name: module.ones((1, 1, 2, 2), dtype=module.float32) - for name in names - } - ) + runner.infer({name: module.ones((1, 1, 2, 2), dtype=module.float32) for name in names}) @pytest.mark.parametrize("module", [torch, np]) def test_error_on_wrong_dtype_feed_dict(self, identity_engine, module): @@ -271,13 +254,9 @@ class TestTrtRunner: with pytest.raises(PolygraphyException, match="incompatible shape."): runner.infer({"x": module.ones((1, 1, 3, 2), dtype=module.float32)}) - @pytest.mark.parametrize( - "use_view", [True, False] - ) # We should be able to use DeviceArray in place of DeviceView + @pytest.mark.parametrize("use_view", [True, False]) # We should be able to use DeviceArray in place of DeviceView def test_device_views(self, use_view, reducable_engine): - with TrtRunner(reducable_engine) as runner, cuda.DeviceArray( - (1,), dtype=np.float32 - ) as x: + with TrtRunner(reducable_engine) as runner, cuda.DeviceArray((1,), dtype=np.float32) as x: x.copy_from(np.ones((1,), dtype=np.float32)) outputs = runner.infer( { @@ -303,71 +282,42 @@ class TestTrtRunner: assert np.all(outputs["y"] == inp) check(runner.infer({"x": inp})) - check( - runner.infer( - { - "x": cuda.DeviceArray( - shape=inp.shape, dtype=inp.dtype - ).copy_from(inp) - } - ) - ) + check(runner.infer({"x": cuda.DeviceArray(shape=inp.shape, dtype=inp.dtype).copy_from(inp)})) torch_outputs = runner.infer({"x": torch.from_numpy(inp)}) check({name: out.numpy() for name, out in torch_outputs.items()}) check(runner.infer({"x": inp})) - @pytest.mark.parametrize( - "use_view", [True, False] - ) # We should be able to use DeviceArray in place of DeviceView + @pytest.mark.parametrize("use_view", [True, False]) # We should be able to use DeviceArray in place of DeviceView def test_device_view_dynamic_shapes(self, use_view): model = ONNX_MODELS["dynamic_identity"] profiles = [ Profile().add("X", (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)), ] - runner = TrtRunner( - EngineFromNetwork( - NetworkFromOnnxBytes(model.loader), CreateConfig(profiles=profiles) - ) - ) + runner = TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader), CreateConfig(profiles=profiles))) with runner, cuda.DeviceArray(shape=(1, 2, 3, 3), dtype=np.float32) as arr: inp = np.random.random_sample(size=(1, 2, 3, 3)).astype(np.float32) arr.copy_from(inp) - outputs = runner.infer( - { - "X": ( - cuda.DeviceView(arr.ptr, arr.shape, arr.dtype) - if use_view - else arr - ) - } - ) + outputs = runner.infer({"X": (cuda.DeviceView(arr.ptr, arr.shape, arr.dtype) if use_view else arr)}) assert np.all(outputs["Y"] == inp) assert outputs["Y"].shape == (1, 2, 3, 3) def test_cannot_use_device_view_shape_tensor(self): model = ONNX_MODELS["empty_tensor_expand"] - with TrtRunner( - EngineFromNetwork(NetworkFromOnnxBytes(model.loader)) - ) as runner, cuda.DeviceArray( + with TrtRunner(EngineFromNetwork(NetworkFromOnnxBytes(model.loader))) as runner, cuda.DeviceArray( shape=(5,), dtype=( np.int32 - if mod.version(trt.__version__) < mod.version("9.0") + if mod.version(trt.__version__) < mod.version("9.0") and not config.USE_TENSORRT_RTX else np.int64 ), ) as arr: - with pytest.raises( - PolygraphyException, match="it must reside in host memory" - ): - runner.infer( - {"data": np.ones((2, 0, 3, 0), dtype=np.float32), "new_shape": arr} - ) + with pytest.raises(PolygraphyException, match="it must reside in host memory"): + runner.infer({"data": np.ones((2, 0, 3, 0), dtype=np.float32), "new_shape": arr}) @pytest.mark.parametrize("hwc_input", [True, False], ids=["hwc_input", "chw_input"]) - @pytest.mark.parametrize( - "hwc_output", [True, False], ids=["hwc_output", "chw_output"] - ) + @pytest.mark.parametrize("hwc_output", [True, False], ids=["hwc_output", "chw_output"]) + @pytest.mark.skipif(config.USE_TENSORRT_RTX, reason="TensorRT-RTX does not support custom I/O format networks") def test_infer_chw_format(self, hwc_input, hwc_output): model = ONNX_MODELS["identity_multi_ch"] inp_shape = model.input_metadata["x"].shape @@ -389,9 +339,7 @@ class TestTrtRunner: outputs = runner.infer({"x": inp}) if hwc_input == hwc_output: # output in CHW/HWC format and similarly shaped assert np.allclose(outputs["y"], inp) - elif ( - not hwc_input and hwc_output - ): # output in HWC format and shaped (N, H, W, C) + elif not hwc_input and hwc_output: # output in HWC format and shaped (N, H, W, C) assert np.allclose(outputs["y"].transpose(0, 3, 1, 2), inp) else: # hwc_input and not hwc_output: output in CHW format and shaped (N, C, H, W) assert np.allclose(outputs["y"].transpose(0, 2, 3, 1), inp) @@ -402,9 +350,7 @@ class TestTrtRunner: with cuda.DeviceArray.raw(shape) as arr: host_buffers = {} stream = cuda.Stream() - host_arr = _get_array_on_cpu( - arr, "test", host_buffers, stream, arr.nbytes, use_torch - ) + host_arr = _get_array_on_cpu(arr, "test", host_buffers, stream, arr.nbytes, use_torch) if use_torch: assert isinstance(host_arr, torch.Tensor) @@ -412,7 +358,7 @@ class TestTrtRunner: assert isinstance(host_arr, np.ndarray) @pytest.mark.skipif( - mod.version(trt.__version__) < mod.version("10.0"), + mod.version(trt.__version__) < mod.version("10.0") and not config.USE_TENSORRT_RTX, reason="Feature not present before 10.0", ) @pytest.mark.parametrize("budget", [None, -2, -1, 0, 0.5, 0.99, 1.0, 1000, np.inf]) @@ -435,3 +381,28 @@ class TestTrtRunner: with TrtRunner(engine, optimization_profile=0, **kwargs) as runner: model.check_runner(runner) + + @pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled") + def test_compute_capabilities_engine_building(self): + """Test compute capabilities integration with engine building""" + model = ONNX_MODELS["identity"] + network_loader = NetworkFromOnnxBytes(model.loader) + + # Test --use-gpu flag + config_loader = CreateConfig(use_gpu=True) + engine = engine_from_network(network_loader, config_loader) + with TrtRunner(engine) as runner: + model.check_runner(runner) + + # Test --compute-capabilities flag + config_loader = CreateConfig(compute_capabilities=[(7, 5), (8, 0), (8, 6)]) + engine = engine_from_network(network_loader, config_loader) + with TrtRunner(engine) as runner: + model.check_runner(runner) + + @pytest.mark.skipif(not config.USE_TENSORRT_RTX, reason="TensorRT-RTX not enabled") + def test_compute_capabilities_mutual_exclusion(self): + """Test that use_gpu and compute_capabilities are mutually exclusive""" + # Test mutual exclusion - should raise an exception + with pytest.raises(PolygraphyException, match="use_gpu and compute_capabilities are mutually exclusive"): + CreateConfig(use_gpu=True, compute_capabilities=[(7, 5)]) diff --git a/tools/Polygraphy/tests/backend/trt/test_util.py b/tools/Polygraphy/tests/backend/trt/test_util.py index 561f3498..bc84b703 100644 --- a/tools/Polygraphy/tests/backend/trt/test_util.py +++ b/tools/Polygraphy/tests/backend/trt/test_util.py @@ -19,12 +19,19 @@ import contextlib from textwrap import dedent import pytest -import tensorrt as trt -from polygraphy import mod -from polygraphy.backend.trt import CreateConfig, Profile, create_network +from polygraphy import config, mod +from polygraphy.backend.trt import Profile, create_network from polygraphy.backend.trt import util as trt_util +# Import CreateConfigRTX conditionally for TensorRT-RTX builds +if config.USE_TENSORRT_RTX: + import tensorrt_rtx as trt + from polygraphy.backend.tensorrt_rtx import CreateConfigRTX as CreateConfig +else: + import tensorrt as trt + from polygraphy.backend.trt import CreateConfig + @pytest.fixture(scope="session") def dummy_network(): @@ -54,14 +61,27 @@ def test_all_layer_types_mapped(layer_class_mapping, layer_type): def adjust_memory_pool_limits_after_8_6(limits): # Adjust tactic DRAM so we can match the output text reliably in update_expected_output. - if mod.version(trt.__version__) >= mod.version("8.6"): + if mod.version(trt.__version__) >= mod.version("8.6") or config.USE_TENSORRT_RTX: limits[trt.MemoryPoolType.TACTIC_DRAM] = 1 << 30 return limits def update_expected_output(expected): - if mod.version(trt.__version__) >= mod.version("8.6"): - if mod.version(trt.__version__) >= mod.version("10.0"): + is_trt_10_plus = ( + mod.version(trt.__version__) >= mod.version("10.0") or + config.USE_TENSORRT_RTX + ) + is_trt_8_6_plus = ( + mod.version(trt.__version__) >= mod.version("8.6") or + config.USE_TENSORRT_RTX + ) + is_trt_8_7_plus = ( + mod.version(trt.__version__) >= mod.version("8.7") or + config.USE_TENSORRT_RTX + ) + + if is_trt_8_6_plus: + if is_trt_10_plus: expected = expected.replace( "MiB]", "MiB, TACTIC_DRAM: 1024.00 MiB, TACTIC_SHARED_MEMORY: 1024.00 MiB]", @@ -70,22 +90,26 @@ def update_expected_output(expected): expected = expected.replace("MiB]", "MiB, TACTIC_DRAM: 1024.00 MiB]") if "Preview Features" not in expected: - if mod.version(trt.__version__) < mod.version("10.0"): + if not is_trt_10_plus: expected = ( dedent(expected).strip() + "\nPreview Features | [FASTER_DYNAMIC_SHAPES_0805, DISABLE_EXTERNAL_TACTIC_SOURCES_FOR_CORE_0805]" ) else: + preview_features = "[PROFILE_SHARING_0806" + if config.USE_TENSORRT_RTX: + preview_features += ", RUNTIME_ACTIVATION_RESIZE_10_10" + preview_features += "]" expected = ( dedent(expected).strip() - + "\nPreview Features | [PROFILE_SHARING_0806]" + + f"\nPreview Features | {preview_features}" ) - if mod.version(trt.__version__) >= mod.version("8.7"): + if is_trt_8_7_plus: # CUBLAS_LT is not longer enabled by default expected = expected.replace("CUBLAS_LT, ", "") - if mod.version(trt.__version__) >= mod.version("10.0"): + if is_trt_10_plus: expected = expected.replace( "EngineCapability.DEFAULT", "EngineCapability.STANDARD" ) @@ -107,12 +131,12 @@ def update_expected_output(expected): ), update_expected_output( """ - Flags | [] + Flags | [{}] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 16.00 MiB] Tactic Sources | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS] Profiling Verbosity | ProfilingVerbosity.DETAILED - """ + """.format("TF32" if config.USE_TENSORRT_RTX else "") ), ), ( @@ -124,12 +148,12 @@ def update_expected_output(expected): ), update_expected_output( """ - Flags | [] + Flags | [{}] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 16.00 MiB] Tactic Sources | [] Profiling Verbosity | ProfilingVerbosity.DETAILED - """ + """.format("TF32" if config.USE_TENSORRT_RTX else "") ), ), ( @@ -140,12 +164,12 @@ def update_expected_output(expected): ), update_expected_output( """ - Flags | [] + Flags | [{}] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 4.00 MiB] Tactic Sources | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS] Profiling Verbosity | ProfilingVerbosity.DETAILED - """ + """.format("TF32" if config.USE_TENSORRT_RTX else "") ), ), ( @@ -153,20 +177,24 @@ def update_expected_output(expected): memory_pool_limits=adjust_memory_pool_limits_after_8_6( {trt.MemoryPoolType.WORKSPACE: 16 << 20} ), - fp16=True, - int8=True, - tf32=True, + **({} if config.USE_TENSORRT_RTX else { + "fp16": True, + "int8": True, + "tf32": True, + }), refittable=True, precision_constraints="obey", ), update_expected_output( """ - Flags | [FP16, INT8, REFIT, TF32, OBEY_PRECISION_CONSTRAINTS] + Flags | [{}REFIT, TF32, OBEY_PRECISION_CONSTRAINTS] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 16.00 MiB] Tactic Sources | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS] Profiling Verbosity | ProfilingVerbosity.DETAILED - """ + """.format( + "" if config.USE_TENSORRT_RTX else "FP16, INT8, ", + ) ), ), ( @@ -181,15 +209,16 @@ def update_expected_output(expected): ), update_expected_output( """ - Flags | [] + Flags | [{}] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 16.00 MiB] Tactic Sources | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS] Profiling Verbosity | ProfilingVerbosity.DETAILED Optimization Profiles | 2 profile(s) - """ + """.format("TF32" if config.USE_TENSORRT_RTX else "") ), ), + ] + ([] if config.USE_TENSORRT_RTX else [ ( CreateConfig( memory_pool_limits=adjust_memory_pool_limits_after_8_6( @@ -208,6 +237,7 @@ def update_expected_output(expected): """ ), ), + ]) + [ ( ( CreateConfig( @@ -218,32 +248,39 @@ def update_expected_output(expected): ), update_expected_output( """ - Flags | [] + Flags | [{}] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 16.00 MiB] Tactic Sources | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS] Profiling Verbosity | ProfilingVerbosity.DETAILED Preview Features | [PROFILE_SHARING_0806] - """ + """.format("TF32" if config.USE_TENSORRT_RTX else "") ), ) - if mod.version(trt.__version__) >= mod.version("10.0") + if mod.version(trt.__version__) >= mod.version("10.0") or config.USE_TENSORRT_RTX else ( CreateConfig( memory_pool_limits=adjust_memory_pool_limits_after_8_6( {trt.MemoryPoolType.WORKSPACE: 16 << 20} ), - preview_features=[trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805], + preview_features=( + [trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03] + if config.USE_TENSORRT_RTX + else [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805] + ), ), update_expected_output( """ - Flags | [] + Flags | [{}] Engine Capability | EngineCapability.DEFAULT Memory Pools | [WORKSPACE: 16.00 MiB] Tactic Sources | [CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS] Profiling Verbosity | ProfilingVerbosity.DETAILED - Preview Features | [FASTER_DYNAMIC_SHAPES_0805] - """ + Preview Features | [{}] + """.format( + "TF32" if config.USE_TENSORRT_RTX else "", + "ALIASED_PLUGIN_IO_10_03" if config.USE_TENSORRT_RTX else "FASTER_DYNAMIC_SHAPES_0805" + ) ), ) ), @@ -252,9 +289,9 @@ def update_expected_output(expected): "default", "tactic-sources", "memory-pool-limits", - "builder-flags", + "builder-flags" + ("-rtx" if config.USE_TENSORRT_RTX else ""), "profiles", - "dla", + ] + ([] if config.USE_TENSORRT_RTX else ["dla"]) + [ "preview-features", ], ) diff --git a/tools/Polygraphy/tests/logger/test_logger.py b/tools/Polygraphy/tests/logger/test_logger.py index 7a3a3572..929398fe 100644 --- a/tools/Polygraphy/tests/logger/test_logger.py +++ b/tools/Polygraphy/tests/logger/test_logger.py @@ -75,8 +75,12 @@ class TestLogger: assert num_times_called == 3 assert num_times_called == 4 - @pytest.mark.serial + @pytest.mark.serial def test_use_python_logging_system(self, tmp_python_log_file): + # Clear log file + with tmp_python_log_file.open("w") as fp: + fp.write("") + logger = Logger(severity=Logger.ULTRA_VERBOSE) logger.use_python_logging_system = True # add custom Polygraphy levels @@ -103,7 +107,16 @@ class TestLogger: with tmp_python_log_file.open() as fp: log_messages = fp.read() - assert log_messages == """\ + # Remove lines containing "pytest_shutil.workspace" + log_messages = "\n".join( + line + for line in log_messages.splitlines() + if "pytest_shutil.workspace" not in line + ) + ("\n" if log_messages.endswith("\n") else "") + + assert ( + log_messages + == """\ ULTRA_VERBOSE:Polygraphy:[U] ultra verbose SUPER_VERBOSE:Polygraphy:[S] super verbose EXTRA_VERBOSE:Polygraphy:[X] extra verbose @@ -115,6 +128,8 @@ WARNING:Polygraphy:[W] warning ERROR:Polygraphy:[E] error CRITICAL:Polygraphy:[!] critical """ + ) + class TestSeverityTrie: @pytest.mark.parametrize( diff --git a/tools/Polygraphy/tests/models/attention.onnx b/tools/Polygraphy/tests/models/attention.onnx new file mode 100644 index 00000000..2c9b966e Binary files /dev/null and b/tools/Polygraphy/tests/models/attention.onnx differ diff --git a/tools/Polygraphy/tests/models/attention_same_qkv.onnx b/tools/Polygraphy/tests/models/attention_same_qkv.onnx new file mode 100644 index 00000000..9fba4701 Binary files /dev/null and b/tools/Polygraphy/tests/models/attention_same_qkv.onnx differ diff --git a/tools/Polygraphy/tests/models/meta.py b/tools/Polygraphy/tests/models/meta.py index 0a81b9cf..185e5378 100644 --- a/tools/Polygraphy/tests/models/meta.py +++ b/tools/Polygraphy/tests/models/meta.py @@ -447,4 +447,19 @@ ONNX_MODELS = { LoaderType=BytesFromPath, check_runner=no_check_implemented, ), + "attention": Model( + path=model_path("attention.onnx"), + LoaderType=BytesFromPath, + check_runner=no_check_implemented, + ), + "multi_attention": Model( + path=model_path("multi_attention.onnx"), + LoaderType=BytesFromPath, + check_runner=no_check_implemented, + ), + "attention_same_qkv": Model( + path=model_path("attention_same_qkv.onnx"), + LoaderType=BytesFromPath, + check_runner=no_check_implemented, + ), } diff --git a/tools/Polygraphy/tests/models/multi_attention.onnx b/tools/Polygraphy/tests/models/multi_attention.onnx new file mode 100644 index 00000000..f06d723c Binary files /dev/null and b/tools/Polygraphy/tests/models/multi_attention.onnx differ diff --git a/tools/Polygraphy/tests/tools/conftest.py b/tools/Polygraphy/tests/tools/conftest.py index 7d17e327..f9f1fe98 100644 --- a/tools/Polygraphy/tests/tools/conftest.py +++ b/tools/Polygraphy/tests/tools/conftest.py @@ -51,9 +51,11 @@ poly_run = make_poly_fixture(["run"]) poly_convert = make_poly_fixture(["convert"]) poly_inspect = make_poly_fixture(["inspect"]) poly_check = make_poly_fixture(["check"]) +poly_multi_device_shard = make_poly_fixture(["multi-device", "shard"]) poly_surgeon = make_poly_fixture(["surgeon"]) poly_surgeon_extract = make_poly_fixture(["surgeon", "extract"]) poly_template = make_poly_fixture(["template"]) +poly_template_shard = make_poly_fixture(["template", "shard-hints"]) poly_debug = make_poly_fixture(["debug"]) poly_data = make_poly_fixture(["data"]) poly_plugin_match = make_poly_fixture(["plugin", "match"]) diff --git a/tools/Polygraphy/tests/tools/test_convert.py b/tools/Polygraphy/tests/tools/test_convert.py index 1d97acb9..19f313de 100644 --- a/tools/Polygraphy/tests/tools/test_convert.py +++ b/tools/Polygraphy/tests/tools/test_convert.py @@ -20,8 +20,10 @@ from textwrap import dedent import onnx import pytest import tensorrt as trt -from polygraphy import mod, util + +from polygraphy import util from polygraphy.backend.common import BytesFromPath +from polygraphy.backend.onnx import gs_from_onnx from polygraphy.backend.trt import EngineFromBytes from tests.models.meta import ONNX_MODELS, TF_MODELS @@ -49,12 +51,15 @@ class TestConvertToOnnx: ) # I/O types should be unchanged model = onnx.load(outmodel.name) - assert model.graph.input[0].type.tensor_type.elem_type == 1 - assert model.graph.node[2].op_type == "Cast" - assert model.graph.node[0].op_type == "Identity" - assert model.graph.node[1].op_type == "Identity" - assert model.graph.node[3].op_type == "Cast" - assert model.graph.output[0].type.tensor_type.elem_type == 1 + graph = gs_from_onnx(model) + graph.toposort() + + assert graph.inputs[0].dtype == "float32" + assert graph.nodes[0].op == "Cast" + assert graph.nodes[1].op == "Identity" + assert graph.nodes[2].op == "Identity" + assert graph.nodes[3].op == "Cast" + assert graph.outputs[0].dtype == "float32" class TestConvertToTrt: diff --git a/tools/Polygraphy/tests/tools/test_inspect.py b/tools/Polygraphy/tests/tools/test_inspect.py index 8237a47d..bd8eeb61 100644 --- a/tools/Polygraphy/tests/tools/test_inspect.py +++ b/tools/Polygraphy/tests/tools/test_inspect.py @@ -775,3 +775,32 @@ class TestInspectSparsity: ipath = ONNX_MODELS[model_name].path status = poly_inspect(["sparsity", ipath]) assert status + +class TestDebugTensors: + @pytest.mark.skipif( + mod.version(trt.__version__) < mod.version("10.13"), + reason="Feature not supported before 10.13", + ) + def test_unfused_debug_tensors(self, poly_run, poly_inspect): + with tempfile.TemporaryDirectory() as outdir: + model = ONNX_MODELS["matmul_2layer"].path + + poly_run( + [ + model, + "--trt", + "--mark-unfused-tensors-as-debug-tensors", + "--save-outputs", + "output.json", + "--save-engine", + "debug_unfused.engine" + ], + cwd=outdir + ) + status = poly_inspect(["model", "debug_unfused.engine", "--show", "layers", "--model-type", "engine", "--combine-tensor-info", "output.json"], cwd=outdir) + + assert status.stdout.count("min") == 5 + assert status.stdout.count("max") == 5 + assert status.stdout.count("avg") == 4 + + diff --git a/tools/Polygraphy/tests/tools/test_shard.py b/tools/Polygraphy/tests/tools/test_shard.py new file mode 100644 index 00000000..10bdfb84 --- /dev/null +++ b/tools/Polygraphy/tests/tools/test_shard.py @@ -0,0 +1,122 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import onnx +from polygraphy.tools.multi_device import ShardHints +from polygraphy import util +from tests.models.meta import ONNX_MODELS + +class TestShard: + + def check_expected_dist_count(self, path, scatter_count, gather_count): + model = onnx.load(path) + + dist_nodes = [node for node in model.graph.node if node.op_type == "DistCollective"] + + def has_collective_operation(node, op): + for attr in node.attribute: + if attr.name == "collective_operation" and attr.s.decode() == op: + return True + return False + + gather_nodes = [node for node in dist_nodes if has_collective_operation(node, "all_gather")] + scatter_nodes = [node for node in dist_nodes if has_collective_operation(node, "reduce_scatter")] + + assert len(gather_nodes) == gather_count + assert len(scatter_nodes) == scatter_count + + def test_multi_attention_head_shard( + self, poly_multi_device_shard, poly_template_shard + ): + # Test for a network with multiple attention heads, shard only affects specified ones + with util.NamedTemporaryFile(suffix=".onnx") as outmodel, util.NamedTemporaryFile(mode='w+', suffix=".json") as hints: + poly_template_shard( + [ + ONNX_MODELS["multi_attention"].path, + "-o", + hints.name + ] + ) + + # Remove the second attention + json = ShardHints.load(hints.name) + json.inputs = json.inputs[:len(json.inputs)//2] + json.outputs = json.outputs[:len(json.outputs)//2] + json.attention_layers = json.attention_layers[:len(json.attention_layers)//2] + json.save(hints.name) + + poly_multi_device_shard( + [ + ONNX_MODELS["multi_attention"].path, + "-o", + outmodel.name, + "-s", + hints.name + ] + ) + + # Should only gather specified kv + self.check_expected_dist_count(outmodel.name, 3, 3) + + def test_shard_same_qkv( + self, poly_multi_device_shard, poly_template_shard + ): + # Test shard doesn't insert multiple all_gathers if any inputs are the same + with util.NamedTemporaryFile(suffix=".onnx") as outmodel, util.NamedTemporaryFile(mode = "w+", suffix=".json") as hints: + poly_template_shard( + [ + ONNX_MODELS["attention_same_qkv"].path, + "-o", + hints.name + ] + ) + poly_multi_device_shard( + [ + ONNX_MODELS["attention_same_qkv"].path, + "-o", + outmodel.name, + "-s", + hints.name + ] + ) + + # One scatter for single input, one gather for qkv (skip gather at end because q was gathered) + self.check_expected_dist_count(outmodel.name, 1, 1) + + def test_shard( + self, poly_multi_device_shard, poly_template_shard + ): + # Test normal sharding of kv for attention + with util.NamedTemporaryFile(suffix=".onnx") as outmodel, util.NamedTemporaryFile(mode = "w+", suffix=".json") as hints: + poly_template_shard( + [ + ONNX_MODELS["attention"].path, + "-o", + hints.name + ] + ) + poly_multi_device_shard( + [ + ONNX_MODELS["attention"].path, + "-o", + outmodel.name, + "-s", + hints.name + ] + ) + + # 3 scatters for each input, 2 gathers for kv, one gather at end + self.check_expected_dist_count(outmodel.name, 3, 3) diff --git a/tools/Polygraphy/tests/util/test_serde.py b/tools/Polygraphy/tests/util/test_serde.py index 89d6d198..689f995e 100644 --- a/tools/Polygraphy/tests/util/test_serde.py +++ b/tools/Polygraphy/tests/util/test_serde.py @@ -89,6 +89,7 @@ class TestDecoder: # to retain backwards compatibility. assert set(Decoder.polygraphy_registered.keys()) == { "__polygraphy_encoded_Algorithm", + "__polygraphy_encoded_AttentionLayerHint", "__polygraphy_encoded_Dummy", "__polygraphy_encoded_FormattedArray", "__polygraphy_encoded_IterationContext", @@ -96,10 +97,13 @@ class TestDecoder: "__polygraphy_encoded_LazyArray", "__polygraphy_encoded_ndarray", "__polygraphy_encoded_RunResults", + "__polygraphy_encoded_ShardHints", + "__polygraphy_encoded_ShardTensor", "__polygraphy_encoded_TacticReplayData", "__polygraphy_encoded_Tensor", "__polygraphy_encoded_TensorInfo", "Algorithm", + "AttentionLayerHint", "Dummy", "FormattedArray", "IterationContext", @@ -108,6 +112,8 @@ class TestDecoder: "LazyNumpyArray", "ndarray", "RunResults", + "ShardHints", + "ShardTensor", "TacticReplayData", "Tensor", "TensorInfo", diff --git a/tools/onnx-graphsurgeon/CHANGELOG.md b/tools/onnx-graphsurgeon/CHANGELOG.md index 00285bfe..7671cf96 100644 --- a/tools/onnx-graphsurgeon/CHANGELOG.md +++ b/tools/onnx-graphsurgeon/CHANGELOG.md @@ -2,6 +2,14 @@ Dates are in YYYY-MM-DD format. +## v0.5.9 (2025-10-28) + +### Fixed + +- Fixed a bug where the pattern matching logic would generate false positives in + cases where there were extra external consumers. + + ## v0.5.8 (2025-04-08) ### Fixed diff --git a/tools/onnx-graphsurgeon/onnx_graphsurgeon/__init__.py b/tools/onnx-graphsurgeon/onnx_graphsurgeon/__init__.py index 676e2825..fb3a252f 100644 --- a/tools/onnx-graphsurgeon/onnx_graphsurgeon/__init__.py +++ b/tools/onnx-graphsurgeon/onnx_graphsurgeon/__init__.py @@ -7,4 +7,4 @@ from onnx_graphsurgeon.ir.node import Node from onnx_graphsurgeon.ir.tensor import Constant, Tensor, Variable from onnx_graphsurgeon.util.exception import OnnxGraphSurgeonException -__version__ = "0.5.8" +__version__ = "0.5.9" diff --git a/tools/onnx-graphsurgeon/onnx_graphsurgeon/graph_pattern/graph_pattern.py b/tools/onnx-graphsurgeon/onnx_graphsurgeon/graph_pattern/graph_pattern.py index 53d0b1a6..21d978c4 100644 --- a/tools/onnx-graphsurgeon/onnx_graphsurgeon/graph_pattern/graph_pattern.py +++ b/tools/onnx-graphsurgeon/onnx_graphsurgeon/graph_pattern/graph_pattern.py @@ -16,6 +16,7 @@ # from typing import Dict, List, Union +import copy from onnx_graphsurgeon.ir.graph import Constant, Graph, Node from onnx_graphsurgeon.logger import G_LOGGER @@ -330,6 +331,76 @@ class GraphPattern: from_inbound, ) if match: + # If we entered this subpattern via an inbound boundary tensor, ensure + # that all internal consumers of the boundary input tensor are covered. + # This supports cases where a single subpattern input fans out to multiple + # internal nodes. + if from_inbound: + # Determine the boundary ONNX tensor corresponding to this inbound edge + tensor_index_for_node = self._get_tensor_index_for_node( + initial_node, from_tensor, is_node_input=True + ) + if tensor_index_for_node >= len(onnx_node.inputs): + return None + boundary_onnx_tensor = onnx_node.inputs[tensor_index_for_node] + + # All internal consumer nodes of the boundary input tensor + internal_consumers = list(self.tensor_outputs.get(from_tensor, [])) + + # Build candidate ONNX consumer nodes for the boundary tensor + candidate_onnx_consumers = list(getattr(boundary_onnx_tensor, "outputs", [])) + used_consumer_ids = set() + + # Mark already mapped consumers as used if applicable + for consumer_node in internal_consumers: + if consumer_node in mapping: + inbound_tensor_index = self._get_tensor_index_for_node( + consumer_node, from_tensor, is_node_input=True + ) + inbound_mapped = self.nodes[consumer_node].get_inbound_or_outbound_onnx_node( + mapping[consumer_node], is_inbound=True, tensor_index=inbound_tensor_index + ) + if inbound_mapped is None: + return None + # Ensure the mapped node is indeed one of the boundary tensor's consumers + if all(c.id != inbound_mapped.id for c in candidate_onnx_consumers): + return None + used_consumer_ids.add(inbound_mapped.id) + + # For remaining internal consumers not yet mapped, try to match them to unused ONNX consumers + for consumer_node in internal_consumers: + if consumer_node in mapping: + continue + matched_this_consumer = False + for c in candidate_onnx_consumers: + if c.id in used_consumer_ids or c.id in mapped_onnx_nodes: + continue + # Try to extend mapping with this branch. Use copies to allow backtracking on failure. + mapping_copy = copy.deepcopy(mapping) + mapped_onnx_nodes_copy = set(mapped_onnx_nodes) + ok = self._match_node( + consumer_node, + c, + from_tensor, + mapping_copy, + mapped_onnx_nodes_copy, + onnx_graph_output_tensors, + from_inbound=True, + ) + if ok: + mapping.clear() + for _k, _v in mapping_copy.items(): + mapping[_k] = _v + mapping.inputs = mapping_copy.inputs + mapping.outputs = mapping_copy.outputs + mapping.constants = mapping_copy.constants + mapped_onnx_nodes = mapped_onnx_nodes_copy + used_consumer_ids.add(c.id) + matched_this_consumer = True + break + if not matched_this_consumer: + return None + return mapping else: return None @@ -443,13 +514,14 @@ class GraphPattern: if onnx_tensor.name in onnx_graph_output_tensors: return False # The pattern tensor is not an output but the onnx tensor is an output tensor of the onnx graph. - # For sub-patterns, each input tensor can only have 1 output node. Otherwise the following test will fail. - if len(self.tensor_outputs[node_output_tensor]) != len(onnx_tensor.outputs): - return False - for output_node, output_onnx_node in zip( - self.tensor_outputs[node_output_tensor], onnx_tensor.outputs - ): - # dfs ends when revisiting a node. We need to check if the edges are matched. + # Flexible consumer matching: each pattern consumer must map to some ONNX consumer, + # but the ONNX tensor may have additional consumers we can ignore. + pattern_consumers = list(self.tensor_outputs[node_output_tensor]) + candidate_consumers = list(getattr(onnx_tensor, "outputs", [])) + used_candidate_ids = set() + + # First validate already mapped consumers + for output_node in pattern_consumers: if output_node in mapping: inbound_tensor_index = self._get_tensor_index_for_node( output_node, node_output_tensor, is_node_input=True @@ -463,22 +535,50 @@ class GraphPattern: ) if ( inbound_onnx_node_of_output_node is None - or inbound_onnx_node_of_output_node.name - != output_onnx_node.name + or all(c.id != inbound_onnx_node_of_output_node.id for c in candidate_consumers) ): return False + used_candidate_ids.add(inbound_onnx_node_of_output_node.id) + + # Then match remaining consumers greedily with backtracking via copies + for output_node in pattern_consumers: + if output_node in mapping: continue - match = self._match_node( - output_node, - output_onnx_node, - node_output_tensor, - mapping, - mapped_onnx_nodes, - onnx_graph_output_tensors, - from_inbound=True, - ) - if not match: + matched_output_consumer = False + for oc in candidate_consumers: + if oc.id in used_candidate_ids or oc.id in mapped_onnx_nodes: + continue + mapping_copy = copy.deepcopy(mapping) + mapped_onnx_nodes_copy = set(mapped_onnx_nodes) + ok = self._match_node( + output_node, + oc, + node_output_tensor, + mapping_copy, + mapped_onnx_nodes_copy, + onnx_graph_output_tensors, + from_inbound=True, + ) + if ok: + # Adopt mapping changes in-place so caller's reference sees updates + mapping.clear() + for _k, _v in mapping_copy.items(): + mapping[_k] = _v + mapping.inputs = mapping_copy.inputs + mapping.outputs = mapping_copy.outputs + mapping.constants = mapping_copy.constants + mapped_onnx_nodes = mapped_onnx_nodes_copy + used_candidate_ids.add(oc.id) + matched_output_consumer = True + break + if not matched_output_consumer: return False + # For non-pattern-output tensors, disallow extra external consumers. + # All consumers of this ONNX tensor must be accounted for by the pattern. + if node_output_tensor not in self.output_tensors: + for oc in candidate_consumers: + if oc.id not in used_candidate_ids: + return False return True def match_all(self, graph: Graph) -> List[PatternMapping]: diff --git a/tools/onnx-graphsurgeon/tests/test_graph_pattern.py b/tools/onnx-graphsurgeon/tests/test_graph_pattern.py index 2238b21d..efef78b3 100644 --- a/tools/onnx-graphsurgeon/tests/test_graph_pattern.py +++ b/tools/onnx-graphsurgeon/tests/test_graph_pattern.py @@ -19,6 +19,8 @@ import os import onnx +import pytest + from onnx_graphsurgeon import GraphPattern, PatternMapping from onnx_graphsurgeon.importers.onnx_importer import import_onnx from onnx_graphsurgeon.logger import G_LOGGER