diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index d7fa631df9750..97534c1b6a70a 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -153,7 +153,6 @@ option(onnxruntime_ENABLE_MICROSOFT_INTERNAL "Use this option to enable/disable option(onnxruntime_USE_VITISAI "Build with Vitis-AI" OFF) option(onnxruntime_USE_TENSORRT "Build with TensorRT support" OFF) option(onnxruntime_USE_TENSORRT_BUILTIN_PARSER "Use TensorRT builtin parser" OFF) -option(onnxruntime_USE_NV "Build with TensorRT support" OFF) option(onnxruntime_ENABLE_LTO "Enable link time optimization" OFF) option(onnxruntime_CROSS_COMPILING "Cross compiling onnx runtime" OFF) option(onnxruntime_GCOV_COVERAGE "Compile with options necessary to run code coverage" OFF) @@ -349,7 +348,6 @@ else() endif() option(onnxruntime_USE_TENSORRT_INTERFACE "Build ONNXRuntime shared lib which is compatible with TensorRT EP interface" OFF) -option(onnxruntime_USE_NV_INTERFACE "Build ONNXRuntime shared lib which is compatible with NV EP interface" OFF) option(onnxruntime_USE_CUDA_INTERFACE "Build ONNXRuntime shared lib which is compatible with Cuda EP interface" OFF) option(onnxruntime_USE_OPENVINO_INTERFACE "Build ONNXRuntime shared lib which is compatible with OpenVINO EP interface" OFF) option(onnxruntime_USE_VITISAI_INTERFACE "Build ONNXRuntime shared lib which is compatible with Vitis-AI EP interface" OFF) @@ -991,15 +989,6 @@ if (onnxruntime_USE_TENSORRT_INTERFACE AND (NOT onnxruntime_USE_TENSORRT)) list(APPEND ORT_PROVIDER_FLAGS -DUSE_TENSORRT_PROVIDER_INTERFACE=1) endif() -if (onnxruntime_USE_NV) - list(APPEND ORT_PROVIDER_FLAGS -DUSE_NV=1) - list(APPEND ONNXRUNTIME_PROVIDER_NAMES nv_tensorrt_rtx) -endif() - -if (onnxruntime_USE_NV_INTERFACE AND (NOT onnxruntime_USE_NV)) - list(APPEND ORT_PROVIDER_FLAGS -DUSE_NV_PROVIDER_INTERFACE=1) -endif() - if (onnxruntime_USE_RKNPU) list(APPEND ORT_PROVIDER_FLAGS -DUSE_RKNPU=1) list(APPEND ONNXRUNTIME_PROVIDER_NAMES rknpu) diff --git a/cmake/deps.txt b/cmake/deps.txt index b31831f7eef7e..55b1e0c0f752d 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -64,7 +64,6 @@ kleidiai;https://github.com/ARM-software/kleidiai/archive/refs/tags/v1.30.0.tar. # kleidiai-qmx is pinned to a specific commit as there are no tagged releases. When an appropriate tagged release becomes available, # this entry will be updated to use refs/tags/ instead of the raw commit hash. kleidiai-qmx;https://github.com/qualcomm/kleidiai/archive/2f10c9a8d32f81ffeeb6d4885a29cc35d2b0da87.zip;5e855730a2d69057a569f43dd7532db3b2d2a05c -vulkan_headers;https://codeload.github.com/KhronosGroup/Vulkan-Headers/tar.gz/refs/tags/v1.4.344;57bc528ef7c4a3f7bfbb59e64a187e3734bd29d8 # cpp_client_telemetry (1DS SDK) for cross-platform telemetry, including vcpkg builds. # This release fixes Windows Debug listener TLS leaks during DLL unload. cpp_client_telemetry;https://github.com/microsoft/cpp_client_telemetry/archive/refs/tags/v3.10.279.1.zip;e97ee8f91c9348551d35873da33a18b4e50cf0db diff --git a/cmake/onnxruntime_framework.cmake b/cmake/onnxruntime_framework.cmake index c875cac5edc79..79105d3730a55 100644 --- a/cmake/onnxruntime_framework.cmake +++ b/cmake/onnxruntime_framework.cmake @@ -70,7 +70,7 @@ endif() if(onnxruntime_ENABLE_INSTRUMENT) target_compile_definitions(onnxruntime_framework PRIVATE ONNXRUNTIME_ENABLE_INSTRUMENT) endif() -if(onnxruntime_USE_TENSORRT OR onnxruntime_USE_NCCL OR onnxruntime_USE_NV) +if(onnxruntime_USE_TENSORRT OR onnxruntime_USE_NCCL) # TODO: for now, core framework depends on CUDA. It should be moved to TensorRT EP # TODO: provider_bridge_ort.cc should not include nccl.h target_include_directories(onnxruntime_framework PRIVATE ${ONNXRUNTIME_ROOT} ${CMAKE_CURRENT_BINARY_DIR} PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 0270be7d87d46..6066df26bc3e2 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -129,10 +129,6 @@ if (onnxruntime_USE_TENSORRT) include(onnxruntime_providers_tensorrt.cmake) endif() -if (onnxruntime_USE_NV) - include(onnxruntime_providers_nv.cmake) -endif() - if (onnxruntime_USE_VITISAI) include(onnxruntime_providers_vitisai.cmake) endif() diff --git a/cmake/onnxruntime_providers_nv.cmake b/cmake/onnxruntime_providers_nv.cmake deleted file mode 100644 index ffb6fd236cf48..0000000000000 --- a/cmake/onnxruntime_providers_nv.cmake +++ /dev/null @@ -1,203 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# Licensed under the MIT License. - if(onnxruntime_CUDA_HOME) - file(TO_CMAKE_PATH ${onnxruntime_CUDA_HOME} CUDAToolkit_ROOT) - endif() - find_package(CUDAToolkit REQUIRED) - enable_language(CUDA) - if(onnxruntime_DISABLE_CONTRIB_OPS) - message( FATAL_ERROR "To compile TensorRT execution provider contrib ops have to be enabled to dump an engine using com.microsoft:EPContext node." ) - endif() - add_definitions(-DUSE_NV=1) - if (onnxruntime_NV_PLACEHOLDER_BUILDER) - add_definitions(-DORT_NV_PLACEHOLDER_BUILDER) - endif() -if (NOT onnxruntime_USE_TENSORRT_BUILTIN_PARSER) - message(FATAL_ERROR "TensorRT RTX can not be used with the open source parser.") -endif () - set(BUILD_LIBRARY_ONLY 1) - add_definitions("-DONNX_ML=1") - add_definitions("-DONNX_NAMESPACE=onnx") - set(CUDA_INCLUDE_DIRS ${CUDAToolkit_INCLUDE_DIRS}) - set(TENSORRT_RTX_ROOT ${onnxruntime_TENSORRT_RTX_HOME}) - set(OLD_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS}) - set(PROTOBUF_LIBRARY ${PROTOBUF_LIB}) - if (WIN32) - set(OLD_CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS}) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4099 /wd4551 /wd4505 /wd4515 /wd4706 /wd4456 /wd4324 /wd4701 /wd4804 /wd4702 /wd4458 /wd4703") - if (CMAKE_BUILD_TYPE STREQUAL "Debug") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4805") - endif() - set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -include algorithm") - set(DISABLED_WARNINGS_FOR_TRT /wd4456) - endif() - if ( CMAKE_COMPILER_IS_GNUCC ) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter -Wno-missing-field-initializers") - endif() - set(CXX_VERSION_DEFINED TRUE) - - find_path(TENSORRT_RTX_INCLUDE_DIR NvInfer.h - HINTS ${TENSORRT_RTX_ROOT} - PATH_SUFFIXES include) - - - file(READ ${TENSORRT_RTX_INCLUDE_DIR}/NvInferVersion.h NVINFER_VER_CONTENT) - string(REGEX MATCH "define TRT_MAJOR_RTX * +([0-9]+)" NV_TRT_MAJOR_RTX "${NVINFER_VER_CONTENT}") - string(REGEX REPLACE "define TRT_MAJOR_RTX * +([0-9]+)" "\\1" NV_TRT_MAJOR_RTX "${NV_TRT_MAJOR_RTX}") - string(REGEX MATCH "define TRT_MINOR_RTX * +([0-9]+)" NV_TRT_MINOR_RTX "${NVINFER_VER_CONTENT}") - string(REGEX REPLACE "define TRT_MINOR_RTX * +([0-9]+)" "\\1" NV_TRT_MINOR_RTX "${NV_TRT_MINOR_RTX}") - math(EXPR NV_TRT_MAJOR_RTX_INT "${NV_TRT_MAJOR_RTX}") - math(EXPR NV_TRT_MINOR_RTX_INT "${NV_TRT_MINOR_RTX}") - - if (NV_TRT_MAJOR_RTX) - MESSAGE(STATUS "NV_TRT_MAJOR_RTX is ${NV_TRT_MAJOR_RTX}") - else() - MESSAGE(STATUS "Can't find NV_TRT_MAJOR_RTX macro") - endif() - - if (WIN32) - set(TRT_RTX_LIB "tensorrt_rtx_${NV_TRT_MAJOR_RTX}_${NV_TRT_MINOR_RTX}") - set(RTX_PARSER_LIB "tensorrt_onnxparser_rtx_${NV_TRT_MAJOR_RTX}_${NV_TRT_MINOR_RTX}") - endif() - - if (NOT TRT_RTX_LIB) - set(TRT_RTX_LIB "tensorrt_rtx") - endif() - - if (NOT RTX_PARSER_LIB) - set(RTX_PARSER_LIB "tensorrt_onnxparser_rtx") - endif() - - MESSAGE(STATUS "Looking for ${TRT_RTX_LIB}") - - find_library(TENSORRT_LIBRARY_INFER ${TRT_RTX_LIB} - HINTS ${TENSORRT_RTX_ROOT} - PATH_SUFFIXES lib lib64 lib/x64) - - if (NOT TENSORRT_LIBRARY_INFER) - MESSAGE(STATUS "Can't find ${TRT_RTX_LIB}") - endif() - - if (onnxruntime_USE_TENSORRT_BUILTIN_PARSER) - MESSAGE(STATUS "Looking for ${RTX_PARSER_LIB}") - - find_library(TENSORRT_LIBRARY_NVONNXPARSER ${RTX_PARSER_LIB} - HINTS ${TENSORRT_RTX_ROOT} - PATH_SUFFIXES lib lib64 lib/x64) - - if (NOT TENSORRT_LIBRARY_NVONNXPARSER) - MESSAGE(STATUS "Can't find ${RTX_PARSER_LIB}") - endif() - - set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_NVONNXPARSER}) - MESSAGE(STATUS "Find TensorRT libs at ${TENSORRT_LIBRARY}") - else() - set(ONNX_USE_LITE_PROTO ON) - - onnxruntime_fetchcontent_declare( - onnx_tensorrt - URL ${DEP_URL_onnx_tensorrt} - URL_HASH SHA1=${DEP_SHA1_onnx_tensorrt} - EXCLUDE_FROM_ALL - ) - if (NOT CUDA_INCLUDE_DIR) - set(CUDA_INCLUDE_DIR ${CUDAToolkit_INCLUDE_DIRS}) # onnx-tensorrt repo needs this variable to build - endif() - # The onnx_tensorrt repo contains a test program, getSupportedAPITest, which doesn't support Windows. It uses - # unistd.h. So we must exclude it from our build. onnxruntime_fetchcontent_makeavailable is for the purpose. - onnxruntime_fetchcontent_makeavailable(onnx_tensorrt) - set(CMAKE_CXX_FLAGS ${OLD_CMAKE_CXX_FLAGS}) - if ( CMAKE_COMPILER_IS_GNUCC ) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-parameter") - endif() - if (WIN32) - set(CMAKE_CUDA_FLAGS ${OLD_CMAKE_CUDA_FLAGS}) - unset(PROTOBUF_LIBRARY) - unset(OLD_CMAKE_CXX_FLAGS) - unset(OLD_CMAKE_CUDA_FLAGS) - set_target_properties(${RTX_PARSER_LIB} PROPERTIES LINK_FLAGS "/ignore:4199") - target_compile_options(nvonnxparser_static PRIVATE /FIio.h /wd4100) - target_compile_options(${RTX_PARSER_LIB} PRIVATE /FIio.h /wd4100) - endif() - # Static libraries are just nvonnxparser_static on all platforms - set(onnxparser_link_libs nvonnxparser_static) - set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER}) - MESSAGE(STATUS "Find TensorRT libs at ${TENSORRT_LIBRARY}") - endif() - - # ${TENSORRT_LIBRARY} is empty if we link nvonnxparser_static. - # nvonnxparser_static is linked against tensorrt libraries in onnx-tensorrt - # See https://github.com/onnx/onnx-tensorrt/blob/8af13d1b106f58df1e98945a5e7c851ddb5f0791/CMakeLists.txt#L121 - # However, starting from TRT 10 GA, nvonnxparser_static doesn't link against tensorrt libraries. - # Therefore, the above code finds ${TENSORRT_LIBRARY_INFER} - set(trt_link_libs ${CMAKE_DL_LIBS} ${TENSORRT_LIBRARY}) - - file(GLOB_RECURSE onnxruntime_providers_nv_tensorrt_rtx_cc_srcs CONFIGURE_DEPENDS - "${ONNXRUNTIME_ROOT}/core/providers/nv_tensorrt_rtx/*.h" - "${ONNXRUNTIME_ROOT}/core/providers/nv_tensorrt_rtx/*.cc" - "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.h" - "${ONNXRUNTIME_ROOT}/core/providers/shared_library/*.cc" - "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_stream_handle.h" - "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_stream_handle.cc" - "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.h" - "${ONNXRUNTIME_ROOT}/core/providers/cuda/cudnn_loader.cc" - "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_graph.h" - "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_graph.cc" - ) - - source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_nv_tensorrt_rtx_cc_srcs}) - onnxruntime_add_shared_library_module(onnxruntime_providers_nv_tensorrt_rtx ${onnxruntime_providers_nv_tensorrt_rtx_cc_srcs}) - onnxruntime_add_include_to_target(onnxruntime_providers_nv_tensorrt_rtx onnxruntime_common) - target_link_libraries(onnxruntime_providers_nv_tensorrt_rtx PRIVATE Eigen3::Eigen onnx flatbuffers::flatbuffers Boost::mp11 safeint_interface Eigen3::Eigen) - add_dependencies(onnxruntime_providers_nv_tensorrt_rtx onnxruntime_providers_shared ${onnxruntime_EXTERNAL_DEPENDENCIES}) - # Make NVML linkage optional: only link CUDA::nvml if the target exists. - set(_nvml_lib "") - if (TARGET CUDA::nvml) - set(_nvml_lib CUDA::nvml) - endif() - if (onnxruntime_USE_TENSORRT_BUILTIN_PARSER) - target_link_libraries(onnxruntime_providers_nv_tensorrt_rtx PRIVATE ${trt_link_libs} ${ONNXRUNTIME_PROVIDERS_SHARED} ${PROTOBUF_LIB} flatbuffers::flatbuffers Boost::mp11 safeint_interface ${ABSEIL_LIBS} ${_nvml_lib} PUBLIC CUDA::cudart CUDA::cuda_driver) - else() - target_link_libraries(onnxruntime_providers_nv_tensorrt_rtx PRIVATE ${onnxparser_link_libs} ${trt_link_libs} ${ONNXRUNTIME_PROVIDERS_SHARED} ${PROTOBUF_LIB} flatbuffers::flatbuffers ${ABSEIL_LIBS} ${_nvml_lib} PUBLIC CUDA::cudart CUDA::cuda_driver) - endif() - target_include_directories(onnxruntime_providers_nv_tensorrt_rtx PRIVATE ${ONNXRUNTIME_ROOT} ${CMAKE_CURRENT_BINARY_DIR} ${TENSORRT_RTX_INCLUDE_DIR} ${onnx_tensorrt_SOURCE_DIR} - PUBLIC ${CUDAToolkit_INCLUDE_DIRS}) - - # ${CMAKE_CURRENT_BINARY_DIR} is so that #include "onnxruntime_config.h" inside tensor_shape.h is found - set_target_properties(onnxruntime_providers_nv_tensorrt_rtx PROPERTIES LINKER_LANGUAGE CUDA) - set_target_properties(onnxruntime_providers_nv_tensorrt_rtx PROPERTIES FOLDER "ONNXRuntime") - target_compile_definitions(onnxruntime_providers_nv_tensorrt_rtx PRIVATE ONNXIFI_BUILD_LIBRARY=1) - if (TARGET CUDA::nvml) - target_compile_definitions(onnxruntime_providers_nv_tensorrt_rtx PRIVATE HAVE_NVML=1) - endif() - target_compile_options(onnxruntime_providers_nv_tensorrt_rtx PRIVATE ${DISABLED_WARNINGS_FOR_TRT}) - if (WIN32) - target_compile_options(onnxruntime_providers_nv_tensorrt_rtx INTERFACE /wd4456) - endif() - # set CUDA_MINIMAL as default for NV provider since we do not have fallback to CUDA - target_compile_definitions(onnxruntime_providers_nv_tensorrt_rtx PRIVATE USE_CUDA_MINIMAL=1) - - # Needed for the provider interface, as it includes training headers when training is enabled - if (onnxruntime_ENABLE_TRAINING_OPS) - target_include_directories(onnxruntime_providers_nv_tensorrt_rtx PRIVATE ${ORTTRAINING_ROOT}) - if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) - onnxruntime_add_include_to_target(onnxruntime_providers_nv_tensorrt_rtx Python::Module) - endif() - endif() - - if(APPLE) - set_property(TARGET onnxruntime_providers_nv_tensorrt_rtx APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker -exported_symbols_list ${ONNXRUNTIME_ROOT}/core/providers/nv_tensorrt_rtx/exported_symbols.lst") - elseif(UNIX) - set_property(TARGET onnxruntime_providers_nv_tensorrt_rtx APPEND_STRING PROPERTY COMPILE_FLAGS "-Wno-deprecated-declarations") - set_property(TARGET onnxruntime_providers_nv_tensorrt_rtx APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker --version-script=${ONNXRUNTIME_ROOT}/core/providers/nv_tensorrt_rtx/version_script.lds -Xlinker --gc-sections") - elseif(WIN32) - set_property(TARGET onnxruntime_providers_nv_tensorrt_rtx APPEND_STRING PROPERTY LINK_FLAGS "-DEF:${ONNXRUNTIME_ROOT}/core/providers/nv_tensorrt_rtx/symbols.def") - else() - message(FATAL_ERROR "onnxruntime_providers_nv_tensorrt_rtx unknown platform, need to specify shared library exports for it") - endif() - - install(TARGETS onnxruntime_providers_nv_tensorrt_rtx - ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} - RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index fb976e6ffbc92..79ca697b1818c 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -1097,36 +1097,6 @@ if (onnxruntime_USE_TENSORRT) ) endif() -if (onnxruntime_USE_NV) - if (WIN32 OR ${CMAKE_SYSTEM_NAME} STREQUAL "Linux") - file(GLOB NV_LIB_FILES LIST_DIRECTORIES false "${TENSORRT_RTX_ROOT}/lib/tensorrt_*.dll" - "${TENSORRT_RTX_ROOT}/bin/tensorrt_*.dll" - "${TENSORRT_RTX_ROOT}/lib/libtensorrt_*.so" - "${TENSORRT_RTX_ROOT}/bin/libtensorrt_*.so") - message(STATUS "NV lib files: " ${NV_LIB_FILES}) - endif() - add_custom_command( - TARGET onnxruntime_pybind11_state POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy - ${NV_LIB_FILES} - $ - $ - $/onnxruntime/capi/ - ) - if (EXISTS "${TENSORRT_RTX_ROOT}/doc/LICENSE.txt") - add_custom_command( - TARGET onnxruntime_pybind11_state POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy "${TENSORRT_RTX_ROOT}/doc/LICENSE.txt" $/onnxruntime/TRT_RTX_LICENSE.txt - ) - endif() - if (EXISTS "${TENSORRT_RTX_ROOT}/doc/Acknowledgements.txt") - add_custom_command( - TARGET onnxruntime_pybind11_state POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy "${TENSORRT_RTX_ROOT}/doc/Acknowledgements.txt" $/onnxruntime/TRT_RTX_Acknowledgements.txt - ) - endif() -endif() - if (onnxruntime_USE_MIGRAPHX) add_custom_command( TARGET onnxruntime_pybind11_state POST_BUILD diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 4e630b5bb45d2..7af7b0399e36b 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -122,7 +122,7 @@ function(AddTest) Threads::Threads) target_compile_definitions(${_UT_TARGET} PRIVATE -DUSE_ONNXRUNTIME_DLL) else() - if(onnxruntime_USE_CUDA OR onnxruntime_USE_NV) + if(onnxruntime_USE_CUDA) #XXX: we should not need to do this. onnxruntime_test_all.exe should not have direct dependency on CUDA DLLs, # otherwise it will impact when CUDA DLLs can be unloaded. target_link_libraries(${_UT_TARGET} PRIVATE CUDA::cudart) @@ -151,11 +151,6 @@ function(AddTest) # used for instantiating placeholder TRT builder to mitigate TRT library load/unload overhead target_include_directories(${_UT_TARGET} PRIVATE ${TENSORRT_INCLUDE_DIR}) endif() - if (onnxruntime_USE_NV) - # used for instantiating placeholder TRT builder to mitigate TRT library load/unload overhead - target_include_directories(${_UT_TARGET} PRIVATE ${NV_INCLUDE_DIR} ${CUDAToolkit_INCLUDE_DIRS}) - endif() - if(MSVC) target_compile_options(${_UT_TARGET} PRIVATE "$<$:SHELL:--compiler-options /utf-8>" "$<$>:/utf-8>") @@ -780,25 +775,6 @@ if(onnxruntime_USE_TENSORRT) list(APPEND onnxruntime_test_providers_libs ${TENSORRT_LIBRARY_INFER}) endif() -if(onnxruntime_USE_NV) - # If an external project (e.g. dawn from Webgpu EP has already added a Vulkan::Headers target we shouldn't try to import another version of the Vulkan headers) - if (NOT TARGET Vulkan::Headers) - onnxruntime_fetchcontent_declare( - vulkan_headers - URL ${DEP_URL_vulkan_headers} - URL_HASH SHA1=${DEP_SHA1_vulkan_headers} - EXCLUDE_FROM_ALL - ) - onnxruntime_fetchcontent_makeavailable(vulkan_headers) - endif() - list(APPEND onnxruntime_test_framework_src_patterns ${TEST_SRC_DIR}/providers/nv_tensorrt_rtx/*) - list(APPEND onnxruntime_test_framework_src_patterns "${ONNXRUNTIME_ROOT}/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h") - list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_nv_tensorrt_rtx onnxruntime_providers_shared) - list(APPEND onnxruntime_test_providers_libs ${TENSORRT_LIBRARY_INFER}) - list(APPEND onnxruntime_test_providers_libs Vulkan::Headers) -endif() - - if(onnxruntime_USE_MIGRAPHX) list(APPEND onnxruntime_test_framework_src_patterns ${TEST_SRC_DIR}/providers/migraphx/*) endif() @@ -1243,7 +1219,7 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR IOS) endif() set(test_all_args) -if (onnxruntime_USE_TENSORRT OR onnxruntime_USE_NV) +if (onnxruntime_USE_TENSORRT) # TRT EP CI takes much longer time when updating to TRT 8.2 # So, we only run trt ep and exclude other eps to reduce CI test time. # @@ -2263,7 +2239,7 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) list(APPEND onnxruntime_perf_test_libs onnxruntime_graph onnxruntime_session ${onnxruntime_providers_target} onnxruntime_framework onnxruntime_util onnxruntime_mlas onnxruntime_optimizer onnxruntime_flatbuffers iconv re2 gtest absl_failure_signal_handler absl_examine_stack absl_flags_parse absl_flags_usage absl_flags_usage_internal) endif() target_link_libraries(onnxruntime_perf_test PRIVATE ${onnxruntime_perf_test_libs} Threads::Threads) - if (onnxruntime_USE_CUDA OR onnxruntime_USE_NV OR onnxruntime_USE_TENSORRT) + if (onnxruntime_USE_CUDA OR onnxruntime_USE_TENSORRT) target_link_libraries(onnxruntime_perf_test PRIVATE CUDA::cudart) endif() if(WIN32) @@ -2272,7 +2248,7 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP) else() target_link_libraries(onnxruntime_perf_test PRIVATE onnx_test_runner_common absl::flags absl::flags_parse ${onnx_test_libs}) # When onnxruntime_BUILD_SHARED_LIB is OFF (the plugin build path), perf test was missing CUDA include directories and CUDA::cudart linkage. - if (onnxruntime_USE_CUDA OR onnxruntime_USE_NV OR onnxruntime_USE_TENSORRT) + if (onnxruntime_USE_CUDA OR onnxruntime_USE_TENSORRT) target_include_directories(onnxruntime_perf_test PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_link_libraries(onnxruntime_perf_test PRIVATE CUDA::cudart) endif() @@ -2469,9 +2445,6 @@ endif() if (onnxruntime_USE_TENSORRT) list(APPEND onnxruntime_shared_lib_test_LIBS ${TENSORRT_LIBRARY_INFER}) endif() - if (onnxruntime_USE_NV) - list(APPEND onnxruntime_shared_lib_test_LIBS ${TENSORRT_LIBRARY_INFER} CUDA::cudart) - endif() if (onnxruntime_USE_DML) list(APPEND onnxruntime_shared_lib_test_LIBS d3d12.lib) endif() @@ -2496,11 +2469,6 @@ endif() target_include_directories(onnxruntime_shared_lib_test PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) target_sources(onnxruntime_shared_lib_test PRIVATE ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/cuda_ops.cu) endif() - if (onnxruntime_USE_NV) - target_include_directories(onnxruntime_shared_lib_test PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) - endif() - - if (CMAKE_SYSTEM_NAME STREQUAL "Android") target_sources(onnxruntime_shared_lib_test PRIVATE "${ONNXRUNTIME_ROOT}/core/platform/android/cxa_demangle.cc" @@ -2804,9 +2772,6 @@ if (NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" AND NOT onnxruntime_CUDA_MINIMAL if (onnxruntime_USE_TENSORRT) list(APPEND onnxruntime_customopregistration_test_LIBS ${TENSORRT_LIBRARY_INFER}) endif() - if (onnxruntime_USE_NV) - list(APPEND onnxruntime_customopregistration_test_LIBS ${TENSORRT_LIBRARY_INFER}) - endif() if (CMAKE_SYSTEM_NAME MATCHES "AIX") list(APPEND onnxruntime_customopregistration_test_LIBS onnxruntime_graph onnxruntime_session ${onnxruntime_providers_target} onnxruntime_lora onnxruntime_framework onnxruntime_util onnxruntime_mlas onnxruntime_optimizer onnxruntime_flatbuffers iconv re2 ${PROTOBUF_LIB} onnx onnx_proto) endif() @@ -3164,7 +3129,7 @@ if (onnxruntime_BUILD_SHARED_LIB AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" list(APPEND onnxruntime_ep_graph_test_LIBS ${CMAKE_DL_LIBS}) endif() - if (onnxruntime_USE_TENSORRT OR onnxruntime_USE_NV) + if (onnxruntime_USE_TENSORRT) # Need this because unittest_main_src defines a global nvinfer1::IBuilder variable. list(APPEND onnxruntime_ep_graph_test_LIBS ${TENSORRT_LIBRARY_INFER}) endif() diff --git a/docs/CUDA_cuDNN_Optional_Design.md b/docs/CUDA_cuDNN_Optional_Design.md index 1d761d31b72b0..d0a6c61388f81 100644 --- a/docs/CUDA_cuDNN_Optional_Design.md +++ b/docs/CUDA_cuDNN_Optional_Design.md @@ -3,10 +3,10 @@ Status: Draft / Proposal Owner: (CUDA EP) Scope: `onnxruntime/core/providers/cuda` (main static CUDA EP), CUDA Plugin EP -(`BUILD_CUDA_EP_AS_PLUGIN`), and the CUDA unit tests. **Out of scope:** TensorRT EP and -NV‑TensorRT‑RTX EP (they create and own their own cuDNN handles and inherently depend on -cuDNN/TensorRT). The existing build-time `USE_CUDA_MINIMAL` path is also out of scope; it is -used by TensorRT / NV‑TensorRT‑RTX integration and should remain available. +(`BUILD_CUDA_EP_AS_PLUGIN`), and the CUDA unit tests. **Out of scope:** TensorRT EP (it creates +and owns its own cuDNN handles and inherently depends on cuDNN/TensorRT). The existing +build-time `USE_CUDA_MINIMAL` path is also out of scope; it is used by TensorRT integration +and should remain available. --- @@ -82,7 +82,7 @@ This is delivered in **three phases**: - `CUDNN_FE_CALL` / `CUDNN_FE_RETURN_IF_ERROR` for the frontend (`cudnn_fe_call.*`). - `CudaErrString` calls `cudnnGetErrorString` (`cuda_call.cc`). - All of the above are already gated by `#ifndef USE_CUDA_MINIMAL` in shared CUDA code. That - build‑time path is used by TensorRT / NV‑TensorRT‑RTX related builds and is a useful + build‑time path is used by TensorRT-related builds and is a useful inventory of cuDNN‑touching code, but it is *not* the CUDA EP runtime behavior we want. ### 2.4 Operators / components that depend on cuDNN @@ -354,9 +354,8 @@ Python wheel packaging is unchanged by this design: cuDNN DLLs are not packed in today, so Phase 1 is not introducing a new "CUDA-minimal" wheel flavor. The runtime loader simply makes the existing package tolerant of environments where cuDNN is absent. -TensorRT / NV‑RTX EPs are untouched and continue to link cuDNN as before. (If both a TRT EP -and the shimmed CUDA EP are in the same process, symbol collision must be avoided — see -§8 Risks.) +The TensorRT EP is untouched and continues to link cuDNN as before. (If both a TRT EP and the +shimmed CUDA EP are in the same process, symbol collision must be avoided — see §8 Risks.) ### 3.6 Applying the same pattern to cuFFT @@ -452,8 +451,8 @@ present, behavior is byte‑for‑byte identical to today. `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING`. This is required so frontend graph code uses `dlsym` / `GetProcAddress` on `cudnn_frontend::cudnn_dlhandle` instead of creating link-time imports for `cudnnBackend*` symbols. - - Keep `USE_CUDA_MINIMAL` working. It is used by TensorRT / NV‑TensorRT‑RTX related - builds and is not replaced by the optional-cuDNN runtime shim. + - Keep `USE_CUDA_MINIMAL` working. It is used by TensorRT-related builds and is not + replaced by the optional-cuDNN runtime shim. Current CMake anchor points: @@ -741,5 +740,5 @@ kernels; CUTLASS (already vendored) for conv/GEMM‑shaped work. from application code before creating the session. - No new "CUDA-minimal" wheel is required for Phase 1. cuDNN DLLs are not packed in the wheel today. -- Keep the existing `USE_CUDA_MINIMAL` build-time path. It is used by RTX/TensorRT-related EP +- Keep the existing `USE_CUDA_MINIMAL` build-time path. It is used by TensorRT-related EP builds and is not replaced by the CUDA EP runtime shim. diff --git a/go/onnxruntime/onnxruntime_c_api.h b/go/onnxruntime/onnxruntime_c_api.h index bfd4136c434d2..4d9c7dd7a2710 100644 --- a/go/onnxruntime/onnxruntime_c_api.h +++ b/go/onnxruntime/onnxruntime_c_api.h @@ -308,6 +308,9 @@ ORT_RUNTIME_CLASS(ThreadingOptions); ORT_RUNTIME_CLASS(ArenaCfg); ORT_RUNTIME_CLASS(PrepackedWeightsContainer); ORT_RUNTIME_CLASS(TensorRTProviderOptionsV2); +/** \deprecated The provider-bridge NvTensorRTRTXExecutionProvider has been removed. Retained for source + * compatibility only. Use the NVIDIA TensorRT RTX plugin EP via OrtApi::RegisterExecutionProviderLibrary instead. + */ ORT_RUNTIME_CLASS(NvTensorRtRtxProviderOptions); ORT_RUNTIME_CLASS(CUDAProviderOptionsV2); ORT_RUNTIME_CLASS(CANNProviderOptions); @@ -1164,7 +1167,7 @@ typedef struct OrtGraphicsInteropConfig { * * Can be used for future extensibility without changing the struct layout. * For example, D3D12 fence sharing flags or provider-specific options like - * onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data + * onnxruntime::tensorrt_rtx::provider_option_names::kExternalComputeQueueDataParamNV_data * for Vulkan interop for the NvTensorRTRTX provider. */ const OrtKeyValuePairs* additional_options; diff --git a/include/onnxruntime/core/graph/constants.h b/include/onnxruntime/core/graph/constants.h index c63c14ad6a32f..72b69ff276b50 100644 --- a/include/onnxruntime/core/graph/constants.h +++ b/include/onnxruntime/core/graph/constants.h @@ -37,6 +37,8 @@ constexpr const char* kDnnlExecutionProvider = "DnnlExecutionProvider"; constexpr const char* kOpenVINOExecutionProvider = "OpenVINOExecutionProvider"; constexpr const char* kVitisAIExecutionProvider = "VitisAIExecutionProvider"; constexpr const char* kTensorrtExecutionProvider = "TensorrtExecutionProvider"; +// Deprecated: the provider-bridge NvTensorRTRTXExecutionProvider has been removed. Retained for source compatibility +// and to identify the NVIDIA TensorRT RTX plugin EP, which registers under the same name. constexpr const char* kNvTensorRTRTXExecutionProvider = "NvTensorRTRTXExecutionProvider"; constexpr const char* kNnapiExecutionProvider = "NnapiExecutionProvider"; constexpr const char* kQnnExecutionProvider = "QNNExecutionProvider"; diff --git a/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h b/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h deleted file mode 100644 index 283c6be0f082b..0000000000000 --- a/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h +++ /dev/null @@ -1,49 +0,0 @@ -#pragma once - -/** - - * @namespace onnxruntime::nv::provider_option_names - * @details The `provider_option_names` namespace contains the following constants: - * - `kDeviceId`: Specifies the GPU device ID to use. - * - `kHasUserComputeStream`: Indicates whether a user-provided compute stream is used. - * - `kUserComputeStream`: Specifies the user-provided compute stream. - * - `kUserAuxStreamArray`: Specifies the user-provided aux stream. - * - `kMaxWorkspaceSize`: Sets the maximum workspace size for GPU memory allocation. - * - 'kMaxSharedMemSize': Sets the maximum amount of shared memory that TensorRT kernels are allowed to use - * - `kLengthAuxStreamArray`: Specifies the length/size of the auxiliary streams array (kUserAuxStreamArray). Also sets the maximum number of auxiliary streams for TensorRT execution. - * - `kDumpSubgraphs`: Enables or disables dumping of subgraphs for debugging. - * - `kDetailedBuildLog`: Enables or disables detailed build logs for debugging. - * - `kProfilesMinShapes`: Specifies the minimum shapes for profiling. - * - `kProfilesMaxShapes`: Specifies the maximum shapes for profiling. - * - `kProfilesOptShapes`: Specifies the optimal shapes for profiling. - * - `kCudaGraphEnable`: Enables or disables CUDA graph optimizations. - * - `kONNXBytestream`: Specifies the ONNX model as a bytestream. - * - `kONNXBytestreamSize`: Specifies the size of the ONNX bytestream. - */ -namespace onnxruntime { -namespace nv { -namespace provider_option_names { -constexpr const char* kDeviceId = "device_id"; -constexpr const char* kHasUserComputeStream = "has_user_compute_stream"; -constexpr const char* kUserComputeStream = "user_compute_stream"; -constexpr const char* kUserAuxStreamArray = "user_aux_stream_array"; -constexpr const char* kMaxWorkspaceSize = "nv_max_workspace_size"; -constexpr const char* kMaxSharedMemSize = "nv_max_shared_mem_size"; -constexpr const char* kLengthAuxStreamArray = "nv_length_aux_stream_array"; -constexpr const char* kDumpSubgraphs = "nv_dump_subgraphs"; -constexpr const char* kDetailedBuildLog = "nv_detailed_build_log"; -constexpr const char* kProfilesMinShapes = "nv_profile_min_shapes"; -constexpr const char* kProfilesMaxShapes = "nv_profile_max_shapes"; -constexpr const char* kProfilesOptShapes = "nv_profile_opt_shapes"; -constexpr const char* kCudaGraphEnable = "enable_cuda_graph"; -constexpr const char* kMultiProfileEnable = "nv_multi_profile_enable"; -constexpr const char* kUseExternalDataInitializer = "nv_use_external_data_initializer"; -constexpr const char* kRuntimeCacheFile = "nv_runtime_cache_path"; -constexpr const char* kExternalComputeQueueDataParamNV_data = "VkExternalComputeQueueDataParamsNV_data"; - -} // namespace provider_option_names -namespace run_option_names { -constexpr const char* kProfileIndex = "nv_profile_index"; -} -} // namespace nv -} // namespace onnxruntime diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index c38ecd3a6794f..dc28672afb522 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -313,6 +313,9 @@ ORT_RUNTIME_CLASS(ThreadingOptions); ORT_RUNTIME_CLASS(ArenaCfg); ORT_RUNTIME_CLASS(PrepackedWeightsContainer); ORT_RUNTIME_CLASS(TensorRTProviderOptionsV2); +/** \deprecated The provider-bridge NvTensorRTRTXExecutionProvider has been removed. Retained for source + * compatibility only. Use the NVIDIA TensorRT RTX plugin EP via OrtApi::RegisterExecutionProviderLibrary instead. + */ ORT_RUNTIME_CLASS(NvTensorRtRtxProviderOptions); ORT_RUNTIME_CLASS(CUDAProviderOptionsV2); ORT_RUNTIME_CLASS(CANNProviderOptions); @@ -1244,7 +1247,7 @@ typedef struct OrtGraphicsInteropConfig { * * Can be used for future extensibility without changing the struct layout. * For example, D3D12 fence sharing flags or provider-specific options like - * onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data + * onnxruntime::tensorrt_rtx::provider_option_names::kExternalComputeQueueDataParamNV_data * for Vulkan interop for the NvTensorRTRTX provider. */ const OrtKeyValuePairs* additional_options; diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 1251fae691b7c..1029498a9b450 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -697,7 +697,7 @@ static const char* const kOrtSessionOptionsQDQMatMulNBitsBlockSize = "session.qd // Enable the DQ->MatMulNBits fusion graph transformer. // "0": disabled (default). "1": enabled. -// This is typically set automatically by InferenceSession when the NvTensorRTRTX EP is registered. +// This is typically enabled automatically when the NvTensorRTRTX plugin EP is registered. static const char* const kOrtSessionOptionsEnableDQMatMulNBitsFusion = "session.enable_dq_matmulnbits_fusion"; // THIS OPTION IS NOT A REGULAR SESSION OPTION SINCE IT CAN BE MODIFIED AT ANY TIME diff --git a/java/build.gradle b/java/build.gradle index 6d4d0835af0c3..38d53c84981fa 100644 --- a/java/build.gradle +++ b/java/build.gradle @@ -195,7 +195,6 @@ test { 'USE_DNNL', 'USE_MIGRAPHX', 'USE_NNAPI', - 'USE_NV', 'USE_OPENVINO', 'USE_QNN', 'USE_RKNPU', diff --git a/onnxruntime/core/providers/cuda/shared_inc/cuda_call.h b/onnxruntime/core/providers/cuda/shared_inc/cuda_call.h index 74300ee76a631..6a56d8f0342e6 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/cuda_call.h +++ b/onnxruntime/core/providers/cuda/shared_inc/cuda_call.h @@ -16,7 +16,6 @@ std::conditional_t CudaCall( const char* file, const int line); #define CUDA_CALL(expr) (::onnxruntime::CudaCall((expr), #expr, "CUDA", cudaSuccess, "", __FILE__, __LINE__)) -#define CU_CALL(expr) (::onnxruntime::CudaCall((expr), #expr, "CUDA", CUDA_SUCCESS, "", __FILE__, __LINE__)) #define CUBLAS_CALL(expr) (::onnxruntime::CudaCall((expr), #expr, "CUBLAS", CUBLAS_STATUS_SUCCESS, "", __FILE__, __LINE__)) #define CURAND_CALL(expr) (::onnxruntime::CudaCall((expr), #expr, "CURAND", CURAND_STATUS_SUCCESS, "", __FILE__, __LINE__)) @@ -26,7 +25,6 @@ std::conditional_t CudaCall( #define CUFFT_CALL(expr) (::onnxruntime::CudaCall((expr), #expr, "CUFFT", CUFFT_SUCCESS, "", __FILE__, __LINE__)) #define CUDA_CALL_THROW(expr) (::onnxruntime::CudaCall((expr), #expr, "CUDA", cudaSuccess, "", __FILE__, __LINE__)) -#define CU_CALL_THROW(expr) (::onnxruntime::CudaCall((expr), #expr, "CUDA", CUDA_SUCCESS, "", __FILE__, __LINE__)) #define CUBLAS_CALL_THROW(expr) (::onnxruntime::CudaCall((expr), #expr, "CUBLAS", CUBLAS_STATUS_SUCCESS, "", __FILE__, __LINE__)) #define CURAND_CALL_THROW(expr) (::onnxruntime::CudaCall((expr), #expr, "CURAND", CURAND_STATUS_SUCCESS, "", __FILE__, __LINE__)) diff --git a/onnxruntime/core/providers/get_execution_providers.cc b/onnxruntime/core/providers/get_execution_providers.cc index 59bfd17a8451a..8a2559dbc1b51 100644 --- a/onnxruntime/core/providers/get_execution_providers.cc +++ b/onnxruntime/core/providers/get_execution_providers.cc @@ -20,14 +20,6 @@ struct ProviderInfo { // kCpuExecutionProvider should always be last constexpr ProviderInfo kProvidersInPriorityOrder[] = { - { - kNvTensorRTRTXExecutionProvider, -#ifdef USE_NV - true, -#else - false, -#endif - }, { kTensorrtExecutionProvider, #ifdef USE_TENSORRT diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/exported_symbols.lst b/onnxruntime/core/providers/nv_tensorrt_rtx/exported_symbols.lst deleted file mode 100644 index f4c41412594af..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/exported_symbols.lst +++ /dev/null @@ -1 +0,0 @@ -_GetProvider diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_allocator.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_allocator.cc deleted file mode 100644 index f8a722d782653..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_allocator.cc +++ /dev/null @@ -1,101 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "nv_allocator.h" -#include "nv_includes.h" -#include "core/providers/cuda/shared_inc/cuda_call.h" -namespace onnxruntime { - -void CUDAAllocator::CheckDevice(bool throw_when_fail) const { -#ifndef NDEBUG - // check device to match at debug build - // if it's expected to change, call cudaSetDevice instead of the check - int current_device; - auto cuda_err = cudaGetDevice(¤t_device); - if (cuda_err == cudaSuccess) { - ORT_ENFORCE(current_device == Info().device.Id()); - } else if (throw_when_fail) { - CUDA_CALL_THROW(cuda_err); - } -#else - ORT_UNUSED_PARAMETER(throw_when_fail); -#endif -} - -void CUDAAllocator::SetDevice(bool throw_when_fail) const { - int current_device; - auto cuda_err = cudaGetDevice(¤t_device); - if (cuda_err == cudaSuccess) { - int allocator_device_id = Info().device.Id(); - if (current_device != allocator_device_id) { - cuda_err = cudaSetDevice(allocator_device_id); - } - } - - if (cuda_err != cudaSuccess && throw_when_fail) { - CUDA_CALL_THROW(cuda_err); - } -} - -void* CUDAAllocator::Alloc(size_t size) { - SetDevice(true); - CheckDevice(true); - void* p = nullptr; - if (size > 0) { - // BFCArena was updated recently to handle the exception and adjust the request size - CUDA_CALL_THROW(cudaMalloc((void**)&p, size)); - } - return p; -} - -void CUDAAllocator::Free(void* p) { - SetDevice(false); - CheckDevice(false); // ignore CUDA failure when free - cudaFree(p); // do not throw error since it's OK for cudaFree to fail during shutdown -} - -void* CUDAExternalAllocator::Alloc(size_t size) { - void* p = nullptr; - if (size > 0) { - p = alloc_(size); - - // review(codemzs): ORT_ENFORCE does not seem appropriate. - ORT_ENFORCE(p != nullptr); - } - - return p; -} - -void CUDAExternalAllocator::Free(void* p) { - free_(p); - std::lock_guard lock(lock_); - auto it = reserved_.find(p); - if (it != reserved_.end()) { - reserved_.erase(it); - if (empty_cache_) empty_cache_(); - } -} - -void* CUDAExternalAllocator::Reserve(size_t size) { - void* p = Alloc(size); - if (!p) return nullptr; - std::lock_guard lock(lock_); - ORT_ENFORCE(reserved_.find(p) == reserved_.end()); - reserved_.insert(p); - return p; -} - -void* CUDAPinnedAllocator::Alloc(size_t size) { - void* p = nullptr; - if (size > 0) { - CUDA_CALL_THROW(cudaMallocHost((void**)&p, size)); - } - return p; -} - -void CUDAPinnedAllocator::Free(void* p) { - CUDA_CALL_THROW(cudaFreeHost(p)); -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_allocator.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_allocator.h deleted file mode 100644 index 3283aea7c33a2..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_allocator.h +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "core/common/inlined_containers.h" -#include "core/framework/allocator.h" -#include - -namespace onnxruntime { - -class CUDAAllocator : public IAllocator { - public: - CUDAAllocator(OrtDevice::DeviceId device_id, const char* name) - : IAllocator(OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, - OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, - device_id), - OrtMemTypeDefault)) {} - void* Alloc(size_t size) override; - void Free(void* p) override; - - private: - void CheckDevice(bool throw_when_fail) const; - void SetDevice(bool throw_when_fail) const; -}; - -class CUDAExternalAllocator : public CUDAAllocator { - typedef void* (*ExternalAlloc)(size_t size); - typedef void (*ExternalFree)(void* p); - typedef void (*ExternalEmptyCache)(); - - public: - CUDAExternalAllocator(OrtDevice::DeviceId device_id, const char* name, void* alloc, void* free, void* empty_cache) - : CUDAAllocator(device_id, name) { - alloc_ = reinterpret_cast(alloc); - free_ = reinterpret_cast(free); - empty_cache_ = reinterpret_cast(empty_cache); - } - - void* Alloc(size_t size) override; - void Free(void* p) override; - void* Reserve(size_t size) override; - - private: - mutable std::mutex lock_; - ExternalAlloc alloc_; - ExternalFree free_; - ExternalEmptyCache empty_cache_; - InlinedHashSet reserved_; -}; - -// TODO: add a default constructor -class CUDAPinnedAllocator : public IAllocator { - public: - CUDAPinnedAllocator(OrtDevice::DeviceId device_id, const char* name) - : IAllocator( - OrtMemoryInfo(name, OrtAllocatorType::OrtDeviceAllocator, - OrtDevice(OrtDevice::GPU, OrtDevice::MemType::HOST_ACCESSIBLE, OrtDevice::VendorIds::NVIDIA, - device_id), - OrtMemTypeCPUOutput)) {} - - void* Alloc(size_t size) override; - void Free(void* p) override; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_cuda_call.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_cuda_call.cc deleted file mode 100644 index 7d92d381617cf..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_cuda_call.cc +++ /dev/null @@ -1,142 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/shared_library/provider_api.h" -#include "core/providers/cuda/shared_inc/cuda_call.h" -#include - -#ifdef _WIN32 -#else // POSIX -#include -#include -#endif - -namespace onnxruntime { - -using namespace common; - -template -const char* CudaErrString(ERRTYPE) { - ORT_NOT_IMPLEMENTED(); -} - -#define CASE_ENUM_TO_STR(x) \ - case x: \ - return #x - -template <> -const char* CudaErrString(cudaError_t x) { - cudaDeviceSynchronize(); - return cudaGetErrorString(x); -} - -template <> -const char* CudaErrString(CUresult x) { - const char* errorStr = NULL; - cuGetErrorString(x, &errorStr); - return errorStr; -} - -#ifndef USE_CUDA_MINIMAL -template <> -const char* CudaErrString(cublasStatus_t e) { - cudaDeviceSynchronize(); - const char* status_string = cublasGetStatusString(e); - return status_string != nullptr ? status_string : "Unknown cuBLAS error status"; -} - -template <> -const char* CudaErrString(curandStatus) { - cudaDeviceSynchronize(); - return "(see curand.h & look for curandStatus or CURAND_STATUS_xxx)"; -} - -template <> -const char* CudaErrString(cudnnStatus_t e) { - cudaDeviceSynchronize(); - return cudnnGetErrorString(e); -} - -template <> -const char* CudaErrString(cufftResult e) { - cudaDeviceSynchronize(); - switch (e) { - CASE_ENUM_TO_STR(CUFFT_SUCCESS); - CASE_ENUM_TO_STR(CUFFT_ALLOC_FAILED); - CASE_ENUM_TO_STR(CUFFT_INVALID_VALUE); - CASE_ENUM_TO_STR(CUFFT_INTERNAL_ERROR); - CASE_ENUM_TO_STR(CUFFT_SETUP_FAILED); - CASE_ENUM_TO_STR(CUFFT_INVALID_SIZE); - default: - return "Unknown cufft error status"; - } -} -#endif - -#ifdef ORT_USE_NCCL -template <> -const char* CudaErrString(ncclResult_t e) { - cudaDeviceSynchronize(); - return ncclGetErrorString(e); -} -#endif - -template -int GetErrorCode(ERRTYPE err) { - return static_cast(err); -} - -template -std::conditional_t CudaCall( - ERRTYPE retCode, const char* exprString, const char* libName, SUCCTYPE successCode, const char* msg, - const char* file, const int line) { - if (retCode != successCode) { - try { -#ifdef _WIN32 - std::string hostname_str = GetEnvironmentVar("COMPUTERNAME"); - if (hostname_str.empty()) { - hostname_str = "?"; - } - const char* hostname = hostname_str.c_str(); -#else - char hostname[HOST_NAME_MAX]; - if (gethostname(hostname, HOST_NAME_MAX) != 0) - strcpy(hostname, "?"); -#endif - int currentCudaDevice = -1; - cudaGetDevice(¤tCudaDevice); - cudaGetLastError(); // clear last CUDA error - static char str[1024]; - snprintf(str, 1024, "%s failure %d: %s ; GPU=%d ; hostname=%s ; file=%s ; line=%d ; expr=%s; %s", - libName, GetErrorCode(retCode), CudaErrString(retCode), currentCudaDevice, - hostname, - file, line, exprString, msg); - if constexpr (THRW) { - // throw an exception with the error info - ORT_THROW(str); - } else { - LOGS_DEFAULT(ERROR) << str; - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, str); - } - } catch (const std::exception& e) { // catch, log, and rethrow since CUDA code sometimes hangs in destruction, - // so we'd never get to see the error - if constexpr (THRW) { - ORT_THROW(e.what()); - } else { - LOGS_DEFAULT(ERROR) << e.what(); - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, e.what()); - } - } - } - if constexpr (!THRW) { - return Status::OK(); - } -} - -template Status CudaCall(cudaError retCode, const char* exprString, const char* libName, cudaError successCode, const char* msg, const char* file, const int line); -template void CudaCall(cudaError retCode, const char* exprString, const char* libName, cudaError successCode, const char* msg, const char* file, const int line); -template Status CudaCall(CUresult retCode, const char* exprString, const char* libName, CUresult successCode, const char* msg, const char* file, const int line); -template void CudaCall(CUresult retCode, const char* exprString, const char* libName, CUresult successCode, const char* msg, const char* file, const int line); - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_data_transfer.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_data_transfer.cc deleted file mode 100644 index d334aeaa86cdf..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_data_transfer.cc +++ /dev/null @@ -1,114 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/shared_library/provider_api.h" - -#include "nv_data_transfer.h" - -#include "core/providers/cuda/shared_inc/cuda_call.h" -#define CUDA_RETURN_IF_ERROR(expr) ORT_RETURN_IF_ERROR(CUDA_CALL(expr)) -namespace onnxruntime { -bool GPUDataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const { - OrtDevice::DeviceType src_type = src_device.Type(); - OrtDevice::DeviceType dst_type = dst_device.Type(); - - // check that only our GPU is involved - if ((src_type == OrtDevice::GPU && src_device.Vendor() != OrtDevice::VendorIds::NVIDIA) || - (dst_type == OrtDevice::GPU && dst_device.Vendor() != OrtDevice::VendorIds::NVIDIA)) { - return false; - } - - // copy must involve a GPU, and be device to device or cpu (exclude other device types) - return (src_type == OrtDevice::GPU || dst_type == OrtDevice::GPU) && - (src_type == OrtDevice::GPU || src_type == OrtDevice::CPU) && - (dst_type == OrtDevice::GPU || dst_type == OrtDevice::CPU); -} - -common::Status GPUDataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const { - size_t bytes = src.SizeInBytes(); - const void* src_data = src.DataRaw(); - void* dst_data = dst.MutableDataRaw(); - - auto& src_device = src.Location().device; - auto& dst_device = dst.Location().device; - - const bool dst_is_gpu_default = dst_device.Type() == OrtDevice::GPU && - dst_device.MemType() == OrtDevice::MemType::DEFAULT; - const bool src_is_gpu_default = src_device.Type() == OrtDevice::GPU && - src_device.MemType() == OrtDevice::MemType::DEFAULT; - - // for the sync version of memcpy, launch to cuda default stream - if (dst_is_gpu_default) { - if (src_is_gpu_default) { - // Copy only if the two addresses are different. - if (dst_data != src_data) { - CUDA_RETURN_IF_ERROR(cudaMemcpy(dst_data, src_data, bytes, cudaMemcpyDeviceToDevice)); - // For device memory to device memory copy, no host-side synchronization is performed by cudaMemcpy. - // see https://docs.nvidia.com/cuda/cuda-runtime-api/api-sync-behavior.html - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(nullptr)); - } - } else { - // copy from other CPU memory to GPU, this is blocking - CUDA_RETURN_IF_ERROR(cudaMemcpy(dst_data, src_data, bytes, cudaMemcpyHostToDevice)); - if (src_device.MemType() != OrtDevice::MemType::HOST_ACCESSIBLE) { - // For cudaMemcpy from pageable host memory to device memory, DMA to final destination may not have completed. - // see https://docs.nvidia.com/cuda/cuda-runtime-api/api-sync-behavior.html - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(nullptr)); - } - } - } else if (src_is_gpu_default) { - // copying from GPU to CPU memory, this is blocking - CUDA_RETURN_IF_ERROR(cudaMemcpy(dst_data, src_data, bytes, cudaMemcpyDeviceToHost)); - } else { - // copying between cpu memory - ORT_ENFORCE(dst_data != src_data); - memcpy(dst_data, src_data, bytes); - } - - return Status::OK(); -} - -common::Status GPUDataTransfer::CopyTensorAsync(const Tensor& src, Tensor& dst, Stream& stream) const { - size_t bytes = src.SizeInBytes(); - const void* src_data = src.DataRaw(); - void* dst_data = dst.MutableDataRaw(); - - auto& src_device = src.Location().device; - auto& dst_device = dst.Location().device; - - const bool dst_is_gpu_default = dst_device.Type() == OrtDevice::GPU && - dst_device.MemType() == OrtDevice::MemType::DEFAULT; - const bool src_is_gpu_default = src_device.Type() == OrtDevice::GPU && - src_device.MemType() == OrtDevice::MemType::DEFAULT; - - if (dst_is_gpu_default) { - if (src_is_gpu_default) { - // copying between GPU, this is non-blocking - if (dst_data != src_data) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_data, src_data, bytes, cudaMemcpyDeviceToDevice, - static_cast(stream.GetHandle()))); - } - } else { - // copy from pinned or non-pinned CPU memory to GPU - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_data, src_data, bytes, cudaMemcpyHostToDevice, - static_cast(stream.GetHandle()))); - } - } else if (src_is_gpu_default) { - // copy from GPU to pinned or non-pinned CPU memory. - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_data, src_data, bytes, cudaMemcpyDeviceToHost, - static_cast(stream.GetHandle()))); - } else { - if (src_device.MemType() == OrtDevice::MemType::HOST_ACCESSIBLE) { - // sync the stream first to make sure the data arrived - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(static_cast(stream.GetHandle()))); - } - - ORT_ENFORCE(dst_data != src_data); - memcpy(dst_data, src_data, bytes); - } - - return Status::OK(); -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_data_transfer.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_data_transfer.h deleted file mode 100644 index 272ea367ac7e4..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_data_transfer.h +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "nv_includes.h" -#include "core/framework/data_transfer.h" - -namespace onnxruntime { - -class GPUDataTransfer : public IDataTransfer { - public: - GPUDataTransfer() = default; - ~GPUDataTransfer() = default; - - bool CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const override; - - // Dumpen MSVC warning about not fully overriding - using IDataTransfer::CopyTensor; - common::Status CopyTensor(const Tensor& src, Tensor& dst) const override; - common::Status CopyTensorAsync(const Tensor& src, Tensor& dst, Stream& stream) const override; -}; - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc deleted file mode 100644 index cc1cecdb7ae5b..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ /dev/null @@ -1,3693 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#include -#include -#include -#include -#include "core/providers/shared_library/provider_api.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_options.h" -#define ORT_API_MANUAL_INIT -#include "core/session/onnxruntime_cxx_api.h" -#include "core/session/onnxruntime_type_conversion.h" -#include "core/common/common.h" -#include "core/common/narrow.h" -#include "core/common/path_utils.h" -#include "core/common/safeint.h" -#include "core/framework/ort_value.h" -#include "nv_execution_provider.h" -#include "nv_execution_provider_utils.h" -#include "nv_execution_provider_custom_ops.h" -#include "nv_allocator.h" -#include "nv_data_transfer.h" -#include "onnx_ctx_model_helper.h" -#include "core/providers/cuda/shared_inc/cuda_call.h" -#include "core/providers/cuda/cuda_graph.h" -#include "core/session/allocator_adapters.h" -#include "cuda_runtime_api.h" -#include "core/common/parse_string.h" -#include -#include -#include -#include -#include -#include -#include -// TODO: find a better way to share this -#include "core/providers/cuda/cuda_stream_handle.h" - -#ifdef _WIN32 -#include -#define LIBTYPE HINSTANCE -#define OPENLIB(libname) LoadLibrary(libname) -#define LIBFUNC(lib, fn) GetProcAddress((lib), (fn)) -#else -#include -#define LIBTYPE void* -#define OPENLIB(libname) dlopen((libname), RTLD_LAZY) -#define LIBFUNC(lib, fn) dlsym((lib), (fn)) -#endif - -#define CUDA_RETURN_IF_ERROR(expr) ORT_RETURN_IF_ERROR(CUDA_CALL(expr)) - -using namespace ONNX_NAMESPACE; -using namespace ::onnxruntime::logging; -namespace { -// Check if cycle exists in the graph after partitioning -bool FindCycleHelper(size_t i, const std::list* adjacency_map, bool visited[], bool* st, std::vector& cycles) { - if (!visited[i]) { - visited[i] = true; - st[i] = true; - for (auto iter = adjacency_map[i].begin(); iter != adjacency_map[i].end(); ++iter) { - if (!visited[*iter] && FindCycleHelper(*iter, adjacency_map, visited, st, cycles)) { - cycles.push_back(*iter); - return true; - } else if (st[*iter]) { - cycles.push_back(*iter); - return true; - } - } - } - st[i] = false; - return false; -} -} // namespace - -namespace google { -namespace protobuf { -void ShutdownProtobufLibrary(); -} -} // namespace google - -struct ShutdownProtobuf { - ~ShutdownProtobuf() { - ::google::protobuf::ShutdownProtobufLibrary(); - } -} g_protobuf; - -namespace onnxruntime { - -// Helper function to check if a data type is supported by input output nodes ofNvTensorRTRTX EP -static bool IsSupportedInputOutputDataType(ONNXTensorElementDataType data_type) { - switch (data_type) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: // kFLOAT - 32-bit floating point - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: // kHALF - IEEE 16-bit floating-point - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: // kBF16 - Brain float 16 - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: // kBOOL - 8-bit boolean - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4: // kINT4 - 4-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: // kINT8 - 8-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: // kUINT8 - 8-bit unsigned integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: // kINT32 - 32-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: // kINT64 - 64-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN: // kFP8 - 8-bit floating point - return true; - default: - return false; - } -} - -// Helper function to check if a data type is supported by NvTensorRTRTX EP -static bool IsSupportedDataType(ONNXTensorElementDataType data_type) { - switch (data_type) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: // kFLOAT - 32-bit floating point - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: // kHALF - IEEE 16-bit floating-point - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: // kBF16 - Brain float 16 - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: // kBOOL - 8-bit boolean - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4: // kINT4 - 4-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: // kINT8 - 8-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: // kUINT8 - 8-bit unsigned integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: // kINT32 - 32-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: // kINT64 - 64-bit signed integer - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN: // kFP8 - 8-bit floating point - case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: // kDOUBLE - 64-bit floating point - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT4E2M1: // kFP4 - 4-bit floating point - return true; - default: - return false; - } -} - -// Helper function to get data type name as string -static std::string GetDataTypeName(ONNXTensorElementDataType data_type) { - switch (data_type) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: - return "FLOAT"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: - return "FLOAT16"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: - return "BFLOAT16"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: - return "BOOL"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4: - return "INT4"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: - return "INT8"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: - return "UINT8"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: - return "INT32"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: - return "INT64"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN: - return "FLOAT8E4M3FN"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: - return "DOUBLE"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING: - return "STRING"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: - return "UINT16"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32: - return "UINT32"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64: - return "UINT64"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: - return "INT16"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64: - return "COMPLEX64"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128: - return "COMPLEX128"; - default: - return "UNKNOWN(" + std::to_string(static_cast(data_type)) + ")"; - } -} - -// Helper function to check if a node has supported data types -static bool CheckNodeDataTypes(const Node* node) { - // Check input data types - for (const auto* input_def : node->InputDefs()) { - if (input_def->Exists()) { - const auto* type_proto = input_def->TypeAsProto(); - if (type_proto && type_proto->has_tensor_type()) { - auto data_type = utils::ToOrtTensorElementDataType(type_proto->tensor_type().elem_type()); - if (!IsSupportedDataType(data_type)) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Node '" << node->Name() - << "' (OpType: " << node->OpType() - << ") has unsupported input data type: " << GetDataTypeName(data_type) - << " for input '" << input_def->Name() << "'"; - return false; - } - } - } - } - - // Check output data types - for (const auto* output_def : node->OutputDefs()) { - if (output_def->Exists()) { - const auto* type_proto = output_def->TypeAsProto(); - if (type_proto && type_proto->has_tensor_type()) { - auto data_type = utils::ToOrtTensorElementDataType(type_proto->tensor_type().elem_type()); - if (!IsSupportedDataType(data_type)) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Node '" << node->Name() - << "' (OpType: " << node->OpType() - << ") has unsupported output data type: " << GetDataTypeName(data_type) - << " for output '" << output_def->Name() << "'"; - return false; - } - } - } - } - - return true; -} - -void* OutputAllocator::reallocateOutputAsync(char const* /*tensorName*/, void* /*currentMemory*/, uint64_t size, - uint64_t /*alignment*/, cudaStream_t /*stream*/) noexcept { - // Some memory allocators return nullptr when allocating zero bytes, but TensorRT requires a non-null ptr - // even for empty tensors, so allocate a dummy byte. - size = std::max(size, static_cast(1)); - if (size > allocated_size) { - alloc_->Free(alloc_, outputPtr); - outputPtr = nullptr; - allocated_size = 0; - outputPtr = alloc_->Alloc(alloc_, size); - if (outputPtr) { - allocated_size = size; - } - } - // if cudaMalloc fails, returns nullptr. - return outputPtr; -} - -void OutputAllocator::notifyShape(char const* /*tensorName*/, nvinfer1::Dims const& dims) noexcept { - output_shapes.clear(); - output_shapes.reserve(dims.nbDims); - for (int i = 0; i < dims.nbDims; i++) { - output_shapes.push_back(dims.d[i]); - } -} - -class Memcpy final : public OpKernel { - public: - Memcpy(const OpKernelInfo& info) : OpKernel(info) {} - - Status Compute(OpKernelContext* ctx) const override { - const auto* X = ctx->Input(0); - ORT_ENFORCE(X != nullptr, "Memcpy: Input tensor is nullptr."); - Tensor* Y = ctx->Output(0, X->Shape()); - ORT_ENFORCE(Y != nullptr, "Memcpy: Failed to allocate output tensor."); - auto* gpu_data_transfer = Info().GetDataTransferManager().GetDataTransfer(X->Location().device, Y->Location().device); - if (!gpu_data_transfer) - return Status(common::ONNXRUNTIME, common::EP_FAIL, "gpu data transfer is missing in TRT EP."); - if (!ctx->GetComputeStream()) - return Status(common::ONNXRUNTIME, common::EP_FAIL, "Compute Stream is missing in TRT MemCpy kernel's context."); - return gpu_data_transfer->CopyTensorAsync(*X, *Y, *(ctx->GetComputeStream())); - } -}; - -template -KernelCreateInfo BuildKernelCreateInfo(); - -ONNX_OPERATOR_KERNEL_EX( - MemcpyFromHost, - kOnnxDomain, - 1, - kNvTensorRTRTXExecutionProvider, - (*KernelDefBuilder::Create()) - .InputMemoryType(OrtMemTypeCPUInput, 0) - .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), - Memcpy); - -ONNX_OPERATOR_KERNEL_EX( - MemcpyToHost, - kOnnxDomain, - 1, - kNvTensorRTRTXExecutionProvider, - (*KernelDefBuilder::Create()) - .OutputMemoryType(OrtMemTypeCPUOutput, 0) - .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), - Memcpy); - -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNvTensorRTRTXExecutionProvider, kOnnxDomain, 1, MemcpyFromHost); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNvTensorRTRTXExecutionProvider, kOnnxDomain, 1, MemcpyToHost); - -static std::shared_ptr s_kernel_registry; - -void InitializeRegistry() { - s_kernel_registry = KernelRegistry::Create(); - - static const BuildKernelCreateInfoFn function_table[] = { - BuildKernelCreateInfo, - BuildKernelCreateInfo, - }; - - for (auto& function_table_entry : function_table) { - ORT_THROW_IF_ERROR(s_kernel_registry->Register(function_table_entry())); - } -} - -void DeleteRegistry() { - s_kernel_registry.reset(); -} - -std::shared_ptr NvExecutionProvider::GetKernelRegistry() const { - return s_kernel_registry; -} - -// Per TensorRT documentation, logger needs to be a singleton. -TensorrtLogger& GetTensorrtLogger(bool verbose_log) { - const auto log_level = verbose_log ? nvinfer1::ILogger::Severity::kVERBOSE : nvinfer1::ILogger::Severity::kWARNING; - static TensorrtLogger trt_logger(log_level); - if (log_level != trt_logger.get_level()) { - trt_logger.set_level(verbose_log ? nvinfer1::ILogger::Severity::kVERBOSE : nvinfer1::ILogger::Severity::kWARNING); - } - return trt_logger; -} - -std::unique_lock NvExecutionProvider::GetApiLock() const { - static std::mutex singleton; - return std::unique_lock(singleton); -} - -/* - * Get the shape of "shape tensor" input - */ -template -Status GetShapeOfShapeTensor(Ort::ConstValue& input_tensor, - void* shape_values, - int shape_size, - cudaStream_t stream) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(shape_values, - input_tensor.GetTensorData(), - shape_size * sizeof(T), - cudaMemcpyDeviceToHost, - stream)); - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); - return Status::OK(); -} - -/* - * Apply TensorRT optimization profile shapes from provider options. - * - * This function supports single/multiple profile(s). - * (Note: An optimization profile describes a range of dimensions for each network input) - * - */ -bool ApplyProfileShapesFromProviderOptions(std::vector& trt_profiles, - nvinfer1::ITensor* input, - std::unordered_map>>& profile_min_shapes, - std::unordered_map>>& profile_max_shapes, - std::unordered_map>>& profile_opt_shapes, - ShapeRangesMap& input_explicit_shape_ranges, - bool& cuda_graph_flag) { - if (trt_profiles.size() == 0) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Number of optimization profiles should be greater than 0, but it's 0."; - return false; - } - - const std::string& input_name = input->getName(); - if (profile_min_shapes.find(input_name) == profile_min_shapes.end()) { - return false; - } - - if (input_explicit_shape_ranges.find(input_name) == input_explicit_shape_ranges.end()) { - std::unordered_map>> inner_map; - input_explicit_shape_ranges[input_name] = inner_map; - } - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Begin to apply profile shapes ..."; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Input tensor name is '" << input_name << "', number of profiles found is " << trt_profiles.size(); - - for (size_t i = 0; i < trt_profiles.size(); i++) { - nvinfer1::Dims dims = input->getDimensions(); - int nb_dims = dims.nbDims; - - auto trt_profile = trt_profiles[i]; - - // Shape tensor - if (input->isShapeTensor()) { - if (cuda_graph_flag) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Shape tensor detected on input '" << input->getName() << "'. Disabling CUDA Graph."; - cuda_graph_flag = false; - } - int shape_size = nb_dims == 0 ? 1 : static_cast(profile_min_shapes[input_name][i].size()); - std::vector shapes_min(shape_size), shapes_opt(shape_size), shapes_max(shape_size); - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] shape size of this shape tensor is " << shape_size; - - for (int j = 0; j < shape_size; j++) { - auto min_value = profile_min_shapes[input_name][i][j]; - auto max_value = profile_max_shapes[input_name][i][j]; - auto opt_value = profile_opt_shapes[input_name][i][j]; - shapes_min[j] = static_cast(min_value); - shapes_max[j] = static_cast(max_value); - shapes_opt[j] = static_cast(opt_value); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] shapes_min.d[" << j << "] is " << shapes_min[j]; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] shapes_max.d[" << j << "] is " << shapes_max[j]; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] shapes_opt.d[" << j << "] is " << shapes_opt[j]; - - if (input_explicit_shape_ranges[input_name].find(j) == input_explicit_shape_ranges[input_name].end()) { - std::vector> profile_vector(trt_profiles.size()); - input_explicit_shape_ranges[input_name][j] = profile_vector; - } - input_explicit_shape_ranges[input_name][static_cast(j)][i].push_back(min_value); - input_explicit_shape_ranges[input_name][static_cast(j)][i].push_back(max_value); - input_explicit_shape_ranges[input_name][static_cast(j)][i].push_back(opt_value); - } - - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, &shapes_min[0], shape_size); - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, &shapes_max[0], shape_size); - trt_profile->setShapeValuesV2(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, &shapes_opt[0], shape_size); - } - // Execution tensor - else { - nvinfer1::Dims dims_min, dims_opt, dims_max; - dims_min.nbDims = nb_dims; - dims_max.nbDims = nb_dims; - dims_opt.nbDims = nb_dims; - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] number of dimension of this execution tensor is " << nb_dims; - - for (int j = 0; j < nb_dims; j++) { - if (dims.d[j] == -1) { - auto min_value = profile_min_shapes[input_name][i][j]; - auto max_value = profile_max_shapes[input_name][i][j]; - auto opt_value = profile_opt_shapes[input_name][i][j]; - dims_min.d[j] = static_cast(min_value); - dims_max.d[j] = static_cast(max_value); - dims_opt.d[j] = static_cast(opt_value); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] dims_min.d[" << j << "] is " << dims_min.d[j]; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] dims_max.d[" << j << "] is " << dims_max.d[j]; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] dims_opt.d[" << j << "] is " << dims_opt.d[j]; - - if (input_explicit_shape_ranges[input_name].find(j) == input_explicit_shape_ranges[input_name].end()) { - std::vector> profile_vector(trt_profiles.size()); - input_explicit_shape_ranges[input_name][j] = profile_vector; - } - input_explicit_shape_ranges[input_name][static_cast(j)][i].push_back(min_value); - input_explicit_shape_ranges[input_name][static_cast(j)][i].push_back(max_value); - input_explicit_shape_ranges[input_name][static_cast(j)][i].push_back(opt_value); - } else { - dims_min.d[j] = dims.d[j]; - dims_max.d[j] = dims.d[j]; - dims_opt.d[j] = dims.d[j]; - } - } - - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, dims_min); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, dims_max); - trt_profile->setDimensions(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, dims_opt); - } - } - return true; -} - -#define CASE_GET_INPUT_TENSOR(DATA_TYPE, SrcT) \ - case DATA_TYPE: { \ - auto input_tensor_ptr = input_tensor.GetTensorData(); \ - if (input_tensor_ptr != nullptr && elem_cnt > 0) { \ - data = const_cast(input_tensor_ptr); \ - } else { \ - scratch_buffers.push_back(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, 1)); \ - data = scratch_buffers.back().get(); \ - } \ - break; \ - } - -#define CASE_GET_OUTPUT_TENSOR(DATA_TYPE, SrcT) \ - case DATA_TYPE: { \ - auto output_tensor_ptr = output_tensor.GetTensorMutableData(); \ - data_ptr = output_tensor_ptr; \ - if (output_tensor_ptr != nullptr && elem_cnt > 0) { \ - buffer = output_tensor_ptr; \ - } else { \ - scratch_buffers.push_back(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, 1)); \ - buffer = scratch_buffers.back().get(); \ - } \ - break; \ - } - -#define CASE_COPY_TENSOR(DATA_TYPE, DstT) \ - case DATA_TYPE: { \ - auto output_tensor_ptr = output_tensor.GetTensorMutableData(); \ - if (output_tensor_ptr != nullptr && elem_cnt > 0) { \ - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor_ptr, allocator->getBuffer(), elem_cnt * sizeof(DstT), cudaMemcpyDeviceToDevice, stream)); \ - } \ - break; \ - } - -/* - * Set Nv executio context input. - * - * There are two types of input tensor: (1) shape tensor and (2) execution tensor. - * The input buffer binding needs to be handled differently. - * - * @param shape_tensor_values holds "shape tensor -> shape values" for the INT32 shape tensor input across this inference run - * @param shape_tensor_values_int64 holds "shape tensor -> shape values" for the INT64 shape tensor input across this inference run - */ -Status BindContextInput(Ort::KernelContext& ctx, - nvinfer1::ICudaEngine* trt_engine, - nvinfer1::IExecutionContext* trt_context, - const char* input_name, - size_t input_index, - std::unordered_map>& shape_tensor_values, - std::unordered_map>& shape_tensor_values_int64, - std::vector>& scratch_buffers, - OrtAllocator* alloc, - cudaStream_t stream, - bool& skip_input_binding_allowed) { - auto input_tensor = ctx.GetInput(input_index); - auto tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); - const auto tensor_shapes = tensor_info.GetShape(); - const auto tensor_type = tensor_info.GetElementType(); - /* - * Return the number of elements specified by the tensor shape (all dimensions multiplied by each other). - * For 0 dimensions, 1 is returned. If any dimension is less than 0, the result is always -1. - * - * Examples:
- * [] = 1
- * [1,3,4] = 12
- * [2,0,4] = 0
- * [-1,3,4] = -1
- */ - const auto elem_cnt = tensor_info.GetElementCount(); - - if (trt_engine->isShapeInferenceIO(input_name)) { - // Bind "shape tensor" input buffer - skip_input_binding_allowed = false; // Shape tensor input binding cannot be skipped - // The shape of the "shape tensor" is either zero dimension (scalar) or 1-dimension - int shape_size = trt_engine->getTensorShape(input_name).nbDims == 0 ? 1 : static_cast(tensor_shapes[0]); - switch (tensor_type) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: { - // get shape tensor value if not present - if (shape_tensor_values.find(input_name) == shape_tensor_values.end()) { - auto input = std::make_unique(shape_size); - auto status = GetShapeOfShapeTensor(input_tensor, input.get(), shape_size, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - shape_tensor_values[input_name].resize(shape_size); - for (int i = 0; i < shape_size; ++i) { - shape_tensor_values[input_name][i] = input[i]; - } - } - - if (!trt_context->setTensorAddress(input_name, &shape_tensor_values[input_name][0])) { - std::string error_input_name = input_name; - std::string error_msg = - "NvTensorRTRTX EP failed to call nvinfer1::IExecutionContext::setTensorAddress() for shape input '" + - error_input_name + "'"; - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, error_msg)); - } - break; - } - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: { - // get shape tensor value if not present - if (shape_tensor_values_int64.find(input_name) == shape_tensor_values_int64.end()) { - auto input = std::make_unique(shape_size); - auto status = GetShapeOfShapeTensor(input_tensor, input.get(), shape_size, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - shape_tensor_values_int64[input_name].resize(shape_size); - for (int i = 0; i < shape_size; ++i) { - shape_tensor_values_int64[input_name][i] = input[i]; - } - } - - if (!trt_context->setTensorAddress(input_name, &shape_tensor_values_int64[input_name][0])) { - std::string error_input_name = input_name; - std::string error_msg = - "NvTensorRTRTX EP failed to call nvinfer1::IExecutionContext::setTensorAddress() for shape input '" + - error_input_name + "'"; - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, error_msg)); - } - break; - } - default: { - std::string error_input_name = input_name; - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "The data type of shape tensor should be INT32 or INT64. Please check the data type of " + error_input_name); - } - } - } else { - // Set shape for input tensor which is execution tensor - nvinfer1::Dims dims = trt_context->getTensorShape(input_name); - int nb_dims = dims.nbDims; - for (int j = 0, end = nb_dims; j < end; ++j) { - dims.d[j] = static_cast(tensor_shapes[j]); - } - if (!trt_context->setInputShape(input_name, dims)) { - std::string error_input_name = input_name; - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP failed to call nvinfer1::IExecutionContext::setInputShape() for input '" + error_input_name + "'")); - } - - // Bind "execution tensor" input buffer - // - // Note: If an engine binding is an empty tensor, it still needs a non-null memory address, and different tensors should have different addresses. - // Therefore, in the case of empty tensor, TRT EP always allocates a dummy byte. - // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#empty-tensors - void* data = nullptr; - switch (tensor_type) { - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, float) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, uint16_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16, uint16_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL, bool) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4, uint8_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8, int8_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8, uint8_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32, int32_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, int64_t) - CASE_GET_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN, uint8_t) - default: { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP input onnx tensor data type: " + std::to_string(tensor_type) + " not supported."); - } - } - trt_context->setTensorAddress(input_name, data); - } - - return Status::OK(); -} - -/* - * Bind Nv executio context output. - * - * Please note that the "data-depedent shape" output needs corresponding allocator provided. - * - * - * param ctx - ORT kernel context - * param trt_context - A pointer to Nv executio context object - * param output_name - Output tensor name - * param output_index - The index of the output to the ORT kernel context - * param output_type - Data type of the output - * param i - Output iteration index - * param output_tensors - Output iteration index to output's ORT value - * param dds_output_allocator_map - DDS output to its allocator - * param scratch_buffer - The allocation buffer created by TRT EP - * param allocator - ORT allocator - * param buffers - It holds all the output values which are binding to TRT's execution context - * - */ -Status BindContextOutput(Ort::KernelContext& ctx, - nvinfer1::IExecutionContext* trt_context, - const char* output_name, - size_t output_index, - size_t output_type, - DDSOutputAllocatorMap& dds_output_allocator_map, - std::vector>& scratch_buffers, - OrtAllocator* alloc, - nvinfer1::Dims& dims, - void*& data_ptr) { - // Get output shape - dims = trt_context->getTensorShape(output_name); - int nb_dims = dims.nbDims; - bool is_DDS = false; - for (int j = 0, end = nb_dims; j < end; ++j) { - // data-dependent shape - if (dims.d[j] == -1) { - is_DDS = true; - break; - } - } - - auto known_DDS = dds_output_allocator_map.find(output_name) != dds_output_allocator_map.end(); - - // If the output tensor has data-dependent shape, TRT EP will provide an IOutputAllocator for enqueueV3 to dynamically allocate memory buffer. - // Once enqueueV3 returns, TRT EP will then bind the output allocation to ORT kernel context output. - // (Please note that we take strategy A mentioned in https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#dynamic-shaped-output, - // which we defer allocation until the size is known and don't call IExecution::setTensorAddress) - // - // Otherwise, if the shape of the output tensor is known prior to the runtime, ORT will pre-allocate memory buffer for the output tensor for enqueueV3. - if (is_DDS || known_DDS) { - if (!known_DDS) { - auto allocatorPtr = std::make_unique(alloc); - trt_context->setOutputAllocator(output_name, allocatorPtr.get()); - dds_output_allocator_map[output_name] = std::move(allocatorPtr); - dims.nbDims = -1; // Set to -1 to indicate that the shape is not known at this point. - data_ptr = nullptr; // Set data_ptr to nullptr for DDS output binding. - } - } else { - auto output_tensor = ctx.GetOutput(output_index, dims.d, nb_dims); - const auto elem_cnt = output_tensor.GetTensorTypeAndShapeInfo().GetElementCount(); - - void* buffer = nullptr; - - switch (output_type) { - // below macros set data_ptr and skip_output_binding_allowed variables - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, float) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, uint16_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16, uint16_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL, bool) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4, uint8_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8, int8_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8, uint8_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32, int32_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, int64_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN, uint8_t) - CASE_GET_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT4E2M1, uint8_t) - default: { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP output tensor data type: " + std::to_string(output_type) + " not supported."); - } - } - trt_context->setTensorAddress(output_name, buffer); - } - - return Status::OK(); -} - -/* - * Bind ORT kernel context Output. - * - * In the case of DDS (data-dependent shape) output, TRT requires a provided allocator to allocate memory during runtime. - * Once the output has been put in the allocation buffer, ORT calls this function to bind the allocation to ORT kernel context output. - * - * Note: Current approach of setting the ORT kernel context output is copying the output data from allocation buffer to ORT context output address which is not optimal, - * we are waiting for ORT core to support "assign" memory address to ORT context output. Some works need to be done in ORT memory planner to be aware of this memory support. - */ -Status BindKernelOutput(Ort::KernelContext& ctx, - DDSOutputAllocatorMap& allocator_map, - char const* output_name, - size_t output_index, - size_t output_type, - cudaStream_t stream) { - auto allocator = allocator_map[output_name].get(); - auto& shape = allocator->getOutputShape(); - auto output_tensor = ctx.GetOutput(output_index, shape); - - /* - * Return the number of elements specified by the tensor shape (all dimensions multiplied by each other). - * For 0 dimensions, 1 is returned. If any dimension is less than 0, the result is always -1. - * - * Examples:
- * [] = 1
- * [1,3,4] = 12
- * [2,0,4] = 0
- * [-1,3,4] = -1
- */ - auto elem_cnt = output_tensor.GetTensorTypeAndShapeInfo().GetElementCount(); - - /* - * Copy output data from allocation buffer to ORT kernel context output location or - * cast (int32 or float) -> (int64 or double) to ORT kernel context output location. - * - * Note: - * 1. If the output tensor is empty tensor (i.e. any of the dimension is 0) which means element count is 0, - * TRT EP does not perform cuda memory copy nor cuda cast to prevent overwriting other location that might belong to other tensors. - * 2. The cudaMemcpyAsync() and cuda::Impl_Cast() (implemented as _UnaryElementWise() in cuda ep) are all async, but we - * don't need to explicitly call cudaStreamSynchronize() after those APIs due to CUDA EP and TRT EP uses same stream, - * and within the same stream, operations are guaranteed to be executed in order. - */ - switch (output_type) { - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, float) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, uint16_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16, uint16_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL, bool) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4, uint8_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8, int8_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8, uint8_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32, int32_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, int64_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN, uint8_t) - CASE_COPY_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT4E2M1, uint8_t) - default: { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP output tensor data type: " + std::to_string(output_type) + " not supported."); - } - } - return Status::OK(); -} - -NvExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, bool has_user_compute_stream, cudaStream_t stream) { - // Only set device if user hasn't provided a compute stream - if (!has_user_compute_stream) { - CUDA_CALL_THROW(cudaSetDevice(device_id)); - (void)stream; - } -} - -NvExecutionProvider::PerThreadContext::~PerThreadContext() { - trt_context_map_.clear(); -} - -void NvExecutionProvider::PerThreadContext::ResetTensorRTContext(std::string fused_node) { - auto it = trt_context_map_.find(fused_node); - if (it != trt_context_map_.end()) { - trt_context_map_[fused_node].reset(); - } -} - -bool NvExecutionProvider::PerThreadContext::UpdateTensorRTContext(std::string fused_node, tensorrt_ptr::unique_pointer_exec_ctx context) { - if (!context) { - context = tensorrt_ptr::unique_pointer_exec_ctx(); - } - trt_context_map_[fused_node] = std::move(context); - - if (trt_context_map_[fused_node]) { - return true; - } - return false; -} - -void NvExecutionProvider::PerThreadContext::DeleteCapturedGraph(CudaGraphAnnotation_t cuda_graph_annotation_id) { - graph_id_to_run_count_.erase(cuda_graph_annotation_id); - cuda_graph_.Reset(); -} - -void NvExecutionProvider::PerThreadContext::ResetWarmupRuns(CudaGraphAnnotation_t cuda_graph_annotation_id) { - if (graph_id_to_run_count_.find(cuda_graph_annotation_id) == graph_id_to_run_count_.end()) { - return; - } - graph_id_to_run_count_[cuda_graph_annotation_id] = 0; -} - -bool NvExecutionProvider::PerThreadContext::IsGraphCaptureAllowed(CudaGraphAnnotation_t cuda_graph_annotation_id) const { - if (!IsGraphCaptureAllowedOnRun(cuda_graph_annotation_id)) { - return false; - } - - // Safe access to map - return false if key doesn't exist yet - auto it = graph_id_to_run_count_.find(cuda_graph_annotation_id); - if (it == graph_id_to_run_count_.end()) { - return false; // Entry doesn't exist yet, not ready for capture - } - - bool allowed = it->second >= min_num_runs_before_cuda_graph_capture_; - if (allowed) { - LOGS_DEFAULT(VERBOSE) << "NvTensorRTRTX EP Graph capture allowed for ID: " << cuda_graph_annotation_id - << ", run count: " << it->second; - } - return allowed; -} - -bool NvExecutionProvider::PerThreadContext::IsGraphCaptureAllowedOnRun(CudaGraphAnnotation_t cuda_graph_annotation_id) const { - return cuda_graph_.IsGraphCaptureAllowedOnRun(cuda_graph_annotation_id); -} - -CudaGraphAnnotation_t NvExecutionProvider::PerThreadContext::GetCudaGraphAnnotationId(const onnxruntime::RunOptions& run_options) const { - // Actual implementation - auto graph_annotation_str = run_options.GetConfigOptions().GetConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation); - CudaGraphAnnotation_t cuda_graph_annotation_id = kCudaGraphAnnotationDefault; - - // Kind of debugging head implementation, can be cleaned and made robust like CUDA EP - if (graph_annotation_str.has_value() && !graph_annotation_str->empty()) { - if (!TryParseStringWithClassicLocale(*graph_annotation_str, cuda_graph_annotation_id)) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Failed to parse cuda graph annotation id: " - << *graph_annotation_str << ", using default: " << kCudaGraphAnnotationDefault; - cuda_graph_annotation_id = kCudaGraphAnnotationDefault; - } - } - return cuda_graph_annotation_id; -} - -void NvExecutionProvider::PerThreadContext::SetCurrentGraphAnnotationId(CudaGraphAnnotation_t cuda_graph_annotation_id) { - current_graph_annotation_id_ = cuda_graph_annotation_id; -} - -CudaGraphAnnotation_t NvExecutionProvider::PerThreadContext::GetCurrentGraphAnnotationId() const { - return current_graph_annotation_id_; -} - -void NvExecutionProvider::PerThreadContext::CaptureBegin(CudaGraphAnnotation_t cuda_graph_annotation_id) { - cuda_graph_.Reset(); - cuda_graph_.CaptureBegin(cuda_graph_annotation_id); -} - -void NvExecutionProvider::PerThreadContext::CaptureEnd(CudaGraphAnnotation_t cuda_graph_annotation_id) { - cuda_graph_.CaptureEnd(cuda_graph_annotation_id); -} - -bool NvExecutionProvider::PerThreadContext::IsGraphCaptured(CudaGraphAnnotation_t cuda_graph_annotation_id) const { - return cuda_graph_.IsGraphCaptured(cuda_graph_annotation_id); -} - -Status NvExecutionProvider::PerThreadContext::ReplayGraph(CudaGraphAnnotation_t cuda_graph_annotation_id, bool sync_status_flag) { - return cuda_graph_.Replay(cuda_graph_annotation_id, sync_status_flag); -} - -void NvExecutionProvider::PerThreadContext::IncrementRegularRunCountBeforeGraphCapture(CudaGraphAnnotation_t cuda_graph_annotation_id) { - graph_id_to_run_count_[cuda_graph_annotation_id]++; -} - -bool NvExecutionProvider::PerThreadContext::IsTensorRTContextInMap(std::string fused_node) { - auto it = trt_context_map_.find(fused_node); - if (it != trt_context_map_.end()) { - return true; - } - return false; -} - -nvinfer1::IExecutionContext& NvExecutionProvider::PerThreadContext::GetTensorRTContext(std::string fused_node) { - auto it = trt_context_map_.find(fused_node); - if (it != trt_context_map_.end()) { - return *(it->second.get()); // dereference shared pointer - } - auto context = tensorrt_ptr::unique_pointer_exec_ctx(); - trt_context_map_[fused_node] = std::move(context); - return *(trt_context_map_[fused_node].get()); // dereference shared pointer -} - -void NvExecutionProvider::ReleasePerThreadContext() const { - const auto& per_thread_context_cache = PerThreadContextCache(); - - auto cached_context_it = per_thread_context_cache->find(this); - ORT_ENFORCE(cached_context_it != per_thread_context_cache->end()); - auto cached_context = cached_context_it->second.lock(); - ORT_ENFORCE(cached_context); - - { - std::lock_guard lock(context_state_.mutex); - context_state_.active_contexts.erase(cached_context); - context_state_.retired_context_pool.push_back(cached_context); - } - - per_thread_context_cache->erase(cached_context_it); -} - -NvExecutionProvider::PerThreadContext& NvExecutionProvider::GetPerThreadContext() const { - const auto& per_thread_context_cache = PerThreadContextCache(); - - // try to use cached context - auto cached_context_it = per_thread_context_cache->find(this); - if (cached_context_it != per_thread_context_cache->end()) { - auto cached_context = cached_context_it->second.lock(); - ORT_ENFORCE(cached_context); - return *cached_context; - } - - // get context and update cache - std::shared_ptr context; - { - std::lock_guard lock(context_state_.mutex); - - // get or create a context - if (context_state_.retired_context_pool.empty()) { - context = std::make_shared(narrow(info_.device_id), - info_.has_user_compute_stream, stream_); - } else { - context = context_state_.retired_context_pool.back(); - context_state_.retired_context_pool.pop_back(); - } - - // insert into active_contexts, should not already be present - const auto active_contexts_insert_result = context_state_.active_contexts.insert(context); - ORT_ENFORCE(active_contexts_insert_result.second); - - // insert into caches_to_update_on_destruction, may already be present - ORT_IGNORE_RETURN_VALUE(context_state_.caches_to_update_on_destruction.insert(per_thread_context_cache)); - } - - per_thread_context_cache->insert(std::make_pair(this, context)); - - return *context; -} - -NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) - : IExecutionProvider{onnxruntime::kNvTensorRTRTXExecutionProvider, - OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, - narrow(info.device_id))}, - info_(info), - device_id_(info.device_id) { - InitProviderOrtApi(); - - if (!info.has_user_compute_stream) { - // If the app is passing in a compute stream, it already has initialized cuda and created a context. - // Calling cudaSetDevice() will set the default context in the current thread - // which may not be compatible with the stream created by the app. - CUDA_CALL_THROW(cudaSetDevice(device_id_)); - } - - cudaDeviceProp prop; - CUDA_CALL_THROW(cudaGetDeviceProperties(&prop, device_id_)); - auto cc = prop.major * 10 + prop.minor; - if (!(cc == 86 || cc == 89 || cc >= 120)) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[NvTensorRTRTX EP] The execution provider only supports RTX devices with compute capabilities 86, 89, 120 and above")); - } - compute_capability_ = GetComputeCapability(prop); - if (info.has_user_compute_stream) { - external_stream_ = true; - stream_ = static_cast(info.user_compute_stream); - } else if (cuda_graph_enable_) { - external_stream_ = false; - CUDA_CALL_THROW(cudaStreamCreate(&stream_)); - } else { - external_stream_ = false; - stream_ = nullptr; // Will be created in compute function - } - - if (info.user_aux_stream_array != nullptr) { - if (info.auxiliary_streams <= 0) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Auxiliary streams must be greater than 0 when using external auxiliary streams")); - } - external_aux_streams_ = true; - aux_streams_ = reinterpret_cast(info.user_aux_stream_array); - } else { - external_aux_streams_ = false; - aux_streams_ = nullptr; - } - - std::string profile_min_shapes, profile_max_shapes, profile_opt_shapes; - - // incase the EP context is dumped the engine cache has to be enabled - auto enable_engine_cache_for_ep_context_model = [this]() { - if (dump_ep_context_model_ && ep_context_embed_mode_ == 0) { - engine_cache_enable_ = true; - } - }; - - max_partition_iterations_ = info.max_partition_iterations; - min_subgraph_size_ = info.min_subgraph_size; - max_workspace_size_ = info.max_workspace_size; - max_shared_mem_size_ = info.max_shared_mem_size; - dump_subgraphs_ = info.dump_subgraphs; - weight_stripped_engine_enable_ = info.weight_stripped_engine_enable; - // make runtime cache path absolute and create directory if it doesn't exist - if (!info.runtime_cache_path.empty()) { - std::filesystem::path p(info.runtime_cache_path); - std::filesystem::path abs_path = std::filesystem::absolute(p); - const auto& env = GetDefaultEnv(); - auto status = env.CreateFolder(abs_path.string()); - if (!status.IsOK()) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] The runtime cache directory could not be created at: " << abs_path - << ". Runtime cache is disabled."; - } else { - runtime_cache_ = abs_path; - } - } - - onnx_model_folder_path_ = info.onnx_model_folder_path; - onnx_model_bytestream_ = info.onnx_bytestream; - onnx_model_bytestream_size_ = info.onnx_bytestream_size; - if ((onnx_model_bytestream_ != nullptr && onnx_model_bytestream_size_ == 0) || - (onnx_model_bytestream_ == nullptr && onnx_model_bytestream_size_ != 0)) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "When providing either 'trt_onnx_bytestream_size' or " - "'trt_onnx_bytestream' both have to be provided")); - } - use_external_data_initializer_ = info.use_external_data_initializer; - onnx_external_data_bytestream_ = info.external_data_bytestream; - onnx_external_data_bytestream_size_ = info.external_data_bytestream_size; - if ((onnx_external_data_bytestream_ != nullptr && onnx_external_data_bytestream_size_ == 0) || - (onnx_external_data_bytestream_ == nullptr && onnx_external_data_bytestream_size_ != 0)) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "When providing either 'onnx_external_data_bytestream_size' or " - "'onnx_external_data_bytestream' both have to be provided")); - } - detailed_build_log_ = info.detailed_build_log; - dump_ep_context_model_ = info.dump_ep_context_model; - ep_context_file_path_ = info.ep_context_file_path; - ep_context_embed_mode_ = info.ep_context_embed_mode; - compile_only_mode_ = info.compile_only_mode; - enable_engine_cache_for_ep_context_model(); - cache_prefix_ = info.engine_cache_prefix; - // use a more global cache if given - engine_decryption_enable_ = info.engine_decryption_enable; - if (engine_decryption_enable_) { - engine_decryption_lib_path_ = info.engine_decryption_lib_path; - } - force_sequential_engine_build_ = info.force_sequential_engine_build; - sparsity_enable_ = info.sparsity_enable; - auxiliary_streams_ = info.auxiliary_streams; - profile_min_shapes = info.profile_min_shapes; - profile_max_shapes = info.profile_max_shapes; - profile_opt_shapes = info.profile_opt_shapes; - - /* - * Parse explicit min/max/opt profile shapes from provider options. - * - * The format of min/max/opt profile shapes is defined as below: - * "input1:dim1xdim2...,input2:dim1xdim2...,...,input1:dim3xdim4...,input2:dim3xdim4...,..." - * - * (Note: if multiple shapes with same input name are specified, TRT EP will consider them as multiple profiles. - * Please refer to ParserProfileShapes() for more details) - * - */ - bool status = true; - if (status) { - status = ParseProfileShapes(profile_min_shapes, profile_min_shapes_); - if (!status) { - profile_min_shapes_.clear(); - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] The format of provider option 'trt_profile_min_shapes' is wrong, please follow the format of 'input1:dim1xdimd2...,input2:dim1xdim2...,...'"; - } - } - - if (status) { - status = ParseProfileShapes(profile_max_shapes, profile_max_shapes_); - if (!status) { - profile_max_shapes_.clear(); - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] The format of provider option 'trt_profile_max_shapes' is wrong, please follow the format of 'input1:dim1xdimd2...,input2:dim1xdim2...,...'"; - } - } - - if (status) { - status = ParseProfileShapes(profile_opt_shapes, profile_opt_shapes_); - if (!status) { - profile_opt_shapes_.clear(); - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] The format of provider option 'trt_profile_opt_shapes' is wrong, please follow the format of 'input1:dim1xdimd2...,input2:dim1xdim2...,...'"; - } - } - - if (status) { - status = ValidateProfileShapes(profile_min_shapes_, profile_max_shapes_, profile_opt_shapes_); - if (!status) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Profile shapes validation failed. Make sure the provider options 'trt_profile_min_shapes', 'trt_profile_max_shapes' and 'trt_profile_opt_shapes' have same input name and number of profile."; - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] TRT EP will implicitly create optimization profiles based on input tensor for you."; - profile_min_shapes_.clear(); - profile_max_shapes_.clear(); - profile_opt_shapes_.clear(); - } - } - - cuda_graph_enable_ = info.cuda_graph_enable; - multi_profile_enable_ = info.multi_profile_enable; - op_types_to_exclude_ = info.op_types_to_exclude; - - // Validate setting - if (max_partition_iterations_ <= 0) { - // LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] TensorRT option nv_max_partition_iterations must be a positive integer value. Set it to 1000"; - max_partition_iterations_ = 1000; - } - if (min_subgraph_size_ <= 0) { - // LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] TensorRT option nv_min_subgraph_size must be a positive integer value. Set it to 1"; - min_subgraph_size_ = 1; - } - - // If ep_context_file_path_ is provided as a directory, create it if it's not existed - if (dump_ep_context_model_ && !ep_context_file_path_.empty() && std::filesystem::path(ep_context_file_path_).extension().empty() && !std::filesystem::is_directory(ep_context_file_path_)) { - if (!std::filesystem::create_directory(ep_context_file_path_)) { - throw std::runtime_error("Failed to create directory " + ep_context_file_path_); - } - } - - // If dump_ep_context_model_ is enabled, TRT EP forces cache_path_ to be the relative path of ep_context_file_path_. - // For example, - // - original cache path = "engine_cache_dir" -> new cache path = "./context_model_dir/engine_cache_dir" - // - original cache path = "" -> new cache path = "./context_model_dir" - // The new cache path will be saved as the "ep_cache_context" node attritue of the EP context node. - // For security reason, it needs to make sure the engine cache is saved inside context model directory. - if (dump_ep_context_model_) { - // TODO(maximilianm) not sure if this is still needed - engine_cache_enable_ = true; - if (path_utils::IsAbsolutePath(cache_path_)) { - LOGS_DEFAULT(ERROR) << "In the case of dumping context model and for security purpose, the trt_engine_cache_path should be set with a relative path, but it is an absolute path: " << cache_path_; - } - if (path_utils::IsRelativePathToParentPath(cache_path_)) { - LOGS_DEFAULT(ERROR) << "In the case of dumping context model and for security purpose, The trt_engine_cache_path has '..', it's not allowed to point outside the directory."; - } - - // Engine cache relative path to context model directory. - // It's used when dumping the "ep_cache_context" node attribute. - engine_cache_relative_path_to_context_model_dir = cache_path_; - - // Make cache_path_ to be the relative path of ep_context_file_path_ - cache_path_ = path_utils::GetDirOrParentPath(ep_context_file_path_).append(cache_path_).string(); - } - - if (engine_decryption_enable_) { - LIBTYPE handle = OPENLIB(engine_decryption_lib_path_.c_str()); - if (handle == nullptr) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP could not open shared library from " + engine_decryption_lib_path_)); - } - engine_decryption_ = (int (*)(const char*, char*, size_t*))LIBFUNC(handle, "decrypt"); - engine_encryption_ = (int (*)(const char*, char*, size_t))LIBFUNC(handle, "encrypt"); - if (engine_decryption_ == nullptr) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP could not find decryption function in shared library from " + engine_decryption_lib_path_)); - } - } - - // cuda graph: - // cudaStreamSynchronize() is not allowed in cuda graph capture. - // - // external stream: - // If user provides "external" cuda stream, only this cuda stream will be used even if multiple threads are running InferenceSession.Run() concurrently. - // So, no need to synchronize different streams after enqueueV3. - if (external_stream_) { - sync_stream_after_enqueue_ = false; - } - - { - auto lock = GetApiLock(); - runtime_ = std::unique_ptr(nvinfer1::createInferRuntime(GetTensorrtLogger(detailed_build_log_))); - } - - trt_version_ = getInferLibVersion(); - CUDA_CALL_THROW(cudaRuntimeGetVersion(&cuda_version_)); - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] TensorRT version is " << trt_version_; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] CUDA version is " << cuda_version_; - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Nv provider options: " - << "device_id: " << device_id_ - << ", nv_max_partition_iterations: " << max_partition_iterations_ - << ", nv_min_subgraph_size: " << min_subgraph_size_ - << ", nv_max_workspace_size: " << max_workspace_size_ - << ", nv_dump_subgraphs: " << dump_subgraphs_ - << ", nv_weight_stripped_engine_enable: " << weight_stripped_engine_enable_ - << ", nv_onnx_model_folder_path: " << onnx_model_folder_path_ - << ", nv_engine_decryption_enable: " << engine_decryption_enable_ - << ", nv_engine_decryption_lib_path: " << engine_decryption_lib_path_ - << ", nv_force_sequential_engine_build: " << force_sequential_engine_build_ - << ", nv_sparsity_enable: " << sparsity_enable_ - << ", nv_auxiliary_streams: " << auxiliary_streams_ - << ", enable_cuda_graph: " << cuda_graph_enable_ - << ", nv_dump_ep_context_model: " << dump_ep_context_model_ - << ", nv_ep_context_file_path: " << ep_context_file_path_ - << ", nv_ep_context_embed_mode: " << ep_context_embed_mode_ - << ", nv_cache_prefix: " << cache_prefix_ - << ", nv_onnx_model_bytestream_size_: " << onnx_model_bytestream_size_ - << ", nv_onnx_external_bytestream_size_: " << onnx_external_data_bytestream_size_ - << ", nv_use_external_data_initializer_: " << use_external_data_initializer_ - << ", nv_op_types_to_exclude: " << op_types_to_exclude_ - << ", nv_runtime_cache_path: " << runtime_cache_; -} - -Status NvExecutionProvider::Sync() const { - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream_)); - return Status::OK(); -} - -NvExecutionProvider::~NvExecutionProvider() { - // clean up thread local context caches - { - std::lock_guard lock(context_state_.mutex); - for (const auto& cache_weak : context_state_.caches_to_update_on_destruction) { - const auto cache = cache_weak.lock(); - if (!cache) continue; - ORT_IGNORE_RETURN_VALUE(cache->erase(this)); - } - } - - if (!external_stream_ && stream_ != nullptr) { - ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaStreamDestroy(stream_))); - } - ReleaseTensorRTCustomOpDomainList(info_.custom_op_domain_list); - - if (alloc_ != nullptr) { - // This code is same as OrtApis::ReleaseAllocator defined in allocator_adapters.cc. - // We can't get api inside destructor so that's why we duplicate the code here. - delete static_cast(alloc_); - } -} - -void NvExecutionProvider::HandleCudaGraphStart(cudaStream_t stream, bool require_io_binding, - CudaGraphAnnotation_t cuda_graph_annotation_id, bool& graph_replay_on_this_run, bool& should_start_capture) { - graph_replay_on_this_run = false; - should_start_capture = false; - - // Case 1: CUDA Graph capture is enabled AND IO binding is required. - // In this case, we force graph re-capture by resetting warmup runs. - // If a graph for this annotation ID already exists, delete it before proceeding. - if (require_io_binding && cuda_graph_enable_) { - GetPerThreadContext().ResetWarmupRuns(cuda_graph_annotation_id); - - if (GetPerThreadContext().IsGraphCaptured(cuda_graph_annotation_id)) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Graph already captured and required_io_binding is true, resetting warmup runs and deleting graph"; - GetPerThreadContext().DeleteCapturedGraph(cuda_graph_annotation_id); - } - // Case 2: CUDA Graph capture is enabled AND IO binding is NOT required - } else if (cuda_graph_enable_ && !require_io_binding) { - // If the graph is not yet captured, increment the regular run counter - if (cuda_graph_annotation_id != kCudaGraphAnnotationSkip && - !GetPerThreadContext().IsGraphCaptured(cuda_graph_annotation_id)) { - GetPerThreadContext().IncrementRegularRunCountBeforeGraphCapture(cuda_graph_annotation_id); - } - - // If capture is allowed and graph not already captured, - // set the stream and begin capture - if (!GetPerThreadContext().IsGraphCaptured(cuda_graph_annotation_id) && - GetPerThreadContext().IsGraphCaptureAllowed(cuda_graph_annotation_id)) { - GetPerThreadContext().SetCudaGraphStream(stream); - GetPerThreadContext().CaptureBegin(cuda_graph_annotation_id); - should_start_capture = true; - } - - // If a graph is already captured for this ID, mark it for replay in this run. - if (GetPerThreadContext().IsGraphCaptured(cuda_graph_annotation_id)) { - graph_replay_on_this_run = true; - } - } -} - -bool NvExecutionProvider::IsGraphCaptureEnabled() const { - // Return false so that ORT's framework does not cache this EP for ORT-managed graph capture/replay. - // NvTensorRTRTX manages CUDA graph capture/replay internally. - return false; -} - -bool NvExecutionProvider::IsGraphCaptured(int graph_annotation_id) const { - // This is hardcoded to always return false because we are not allowing the ORT framework to have the CUDA graph control. - (void)graph_annotation_id; - return false; -} - -Status NvExecutionProvider::ReplayGraph(int graph_annotation_id, bool /*sync*/) { - // The sync parameter is ignored: NV TRT RTX EP manages its own CUDA graph lifecycle and always replays synchronously. - // This is hardcoded to always return OK because we are not allowing the ORT framework to have the CUDA graph control. - (void)graph_annotation_id; - return Status::OK(); -} - -Status NvExecutionProvider::OnRunStart(const onnxruntime::RunOptions& run_options) { - if (cuda_graph_enable_) { - CudaGraphAnnotation_t cuda_graph_annotation_id = GetPerThreadContext().GetCudaGraphAnnotationId(run_options); - GetPerThreadContext().SetCurrentGraphAnnotationId(cuda_graph_annotation_id); - } - - if (multi_profile_enable_ == true) { - auto graph_annotation_str = - run_options.GetConfigOptions().GetConfigEntry(nv::run_option_names::kProfileIndex); - TryParseStringWithClassicLocale(*graph_annotation_str, nv_profile_index_); - } - return Status::OK(); -} - -Status NvExecutionProvider::OnRunEnd(bool sync_stream, const onnxruntime::RunOptions& run_options) { - (void)run_options; - - if (sync_stream && external_stream_) { - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream_)); - } - return Status::OK(); -} - -std::vector NvExecutionProvider::CreatePreferredAllocators() { - OrtArenaCfg arena_cfg(0, static_cast(ArenaExtendStrategy::kNextPowerOfTwo), - -1, -1, -1, -1); - AllocatorCreationInfo default_memory_info( - [](OrtDevice::DeviceId device_id) { return std::make_unique(device_id, CUDA); }, - narrow(device_id_), - true, - arena_cfg, - // make it stream aware - true); - - AllocatorCreationInfo pinned_allocator_info( - [](OrtDevice::DeviceId device_id) { - return std::make_unique(device_id, CUDA_PINNED); - }, - narrow(device_id_)); - - return std::vector{CreateAllocator(default_memory_info), CreateAllocator(pinned_allocator_info)}; -} - -std::unique_ptr NvExecutionProvider::GetDataTransfer() const { - return std::make_unique(); -} - -// Get the pointer to the IBuilder instance. -// Note: This function is not thread safe. Calls to this function from different threads must be serialized -// even though it doesn't make sense to have multiple threads initializing the same inference session. -nvinfer1::IBuilder* NvExecutionProvider::GetBuilder(TensorrtLogger& trt_logger) const { - if (!builder_) { - { - auto lock = GetApiLock(); - builder_ = std::unique_ptr(nvinfer1::createInferBuilder(trt_logger)); - unsigned int num_threads = std::thread::hardware_concurrency(); - builder_->setMaxThreads(num_threads / 2); - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Set threads that the builder can use to:" << builder_->getMaxThreads(); - } - } - return builder_.get(); -} - -void NvExecutionProvider::GetCustomOpDomainList(std::vector& custom_op_domain_list) const { - auto status = CreateTensorRTCustomOpDomainList(custom_op_domain_list, info_.extra_plugin_lib_paths); - if (status != Status::OK()) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Failed to get TRT plugins from TRT plugin registration."; - } -} - -// Check the graph is the subgraph of control flow op -bool NvExecutionProvider::IsSubGraphOfControlFlowOp(const GraphViewer& graph) const { - if (graph.IsSubgraph()) { - const auto& node = graph.ParentNode(); - if (control_flow_op_set_.find(node->OpType()) != control_flow_op_set_.end()) { - return true; - } - } - return false; -} - -// Check whether all the nodes of the graph are assigned to specific ep -bool NvExecutionProvider::AllNodesAssignedToSpecificEP(const GraphViewer& graph, const std::string& provider_type) const { - const int number_of_ort_nodes = graph.NumberOfNodes(); - std::vector nodes_vector(number_of_ort_nodes); - std::iota(std::begin(nodes_vector), std::end(nodes_vector), 0); - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - for (const auto& index : nodes_vector) { - const auto& node = graph.GetNode(node_index[index]); - if (node->GetExecutionProviderType() != provider_type) { - return false; - } - } - - return number_of_ort_nodes != 0; -} - -// Check whether all the nodes of subgraph are supported -bool NvExecutionProvider::IsSubGraphFullySupported(SubGraphCollection_t supported_nodes_vector, const int number_of_ort_nodes) const { - int number_of_trt_nodes = 0; - for (const auto& group : supported_nodes_vector) { - if (!group.first.empty()) { - number_of_trt_nodes += static_cast(group.first.size()); - } - } - - return number_of_trt_nodes == number_of_ort_nodes; -} - -std::unique_ptr NvExecutionProvider::GetSubGraph(SubGraph_t graph_nodes_index, const GraphViewer& graph, const HashValue& model_hash, int subgraph_index) const { - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - std::unordered_set node_set; - node_set.reserve(graph_nodes_index.first.size()); - for (const auto& index : graph_nodes_index.first) { - node_set.insert(node_index[index]); - } - - // Get parent graph output names - std::unordered_set graph_output_names; - for (const auto* output_arg : graph.GetOutputs()) { - graph_output_names.insert(output_arg->Name()); - } - - // Find inputs and outputs of the subgraph - - std::unique_ptr sub_graph = onnxruntime::IndexedSubGraph::Create(); - std::unordered_map original_inputs; - - // These maps store the inputs and outputs of the subgraph. - // Please note that the inputs and outputs of the maps will be dynamically updated during node iteration - // to determine the final inputs and outputs of the subgraph. - std::unordered_map fused_inputs, fused_outputs; - - // This map stores the node's output that will be consumed by another node outside of this subgraph. - // So the node's output should be put into the subgraph's output list. - std::unordered_map fused_outputs_to_add; - - // This map stores the node's output that is original graph's output. - // So the node's output should be put into the subgraph's output list. - std::unordered_map graph_outputs_to_add; - - std::unordered_set erased; - - // This is the relative ordering that ensures node's input or output being added to the 'fused_inputs', - // 'fused_outputs', 'fused_outputs_to_add' and 'graph_outputs_to_add' maps is associated with a relative order index. - // Items added earlier receive a smaller order index than items added later. - // When constructing the final sub_graph's input or output lists, entries with smaller - // order indices will appear before those with larger indices. - int input_order = 0; - int output_order = 0; - - std::vector initializers; - for (const auto& index : graph_nodes_index.first) { - sub_graph->Nodes().push_back(node_index[index]); - const auto& node = graph.GetNode(node_index[index]); - for (const auto& input : node->InputDefs()) { - if (graph.IsConstantInitializer(input->Name(), true)) { - initializers.push_back(input->Name()); - continue; - } - const auto& it = fused_outputs.find(input); - if (it != fused_outputs.end()) { - fused_outputs.erase(it); - erased.insert(input); - } else if (erased.find(input) == erased.end()) { - // Only when input is neither in output list nor erased list, add the input to input list - fused_inputs.insert({input, input_order++}); - } - } - - for (const auto& input : node->ImplicitInputDefs()) { - if (graph.IsConstantInitializer(input->Name(), true)) { - initializers.push_back(input->Name()); - continue; - } - const auto& it = fused_outputs.find(input); - if (it != fused_outputs.end()) { - fused_outputs.erase(it); - erased.insert(input); - } else if (erased.find(input) == erased.end()) { - // Only when input is neither in output list nor erased list, add the input to input list - fused_inputs.insert({input, input_order++}); - } - } - - // For output searching, there are two special cases, - // One is, if node's OutputEdges are more than its outputs, meaning certain output is used more than once, - // if the output is connected to nodes that don't belong to the subgraph, the output need to be added - // to the output list - // The other one is, if subgraph's node output is parent graph's output. the node output should - // be also added to the subgraph's output list - if (node->GetOutputEdgesCount() > node->OutputDefs().size()) { - for (auto it = node->OutputEdgesBegin(), end = node->OutputEdgesEnd(); it != end; ++it) { - const auto& node_idx = it->GetNode().Index(); - const onnxruntime::NodeArg* output; - // The dst_arg_index from GetDstArgIndex() could be the index for explicit/implicit input defs of the node. - // We need to get the correct input index accordingly. (See Graph::BuildConnections() in graph.cc for more details) - if (it->GetDstArgIndex() < static_cast(it->GetNode().InputDefs().size())) { - output = (it->GetNode()).InputDefs()[it->GetDstArgIndex()]; - } else { - output = (it->GetNode()).ImplicitInputDefs()[it->GetDstArgIndex() - static_cast(it->GetNode().InputDefs().size())]; - } - - if (node_set.find(node_idx) == node_set.end()) { - // This output will be consumed by another node outside of this subgraph. - // So the output should be put into the subgraph's output list. - fused_outputs_to_add.insert({output, output_order++}); - } - } - } - - for (const auto& output : node->OutputDefs()) { - const auto& it = fused_inputs.find(output); - if (it != fused_inputs.end()) { - fused_inputs.erase(it); - erased.insert(output); - } else if (erased.find(output) == erased.end()) { - if (graph.GetGraph().GetConsumerNodes(output->Name()).size() > 0) { - // Only when output is neither in input list nor erased list, - // and the output is consumed by another node, add the output to output list - fused_outputs.insert({output, output_order++}); - } - } - - if (graph_output_names.find(output->Name()) != graph_output_names.end()) { - // This output is the graph's output. - // So the output should be put into the subgraph's output list. - graph_outputs_to_add.insert({output, output_order++}); - } - } - } - - fused_outputs.insert(fused_outputs_to_add.begin(), fused_outputs_to_add.end()); - fused_outputs.insert(graph_outputs_to_add.begin(), graph_outputs_to_add.end()); - - std::multimap inputs, outputs; - - // Get the input order of the original graph - int order = 0; - for (const auto* input : graph.GetInputs()) { - original_inputs[input] = order++; - } - - // input order needs to be consistent with original graph's input order - for (auto it = fused_inputs.begin(), end = fused_inputs.end(); it != end; ++it) { - const auto& iter = original_inputs.find(it->first); - if (iter != original_inputs.end()) { - inputs.insert(std::pair(iter->second, iter->first)); - } else { - inputs.insert(std::pair(it->second, it->first)); - } - } - - // Sort outputs by the order they were added - for (auto it = fused_outputs.begin(), end = fused_outputs.end(); it != end; ++it) { - outputs.insert(std::pair(it->second, it->first)); - } - - // Generate unique kernel name for TRT subgraph - std::string subgraph_id = std::to_string(model_hash) + "_" + std::to_string(subgraph_index); - auto meta_def = IndexedSubGraph_MetaDef::Create(); - const std::string graph_type = graph.IsSubgraph() ? "subgraph" : "graph"; - meta_def->name() = "TRTKernel_" + graph_type + "_" + graph.Name() + "_" + subgraph_id; - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] TensorRT subgraph MetaDef name " + meta_def->name(); - - // Assign inputs and outputs to subgraph's meta_def - for (const auto& input : inputs) { - if (input.second->Exists()) { - meta_def->inputs().push_back(input.second->Name()); - } - } - - for (const auto& initializer : initializers) { - meta_def->constant_initializers().push_back(initializer); - } - - for (const auto& output : outputs) { - if (output.second->Exists()) { - meta_def->outputs().push_back(output.second->Name()); - } - } - - meta_def->domain() = kMSDomain; - meta_def->since_version() = 1; - sub_graph->SetMetaDef(std::move(meta_def)); - - return sub_graph; -} - -SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t nodes_vector_input, int iterations, const int max_iterations, - const GraphViewer& graph, bool* early_termination) const { - // Return if iterations are exceeding predefined number - SubGraphCollection_t nodes_list_output; - if (iterations > max_iterations) { - *early_termination = true; - return nodes_list_output; - } - - // Get parent graph output names - std::unordered_set graph_output_names; - for (const auto* output_arg : graph.GetOutputs()) { - graph_output_names.insert(output_arg->Name()); - } - - iterations++; - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - for (const auto& group : nodes_vector_input) { - // Construct subgraph - if (!group.first.empty()) { - if (group.second) { - nodes_list_output.push_back(group); - } else { - auto model_build = graph.CreateModel(*GetLogger()); - auto& graph_build = model_build->MainGraph(); - bool has_control_flow_op = false; - - // Add node and node args - // If node output is also parent graph output, the output will be added to the - // subgraph's output list - // - // Initializers that refer to a memory location in OrtValue - // can not be handled by TRT (unlike those that are on disk). - // This prevents us from sharing the data and we have to make a copy here. - constexpr const bool load_initializers_inline_true = true; - std::vector subgraph_output_names; - for (const auto& index : group.first) { - const auto& node = graph.GetNode(node_index[index]); - std::vector inputs, outputs; - for (auto input : node->InputDefs()) { - auto& n_input = graph_build.GetOrCreateNodeArg(input->Name(), input->TypeAsProto()); - inputs.push_back(&n_input); - graph_utils::MakeInitializerCopyIfNotExist(graph.GetGraph(), graph_build, input->Name(), - load_initializers_inline_true); - } - - for (auto input : node->ImplicitInputDefs()) { - graph_utils::MakeInitializerCopyIfNotExist(graph.GetGraph(), graph_build, input->Name(), - load_initializers_inline_true); - } - - for (auto output : node->OutputDefs()) { - auto& n_output = graph_build.GetOrCreateNodeArg(output->Name(), output->TypeAsProto()); - outputs.push_back(&n_output); - const auto name = output->Name(); - if (graph_output_names.find(name) != graph_output_names.end()) { - subgraph_output_names.push_back(name); - } - } - - if (control_flow_op_set_.find(node->OpType()) != control_flow_op_set_.end()) { - has_control_flow_op = true; - } - - // If the node has subgraph, it's possible that the ORT graph of that subgraph and the GraphProto in the node attributes are not in sync because of graph optimization. - // Therefore, we need to force GraphProto attributes to be updated in order to get the valid GraphProto. - if (node->GetAttributes().size() > 0) { - auto node_proto = ONNX_NAMESPACE::NodeProto::Create(); - // we need to update any GraphProto attributes for subgraphs so that any changes made by things - // such as the optimizers are captured. otherwise we can end up saving an invalid graph. - node->ToProto(*node_proto, /* update_subgraphs */ true); - const int num_attributes = node_proto->attribute_size(); - auto node_attributes = ONNX_NAMESPACE::NodeAttributes::Create(); - node_attributes->reserve(num_attributes); - - for (int i = 0; i < num_attributes; ++i) { - auto& attr = node_proto->attribute(i); - node_attributes->emplace(attr.name(), attr); - } - - // The GraphProto attributes are the updated ones. - graph_build.AddNode(node->Name(), node->OpType(), node->Description(), inputs, outputs, node_attributes.get(), node->Domain()); - } else { - // The GraphProto attributes are the original ones. - graph_build.AddNode(node->Name(), node->OpType(), node->Description(), inputs, outputs, &node->GetAttributes(), node->Domain()); - } - } - - // Only if the newly built graph has control flow op as well as it has parent node, - // it needs to handle outer scope values before calling graph.Resolve(). - if (has_control_flow_op && graph.ParentNode()) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Handle outer scope values for the subgraph " << graph_build.Name(); - BuildSubGraphContext(graph_build); - SetGraphOuterScopeValuesAndInputs(graph_build, graph.GetGraph()); - SetAllGraphInputs(graph_build); - } - - ORT_THROW_IF_ERROR(graph_build.Resolve()); - - // Add parent graph output to the subgraph - int i = 0; - std::vector subgraph_outputs; - subgraph_outputs.resize(subgraph_output_names.size()); - for (auto& name : subgraph_output_names) { - auto output_arg = graph.GetNodeArg(name); - auto& subgraph_output_arg = graph_build.GetOrCreateNodeArg(output_arg->Name(), output_arg->TypeAsProto()); - subgraph_outputs[i] = &subgraph_output_arg; - ++i; - } - auto& graph_build_outputs = graph_build.GetOutputs(); - subgraph_outputs.insert(subgraph_outputs.begin(), graph_build_outputs.begin(), graph_build_outputs.end()); - graph_build.SetOutputs(graph_build_outputs); - ORT_ENFORCE(graph_build.Resolve().IsOK()); - - // Check if input tensors have shapes - if (iterations > 1) { - auto graph_inputs = graph_build.GetInputs(); - for (auto input_arg : graph_inputs) { - bool has_dim_value_or_param = true; - auto input_shape = input_arg->Shape(); - if (input_shape != nullptr) { - auto dim_size = input_shape->dim_size(); - for (int i = 0; i < dim_size; ++i) { - auto& dim = input_shape->dim(i); - if (!dim.has_dim_value() && !dim.has_dim_param()) { - has_dim_value_or_param = false; - break; - } - } - } - - if (input_shape == nullptr || !has_dim_value_or_param) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "TensorRT input: " + input_arg->Name() + " has no shape specified. " + - "Please run shape inference on the onnx model first. Details can be found in " + - "https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#shape-inference-for-tensorrt-subgraphs")); - } - } - } - - // Serialize modelproto to string - auto graph_viewer = graph_build.CreateGraphViewer(); - auto model = graph_viewer->CreateModel(*GetLogger()); - auto model_proto = model->ToProto(); - - // save user provided external data in memory instead of writing to ModelProto - // needed for models > 2GB - std::vector userWeights; - if (use_external_data_initializer_) { - const auto& allInitializers = graph_viewer->GetAllInitializedTensors(); - userWeights.reserve(allInitializers.size()); - for (const auto& [name, tp] : allInitializers) { - if (utils::HasRawData(*tp)) { - // Keep inits in memory instead of writing to ModelProto. - userWeights.emplace_back(name, tp->raw_data().data(), tp->raw_data().size()); - } else if (utils::HasExternalDataInMemory(*tp)) { - // the initializer was marked as external data by the ORT graph at load time since it was provided in memory - if (OrtValue v; graph_viewer->GetOrtValueInitializer(name, v)) { - Ort::ConstValue initializer_value{&v}; - const size_t size = initializer_value.GetTensorSizeInBytes(); - const void* ptr = initializer_value.GetTensorRawData(); - userWeights.emplace_back(name, ptr, size); - } else { - // only copy and take ownership of the data if none of the above conditions are met - std::unique_ptr full_init; - ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tp, full_init)); - userWeights.emplace_back(name, full_init->raw_data()); - } - } - } - } - - // ORT's default topological sort is using reversed DFS. - // When creating model proto from graph viewer, let ORT use priority-based topological sort based on node index. - // The reason is, in some cases, for example ResNet50, using default topological sort will end up with generating - // the model proto that has different node ordering compared to original onnx model. - graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, !use_external_data_initializer_ /*include raw initializers*/); - model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); - - std::string string_buf; - model_proto->SerializeToString(string_buf); - - if (dump_subgraphs_) { - // Dump TensorRT subgraph for debugging - std::fstream dump("NvExecutionProvider_TRT_Subgraph.onnx", std::ios::out | std::ios::trunc | std::ios::binary); - model_proto->SerializeToOstream(dump); - } - - // Get supported node list recursively - SubGraphCollection_t parser_nodes_list; - TensorrtLogger& trt_logger = GetTensorrtLogger(detailed_build_log_); - auto trt_builder = GetBuilder(trt_logger); -#if TRT_MAJOR_RTX >= 2 || (TRT_MAJOR_RTX == 1 && ((TRT_MINOR_RTX == 5 && TRT_BUILD_RTX >= 97) || TRT_MINOR_RTX >= 6)) - // kSTRONGLY_TYPED == 0 => bit flag value is 1U. Use literal to avoid deprecated-enum warning (deprecated since 1.5.0.97). - constexpr uint32_t network_flags = 1U; -#else - auto network_flags = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kSTRONGLY_TYPED); -#endif - auto trt_network = std::unique_ptr(trt_builder->createNetworkV2(network_flags)); - - bool is_model_supported = false; - - // limit the scope of trt_parser so that model gets unloaded from memory asap - { - auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); - - if (use_external_data_initializer_) { -#if TRT_MAJOR_RTX > 1 || TRT_MINOR_RTX >= 1 - trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - for (auto const& userWeight : userWeights) { - trt_parser->loadInitializer(userWeight.Name(), userWeight.Data(), userWeight.Size()); - } - is_model_supported = trt_parser->parseModelProto(); -#else - ORT_THROW("'nv_use_external_data_initializer' is only supported on TensorRT RTX 1.1.x.x and above."); -#endif - } else { - is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); - } - - // Note: Calling getNbSubgraphs or getSubgraphNodes before calling supportsModelV2 results in undefined behavior. - auto num_subgraphs = trt_parser->getNbSubgraphs(); - parser_nodes_list.reserve(num_subgraphs); - - for (int64_t i = 0; i < num_subgraphs; ++i) { - int64_t subgraph_len = 0; - int64_t* nodes = trt_parser->getSubgraphNodes(i, subgraph_len); - parser_nodes_list.emplace_back(); - parser_nodes_list.back().first.reserve(subgraph_len); - for (int64_t j = 0; j < subgraph_len; ++j) { - parser_nodes_list.back().first.push_back(nodes[j]); - } - parser_nodes_list.back().second = is_model_supported ? true : false; - } - } - SubGraphCollection_t next_nodes_list; - const std::vector& subgraph_node_index = graph_viewer->GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - next_nodes_list = GetSupportedList(parser_nodes_list, iterations, max_iterations, *graph_viewer, early_termination); - for (size_t i = 0, end = next_nodes_list.size(); i < end; ++i) { - for (size_t j = 0, end = next_nodes_list[i].first.size(); j < end; ++j) { - /* - * Convert the supported node list returning from onnx-tensorrt parser to the node list recognized by ORT TRT. - * - * TRT EP reconstructs the graph based on the nodes in group.first and feeds this graph (converts to model proto and to string buffer) to onnx-tensorrt parser. - * The node index in the list returning from onnx-tensorrt parser might not be the same as the node index in group.first. Therefore, TRT EP needs a node index mapping table here. - * - * The order of iterating the nodes in group.first and calling graph_build.AddNode() determines the node order in the newly constructed graph (see Graph::AllocateNode() in graph.cc), - * however, once the graph is converted to model proto, the node proto order in model proto (ex: onnx-tensorrt calls model.graph().node() to iterate NodeProto in ModelProto) is decided by topo sort. - * - * The topo sort list (i.e. subgraph_node_index) acts as the node index mapping table: - * subgraph_node_index[node index from onnx-tensorrt parser] = index in group.first - * - * In the past, TRT EP uses ORT's default reversed DFS topo sort which might end up with the sorting result not sequence of 0, 1, ... n-1, ex: the subgraph_node_index = [0,2,1,3,4]. - * With the change of using ORT's priority-based topo sort (node with lower node index outputs first) the sorting result is the sequence of 0, 1, ... n-1 for most of the cases, - * therefore subgraph_node_index as a mapping table is not needed anymore. - * - * TODO: Remove the subgraph_node_index - */ - next_nodes_list[i].first[j] = group.first[subgraph_node_index[next_nodes_list[i].first[j]]]; - } - nodes_list_output.push_back(next_nodes_list[i]); - } - } - } - } - return nodes_list_output; -} - -// Detect and remove cycles from supported node list -bool NvExecutionProvider::DetectTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const GraphViewer& graph, const HashValue& model_hash, bool remove_cycles) const { - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - bool trt_cycle = true, cycle_detected = false; - while (trt_cycle) { - trt_cycle = false; - std::unordered_map node_to_index_map; - std::unordered_map index_to_node_map; - std::unordered_map> input_to_nodes_map, node_to_outputs_map; - std::unordered_set non_trt_node_index(node_index.begin(), node_index.end()); - size_t id = 0; - int subgraph_index = 0; - for (const auto& group : supported_nodes_vector) { - if (!group.first.empty()) { - // Construct subgraph from node list - std::unique_ptr sub_graph = GetSubGraph(group, graph, model_hash, subgraph_index); - - // Create node to inputs/outputs/index maps - const auto& meta_def = sub_graph->GetMetaDef(); - const std::string node_name = meta_def->name(); - if (node_to_index_map.find(node_name) == node_to_index_map.end()) { - index_to_node_map[id] = node_name; - node_to_index_map[node_name] = id++; - } - - if (meta_def != nullptr) { - for (const auto& input : meta_def->inputs()) { - input_to_nodes_map[input].insert(node_name); - } - for (const auto& output : meta_def->outputs()) { - node_to_outputs_map[node_name].insert(output); - } - } - - // Remove TensorRT nodes from node index list - for (const auto& index : group.first) { - non_trt_node_index.erase(node_index[index]); - } - subgraph_index++; - } - } - - // Add non TensorRT nodes to the maps - for (const auto& index : non_trt_node_index) { - const auto& node = graph.GetNode(index); - const std::string node_name = node->Name(); - if (node_to_index_map.find(node_name) == node_to_index_map.end()) { - index_to_node_map[id] = node_name; - node_to_index_map[node_name] = id++; - } - - for (const auto& input : node->InputDefs()) { - input_to_nodes_map[input->Name()].insert(node_name); - } - - for (const auto& input : node->ImplicitInputDefs()) { - input_to_nodes_map[input->Name()].insert(node_name); - } - - for (const auto& output : node->OutputDefs()) { - node_to_outputs_map[node_name].insert(output->Name()); - } - } - - // Create adjacency list - size_t graph_size = node_to_index_map.size(); - std::list* adjacency_map = new std::list[graph_size]; - for (const auto& node : node_to_outputs_map) { - for (auto iter = node.second.begin(); iter != node.second.end(); ++iter) { - const auto& loc = input_to_nodes_map.find(*iter); - if (loc != input_to_nodes_map.end()) { - size_t parent_node_index = node_to_index_map.find(node.first)->second; - for (auto child_node : loc->second) { - size_t child_node_index = node_to_index_map.find(child_node)->second; - adjacency_map[parent_node_index].push_back(child_node_index); - } - } - } - } - - // Check cycle in the graph - bool* visited = new bool[graph_size]; - bool* st = new bool[graph_size]; - for (size_t i = 0; i < graph_size; ++i) { - visited[i] = false; - st[i] = false; - } - - std::vector cycles; - bool has_cycle = false; - for (size_t i = 0; i < graph_size; ++i) { - if (FindCycleHelper(i, adjacency_map, visited, st, cycles)) { - has_cycle = true; - cycle_detected = true; - break; - } - } - - // Remove TensorRT subgraph from the supported node list if it's part of the cycle - if (has_cycle && remove_cycles) { - for (size_t i = 0; i < cycles.size(); ++i) { - auto loc = index_to_node_map.find(cycles[i]); - if (loc != index_to_node_map.end() && loc->second.find("TRTKernel") != std::string::npos) { - supported_nodes_vector.erase(supported_nodes_vector.begin() + cycles[i]); - trt_cycle = true; - break; - } - } - } - - delete[] adjacency_map; - delete[] visited; - delete[] st; - } - return cycle_detected; -} - -std::vector> -NvExecutionProvider::GetCapability(const GraphViewer& graph, - const IKernelLookup& /*kernel_lookup*/, - const GraphOptimizerRegistry& /*graph_optimizer_registry*/, - IResourceAccountant* /* resource_accountant */) const { - // Construct subgraph capability from node list - std::vector> result; - // Get ModelPath - const auto& path_string = graph.ModelPath().string(); -#ifdef _WIN32 - strncpy_s(model_path_, path_string.c_str(), sizeof(model_path_) - 1); -#else - strncpy(model_path_, path_string.c_str(), sizeof(model_path_) - 1); -#endif - model_path_[sizeof(model_path_) - 1] = '\0'; - - // Early return if the model has unsupported input/output data types - for (const auto* input : graph.GetInputs()) { - const auto* tp = input->TypeAsProto(); - if (tp && tp->has_tensor_type()) { - auto data_type = utils::ToOrtTensorElementDataType(tp->tensor_type().elem_type()); - if (!IsSupportedInputOutputDataType(data_type)) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Unsupported data type " << GetDataTypeName(data_type) << " for input node: " << input->Name(); - return result; - } - } - } - for (const auto* output : graph.GetOutputs()) { - const auto* tp = output->TypeAsProto(); - if (tp && tp->has_tensor_type()) { - auto data_type = utils::ToOrtTensorElementDataType(tp->tensor_type().elem_type()); - if (!IsSupportedInputOutputDataType(data_type)) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Unsupported data type " << GetDataTypeName(data_type) << " for output node: " << output->Name(); - return result; - } - } - } - - const int number_of_ort_nodes = graph.NumberOfNodes(); - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - - // Generate unique kernel name for TRT graph - HashValue model_hash = TRTGenerateId(graph, std::to_string(trt_version_), std::to_string(cuda_version_)); - - // If there are "EPContext" contrib op nodes, it means TRT EP can fetch the precompiled engine info from the node and - // load the engine directly without having to go through the processes of graph proto reconstruction, calling TRT - // parser and engine compilation. So, simply return subgraphs consists of single ep context nodes here. - int subgraph_idx = 0; - for (size_t node_idx : node_index) { - const auto& node = graph.GetNode(node_idx); - const bool is_context_node = node && !node->OpType().empty() && node->OpType() == EPCONTEXT_OP; - if (is_context_node) { - // Only claim EPContext nodes that belong to this EP. - // If the source attribute is present and doesn't match, skip the node. - const auto& attrs = node->GetAttributes(); - if (attrs.count(SOURCE) > 0 && - attrs.at(SOURCE).s() != kNvTensorRTRTXExecutionProvider) { - continue; - } - SubGraph_t supported_node_vector(std::make_pair(std::vector{node_idx}, true)); - std::unique_ptr sub_graph = GetSubGraph(supported_node_vector, graph, model_hash, subgraph_idx++); - - result.push_back(ComputeCapability::Create(std::move(sub_graph))); - } - } - // return early if context nodes where found - if (!result.empty()) { - return result; - } - - // For regular ONNX nodes, get supported node list from TensorRT parser - - std::vector nodes_vector(number_of_ort_nodes); - std::iota(std::begin(nodes_vector), std::end(nodes_vector), 0); - - auto get_exclude_ops_set = [&](std::string node_list_to_exclude) -> std::set { - std::set set; - if (!node_list_to_exclude.empty()) { - std::stringstream node_list(node_list_to_exclude); - std::string node; - while (std::getline(node_list, node, ',')) { - set.insert(node); - } - } - return set; - }; - - auto exclude_ops_set = get_exclude_ops_set(op_types_to_exclude_); - - SubGraphCollection_t parser_nodes_vector, supported_nodes_vector; - bool new_subgraph = true; - - /* Iterate all the nodes and exclude the node if: - * 1. It's a control flow op and its subgraph(s) is not fully TRT eligible. - * 2. It's a DDS op. - * 3. It has unsupported data types. - */ - for (const auto& index : nodes_vector) { - const auto& node = graph.GetNode(node_index[index]); - bool supported_node = true; - - /* If current node is control flow op, we take different approach based on following four cases: - * - * (1) control flow op is supported by TRT, and its subgraphs are all supported by TRT. Assign this node to TRT. - * (2) control flow op is supported by TRT, but not all its subgraphs supported by TRT. Don't assign this node to TRT. - * (3) control flow op is not supported by TRT, but its subgraphs all supported by TRT. Don't assign this node to TRT. - * (4) control flow op is not supported by TRT, and not all its subgraphs supported by TRT. Don't assign this node to TRT. - * - * For cases 2, 3, 4, even though the control flow op is not assigned to TRT, any portion of its subgraphs that can run in TRT will be still fused and assigned to TRT EP. - */ - if (control_flow_op_set_.find(node->OpType()) != control_flow_op_set_.end()) { - auto supported_control_flow_op = [&](const Node* node) { - auto sub_graphs = node->GetSubgraphs(); - if (sub_graphs.size() != 0) { - for (auto sub_graph : sub_graphs) { - // TRT EP should consider the empty subgraph is fully supported by TRT. - if (sub_graph->CreateGraphViewer()->NumberOfNodes() == 0) { - continue; - } - if (!AllNodesAssignedToSpecificEP(*(sub_graph->CreateGraphViewer()), kNvTensorRTRTXExecutionProvider)) { - // if not all its subgraphs are supported, we need to exclude this control flow op - return false; - } - } - } - return true; - }; - supported_node = supported_control_flow_op(node); - } - - // Exclude any ops, if applicable - if (exclude_ops_set.find(node->OpType()) != exclude_ops_set.end()) { - supported_node = false; - } - - // Check data types and print warnings for unsupported types - if (supported_node) { - if (!CheckNodeDataTypes(node)) { - supported_node = false; - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Node '" << node->Name() - << "' (OpType: " << node->OpType() - << ") excluded due to unsupported data types"; - } - } - - if (supported_node) { - if (new_subgraph) { - parser_nodes_vector.emplace_back(); - // Mark all new graphs as "UnKnown" which will later be parsed by TRT parser - parser_nodes_vector.back().second = false; - new_subgraph = false; - } - parser_nodes_vector.back().first.emplace_back(index); - } else { - new_subgraph = true; - } - } - - bool early_termination = false; - supported_nodes_vector = GetSupportedList(parser_nodes_vector, 0, max_partition_iterations_, graph, &early_termination); - if (early_termination) { - supported_nodes_vector.clear(); - } - - // Remove subgraphs if its size is less than the predefined minimal size - for (auto it = supported_nodes_vector.begin(); it != supported_nodes_vector.end();) { - const size_t subgraph_size = it->first.size(); - if (subgraph_size < min_subgraph_size_) { - it = supported_nodes_vector.erase(it); - } else { - ++it; - } - } - - // Detect and remove cycles from supported node list - DetectTensorRTGraphCycles(supported_nodes_vector, graph, model_hash); - - // Consolidate supported node list - if (supported_nodes_vector.size() > 1) { - nodes_vector.clear(); - for (const auto& group : supported_nodes_vector) { - if (!group.first.empty()) { - nodes_vector.insert(nodes_vector.end(), group.first.begin(), group.first.end()); - } - } - SubGraphCollection_t consolidated_supported_nodes_vector = {{nodes_vector, true}}; - if (DetectTensorRTGraphCycles(consolidated_supported_nodes_vector, graph, model_hash, false)) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] TensorRT nodes are not consolidated because graph will have cycles after consolidation"; - } else { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] TensorRT nodes are consolidated into one subgraph"; - supported_nodes_vector = consolidated_supported_nodes_vector; - } - } - - // Handle the case where the graph is subgraph of control flow op. - // The purpose is to make control flow op as well as its subgraphs run on TRT. - // Here we need to check whether subgraph is fully supported by TRT and don't fuse the nodes of the subgraph until control flow op level. - if (IsSubGraphOfControlFlowOp(graph) && IsSubGraphFullySupported(supported_nodes_vector, number_of_ort_nodes)) { - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); - bool all_subgraphs_are_supported = true; - - // "If" control flow op has two subgraph bodies, "then" body and "else" body respectively. - // Check its parent node's another subgraph to see whether that subgraph is also fully supported by TRT. - if (graph.ParentNode()->OpType() == "If") { - all_subgraphs_are_supported = false; - SubGraphCollection_t subgraph_supported_nodes_vector; - auto sub_graphs = graph.ParentNode()->GetSubgraphs(); - for (auto sub_graph : sub_graphs) { - if (sub_graph.get() != &graph.GetGraph()) { - auto sub_graph_viewer = sub_graph->CreateGraphViewer(); - const int number_of_ort_subgraph_nodes = sub_graph_viewer->NumberOfNodes(); - std::vector subgraph_nodes_vector(number_of_ort_subgraph_nodes); - std::iota(std::begin(subgraph_nodes_vector), std::end(subgraph_nodes_vector), 0); - SubGraphCollection_t parser_subgraph_nodes_vector = {{subgraph_nodes_vector, false}}; - bool subgraph_early_termination = false; - - // Another subgraph of "If" control flow op has no nodes. - // In this case, TRT EP should consider this empty subgraph is fully supported by TRT. - if (sub_graph_viewer->NumberOfNodes() == 0) { - all_subgraphs_are_supported = true; - break; - } - // Another subgraph of "If" control flow op has been parsed by GetCapability before and all subgraph's nodes assigned to TRT EP. - else if (AllNodesAssignedToSpecificEP(*sub_graph_viewer, kNvTensorRTRTXExecutionProvider)) { - all_subgraphs_are_supported = true; - break; - } - // Another subgraph of "If" control flow has been parsed by GetCapability and not all subgraph's nodes assigned to TRT EP. - // (Note: GetExecutionProviderType() returns "" meaning node has not yet been assigned to any EPs) - else if (!AllNodesAssignedToSpecificEP(*sub_graph_viewer, "")) { - all_subgraphs_are_supported = false; - break; - } - - // Another subgraph of "If" control flow has not yet been parsed by GetCapability. - subgraph_supported_nodes_vector = GetSupportedList(parser_subgraph_nodes_vector, 0, max_partition_iterations_, *sub_graph_viewer, &subgraph_early_termination); - all_subgraphs_are_supported = IsSubGraphFullySupported(subgraph_supported_nodes_vector, number_of_ort_subgraph_nodes); - break; - } - } - } - - if (all_subgraphs_are_supported) { - // We want the subgraph nodes to be assigned to TRT EP but don't want them to be fused until later at the control flow op level. - // Simply request the subgraph nodes with a single ComputeCapability for each with no MetaDef (i.e. what the default implementation for IExecutionProvider::GetCapability does). - for (const auto& group : supported_nodes_vector) { - if (!group.first.empty()) { - for (const auto& index : group.first) { - std::unique_ptr sub_graph = onnxruntime::IndexedSubGraph::Create(); - sub_graph->Nodes().push_back(node_index[index]); - result.push_back(ComputeCapability::Create(std::move(sub_graph))); - } - } - } - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Whole graph will run on Nv execution provider"; - - // The context map is only used during EP compile time, release it to save memory space. - subgraph_context_map_.clear(); - return result; - } - } - - int number_of_trt_nodes = 0, subgraph_index = 0; - for (const auto& group : supported_nodes_vector) { - if (!group.first.empty()) { - std::unique_ptr sub_graph = GetSubGraph(group, graph, model_hash, subgraph_index); - result.push_back(ComputeCapability::Create(std::move(sub_graph))); - number_of_trt_nodes += static_cast(group.first.size()); - subgraph_index++; - } - } - - const size_t number_of_subgraphs = supported_nodes_vector.size(); - if (number_of_trt_nodes == 0) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] No graph will run on Nv execution provider"; - } else if (number_of_trt_nodes == number_of_ort_nodes) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Whole graph will run on Nv execution provider"; - } else { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Graph is partitioned and number of subgraphs running on Nv executio provider is " << number_of_subgraphs; - } - - // The context map is only used during EP compile time, release it to save memory space. - subgraph_context_map_.clear(); - return result; -} - -/** - * Refit the weight-stripped engine - */ -common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, - std::string& onnx_model_folder_path, - const void* onnx_model_bytestream, - size_t onnx_model_bytestream_size, - const void* onnx_external_data_bytestream, - size_t onnx_external_data_bytestream_size, - nvinfer1::ICudaEngine* trt_engine, - bool detailed_build_log) { - bool refit_from_file = onnx_model_bytestream == nullptr && onnx_model_bytestream_size == 0; - bool refit_with_external_data = onnx_external_data_bytestream != nullptr && onnx_external_data_bytestream_size != 0; - bool refit_complete = false; - std::filesystem::path onnx_model_path{onnx_model_folder_path}; - if (refit_from_file) { - if (!onnx_model_filename.empty()) { - onnx_model_path.append(onnx_model_filename); - } - if (onnx_model_path.empty()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "The ONNX model was not provided as path. " - "Please use provide an ONNX bytestream to enable refitting the weightless engine."); - } else { - // Validate that the ONNX model path does not escape the model directory. - if (!onnx_model_filename.empty()) { - ORT_RETURN_IF_ERROR(utils::ValidateExternalDataPathFromDir( - std::filesystem::path(onnx_model_folder_path), std::filesystem::path(onnx_model_filename))); - } - - if (!(std::filesystem::exists(onnx_model_path) && std::filesystem::is_regular_file(onnx_model_path))) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "The ONNX model " + onnx_model_path.string() + - " does not exist."); - } - } - } - - // weight-stripped engine refit logic - TensorrtLogger& trt_logger = GetTensorrtLogger(detailed_build_log); - auto refitter = std::unique_ptr(nvinfer1::createInferRefitter(*trt_engine, trt_logger)); - auto parser_refitter = std::unique_ptr( - nvonnxparser::createParserRefitter(*refitter, trt_logger)); - - // New refit APIs - if (refit_with_external_data) { -#if TRT_MAJOR_RTX > 1 || TRT_MINOR_RTX >= 1 - // A valid model bytestream must be passed. - if (refit_from_file) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's refit with external data must be called with a valid ONNX model bytestream"); - } - - if (!parser_refitter->loadModelProto(onnx_model_bytestream, onnx_model_bytestream_size, nullptr)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's IParserRefitter could not load model from provided onnx_model_bytestream"); - } - - // Extract weight information from the Refitter. - int required_weights = refitter->getAllWeights(0, nullptr); - std::vector refit_names_prealocated(required_weights); - refitter->getAllWeights(required_weights, refit_names_prealocated.data()); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refitter requires " << required_weights << " weights"; - std::unordered_set refit_names(std::make_move_iterator(refit_names_prealocated.begin()), - std::make_move_iterator(refit_names_prealocated.end())); - - // Vectors to keep track of data pointers. - std::vector names; - names.reserve(required_weights); - std::vector bytes; - bytes.reserve(required_weights); - std::vector sizes; - sizes.reserve(required_weights); - - auto onnx_model = ModelProto::Create(); - TensorProtos* allInitializers_byte_stream; - - // Reconstruct onnx model view. - const auto onnx_model_view = std::string((const char*)onnx_model_bytestream, - onnx_model_bytestream_size); - if (!onnx_model->ParseFromString(onnx_model_view)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "The provided ONNX bytestream to refit could not be parsed."); - } - - // Extract graph and initializer information. - auto const& graph = onnx_model->mutable_graph(); - allInitializers_byte_stream = graph->mutable_initializer(); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Initializers that were found " << allInitializers_byte_stream->size(); - - // Loop through all initializers - int missing_initializer_data = 0; - for (int initializer_idx = 0; initializer_idx < allInitializers_byte_stream->size(); ++initializer_idx) { - auto& proto = allInitializers_byte_stream->at(initializer_idx); - auto& proto_name = proto.name(); - if (refit_names.find(proto_name) != refit_names.end()) { - if (proto.has_data_location()) { - if (proto.data_location() == TensorProto_DataLocation_EXTERNAL) { - // Default values for reading into external_data blob. - int64_t offset = 0; - size_t length = 0; - auto external_data = proto.mutable_external_data(); - const std::string kOffset = "offset", kLength = "length"; - for (int entry_idx = 0; entry_idx < external_data->size(); ++entry_idx) { - auto current_key = external_data->at(entry_idx).mutable_key(); - auto current_value = external_data->at(entry_idx).mutable_value(); - if (*current_key == kOffset && !current_value->empty()) { - offset = std::stoll(*current_value); - } else if (*current_key == kLength && !current_value->empty()) { - length = std::stoul(*current_value); - } - } - names.push_back(proto.name()); - bytes.push_back(static_cast(onnx_external_data_bytestream) + offset); - sizes.push_back(length); - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[NvTensorRTRTX EP] Proto: " + proto_name + " expected to have external datalocation, but default datalocation was provided instead."); - } - } else if (proto.has_raw_data()) { - auto& raw_data = proto.raw_data(); - names.push_back(proto.name()); - bytes.push_back(raw_data.c_str()); - sizes.push_back(raw_data.size()); - } else { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Proto: " + proto_name + " has no raw nor external data."; - ++missing_initializer_data; - } - } else { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Initializer with name: " << proto_name << " was not marked as refittable"; - } - } - if (missing_initializer_data) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[NvTensorRTRTX EP] RefitEngine is missing " + std::to_string(missing_initializer_data) + " initializers."); - } - - // Load extracted initializers into the parser - if (!names.empty()) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Number of initializers submitted to refitter " << names.size(); - for (size_t i = 0; i < names.size(); i++) { - bool refloadInit = parser_refitter->loadInitializer(names[i].c_str(), bytes[i], sizes[i]); - if (!refloadInit) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); - } - } - } - // Perform refit. - if (!parser_refitter->refitModelProto()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); - } - refit_complete = true; -#else - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Refit with external data is only supported on TensorRT RTX 1.1.x.x and above."); -#endif - } - - // If new refit flow was not completed, then fallback to refit_from_file. - if (!refit_complete) { - if (refit_from_file) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refitting from file on disk: " << onnx_model_path.string(); - if (!parser_refitter->refitFromFile(onnx_model_path.string().c_str())) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); - } - } else { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refitting from byte array"; - if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); - } - } - } - if (refitter->refitCudaEngine()) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully refitted the weight-stripped engine."; - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); - } - - return Status::OK(); -} - -common::Status NvExecutionProvider::Compile(const std::vector& fused_nodes_and_graphs, - std::vector& node_compute_funcs) { - for (auto& fused_node_graph : fused_nodes_and_graphs) { - const GraphViewer& graph_body_viewer = fused_node_graph.filtered_graph; - const Node& fused_node = fused_node_graph.fused_node; - // Build map from input name to its index in input definitions - std::unordered_map input_map; - const auto& input_defs = fused_node.InputDefs(); - input_map.reserve(input_defs.size()); - for (size_t i = 0, end = input_defs.size(); i < end; ++i) { - input_map[input_defs[i]->Name()] = i; - } - - // Build map from output name to its index in output definitions - std::unordered_map output_map; - const auto& output_defs = fused_node.OutputDefs(); - output_map.reserve(output_defs.size()); - for (size_t i = 0, end = output_defs.size(); i < end; ++i) { - output_map[output_defs[i]->Name()] = i; - } - - Status status; - size_t node_idx = 0; - if (GraphHasCtxNode(graph_body_viewer, node_idx)) { - status = CreateNodeComputeInfoFromPrecompiledEngine(graph_body_viewer, - node_idx, - fused_node, - input_map, - output_map, - node_compute_funcs); - } else { - status = CreateNodeComputeInfoFromGraph(graph_body_viewer, fused_node, input_map, output_map, node_compute_funcs); - } - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, status.ErrorMessage()); - } - } - return Status::OK(); -} - -/** - * @brief Determines whether I/O binding is required for TensorRT execution. - * - * This function optimizes TensorRT inference performance by determining when tensor - * input/output binding operations can be skipped. Binding is an expensive operation - * that involves setting up tensor pointers in the TensorRT execution context, so - * avoiding unnecessary rebinding can significantly improve inference throughput. - * - * The function implements a three-tier decision logic: - * 1. First run: Always requires binding to establish initial tensor mappings - * 2. Subsequent runs with optimization allowed: Only rebind if tensors have changed - * 3. Subsequent runs without optimization: Always rebind for safety - * - * @tparam TRTState The TensorRT state type (TensorrtFuncState or TensorrtShortFuncState) - * @param trt_state Pointer to the TensorRT execution state containing tensor cache - * and configuration flags - * @param ctx ONNX Runtime kernel context providing access to current input tensors - * - * @return true if I/O binding is required (tensors changed or safety conditions apply), - * false if binding can be safely skipped (optimization enabled and tensors unchanged) - * - * @note This function modifies trt_state by: - * - Setting is_first_run to false after first execution - * - Caching current tensor parameters in input_tensors vector - * - Updating cached tensors when changes are detected - * - * @warning The skip_io_binding_allowed flag must be carefully managed as incorrect - * usage can lead to inference with stale tensor bindings and incorrect results. - */ -template -static bool IsIOBindingRequired(TRTState* const trt_state, const Ort::KernelContext& ctx) { - // Check if input tensors have changed since the last run - // If so, we need to bind input tensors again - bool require_io_binding = false; - - if (trt_state->is_first_run) { - // If this is the first run, we always bind input tensors - require_io_binding = true; - auto input_tensor_count = ctx.GetInputCount(); - auto output_tensor_count = ctx.GetOutputCount(); - trt_state->input_tensors.resize(input_tensor_count); - trt_state->output_tensors.resize(output_tensor_count); - for (size_t input_index = 0; input_index < input_tensor_count; ++input_index) { - const auto& input_tensor = ctx.GetInput(input_index); - const auto& tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); - - trt_state->input_tensors[input_index] = TensorParams{input_tensor.GetTensorRawData(), tensor_info.GetShape()}; - } - trt_state->is_first_run = false; - } else if (trt_state->skip_io_binding_allowed) { - // If skip_io_binding_allowed is true, we can skip binding if input tensors are the same as before - auto input_tensor_count = ctx.GetInputCount(); - for (size_t input_index = 0; input_index < input_tensor_count; ++input_index) { - const auto& input_tensor = ctx.GetInput(input_index); - const auto& tensor_info = input_tensor.GetTensorTypeAndShapeInfo(); - - TensorParams ip_tensor{input_tensor.GetTensorRawData(), tensor_info.GetShape()}; - - if (ip_tensor != trt_state->input_tensors[input_index]) { - require_io_binding = true; - trt_state->input_tensors[input_index] = ip_tensor; - } - } - } else { - // If this is not the first run and skip_io_binding_allowed is false, we need to bind input tensors - require_io_binding = true; - } - - if (!require_io_binding) { - // no need to bind inputs, but check outputs as well - auto output_tensor_count = ctx.GetOutputCount(); - - for (size_t output_index = 0; output_index < output_tensor_count; ++output_index) { - const auto& prev_output_tensor = trt_state->output_tensors[output_index]; - - if (prev_output_tensor.dims.nbDims != -1) { - const auto& new_output_tensor = ctx.GetOutput(output_index, prev_output_tensor.dims.d, prev_output_tensor.dims.nbDims); - - // different output tensor data means we need to bind outputs again - if (prev_output_tensor.data != new_output_tensor.GetTensorRawData()) { - require_io_binding = true; - break; - } - } - } - } - - return require_io_binding; -} - -const InlinedVector NvExecutionProvider::GetEpContextNodes() const { - InlinedVector ep_context_nodes; - if (ep_context_model_) { - for (auto* node : ep_context_model_->MainGraph().Nodes()) { - ep_context_nodes.push_back(node); - } - } - return ep_context_nodes; -} - -std::string NvExecutionProvider::GetCompiledModelCompatibilityInfo( - const onnxruntime::GraphViewer& graph_viewer) const { - ORT_UNUSED_PARAMETER(graph_viewer); - - // Protect read access to engine_headers_ for thread safety - auto lock = GetApiLock(); - - // Compatibility info is only supported when there is exactly one engine. - // If multiple EPContext nodes/engines exist, return empty so validation is not applicable. - if (engine_headers_.size() > 1) { - return std::string(); - } - - // If we have stored engine headers, return the first one found - // (typically there's only one per EP context) - if (!engine_headers_.empty()) { - return engine_headers_.begin()->second; - } - - // No headers available - validation not supported for this model - return std::string(); -} - -common::Status NvExecutionProvider::ValidateCompiledModelCompatibilityInfo( - const std::string& compatibility_info, - OrtCompiledModelCompatibility& model_compatibility) const { - // If no compatibility info provided, validation not applicable - if (compatibility_info.empty()) { - model_compatibility = OrtCompiledModelCompatibility_EP_NOT_APPLICABLE; - return Status::OK(); - } - - // Decode hex string to binary - std::vector engine_header; - try { - engine_header = HexStringToBinary(compatibility_info); - } catch (const std::exception& ex) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Failed to decode engine header: " << ex.what(); - model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - return Status::OK(); - } - - // Use TensorRT RTX's getEngineValidity to check compatibility - uint64_t diagnostics = 0; - nvinfer1::EngineValidity validity = runtime_->getEngineValidity( - engine_header.data(), - engine_header.size(), - &diagnostics); - - // Map TensorRT RTX validity to ORT compatibility status - switch (validity) { - case nvinfer1::EngineValidity::kVALID: - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Engine is fully compatible with this system"; - model_compatibility = OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL; - break; - - case nvinfer1::EngineValidity::kSUBOPTIMAL: - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine is compatible but recompilation recommended " - << "(diagnostics: 0x" << std::hex << diagnostics << std::dec << ")"; - model_compatibility = OrtCompiledModelCompatibility_EP_SUPPORTED_PREFER_RECOMPILATION; - break; - - case nvinfer1::EngineValidity::kINVALID: - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine is incompatible with this system " - << "(diagnostics: 0x" << std::hex << diagnostics << std::dec << ")"; - model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - break; - - default: - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Unknown TensorRT validity status: " - << static_cast(validity); - model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - break; - } - - return Status::OK(); -} - -Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& graph_body_viewer, - const Node& fused_node, - std::unordered_map& input_map, - std::unordered_map& output_map, - std::vector& node_compute_funcs) { - // Reconstruct graph proto from fused node's function body - auto model = graph_body_viewer.CreateModel(*GetLogger()); - auto model_proto = model->ToProto(); - - // exclude weights if external - std::vector userWeights; - if (use_external_data_initializer_) { - const InitializedTensorSet& allInitializers = graph_body_viewer.GetAllInitializedTensors(); - userWeights.reserve(allInitializers.size()); - for (const auto& [name, tp] : allInitializers) { - if (utils::HasRawData(*tp)) { - userWeights.emplace_back(name, tp->raw_data().data(), tp->raw_data().size()); - } else if (utils::HasExternalDataInMemory(*tp)) { - // the initializer was marked as external data by the ORT graph at load time since it was provided in memory - if (OrtValue v; graph_body_viewer.GetOrtValueInitializer(name, v)) { - Ort::ConstValue initializer_value{&v}; - const size_t size = initializer_value.GetTensorSizeInBytes(); - const void* ptr = initializer_value.GetTensorRawData(); - userWeights.emplace_back(name, ptr, size); - } else { - // only copy and take ownership of the data if none of the above conditions are met - std::unique_ptr full_init; - ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tp, full_init)); - userWeights.emplace_back(name, full_init->raw_data()); - } - } - } - } - // ORT's default topological sort is using reversed DFS. - // When creating model proto from graph viewer, let ORT use priority-based topological sort based on node index. - // The reason is, in some cases, for example ResNet50, using default topological sort will end up with generating - // the model proto that has different node ordering compared to original onnx model. - graph_body_viewer.ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, !use_external_data_initializer_ /*include raw initializers*/); - model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); - std::string string_buf; - model_proto->SerializeToString(string_buf); - - if (dump_subgraphs_) { - // Dump TensorRT subgraphs - std::fstream dump(fused_node.Name() + ".onnx", std::ios::out | std::ios::trunc | std::ios::binary); - model_proto->SerializeToOstream(dump); - } - - TensorrtLogger& trt_logger = GetTensorrtLogger(detailed_build_log_); - auto trt_builder = GetBuilder(trt_logger); -#if TRT_MAJOR_RTX >= 2 || (TRT_MAJOR_RTX == 1 && ((TRT_MINOR_RTX == 5 && TRT_BUILD_RTX >= 97) || TRT_MINOR_RTX >= 6)) - // kSTRONGLY_TYPED == 0 => bit flag value is 1U. Use literal to avoid deprecated-enum warning (deprecated since 1.5.0.97). - constexpr uint32_t network_flags = 1U; -#else - auto network_flags = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kSTRONGLY_TYPED); -#endif - auto trt_network = std::unique_ptr(trt_builder->createNetworkV2(network_flags)); - auto trt_config = std::unique_ptr(trt_builder->createBuilderConfig()); - auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); - - if (use_external_data_initializer_) { -#if TRT_MAJOR_RTX > 1 || TRT_MINOR_RTX >= 1 - trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - for (auto const& userWeight : userWeights) { - trt_parser->loadInitializer(userWeight.Name(), userWeight.Data(), userWeight.Size()); - } - trt_parser->parseModelProto(); -#else - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "'nv_use_external_data_initializer' is only supported on TensorRT RTX 1.1.x.x and above."); -#endif - } else { - trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); - } - - if (max_workspace_size_ > 0) { - trt_config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_size_); - } - if (max_shared_mem_size_ > 0) { - trt_config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kTACTIC_SHARED_MEMORY, max_shared_mem_size_); - } - - // Set compute capability to kCURRENT by default - // Must set the number of compute capabilities before setting the capability itself - constexpr int kDefaultNumComputeCapabilities = 1; - if (trt_config->getNbComputeCapabilities() == 0) { - trt_config->setNbComputeCapabilities(kDefaultNumComputeCapabilities); - } - trt_config->setComputeCapability(nvinfer1::ComputeCapability::kCURRENT, 0); - - int num_inputs = trt_network->getNbInputs(); - int num_outputs = trt_network->getNbOutputs(); - std::unordered_map input_indexes(num_inputs); - std::unordered_map output_indexes(num_outputs); - std::unordered_map output_types(num_outputs); - - /* - * Initialize shape range for each dynamic shape input tensor: - * 1) If user explicitly specifies optimization profiles via provider options, TRT EP will create those profiles during EP compile time. - * It won't make adjustment for profile values during EP compute time. - * - * 2) If no explicit optimization profiles provided by user, TRT EP will firstly set min/max/opt shape to [INT_MAX, INT_MIN, INT_MIN]. - * Later in EP compute time, the shape will be adjusted to [min_input_value, max_input_value, max_input_value] based on input tensor value. - * - * - * Once the TRT profiles are created: - * 1) If all the dynamic shape input tensors have associated profiles explicitly provided by user, those profiles will be applied to TRT builder config - * and the engine will be built at EP compile time. - * - * 2) As long as one of the dynamic shape input tensors has no explicitly associated profile, TRT EP will create default shape as described above, - * and all the profiles won't be applied and engine won't be built until EP compute time. - */ - bool has_explicit_profile = false; - bool has_implicit_profile = false; - int num_profiles = 0; - std::vector trt_profiles; - - // Following c++ map data structure is used to help serialize/deserialize profiles where it saves dynamic shape dimension(s) and min/max/opt values for dynamic shape input tensor. - // - // (1) Single profile case: - // For example, assume tensor_a has two dynamic shape dimensions: dim_0 and dim_2, and tensor_b - // has one dynamic shape dimension: dim_1. The data will be: - // { - // tensor_a: { - // dim_0: [[min_shape, max_shape, opt_shape]], - // dim_2: [[min_shape, max_shape, opt_shape]] - // }, - // tensor_b: { - // dim_1: [[min_shape, max_shape, opt_shape]] - // } - // } - // - // (2) Multiple profiles case: - // For example, assume tensor_a has one dynamic shap dimension: dim 0, and tensor_b has one dynamic shape dimension: dim_1, - // and both of the tensors have two profiles. The data will be: - // { - // tensor_a: { - // dim_0: [[min_shape_0, max_shape_0, opt_shape_0], [min_shape_1, max_shape_1, opt_shape_1]] - // }, - // tensor_b: { - // dim_1: [[min_shape_2, max_shape_2, opt_shape_2], [min_shape_3, max_shape_3, opt_shape_3]] - // } - // } - ShapeRangesMap input_explicit_shape_ranges; - ShapeRangesMap input_implicit_shape_ranges; - - bool has_dynamic_shape = false; // True if input tensor has dynamic shape and no explicit profile is specified, otherwise false - if ((!profile_min_shapes_.empty()) && (!profile_max_shapes_.empty()) && (!profile_opt_shapes_.empty())) { - has_explicit_profile = true; - has_dynamic_shape = true; - num_profiles = GetNumProfiles(profile_min_shapes_); - for (int i = 0; i < num_profiles; i++) { - trt_profiles.push_back(trt_builder->createOptimizationProfile()); - } - } else { - for (unsigned int i = 0, end = num_inputs; i < end; ++i) { - auto input = trt_network->getInput(i); - has_dynamic_shape |= checkTrtTensorIsDynamic(input); - } - if (has_dynamic_shape) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] No explicit optimization profile was specified. " - "We will assume a single profile with fully dynamic range. " - "This feature is experimental and may change in the future." - "If you plan to use this model as fixed shape we recommend using a free dimension override: " - "https://onnxruntime.ai/docs/tutorials/web/env-flags-and-session-options.html#freedimensionoverrides."; - trt_profiles.push_back(trt_builder->createOptimizationProfile()); - } - } - if (has_dynamic_shape) { - // Iterate all input tensors to check dynamic shape - for (unsigned int i = 0, end = num_inputs; i < end; ++i) { - auto input = trt_network->getInput(i); - const std::string& input_name = input->getName(); - nvinfer1::Dims dims = input->getDimensions(); - - // Apply explicit optimization profiles provided by user - bool apply_profile = false; - bool tensor_has_profile = profile_min_shapes_.find(input_name) != profile_min_shapes_.end() && - profile_opt_shapes_.find(input_name) != profile_opt_shapes_.end() && - profile_max_shapes_.find(input_name) != profile_max_shapes_.end(); - if (has_explicit_profile && tensor_has_profile) { - apply_profile = ApplyProfileShapesFromProviderOptions(trt_profiles, input, profile_min_shapes_, profile_max_shapes_, profile_opt_shapes_, input_explicit_shape_ranges, cuda_graph_enable_); - } else { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Creating implicit profile for tensor " << input_name; - profile_min_shapes_[input_name] = std::vector>{{}}; - profile_min_shapes_[input_name][0].resize(dims.nbDims); - profile_opt_shapes_[input_name] = std::vector>{{}}; - profile_opt_shapes_[input_name][0].resize(dims.nbDims); - profile_max_shapes_[input_name] = std::vector>{{}}; - profile_max_shapes_[input_name][0].resize(dims.nbDims); - for (int idx_dim = 0; idx_dim < dims.nbDims; ++idx_dim) { - auto dim_value = dims.d[idx_dim]; - if (dim_value == -1) { - has_implicit_profile = true; - // TODO(maximilianm) this is needed until we have a wildcard in the API to support dynamic shapes - profile_min_shapes_[input_name][0][idx_dim] = 0; - // TODO(maximilianm) This can be buggy since shape inference can failt with 1 being used as optimal shape - // [2025-04-04 15:41:58 ERROR] IBuilder::buildSerializedNetwork: Error Code 4: Internal Error (kOPT values for profile 0 violate shape constraints: - // /conv1/Conv: spatial dimension of convolution/deconvolution output cannot be negative (build-time output dimension of axis 2 is - // (+ (CEIL_DIV (+ h -6) 2) 1)) Condition '<' violated: 2 >= 1.) - profile_opt_shapes_[input_name][0][idx_dim] = 1; - profile_max_shapes_[input_name][0][idx_dim] = std::numeric_limits::max(); - } else { - profile_min_shapes_[input_name][0][idx_dim] = dim_value; - profile_opt_shapes_[input_name][0][idx_dim] = dim_value; - profile_max_shapes_[input_name][0][idx_dim] = dim_value; - } - } - apply_profile = ApplyProfileShapesFromProviderOptions(trt_profiles, input, profile_min_shapes_, profile_max_shapes_, profile_opt_shapes_, input_explicit_shape_ranges, cuda_graph_enable_); - } - if (!apply_profile) { - std::ostringstream msg; - msg << "Optimization profile could not be applied for tensor:\n"; - msg << input_name; - msg << "\n["; - for (int idx_dim = 0; idx_dim < dims.nbDims; ++idx_dim) { - msg << dims.d[idx_dim] << ","; - } - msg << "]"; - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, msg.str()); - } - } - - // Set explicit profiles in TRT config if all dynamic shape inputs have associated profiles provided by user - if (has_explicit_profile || has_implicit_profile) { - // TRT EP has a constraint here. - // Users need to provide all the dynamic shape inputs with associated profiles if they want to explicitly specify profiles through provider options. - for (auto trt_profile : trt_profiles) { - trt_config->addOptimizationProfile(trt_profile); - } - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "No explicit or implicit shapes were provided for dynamic shape inputs."); - ; - } - } - - // enable sparse weights - if (sparsity_enable_) { - trt_config->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Sparse weights are allowed"; - } - - // limit auxiliary streams - if (auxiliary_streams_ >= 0) { - trt_config->setMaxAuxStreams(auxiliary_streams_); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Auxiliary streams are se to " << auxiliary_streams_; - } - - if (weight_stripped_engine_enable_) { - trt_config->setFlag(nvinfer1::BuilderFlag::kSTRIP_PLAN); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] STRIP_PLAN is enabled"; - trt_config->setFlag(nvinfer1::BuilderFlag::kREFIT_IDENTICAL); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] REFIT_IDENTICAL is enabled"; - } - - // Build TRT engine (if needed) and load TRT engine if: - // (1) Graph has no dynamic shape input - // (2) All the dynamic shape inputs have associated explicit profiles specified by user - // - // Otherwise engine will be handled at inference time. - std::unique_ptr trt_engine; - tensorrt_ptr::unique_pointer_exec_ctx trt_context; - std::unique_ptr trt_runtime_cache; - std::unique_ptr trt_runtime_config; - std::string runtime_cache_file = ""; - - // Generate file name for dumping ep context model - if (dump_ep_context_model_ && ctx_model_path_.empty()) { - ctx_model_path_ = path_utils::GetPathWithStemSuffix(ep_context_file_path_, model_path_, "_ctx.onnx"); - } - { - auto lock = GetApiLock(); - // Build engine - std::chrono::steady_clock::time_point engine_build_start; - if (detailed_build_log_) { - engine_build_start = std::chrono::steady_clock::now(); - } - std::unique_ptr serialized_engine{trt_builder->buildSerializedNetwork(*trt_network, *trt_config)}; - if (serialized_engine == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP failed to create engine from network for fused node: " + fused_node.Name()); - } - - // Capture engine header (first 64 bytes) for compatibility validation - if (serialized_engine->size() >= kTensorRTEngineHeaderSize) { - std::string engine_header_hex = BinaryToHexString( - serialized_engine->data(), - kTensorRTEngineHeaderSize); - engine_headers_[fused_node.Name()] = engine_header_hex; - } else { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine too small to capture header for validation: " - << serialized_engine->size() << " bytes"; - } - - // In compile-only mode the session will not be used for inference. Skip deserialization - // and GPU context creation — the serialized engine is already saved as an EP context node. - if (!compile_only_mode_) { - trt_engine = std::unique_ptr(runtime_->deserializeCudaEngine(serialized_engine->data(), serialized_engine->size())); - if (trt_engine == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP failed to deserialize engine for fused node: " + fused_node.Name()); - } - - trt_runtime_config = std::unique_ptr(trt_engine->createRuntimeConfig()); - if (trt_runtime_config && cuda_graph_enable_) { - trt_runtime_config->setDynamicShapesKernelSpecializationStrategy(nvinfer1::DynamicShapesKernelSpecializationStrategy::kEAGER); -#if TRT_MAJOR_RTX > 1 || (TRT_MAJOR_RTX == 1 && TRT_MINOR_RTX >= 3) - auto cuda_strategy_flag = trt_runtime_config->setCudaGraphStrategy(nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE); - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] CUDA graph strategy with RTX Graph capture enabled : " << cuda_strategy_flag; -#else - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] CUDA graph is enabled but RTX Graph capture is not available. " - << "The current TRT RTX version does not support RTX Graph. " - << "Please upgrade to TRT RTX >= 1.3 to use RTX Graph capture feature for optimal CUDA graph performance."; -#endif - } - trt_runtime_config->setExecutionContextAllocationStrategy(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED); - if (!runtime_cache_.empty()) { - runtime_cache_file = (runtime_cache_ / fused_node.Name()).string(); - trt_runtime_cache = std::unique_ptr(trt_runtime_config->createRuntimeCache()); - auto cache_data = file_utils::ReadFile(runtime_cache_file); - if (!trt_runtime_cache->deserialize(cache_data.data(), cache_data.size())) { - trt_runtime_cache = std::unique_ptr(trt_runtime_config->createRuntimeCache()); - LOGS_DEFAULT(INFO) << "TensorRT RTX failed to deserialize the runtime cache, will overwrite with new one" << std::endl; - } - if (!trt_runtime_config->setRuntimeCache(*trt_runtime_cache)) { - LOGS_DEFAULT(INFO) << "TensorRT RTX failed to set the runtime cache" << std::endl; - } - } - } - - if (detailed_build_log_) { - auto engine_build_stop = std::chrono::steady_clock::now(); - LOGS_DEFAULT(INFO) << "TensorRT engine build for " << fused_node.Name() << " took: " << std::chrono::duration_cast(engine_build_stop - engine_build_start).count() << "ms" << std::endl; - } - // dump EP context node model - if (dump_ep_context_model_) { - // "ep_cache_context" node attribute should be a relative path to context model directory - - std::string cache_path = ""; - // Customize cache prefix if assigned - if (!cache_prefix_.empty()) { - // Generate cache suffix in case user would like to customize cache prefix - cache_path = GetCachePath(cache_path_, cache_prefix_) + fused_node.Name() + ".engine"; - ; - } else { - cache_path = GetCachePath(cache_path_, fused_node.Name()) + ".engine"; - ; - } - // NV TRT EP per default generates hardware compatible engines for any RTX device with compute capability > 80 - std::string compute_capability_hw_compat = "80+"; - if (!ep_context_model_) { - ep_context_model_ = Model::Create("nv_trt_rtx_ep_context_model", false, *GetLogger()); - } - - auto status = CreateCtxNode(graph_body_viewer, - ep_context_model_->MainGraph(), - cache_path, - reinterpret_cast(serialized_engine->data()), - serialized_engine->size(), - ep_context_embed_mode_, - compute_capability_hw_compat, - model_path_, - fused_node.Name(), - trt_version_); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - } - - // In compile-only mode the engine was built and saved but GPU deserialization was skipped. - // Return a stub compute function — it will never be called since the session is - // destroyed immediately after compilation without running any inference. - if (compile_only_mode_) { - NodeComputeInfo stub_compute_info; - stub_compute_info.create_state_func = [](ComputeContext*, FunctionState* state) -> int { - *state = nullptr; - return 0; - }; - stub_compute_info.release_state_func = [](FunctionState) {}; - stub_compute_info.compute_func = [](FunctionState, const OrtApi*, OrtKernelContext*) -> Status { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "NvTensorRTRTX EP: inference is not available in a compile-only session"); - }; - node_compute_funcs.push_back(std::move(stub_compute_info)); - return Status::OK(); - } - - if (weight_stripped_engine_refit_) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refit engine from main ONNX file after engine build"; - std::filesystem::path model_fs_path(model_path_); - std::string refit_folder = onnx_model_folder_path_.empty() - ? model_fs_path.parent_path().string() - : onnx_model_folder_path_; - auto status = RefitEngine(model_fs_path.filename().string(), - refit_folder, - onnx_model_bytestream_, - onnx_model_bytestream_size_, - onnx_external_data_bytestream_, - onnx_external_data_bytestream_size_, - trt_engine.get(), - detailed_build_log_); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - - // Build context - // Note: Creating an execution context from an engine is thread safe per TRT doc - // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - trt_context = tensorrt_ptr::unique_pointer_exec_ctx( - trt_engine->createExecutionContext(trt_runtime_config.get()), - tensorrt_ptr::IExecutionContextDeleter(runtime_cache_file, std::move(trt_runtime_cache))); - if (!trt_context) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP could not build execution context for fused node: " + fused_node.Name()); - } - - bool is_dynamic_shape_context = false; - // Create input to index map - for (int i = 0; i < num_inputs; ++i) { - auto input = trt_network->getInput(i); - const std::string& input_name = input->getName(); - is_dynamic_shape_context |= checkTrtDimIsDynamic(trt_engine->getTensorShape(input_name.c_str())); - const auto& iter = input_map.find(input_name); - if (iter != input_map.end()) { - input_indexes[input_name] = iter->second; - } - } - - // Create output to index and type maps - const auto& graph_output = model_proto->graph().output(); - for (int i = 0; i < num_outputs; ++i) { - const std::string& output_name = trt_network->getOutput(i)->getName(); - const auto& iter = output_map.find(output_name); - if (iter != output_map.end()) { - output_indexes[output_name] = iter->second; - } - const auto& tensor_type = graph_output[i].type().tensor_type(); - output_types[output_name] = tensor_type.elem_type(); - } - - // Save TRT engine, other TRT objects and input/output info to map - engines_.emplace(fused_node.Name(), std::move(trt_engine)); - contexts_.emplace(fused_node.Name(), std::move(trt_context)); - networks_.emplace(fused_node.Name(), std::move(trt_network)); - input_info_[fused_node.Name()].push_back(input_indexes); - output_info_[fused_node.Name()].push_back(output_indexes); - output_info_[fused_node.Name()].push_back(output_types); - input_shape_ranges_[fused_node.Name()] = input_implicit_shape_ranges; - profiles_.emplace(fused_node.Name(), std::move(trt_profiles)); - - // Create function state - // TODO: remove default capture - NodeComputeInfo compute_info; - compute_info.create_state_func = [=](ComputeContext* context, FunctionState* state) { - std::unique_ptr p = std::make_unique(); - *p = {context->allocate_func, context->release_func, context->allocator_handle, context->node_name, builder_.get(), - &engines_[context->node_name], &contexts_[context->node_name], - &networks_[context->node_name], input_info_[context->node_name], output_info_[context->node_name], - input_shape_ranges_[context->node_name], &tensorrt_mu_, - engine_cache_enable_, cache_path_, - runtime_.get(), profiles_[context->node_name], - engine_decryption_enable_, engine_decryption_, engine_encryption_, - detailed_build_log_, sparsity_enable_, - auxiliary_streams_, cuda_graph_enable_, is_dynamic_shape_context, cache_prefix_}; - *state = p.release(); - return 0; - }; - - // Release function state - compute_info.release_state_func = [](FunctionState state) { - delete static_cast(state); - }; - - // Create compute function - compute_info.compute_func = [this](FunctionState state, const OrtApi* api, OrtKernelContext* context) { - Ort::KernelContext ctx(context); - - TensorrtFuncState* trt_state = reinterpret_cast(state); - - // The whole compute_function should be considered the critical section where multiple threads may update kernel function state, access one builder, create/serialize/save engine, - // save profile and serialize/save timing cache. Therefore, those operations should be synchronized across different threads when ORT is using multithreading. - // More details here, https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - std::lock_guard lock(*(trt_state->tensorrt_mu_ptr)); - const std::unordered_map& input_indexes = (trt_state->input_info)[0]; - const std::unordered_map& output_indexes = (trt_state->output_info)[0]; - const std::unordered_map& output_types = (trt_state->output_info)[1]; - auto fused_node_name = trt_state->fused_node_name; - - std::unordered_map> shape_tensor_values; // This map holds "shape tensor -> shape values" for the shape tensor input across this inference run - std::unordered_map> shape_tensor_values_int64; // same as above but for int64 shape tensor input - auto& dds_output_allocator_map = this->dds_output_allocator_maps_[fused_node_name]; - auto trt_engine = trt_state->engine->get(); - auto trt_context = trt_state->context->get(); - auto trt_profiles = trt_state->profiles; - std::unordered_set input_names; - - if (alloc_ == nullptr) { - OrtDevice device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, - narrow(device_id_)); - OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, device); - Ort::ThrowOnError(api->KernelContext_GetAllocator(context, &mem_info, &alloc_)); - } - OrtAllocator* alloc = alloc_; - - cudaStream_t stream; - if (stream_ != nullptr) { - // Use our existing stream (either user's or our early-created) - stream = stream_; - } else { - // Create stream now (lazy creation case) - void* cuda_stream; - Ort::ThrowOnError(api->KernelContext_GetGPUComputeStream(context, &cuda_stream)); - stream = static_cast(cuda_stream); - stream_ = stream; - } - - if (multi_profile_enable_ == true) { - if (!trt_context->setOptimizationProfileAsync(nv_profile_index_, stream)) - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "NvTensorRTRTX EP select an optimization profile for the current context failed"); - } - - // Set auxiliary stream if provided by user - if (external_aux_streams_ && aux_streams_ != nullptr) { - trt_context->setAuxStreams(aux_streams_, (int32_t)auxiliary_streams_); - } - - // Check before using trt_engine - if (trt_engine == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "No engine is found."); - } - - bool require_io_binding = IsIOBindingRequired(trt_state, ctx); - - // Get input and output binding names - int total_bindings = trt_engine->getNbIOTensors(); - std::vector input_binding_names, output_binding_names; - for (int i = 0, end = total_bindings; i < end; ++i) { - auto const& name = trt_engine->getIOTensorName(i); - auto const& mode = trt_engine->getTensorIOMode(name); - if (mode == nvinfer1::TensorIOMode::kINPUT) { - input_binding_names.push_back(name); - } else { - output_binding_names.push_back(name); - } - } - - /* - * Set input shapes and bind input buffers - */ - auto& scratch_buffers = trt_state->scratch_buffers; - if (require_io_binding) { - scratch_buffers.clear(); - bool skip_input_binding_allowed = true; - for (size_t i = 0, end = input_binding_names.size(); i < end; ++i) { - char const* input_name = input_binding_names[i]; - - size_t input_index = 0; - const auto iter = input_indexes.find(input_name); - if (iter != input_indexes.end()) { - input_index = iter->second; - } - - auto status = BindContextInput(ctx, trt_engine, trt_context, input_name, input_index, shape_tensor_values, shape_tensor_values_int64, scratch_buffers, alloc, stream, skip_input_binding_allowed); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - trt_state->skip_io_binding_allowed = skip_input_binding_allowed; - } - - /* - * Set output shapes and bind output buffers - */ - if (require_io_binding) { - for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { - char const* output_name = output_binding_names[i]; - - size_t output_index = 0; - const auto& index_iter = output_indexes.find(output_name); - if (index_iter != output_indexes.end()) { - output_index = index_iter->second; - } - - size_t output_type = 0; - const auto type_iter = output_types.find(output_name); - if (type_iter != output_types.end()) { - output_type = type_iter->second; - } - - nvinfer1::Dims dims; - void* data_ptr = nullptr; - - Status status = BindContextOutput(ctx, trt_context, output_name, output_index, output_type, - dds_output_allocator_map, scratch_buffers, alloc, dims, data_ptr); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - - trt_state->output_tensors[output_index] = TensorParams{data_ptr, dims}; - } - } - - // Set execution context memory - if (require_io_binding) { - size_t mem_size = trt_engine->getDeviceMemorySizeV2(); - if (trt_state->is_dynamic_shape) { - mem_size = trt_context->updateDeviceMemorySizeForShapes(); - } - if (trt_state->context_memory_size != mem_size) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] A new context memory was allocated with size " << mem_size; - trt_state->context_memory = IAllocator::MakeUniquePtrFromOrtAllocator(alloc, mem_size, true /*use_reserve*/); - trt_state->context_memory_size = mem_size; - trt_context->setDeviceMemoryV2(trt_state->context_memory.get(), mem_size); - } - } - - if (!trt_context->enqueueV3(stream)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "NvTensorRTRTX EP execution context enqueue failed."); - } - - /* - * Given that InferenceSession::Run() is guaranteed to be thread-safe meaning multiple threads can call this function concurrently, - * TRT EP needs to carefully take care of concurrency here, if not, following concurrent issue might happen: - * - * It's suggested that to perform inference concurrently in multiple streams, use one trt execution context per stream. - * In the design of TRT EP (Not apply per-thread context implementation) and if multiple threads are calling InferenceSession::Run() concurrently, - * the trt execution context instance is shared by all the threads and each thread aquires different stream from ORT. - * So TRT EP will end up having one trt execution context using multiple streams which is not suggested. - * But, since the whole compute_func() is protected by the lock and if cudaStreamSynchronize() is enforced here, one trt execution context per stream - * is guaranteed. - * - * Therefore, TRT EP needs to call cudaStreamSynchronize() which means to wait until stream has completed all operations to prevent the concurrent issue mentioned above. - * However, if cuda graph is enabled, TRT EP won't call cudaStreamSynchronize() since it's not allowed during graph capture. - */ - - if (sync_stream_after_enqueue_) { - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); - } - // Assign TRT output back to ORT output - // (1) Bind TRT DDS output to ORT kernel context output. (It needs to wait until enqueueV3 is finished) - // (2) Cast TRT INT32 output to ORT INT64 output or TRT double output to float output - for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { - char const* output_name = output_binding_names[i]; - - size_t output_type = 0; - const auto& iter = output_types.find(output_name); - if (iter != output_types.end()) { - output_type = iter->second; - } - - if (dds_output_allocator_map.find(output_name) != dds_output_allocator_map.end()) { - size_t output_index = 0; - const auto& index_iter = output_indexes.find(output_name); - if (index_iter != output_indexes.end()) { - output_index = index_iter->second; - } - auto status = BindKernelOutput(ctx, dds_output_allocator_map, output_name, output_index, output_type, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, status.ErrorMessage()); - } - } - } - - return Status::OK(); - }; - - node_compute_funcs.push_back(compute_info); - return Status::OK(); -} - -Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const GraphViewer& graph_body_viewer, - size_t node_idx, - const Node& fused_node, - std::unordered_map& input_map, - std::unordered_map& output_map, - std::vector& node_compute_funcs) { - // Compile-only sessions never run inference. The input is already an EPContext model, - // so no engine save is needed either — skip deserialization and execution context creation - // and register a stub compute function that will not be invoked. - if (compile_only_mode_) { - NodeComputeInfo stub_compute_info; - stub_compute_info.create_state_func = [](ComputeContext*, FunctionState* state) -> int { - *state = nullptr; - return 0; - }; - stub_compute_info.release_state_func = [](FunctionState) {}; - stub_compute_info.compute_func = [](FunctionState, const OrtApi*, OrtKernelContext*) -> Status { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "NvTensorRTRTX EP: inference is not available in a compile-only session"); - }; - node_compute_funcs.push_back(std::move(stub_compute_info)); - return Status::OK(); - } - - std::unique_ptr trt_engine; - tensorrt_ptr::unique_pointer_exec_ctx trt_context; - std::unordered_map input_indexes; // TRT engine input name -> ORT kernel context input index - std::unordered_map output_indexes; // TRT engine output name -> ORT kernel context output index - std::unordered_map output_types; // TRT engine output name -> ORT output tensor type - - // Get engine binary data and deserialize it - auto trt_cache_model_handler = TensorRTCacheModelHandler(&trt_engine, - runtime_.get(), - model_path_, - compute_capability_, - weight_stripped_engine_enable_, - onnx_model_folder_path_, - onnx_model_bytestream_, - onnx_model_bytestream_size_, - onnx_external_data_bytestream_, - onnx_external_data_bytestream_size_, - detailed_build_log_); - auto status = trt_cache_model_handler.GetEpContextFromGraph(*graph_body_viewer.GetNode(node_idx)); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - - std::unique_ptr trt_runtime_cache; - auto trt_runtime_config = std::unique_ptr(trt_engine->createRuntimeConfig()); - if (trt_runtime_config && cuda_graph_enable_) { - trt_runtime_config->setDynamicShapesKernelSpecializationStrategy(nvinfer1::DynamicShapesKernelSpecializationStrategy::kEAGER); -#if TRT_MAJOR_RTX > 1 || (TRT_MAJOR_RTX == 1 && TRT_MINOR_RTX >= 3) - auto cuda_strategy_flag = trt_runtime_config->setCudaGraphStrategy(nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE); - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] CUDA graph strategy with RTX Graph capture enabled : " << cuda_strategy_flag; -#else - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] CUDA graph is enabled but RTX Graph capture is not available. " - << "The current TRT RTX version does not support RTX Graph. " - << "Please upgrade to TRT RTX >= 1.3 to use RTX Graph capture feature for optimal CUDA graph performance."; -#endif - } - trt_runtime_config->setExecutionContextAllocationStrategy(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED); - std::string runtime_cache_file = ""; - if (!runtime_cache_.empty()) { - runtime_cache_file = (runtime_cache_ / graph_body_viewer.GetNode(node_idx)->Name()).string(); - trt_runtime_cache = std::unique_ptr(trt_runtime_config->createRuntimeCache()); - auto cache_data = file_utils::ReadFile(runtime_cache_file); - if (!trt_runtime_cache->deserialize(cache_data.data(), cache_data.size())) { - trt_runtime_cache = std::unique_ptr(trt_runtime_config->createRuntimeCache()); - LOGS_DEFAULT(INFO) << "TensorRT RTX failed to deserialize the runtime cache, will overwrite with new one" << std::endl; - } - if (!trt_runtime_config->setRuntimeCache(*trt_runtime_cache)) { - LOGS_DEFAULT(INFO) << "TensorRT RTX failed to set the runtime cache" << std::endl; - } - } - - // Build context - // - // Note: Creating an execution context from an engine is thread safe per TRT doc - // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - trt_context = tensorrt_ptr::unique_pointer_exec_ctx( - trt_engine->createExecutionContext(trt_runtime_config.get()), - tensorrt_ptr::IExecutionContextDeleter(runtime_cache_file, std::move(trt_runtime_cache))); - if (!trt_context) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP could not build execution context for fused node: " + fused_node.Name()); - } - - bool is_dynamic_shape_context = false; - // Create input/output to index maps - for (int32_t i = 0; i < trt_engine->getNbIOTensors(); ++i) { - auto const& name = trt_engine->getIOTensorName(i); - auto const& mode = trt_engine->getTensorIOMode(name); - if (mode == nvinfer1::TensorIOMode::kINPUT) { - is_dynamic_shape_context |= checkTrtDimIsDynamic(trt_engine->getTensorShape(name)); - const auto& iter = input_map.find(name); - if (iter != input_map.end()) { - input_indexes[name] = iter->second; - } - } else { - const auto& iter = output_map.find(name); - if (iter != output_map.end()) { - output_indexes[name] = iter->second; - } - } - } - - // Create output to type map - for (auto node_arg : graph_body_viewer.GetOutputs()) { - auto output_name = node_arg->Name(); - auto& type = node_arg->TypeAsProto()->tensor_type(); - output_types[output_name] = type.elem_type(); - } - - // Save TRT engine, TRT context and input/output info to map - engines_.emplace(fused_node.Name(), std::move(trt_engine)); - contexts_.emplace(fused_node.Name(), std::move(trt_context)); - input_info_[fused_node.Name()].push_back(input_indexes); - output_info_[fused_node.Name()].push_back(output_indexes); - output_info_[fused_node.Name()].push_back(output_types); - - // Create function state - // TODO: remove default capture - NodeComputeInfo compute_info; - compute_info.create_state_func = [=](ComputeContext* context, FunctionState* state) { - std::unique_ptr p = std::make_unique(); - *p = {context->allocate_func, - context->release_func, - context->allocator_handle, - context->node_name, - &engines_[context->node_name], - &contexts_[context->node_name], - input_info_[context->node_name], - output_info_[context->node_name], - &tensorrt_mu_, - is_dynamic_shape_context}; - *state = p.release(); - return 0; - }; - - // Release function state - compute_info.release_state_func = [](FunctionState state) { - delete static_cast(state); - }; - - // Create compute function - compute_info.compute_func = [this](FunctionState state, const OrtApi* api, OrtKernelContext* context) { - Ort::KernelContext ctx(context); - - TensorrtShortFuncState* trt_state = reinterpret_cast(state); - - // The whole compute_function should be considered the critical section. - // More details here, https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - std::lock_guard lock(*(trt_state->tensorrt_mu_ptr)); - - const std::unordered_map& input_indexes = (trt_state->input_info)[0]; - const std::unordered_map& output_indexes = (trt_state->output_info)[0]; - const std::unordered_map& output_types = (trt_state->output_info)[1]; - auto fused_node_name = trt_state->fused_node_name; - auto& dds_output_allocator_map = this->dds_output_allocator_maps_[fused_node_name]; - auto trt_engine = trt_state->engine->get(); - auto trt_context = trt_state->context->get(); - std::unordered_map> shape_tensor_values; // This map holds "shape tensor -> shape values" for the shape tensor input across this inference run - std::unordered_map> shape_tensor_values_int64; // same as above but for int64 shape tensor input - - if (alloc_ == nullptr) { - OrtDevice device(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NVIDIA, - narrow(device_id_)); - OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, device); - Ort::ThrowOnError(api->KernelContext_GetAllocator(context, &mem_info, &alloc_)); - } - OrtAllocator* alloc = alloc_; - - cudaStream_t stream; - if (stream_ != nullptr) { - // Use our existing stream (either user's or our early-created) - stream = stream_; - } else { - // Create stream now (lazy creation case) - void* cuda_stream; - Ort::ThrowOnError(api->KernelContext_GetGPUComputeStream(context, &cuda_stream)); - stream = static_cast(cuda_stream); - } - - // Check before using trt_engine - if (trt_engine == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "No engine is found."); - } - - bool require_io_binding = IsIOBindingRequired(trt_state, ctx); - - // Get input and output binding names - int total_bindings = trt_engine->getNbIOTensors(); - std::vector input_binding_names, output_binding_names; - for (int i = 0, end = total_bindings; i < end; ++i) { - auto const& name = trt_engine->getIOTensorName(i); - auto const& mode = trt_engine->getTensorIOMode(name); - if (mode == nvinfer1::TensorIOMode::kINPUT) { - input_binding_names.push_back(name); - } else { - output_binding_names.push_back(name); - } - } - - /* - * Set input shapes and bind input buffers - */ - auto& scratch_buffers = trt_state->scratch_buffers; - if (require_io_binding) { - scratch_buffers.clear(); - bool skip_input_binding_allowed = true; - for (size_t i = 0, end = input_binding_names.size(); i < end; ++i) { - char const* input_name = input_binding_names[i]; - - size_t input_index = 0; - const auto iter = input_indexes.find(input_name); - if (iter != input_indexes.end()) { - input_index = iter->second; - } - - Status status = BindContextInput(ctx, trt_engine, trt_context, input_name, input_index, shape_tensor_values, shape_tensor_values_int64, scratch_buffers, alloc, stream, skip_input_binding_allowed); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - trt_state->skip_io_binding_allowed = skip_input_binding_allowed; - } - - /* - * Set output shapes and bind output buffers - */ - if (require_io_binding) { - for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { - char const* output_name = output_binding_names[i]; - - size_t output_index = 0; - const auto& index_iter = output_indexes.find(output_name); - if (index_iter != output_indexes.end()) { - output_index = index_iter->second; - } - - size_t output_type = 0; - const auto type_iter = output_types.find(output_name); - if (type_iter != output_types.end()) { - output_type = type_iter->second; - } - - nvinfer1::Dims dims; - void* data_ptr = nullptr; - - Status status = BindContextOutput(ctx, trt_context, output_name, output_index, output_type, - dds_output_allocator_map, scratch_buffers, alloc, dims, data_ptr); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - - trt_state->output_tensors[output_index] = TensorParams{data_ptr, dims}; - } - } - - // Set execution context memory - if (require_io_binding) { - size_t mem_size = trt_engine->getDeviceMemorySizeV2(); - if (trt_state->is_dynamic_shape) { - mem_size = trt_context->updateDeviceMemorySizeForShapes(); - } - if (trt_state->context_memory_size != mem_size) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] A new context memory was allocated with size " << mem_size; - trt_state->context_memory = IAllocator::MakeUniquePtrFromOrtAllocator(alloc, mem_size, true /*use_reserve*/); - // trt_state->context_memory = IAllocator::MakeUniquePtr(alloc, mem_size, false /*use_reserve*/, stream); - trt_state->context_memory_size = mem_size; - trt_context->setDeviceMemoryV2(trt_state->context_memory.get(), mem_size); - } - } - - // Set auxiliary stream if provided by user - if (external_aux_streams_ && aux_streams_ != nullptr) { - trt_context->setAuxStreams(aux_streams_, (int32_t)auxiliary_streams_); - } - - if (!trt_context->enqueueV3(stream)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "NvTensorRTRTX EP execution context enqueue failed."); - } - - /* - * Given that InferenceSession::Run() is guaranteed to be thread-safe meaning multiple threads can call this function concurrently, - * TRT EP needs to carefully take care of concurrency here, if not, following concurrent issue might happen: - * - * It's suggested that to perform inference concurrently in multiple streams, use one trt execution context per stream. - * In the design of TRT EP (Not apply per-thread context implementation) and if multiple threads are calling InferenceSession::Run() concurrently, - * the trt execution context instance is shared by all the threads and each thread aquires different stream from ORT. - * So TRT EP will end up having one trt execution context using multiple streams which is not suggested. - * But, since the whole compute_func() is protected by the lock and if cudaStreamSynchronize() is enforced here, one trt execution context per stream - * is guaranteed. - * - * Therefore, TRT EP needs to call cudaStreamSynchronize() which means to wait until stream has completed all operations to prevent the concurrent issue mentioned above. - * However, if cuda graph is enabled, TRT EP won't call cudaStreamSynchronize() since it's not allowed during graph capture. - */ - - if (sync_stream_after_enqueue_) { - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); - } - // Assign TRT output back to ORT output - // (1) Bind TRT DDS output to ORT kernel context output. (It needs to wait until enqueueV3 is finished) - // (2) Cast TRT INT32 output to ORT INT64 output or TRT double output to float output - for (size_t i = 0, end = output_binding_names.size(); i < end; ++i) { - char const* output_name = output_binding_names[i]; - - size_t output_type = 0; - const auto& iter = output_types.find(output_name); - if (iter != output_types.end()) { - output_type = iter->second; - } - - if (dds_output_allocator_map.find(output_name) != dds_output_allocator_map.end()) { - size_t output_index = 0; - const auto& index_iter = output_indexes.find(output_name); - if (index_iter != output_indexes.end()) { - output_index = index_iter->second; - } - auto status = BindKernelOutput(ctx, dds_output_allocator_map, output_name, output_index, output_type, stream); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, status.ErrorMessage()); - } - } - } - - return Status::OK(); - }; - - node_compute_funcs.push_back(compute_info); - return Status::OK(); -} - -void NvExecutionProvider::RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const { - auto allocator = allocators[GetOrtDeviceByMemType(OrtMemTypeCPU)]; - RegisterCudaStreamHandles(stream_handle_registry, - OrtDevice::GPU, - allocator, - true /* release_cpu_buffer_on_cuda_stream */, - stream_, - external_stream_ /* use_existing_stream */, - external_cudnn_handle_, - external_cublas_handle_, - {}); -} - -OrtDevice NvExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) const { - if (mem_type == OrtMemTypeCPUInput) - return OrtDevice(); - if (mem_type == OrtMemTypeCPUOutput) - return OrtDevice(OrtDevice::GPU, OrtDevice::MemType::HOST_ACCESSIBLE, OrtDevice::VendorIds::NVIDIA, - default_device_.Id()); - return default_device_; -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h deleted file mode 100644 index 89b18a08fff3c..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ /dev/null @@ -1,665 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include -#include -#include -#ifndef USE_CUDA_MINIMAL -#include -#else -typedef void* cudnnHandle_t; -typedef void* cublasHandle_t; -typedef void* cudnnStatus_t; -#endif -#include "core/providers/nv_tensorrt_rtx/nv_includes.h" -#include "core/session/onnxruntime_run_options_config_keys.h" -#include -#include "core/providers/cuda/cuda_graph.h" -#include "nv_execution_provider_info.h" -#include "core/providers/nv_tensorrt_rtx/nv_file_utils.h" - -namespace onnxruntime { - -class TensorrtLogger : public nvinfer1::ILogger { - nvinfer1::ILogger::Severity verbosity_; - - public: - TensorrtLogger(Severity verbosity = Severity::kWARNING) - : verbosity_(verbosity) {} - void log(Severity severity, const char* msg) noexcept override { - if (severity <= verbosity_) { - time_t rawtime = std::time(0); - struct tm stm; -#ifdef _MSC_VER - gmtime_s(&stm, &rawtime); -#else - gmtime_r(&rawtime, &stm); -#endif - char buf[256]; - strftime(&buf[0], 256, - "%Y-%m-%d %H:%M:%S", - &stm); - const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? " BUG" : severity == Severity::kERROR ? " ERROR" - : severity == Severity::kWARNING ? "WARNING" - : severity == Severity::kINFO ? " INFO" - : "UNKNOWN"); - if (severity <= Severity::kERROR) { - LOGS_DEFAULT(ERROR) << "[" << buf << " " << sevstr << "] " << msg; - } else { - LOGS_DEFAULT(WARNING) << "[" << buf << " " << sevstr << "] " << msg; - } - } - } - void set_level(Severity verbosity) { - verbosity_ = verbosity; - } - Severity get_level() const { - return verbosity_; - } -}; - -namespace tensorrt_ptr { -/* - * custom deleter that will dump the optimized runtime cache when the execution context is destructed - */ -struct IExecutionContextDeleter { - IExecutionContextDeleter() = default; - IExecutionContextDeleter(const std::string& runtime_cache_path, std::unique_ptr&& runtime_cache) : runtime_cache_path_(runtime_cache_path), runtime_cache_(std::move(runtime_cache)) {}; - void operator()(nvinfer1::IExecutionContext* context) { - if (context != nullptr) { - if (!runtime_cache_path_.empty()) { - auto serialized_cache_data = std::unique_ptr(runtime_cache_->serialize()); - file_utils::WriteFile(runtime_cache_path_, serialized_cache_data->data(), serialized_cache_data->size()); - } - delete context; - } - } - - private: - std::string runtime_cache_path_; - std::unique_ptr runtime_cache_; -}; - -struct TensorrtInferDeleter { - template - void operator()(T* obj) const { - if (obj) { - delete obj; - } - } -}; - -template -using unique_pointer = std::unique_ptr; -using unique_pointer_exec_ctx = std::unique_ptr; -}; // namespace tensorrt_ptr - -// -// Class to allocate memory for outputs with data-dependent shapes. The sizes of those are unknown so pre-allocation is -// not possible. -// -class OutputAllocator : public nvinfer1::IOutputAllocator { - public: - OutputAllocator() = delete; - OutputAllocator(OrtAllocator* allocator) : alloc_(allocator) {}; - - void* reallocateOutputAsync(char const* tensorName, void* currentMemory, uint64_t size, uint64_t alignment, cudaStream_t stream) noexcept override; - - void notifyShape(char const* tensorName, nvinfer1::Dims const& dims) noexcept override; - - void* getBuffer() { - return outputPtr; - } - - std::vector& getOutputShape() { - return output_shapes; - } - - uint64_t getSize() { - return allocated_size; - } - - ~OutputAllocator() override { - alloc_->Free(alloc_, outputPtr); - } - - private: - OrtAllocator* alloc_; - void* outputPtr{nullptr}; - uint64_t allocated_size = 0; - std::vector output_shapes; -}; - -/* - * This map saves the dimension range of the shape of the shape tensor or execution tensor: - * tensor name -> ( dimension -> [min, max, opt] ) - */ -using ShapeRangesMap = std::unordered_map>>>; - -// SubGraph_t and SubGraphCollection_t were removed from NvOnnxParser.h starting in TRT-RTX 1.5.0.99. -// Define them here so the provider owns these ORT-internal types for any SDK that no longer ships them. -#if TRT_MAJOR_RTX >= 2 || (TRT_MAJOR_RTX == 1 && ((TRT_MINOR_RTX == 5 && TRT_BUILD_RTX >= 99) || TRT_MINOR_RTX >= 6)) -using SubGraph_t = std::pair, bool>; -using SubGraphCollection_t = std::vector; -#endif - -/** - * @brief Container for tensor data and their shape. - * - */ -struct TensorParams { - const void* data{nullptr}; - nvinfer1::Dims dims; - - TensorParams() = default; - - TensorParams(const void* data_ptr, const std::vector& shape) { - // Initialize data and dims from the Ort::ConstValue - data = data_ptr; - - dims.nbDims = static_cast(shape.size()); - for (int i = 0; i < dims.nbDims; ++i) { - dims.d[i] = static_cast(shape[i]); - } - } - - TensorParams(const void* data_ptr, nvinfer1::Dims& shape) { - // Initialize data and dims from the Ort::ConstValue - data = data_ptr; - - dims = shape; - } - - bool operator!=(const TensorParams& other) const { - if (data != other.data || dims.nbDims != other.dims.nbDims) - return true; - - for (int i = 0; i < dims.nbDims; ++i) { - if (dims.d[i] != other.dims.d[i]) - return true; - } - return false; - } -}; - -// Data structure to hold user weights when ModelProtos are serialized with external data -class TensorrtUserWeights { - public: - TensorrtUserWeights(const std::string& name, const std::string& data) : name_(name), - data_cpy_(data) { - }; - - TensorrtUserWeights(const std::string& name, const void* data, size_t size) : name_(name), data_(data), size_(size) { - }; - - const char* Name() const { - return name_.c_str(); - }; - - const void* Data() const { - if (!data_cpy_.empty()) { - return data_cpy_.data(); - } - return data_; - } - - int64_t Size() const { - if (!data_cpy_.empty()) { - return static_cast(data_cpy_.size()); - } - return static_cast(size_); - } - - private: - std::string name_{}; - std::string data_cpy_{}; - void const* data_; - size_t size_; -}; - -// Information to construct kernel function state. -struct TensorrtFuncState { - AllocateFunc test_allocate_func = nullptr; - DestroyFunc test_release_func = nullptr; - AllocatorHandle allocator = nullptr; - std::string fused_node_name; - nvinfer1::IBuilder* builder; - std::unique_ptr* engine = nullptr; - tensorrt_ptr::unique_pointer_exec_ctx* context = nullptr; - std::unique_ptr* network = nullptr; - std::vector> input_info; - std::vector> output_info; - std::unordered_map>>> input_shape_ranges; - std::mutex* tensorrt_mu_ptr = nullptr; - bool engine_cache_enable = false; - std::string engine_cache_path; - nvinfer1::IRuntime* runtime = nullptr; - std::vector profiles; - bool engine_decryption_enable = false; - int (*engine_decryption)(const char*, char*, size_t*) = nullptr; - int (*engine_encryption)(const char*, char*, size_t) = nullptr; - bool detailed_build_log = false; - bool sparsity_enable = false; - int auxiliary_streams = -1; - bool cuda_graph_enable = 0; - bool is_dynamic_shape = false; - std::string cache_prefix; - std::string cache_suffix; - // runtime parameters - std::vector> scratch_buffers; - std::vector input_tensors; - std::vector output_tensors; - bool is_first_run = true; // Indicates if this is the first run of the engine - bool skip_io_binding_allowed = false; // Indicates if input/output binding can be skipped - IAllocatorUniquePtr context_memory = nullptr; - size_t context_memory_size = 0; -}; - -// Minimum information to construct kernel function state for direct engine load code path -struct TensorrtShortFuncState { - AllocateFunc test_allocate_func = nullptr; - DestroyFunc test_release_func = nullptr; - AllocatorHandle allocator = nullptr; - std::string fused_node_name; - std::unique_ptr* engine = nullptr; - tensorrt_ptr::unique_pointer_exec_ctx* context = nullptr; - std::vector> input_info; - std::vector> output_info; - std::mutex* tensorrt_mu_ptr = nullptr; - bool is_dynamic_shape = false; - // runtime parameters - std::vector> scratch_buffers; - std::vector input_tensors; - std::vector output_tensors; - bool is_first_run = true; // Indicates if this is the first run of the engine - bool skip_io_binding_allowed = false; // Indicates if input/output binding can be skipped - IAllocatorUniquePtr context_memory = nullptr; - size_t context_memory_size = 0; -}; - -// Holds important information for building valid ORT graph. -struct SubGraphContext { - std::unordered_set output_args; - std::unordered_map inputs_and_initializers; - std::unordered_map manually_added_graph_inputs; -}; - -using SubGraphContextMap = std::unordered_map>; -using DDSOutputAllocatorMap = std::unordered_map>; -std::string GetWeightRefittedEnginePath(std::string engine_cache_path); - -// Logical device representation. -class NvExecutionProvider : public IExecutionProvider { - public: - explicit NvExecutionProvider(const NvExecutionProviderInfo& info); - // TODO: we might want to transition to this, it allows for an easier option specification: - // explicit NvExecutionProvider(const ProviderOptions& provider_options_map, const ConfigOptions* config_options); - virtual ~NvExecutionProvider(); - - cublasHandle_t PerThreadDefaultCublasHandle() { - return GetPerThreadContext().CublasHandle(); - } - - cudnnHandle_t PerThreadDefaultCudnnHandle() { - return GetPerThreadContext().CudnnHandle(); - } - - virtual std::shared_ptr GetKernelRegistry() const override; - std::unique_ptr GetDataTransfer() const override; - - std::vector> - GetCapability(const GraphViewer& graph, - const IKernelLookup& /*kernel_lookup*/, - const GraphOptimizerRegistry& graph_optimizer_registry, - IResourceAccountant* /* resource_accountant */) const override; - - int GetDeviceId() const override { return device_id_; } - Status Sync() const override; - - common::Status Compile(const std::vector& fused_nodes_and_graphs, - std::vector& node_compute_funcs) override; - - Status OnRunStart(const onnxruntime::RunOptions& run_options) override; - Status OnRunEnd(bool sync_stream, const onnxruntime::RunOptions& run_options) override; - - ProviderOptions GetProviderOptions() const override { - return NvExecutionProviderInfo::ToProviderOptions(info_); - } - - void RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const override; - - void GetCustomOpDomainList(std::vector& custom_op_domain_list) const override; - - OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override; - - std::vector CreatePreferredAllocators() override; - - // CUDA Graph support - bool IsGraphCaptureEnabled() const override; - bool IsGraphCaptured(int graph_annotation_id) const override; - Status ReplayGraph(int graph_annotation_id, bool sync = true) override; - void HandleCudaGraphStart(cudaStream_t stream, bool require_io_binding, CudaGraphAnnotation_t cuda_graph_annotation_id, bool& graph_replay_on_this_run, bool& should_start_capture); - - static common::Status RefitEngine(std::string onnx_model_filename, - std::string& onnx_model_folder_path, - const void* onnx_model_bytestream, - size_t onnx_model_bytestream_size, - const void* onnx_external_data_bytestream, - size_t onnx_external_data_bytestream_size, - nvinfer1::ICudaEngine* trt_engine, - bool detailed_build_log); - - const InlinedVector GetEpContextNodes() const override; - - // Engine compatibility validation methods - std::string GetCompiledModelCompatibilityInfo(const onnxruntime::GraphViewer& graph_viewer) const override; - - common::Status ValidateCompiledModelCompatibilityInfo( - const std::string& compatibility_info, - OrtCompiledModelCompatibility& model_compatibility) const override; - - private: - mutable NvExecutionProviderInfo info_; - bool external_stream_ = false; - cudaStream_t stream_ = nullptr; - bool external_aux_streams_ = false; - cudaStream_t* aux_streams_ = nullptr; - int max_partition_iterations_ = 1000; - size_t min_subgraph_size_ = 1; - size_t max_workspace_size_ = 0; - size_t max_shared_mem_size_ = 0; - bool force_sequential_engine_build_ = false; - bool dump_subgraphs_ = false; - bool engine_cache_enable_ = false; - bool weight_stripped_engine_enable_ = false; - bool weight_stripped_engine_refit_ = false; - std::string onnx_model_folder_path_; - const void* onnx_model_bytestream_; - size_t onnx_model_bytestream_size_; - bool use_external_data_initializer_ = false; - const void* onnx_external_data_bytestream_ = nullptr; - size_t onnx_external_data_bytestream_size_ = 0; - bool sparsity_enable_ = false; - int auxiliary_streams_ = -1; - std::string cache_path_, engine_decryption_lib_path_; - std::unique_ptr runtime_ = nullptr; - std::mutex tensorrt_mu_; - int device_id_; - std::string compute_capability_; - size_t max_ctx_mem_size_ = 0; - mutable char model_path_[4096] = {}; // Reserved for max path length - bool engine_decryption_enable_ = false; - int (*engine_decryption_)(const char*, char*, size_t*) = nullptr; - int (*engine_encryption_)(const char*, char*, size_t) = nullptr; - bool detailed_build_log_ = false; - bool cuda_graph_enable_ = false; - bool multi_profile_enable_ = false; - std::filesystem::path runtime_cache_; - std::string cache_prefix_; - std::string op_types_to_exclude_; - int nv_profile_index_ = 0; - std::unique_ptr ep_context_model_; - - // The format is as for TENSORRT_VERSION: (MAJOR * 100 + MINOR) * 100 + PATCH - int32_t trt_version_; - int32_t cuda_version_; - - // The OrtAllocator object will be get during ep compute time - // and should be kept for the lifetime of TRT EP object. - OrtAllocator* alloc_ = nullptr; - - // For create/dump EP context node model - bool dump_ep_context_model_ = false; - // Set when the EP is instantiated by OrtCompileAPI::CompileModel(). Causes - // CreateNodeComputeInfoFromGraph to skip GPU deserialization and context creation. - bool compile_only_mode_ = false; - std::string ep_context_file_path_; - int ep_context_embed_mode_ = 0; - std::string ctx_model_path_; - std::string engine_cache_relative_path_to_context_model_dir; - - std::unordered_set control_flow_op_set_ = {"If", "Loop", "Scan"}; - mutable std::unordered_map> subgraph_context_map_; - - mutable std::unique_ptr builder_; - - // Following maps that hold TRT objects will be accessible by different threads if ORT is using multithreading. - // In general, TensorRT objects are not thread safe; accesses to an object from different threads must be serialized by the client. - // But there are still some thread safe operations, please see here https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - // For those non thread safe operations, TRT EP uses (1) lock_guard or (2) PerThreadContext to make sure synchronization. - std::unordered_map> engines_; - std::unordered_map contexts_; - std::unordered_map> builders_; - std::unordered_map> networks_; - std::unordered_map>> input_info_; - std::unordered_map>> output_info_; - std::unordered_map>> profile_min_shapes_; - std::unordered_map>> profile_max_shapes_; - std::unordered_map>> profile_opt_shapes_; - std::unordered_map input_shape_ranges_; // The profile shape ranges that the engine is built with - std::unordered_map> profiles_; - std::unordered_map dds_output_allocator_maps_; - - // Storage for engine headers (64 bytes) for compatibility validation - // Maps fused_node_name -> hex-encoded engine header - mutable std::unordered_map engine_headers_; - - // for external stream, we need to create its cudnn/cublass handle before cuda EP enable cuda graph capture - cudnnHandle_t external_cudnn_handle_ = nullptr; - cublasHandle_t external_cublas_handle_ = nullptr; - - // Call cudaStreamSynchronize() after TRT enqueueV3() - mutable bool sync_stream_after_enqueue_ = true; - - // [Note] We don't use PerThreadContext for now since it has issue with multithreading - // - // TRT or CUDA objects that must be maintained on a per thread basis will be put under this PerThreadContext data structure. - // For example, TensorRT execution context and CUDA graph are the ones to be put here. - class PerThreadContext final { - public: - PerThreadContext(OrtDevice::DeviceId device_id, bool has_user_compute_stream, cudaStream_t stream); - ~PerThreadContext(); - - cublasHandle_t CublasHandle() const { - return external_cublas_handle_; - } - - cudnnHandle_t CudnnHandle() const { - return external_cudnn_handle_; - } - - bool IsTensorRTContextInMap(std::string fused_node); - nvinfer1::IExecutionContext& GetTensorRTContext(std::string fused_node); - bool UpdateTensorRTContext(std::string fused_node, tensorrt_ptr::unique_pointer_exec_ctx context); - void ResetTensorRTContext(std::string fused_node); - - // CUDA Graph management - void SetCudaGraphStream(cudaStream_t stream) { cuda_graph_.SetStream(stream); } - bool IsGraphCaptureAllowed(CudaGraphAnnotation_t cuda_graph_annotation_id) const; - bool IsGraphCaptureAllowedOnRun(CudaGraphAnnotation_t cuda_graph_annotation_id) const; - CudaGraphAnnotation_t GetCudaGraphAnnotationId(const onnxruntime::RunOptions& run_options) const; - void SetCurrentGraphAnnotationId(CudaGraphAnnotation_t cuda_graph_annotation_id); - CudaGraphAnnotation_t GetCurrentGraphAnnotationId() const; - void CaptureBegin(CudaGraphAnnotation_t cuda_graph_annotation_id); - void CaptureEnd(CudaGraphAnnotation_t cuda_graph_annotation_id); - bool IsGraphCaptured(CudaGraphAnnotation_t cuda_graph_annotation_id) const; - Status ReplayGraph(CudaGraphAnnotation_t cuda_graph_annotation_id, bool sync_status_flag); - void IncrementRegularRunCountBeforeGraphCapture(CudaGraphAnnotation_t cuda_graph_annotation_id); - void ResetWarmupRuns(CudaGraphAnnotation_t cuda_graph_annotation_id); - void DeleteCapturedGraph(CudaGraphAnnotation_t cuda_graph_annotation_id); - - private: - cudnnHandle_t external_cudnn_handle_ = nullptr; - cublasHandle_t external_cublas_handle_ = nullptr; - - // Maintaining execution context on a per thread basis is suggested by TRT doc. - // Also, for enqueueV2() in execution context, to perform inference concurrently in multiple streams, use one execution context per stream. - // ORT multi-streams feature uses one stream for one thread, therefore maintaining execution context on a per thread basis is necessary for TRT EP, - // otherwise it may result in undefined behavior or synchronization issues. - // - // See more details here: - // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading - // https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_execution_context.html#a63cd95430852038ce864e17c670e0b36 - std::unordered_map trt_context_map_; - - // The profile shape ranges for the engine that the execution context maintained by the PerThreadContext is built with. - // TRT EP needs this info to determine whether to rebuild the execution context. - std::unordered_map input_shape_ranges_; - - // Cuda graph with multi threads will be supported in the future, so cuda_graph_ is put under PerThreadContext. - // ORT TRT only supports CUDA graph when whole model is supported by TRT, so simply maintaining a CUDAGraph instance is enough (no need to maintain one CUDAGraph instance per TRT subgraph) - CUDAGraph cuda_graph_; - // Map of graph id to regular_run_count_before_graph_capture - std::unordered_map graph_id_to_run_count_; - bool is_graph_captured_ = false; - int regular_run_count_before_graph_capture_ = 0; - // Current graph annotation ID for this run - CudaGraphAnnotation_t current_graph_annotation_id_ = kCudaGraphAnnotationDefault; - // There is chance (currently only happens in CUDA EP) that the second regular run allocates GPU memory for causes like: - // (1) memory pattern is enabled. (2) arena allocation for stream. - // Since no GPU memory allocation is allowed during graph capturing, we need at least two regular runs - // to allocate enough memory in Arena before graph capturing. - const int min_num_runs_before_cuda_graph_capture_ = 2; // required min regular runs before graph capture for the necessary memory allocations. - // https://github.com/NVIDIA/TensorRT/blob/main/samples/common/sampleInference.cpp#L1258-L1291 Based on the trtexec code - }; - - using PerThreadContextMap = std::unordered_map>; - // thread local PerThreadContext cache - - struct ContextCacheHolder { - ContextCacheHolder() { - // Keep a weak pointer to the object, if the weak pointer can be locked, then the shared pointer is still around, so we can reset it - RunOnUnload([&, weak_p_ = std::weak_ptr(p)] { - if (auto lock = weak_p_.lock()) { - p.reset(); - } - }); - } - - std::shared_ptr p = std::make_shared(); - }; - - static const std::shared_ptr& PerThreadContextCache() { - thread_local const ContextCacheHolder per_thread_context_cache; - return per_thread_context_cache.p; - } - - struct PerThreadContextState { - // contexts that are currently active - std::set, std::owner_less>> active_contexts; - // contexts available for reuse - std::vector> retired_context_pool; - // weak references to thread local caches from which this NvExecutionProvider instance's entry should be removed - // upon destruction - std::set, std::owner_less>> - caches_to_update_on_destruction; - // synchronizes access to PerThreadContextState members - std::mutex mutex; - }; - - // The execution provider maintains the PerThreadContexts in this structure. - // Synchronization is required to update the contained structures. - // On the other hand, access to an individual PerThreadContext is assumed to be from a single thread at a time, - // so synchronization is not required for that. - mutable PerThreadContextState context_state_; - - PerThreadContext& GetPerThreadContext() const; - void ReleasePerThreadContext() const; - - /**Get IndexedSubGraph based on node list of the subgraph*/ - std::unique_ptr GetSubGraph(SubGraph_t graph_nodes_index, - const GraphViewer& graph, const HashValue& model_hash, int subgraph_index) const; - - /** - Get TensorRT supported node lists by calling Onnx-TensorRT parser recursively. Since each time the parser - can only detect first unsupported node failure, it needs to wait for Onnxruntime to partition the graph - and then detect next failure again. If there are too many iterations, which means many nodes in the graph - are not supported by TensorRT, the process will be terminated and the whole graph is simply assigned to - other execution provider. - */ - SubGraphCollection_t GetSupportedList(SubGraphCollection_t supported_nodes_list, int iterations, const int max_iterations, - const GraphViewer& graph, bool* early_termination) const; - - bool DetectTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const GraphViewer& graph, const HashValue& model_hash, bool remove_cycles = true) const; - - /** - Get a unique_lock object to control the concurrency behavior. - Every api call not in the thread-safe operations(https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading) - should be protected by a lock when invoked by multiple threads concurrently. - */ - std::unique_lock GetApiLock() const; - - /**Check the graph is the subgraph of control flow op*/ - bool IsSubGraphOfControlFlowOp(const GraphViewer& graph) const; - - /**Check whether all the nodes of the graph are assigned to specific ep*/ - bool AllNodesAssignedToSpecificEP(const GraphViewer& graph, const std::string& provider_type) const; - - /**Check whether all the nodes of subgraph are supported*/ - bool IsSubGraphFullySupported(SubGraphCollection_t supported_nodes_vector, const int number_of_ort_nodes) const; - - /** - * Set inputs, initializers and outputs for all subgraphs during NvExecutionProvider::GetSupportedList() - * and save those information in subgraph context data structure. It's useful for building a valid graph and - * make Graph::Resolve() happy especially when dealing with nested control-flow op graph. - */ - void BuildSubGraphContext(const Graph& build_graph) const; - - /** - * Set outer scope values for subgraphs and add thoes values as top-level graph's inputs if needed. - */ - void SetGraphOuterScopeValuesAndInputs(Graph& build_graph, const Graph& graph) const; - - /** - * If ORT TRT manually sets graph input in NvExecutionProvider::SetGraphOuterScopeValuesAndInputs(), - * we have to manully set all the graph inputs in order to pass Graph::Resolve(). - */ - void SetAllGraphInputs(Graph& graph) const; - - /** - * The newly-built graph has not yet being resolved by Graph::Resolve(), so we can't leverage - * Graph::ResolveContext::IsInputInitializerOrOutput(). We have to implement this fuction again. - */ - bool IsInputInitializerOrOutput(const Graph& graph, const std::string& name, bool check_ancestors) const; - - /** - * The newly-built graph has not yet being resolved by Graph::Resolve(), so we can't leverage - * Graph::ResolveContext::IsOuterScopeValue(). We have to implement this fuction again. - */ - bool IsOuterScopeValue(const Graph& graph, const std::string& name) const; - - /** - * The newly-built graph has not yet being resolved by Graph::Resolve(), so we can't leverage - * Graph::ResolveContext::IsLocalValue(). We have to implement this fuction again. - */ - bool IsLocalValue(const Graph& graph, const std::string& name) const; - - /** - * Create a vector of NodeComputeInfo instances directly from "TRT engine" wrapped onnx model without - * going through the time-consuming processes of model parsing and engine building. - */ - Status CreateNodeComputeInfoFromPrecompiledEngine(const GraphViewer& graph_body_viewer, - size_t node_idx, - const Node& fused_node, - std::unordered_map& input_map, - std::unordered_map& output_map, - std::vector& node_compute_funcs); - - /** - * Create a vector of NodeComputeInfo instances from graph. - */ - Status CreateNodeComputeInfoFromGraph(const GraphViewer& graph_body_viewer, - const Node& fused_node, - std::unordered_map& input_map, - std::unordered_map& output_map, - std::vector& node_compute_funcs); - - /** - * Get the pointer to the IBuilder instance. - * This function only creates the instance at the first time it's being called." - */ - nvinfer1::IBuilder* GetBuilder(TensorrtLogger& trt_logger) const; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.cc deleted file mode 100644 index 78f62e14d4aa2..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.cc +++ /dev/null @@ -1,186 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include - -#include "core/framework/provider_options.h" -#include "nv_execution_provider_custom_ops.h" -#include "nv_execution_provider.h" -#include "nv_platform_utils.h" - -namespace onnxruntime { -extern TensorrtLogger& GetTensorrtLogger(bool verbose); - -/* - * Create custom op domain list for TRT plugins. - * - * Here, we collect all registered TRT plugins from TRT registry and create custom ops with "trt.plugins" domain. - * Additionally, if users specify extra plugin libraries, TRT EP will load them at runtime which will register those - * plugins to TRT plugin registry and later TRT EP can get them as well. - * - * There are several TRT plugins registered as onnx schema op through contrib op with ONNX domain in the past, - * for example, EfficientNMS_TRT, MultilevelCropAndResize_TRT, PyramidROIAlign_TRT and DisentangledAttention_TRT. - * In order not to break the old models using those TRT plugins which were registered with ONNX domain and maintain - * backward compatible, we need to keep those legacy TRT plugins registered with ONNX domain with contrib ops. - * - * Note: Current TRT plugin doesn't have APIs to get number of inputs/outputs of the plugin. - * So, TensorRTCustomOp uses variadic inputs/outputs to pass ONNX graph validation. - */ -common::Status CreateTensorRTCustomOpDomainList(std::vector& domain_list, const std::string extra_plugin_lib_paths) { - // Domain for TRT plugin custom ops (domain name: "trt.plugins"). Owns the OrtCustomOpDomain object. - // Raw pointers from .get() are handed out to callers via domain_list and may be held by InferenceSession. - static std::unique_ptr custom_op_domain = std::make_unique(); - - // Owns the TensorRTCustomOp objects for TRT plugins. Raw pointers are stored in custom_op_domain->custom_ops_. - static std::vector> created_custom_op_list; - - // Domain for native custom ops (domain name: "trt"). Owns the OrtCustomOpDomain object. - // Raw pointers from .get() are handed out to callers via domain_list and may be held by InferenceSession. - static std::unique_ptr native_custom_op_domain = std::make_unique(); - - // Owns the TensorRTCustomOp objects for native custom ops. Raw pointers are stored in native_custom_op_domain->custom_ops_. - // Non-empty list indicates native custom ops have been registered (used to avoid re-registration on subsequent calls). - static std::vector> native_custom_op_list; - - // Protects concurrent access to all the above static members. - static std::mutex mutex; - std::lock_guard lock(mutex); - - // Add already-initialized native ops to domain list - if (!native_custom_op_list.empty()) { - domain_list.push_back(native_custom_op_domain.get()); - } - - if (custom_op_domain->domain_ != "" && custom_op_domain->custom_ops_.size() > 0) { - domain_list.push_back(custom_op_domain.get()); - return Status::OK(); - } - - // Load any extra TRT plugin library if any. - // When the TRT plugin library is loaded, the global static object is created and the plugin is registered to TRT registry. - // This is done through macro, for example, REGISTER_TENSORRT_PLUGIN(VisionTransformerPluginCreator). - // extra_plugin_lib_paths has the format of "path_1;path_2....;path_n" - static bool is_loaded = false; - if (!extra_plugin_lib_paths.empty() && !is_loaded) { - std::stringstream extra_plugin_libs(extra_plugin_lib_paths); - std::string lib; - while (std::getline(extra_plugin_libs, lib, ';')) { - auto status = LoadDynamicLibrary(ToPathString(lib)); - if (status == Status::OK()) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully load " << lib; - } else { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP]" << status.ToString(); - } - } - is_loaded = true; - } - - // Load external tensorrt_plugins library from EP directory (same location as nvinfer_plugin) - // This library contains GroupQueryAttention and RotaryEmbedding plugins for transformer models - try { - const auto& env = onnxruntime::GetDefaultEnv(); - auto external_plugin_path = GetEPLibraryDirectory() + - PathString(LIBRARY_PREFIX ORT_TSTR("tensorrt_plugins") LIBRARY_EXTENSION); - void* external_plugin_handle = nullptr; - auto status = env.LoadDynamicLibrary(external_plugin_path, false, &external_plugin_handle); - if (status.IsOK()) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] External plugins loaded: tensorrt_plugins (GQA + RotaryEmbedding)"; - } else { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] tensorrt_plugins library not found in EP library path (optional)"; - } - } catch (const std::exception& e) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] tensorrt_plugins library not available: " << e.what(); - } - - try { - // Get all registered TRT plugins from registry - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Getting all registered TRT plugins from TRT plugin registry ..."; - TensorrtLogger trt_logger = GetTensorrtLogger(false); - void* library_handle = nullptr; - const auto& env = onnxruntime::GetDefaultEnv(); - auto full_path = env.GetRuntimePath() + - PathString(LIBRARY_PREFIX ORT_TSTR("nvinfer_plugin") LIBRARY_EXTENSION); - ORT_THROW_IF_ERROR(env.LoadDynamicLibrary(full_path, false, &library_handle)); - - bool (*dyn_initLibNvInferPlugins)(void* logger, char const* libNamespace); - ORT_THROW_IF_ERROR(env.GetSymbolFromLibrary(library_handle, "initLibNvInferPlugins", (void**)&dyn_initLibNvInferPlugins)); - dyn_initLibNvInferPlugins(&trt_logger, ""); - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Default plugins successfully loaded."; - -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4996) // Ignore warning C4996: 'nvinfer1::*' was declared deprecated -#endif - } catch (const std::exception&) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Default plugin library is not on the path and is therefore ignored"; - } - try { - // getAllCreators() is the TRT 10+ replacement for the removed getPluginCreatorList(). - // It returns IPluginCreatorInterface* const*; cast each entry to the deprecated-but-present - // IPluginCreator* to access getPluginName()/getPluginVersion(). - int32_t num_plugin_creator = 0; - auto plugin_creators = getPluginRegistry()->getAllCreators(&num_plugin_creator); - std::unordered_set registered_plugin_names; - - for (int32_t i = 0; i < num_plugin_creator; i++) { - auto* plugin_creator = dynamic_cast(plugin_creators[i]); - if (!plugin_creator) continue; // skip IPluginCreatorV3One entries that lack these accessors - std::string plugin_name(plugin_creator->getPluginName()); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] " << plugin_name << ", version : " << plugin_creator->getPluginVersion(); - - // plugin has different versions and we only register once - if (registered_plugin_names.find(plugin_name) != registered_plugin_names.end()) { - continue; - } - - created_custom_op_list.push_back(std::make_unique(onnxruntime::kNvTensorRTRTXExecutionProvider, nullptr)); // Make sure TensorRTCustomOp object won't be cleaned up - created_custom_op_list.back().get()->SetName(plugin_creator->getPluginName()); - custom_op_domain->custom_ops_.push_back(created_custom_op_list.back().get()); - registered_plugin_names.insert(plugin_name); - } - -#if defined(_MSC_VER) -#pragma warning(pop) -#endif - - custom_op_domain->domain_ = "trt.plugins"; - domain_list.push_back(custom_op_domain.get()); - } catch (const std::exception&) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Failed to get TRT plugins from TRT plugin registration. Therefore, TRT EP can't create custom ops for TRT plugins"; - } - - // Register native custom ops (register these independent of TRT plugin library availability) - if (native_custom_op_list.empty()) { - const char* native_custom_ops_names[] = {"TRT_FP4DynamicQuantize", "TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"}; - size_t num_native_custom_ops = std::size(native_custom_ops_names); - - for (size_t i = 0; i < num_native_custom_ops; i++) { - native_custom_op_list.push_back(std::make_unique(onnxruntime::kNvTensorRTRTXExecutionProvider, nullptr)); - native_custom_op_list.back()->SetName(native_custom_ops_names[i]); - native_custom_op_domain->custom_ops_.push_back(native_custom_op_list.back().get()); - } - - native_custom_op_domain->domain_ = "trt"; - domain_list.push_back(native_custom_op_domain.get()); - } - - return Status::OK(); -} - -void ReleaseTensorRTCustomOpDomain(OrtCustomOpDomain* domain) { - (void)domain; // Suppress unused parameter warning - // The domain and its custom ops are owned by static unique_ptrs in CreateTensorRTCustomOpDomainList(). - // Callers receive raw pointers via .get(). - // 1. Manually deleting them would cause a double-free when the static unique_ptrs are destroyed at program exit. - // 2. Resetting the static unique_ptrs is also unsafe because other EP instances or InferenceSession objects - // may still hold raw pointers to these same objects (handed out via domain_list). - // The static objects would be shared across EP instances and would persist for the program lifetime. -} - -void ReleaseTensorRTCustomOpDomainList(std::vector& custom_op_domain_list) { - // Only clear the reference vector, don't delete the static domain objects. - custom_op_domain_list.clear(); -} - -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.h deleted file mode 100644 index 81c0d49239ec8..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.h +++ /dev/null @@ -1,86 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#define ORT_API_MANUAL_INIT -#include "core/session/onnxruntime_c_api.h" -#include "core/session/onnxruntime_cxx_api.h" -#include "core/providers/shared_library/provider_api.h" -#include "nv_execution_provider_info.h" - -using namespace onnxruntime; - -namespace onnxruntime { - -common::Status LoadDynamicLibrary(onnxruntime::PathString library_name); -common::Status CreateTensorRTCustomOpDomainList(std::vector& domain_list, - const std::string extra_plugin_lib_paths); -common::Status CreateTensorRTCustomOpDomainList(NvExecutionProviderInfo& info); -void ReleaseTensorRTCustomOpDomain(OrtCustomOpDomain* domain); -void ReleaseTensorRTCustomOpDomainList(std::vector& custom_op_domain_list); - -struct TensorRTCustomKernel { - TensorRTCustomKernel(const OrtKernelInfo* /*info*/, void* compute_stream) - : compute_stream_(compute_stream) { - } - - void Compute(OrtKernelContext* /*context*/) { - // The implementation is in TensorRT plugin. No need to implement it here. - }; - - private: - void* compute_stream_; -}; - -struct TensorRTCustomOp : Ort::CustomOpBase { - explicit TensorRTCustomOp(const char* provider, void* compute_stream) : provider_(provider), - compute_stream_(compute_stream) { - } - - void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { - return new TensorRTCustomKernel(info, compute_stream_); - }; - - const char* GetName() const { return name_; }; - - void SetName(const char* name) { name_ = name; }; - - const char* GetExecutionProviderType() const { return provider_; }; - - size_t GetInputTypeCount() const { return num_inputs_; }; - - void SetInputTypeCount(size_t num) { num_inputs_ = num; }; - - ONNXTensorElementDataType GetInputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; }; - - OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t) const { - return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC; - }; - - size_t GetOutputTypeCount() const { return num_outputs_; }; - - void SetOutputTypeCount(size_t num) { num_outputs_ = num; }; - - ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; }; - - OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t) const { - return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC; - }; - - bool GetVariadicInputHomogeneity() const { - return false; // heterogenous - } - - bool GetVariadicOutputHomogeneity() const { - return false; // heterogeneous - } - - private: - const char* provider_{onnxruntime::kNvTensorRTRTXExecutionProvider}; - void* compute_stream_; - const char* name_; - size_t num_inputs_ = 1; // set to 1 to match with default min_arity for variadic input - size_t num_outputs_ = 1; // set to 1 to match with default min_arity for variadic output -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_helper.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_helper.cc deleted file mode 100644 index 8728558006fc5..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_helper.cc +++ /dev/null @@ -1,262 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/shared_library/provider_api.h" -#include "nv_execution_provider.h" -#include "core/framework/murmurhash3.h" -#include - -namespace onnxruntime { - -namespace { -// Get unique graph name based on graph's name and all nodes' name -std::string GetUniqueGraphName(const Graph& graph) { - HashValue model_hash = 0; - uint32_t hash[4] = {0, 0, 0, 0}; - - auto hash_str = [&hash](const std::string& str) { - MurmurHash3::x86_128(str.data(), gsl::narrow_cast(str.size()), hash[0], &hash); - }; - - // Hash all nodes' name - for (int i = 0; i < graph.MaxNodeIndex(); ++i) { - auto node = graph.GetNode(i); - if (node == nullptr) { - continue; - } - hash_str(node->Name()); - } - - model_hash = hash[0] | (uint64_t(hash[1]) << 32); - - return graph.Name() + "_" + std::to_string(model_hash); -} -} // namespace - -// The newly-built graph has not yet being resolved by Graph::Resolve(), so we can't leverage -// Graph::ResolveContext::IsInputInitializerOrOutput(). We have to implement this fuction again. -bool NvExecutionProvider::IsInputInitializerOrOutput(const Graph& graph, - const std::string& name, - bool check_ancestors) const { - const Graph* parent_graph = nullptr; - return IsLocalValue(graph, name) || - (check_ancestors && (parent_graph = graph.ParentGraph()) != nullptr && - IsInputInitializerOrOutput(*parent_graph, name, check_ancestors)); -} - -// The newly-built graph has not yet being resolved by Graph::Resolve(), so we can't leverage -// Graph::ResolveContext::IsOuterScopeValue(). We have to implement this function again. -bool NvExecutionProvider::IsOuterScopeValue(const Graph& graph, - const std::string& name) const { - const Graph* parent_graph = nullptr; - return (parent_graph = graph.ParentGraph()) != nullptr && - IsInputInitializerOrOutput(*parent_graph, name, true); -} - -// The newly-built graph has not yet being resolved by Graph::Resolve(), so we can't leverage -// Graph::ResolveContext::IsLocalValue(). We have to implement this function again. -bool NvExecutionProvider::IsLocalValue(const Graph& graph, - const std::string& name) const { - std::string unique_graph_name = GetUniqueGraphName(graph); - if (subgraph_context_map_.find(unique_graph_name) == subgraph_context_map_.end()) { - return false; - } - SubGraphContext* context = subgraph_context_map_.at(unique_graph_name).get(); - return context->output_args.find(name) != context->output_args.cend() || - context->inputs_and_initializers.find(name) != context->inputs_and_initializers.cend(); -} - -/** - * Set inputs, initializers and outputs for all subgraphs during NvExecutionProvider::GetSupportedList() - * and save those information in subgraph context data structure. It's useful for building a valid graph and - * make Graph::Resolve() happy especially when dealing with nested control-flow op graph. - */ -void NvExecutionProvider::BuildSubGraphContext(const Graph& graph) const { - // Iterate all the nodes and recurse into inner most subgraph first - for (int i = 0; i < graph.MaxNodeIndex(); ++i) { - auto node = graph.GetNode(i); - if (node == nullptr) { - continue; - } - - auto subgraph_map = node->GetAttributeNameToSubgraphMap(); - for (auto& entry : subgraph_map) { - const Graph* subgraph = entry.second; - BuildSubGraphContext(*subgraph); - } - } - - std::string unique_graph_name = GetUniqueGraphName(graph); - - // Subgraph context has been built before, no need to do it again - if (subgraph_context_map_.find(unique_graph_name) != subgraph_context_map_.end()) { - return; - } - - subgraph_context_map_.emplace(unique_graph_name, std::make_unique()); - SubGraphContext* context = subgraph_context_map_.at(unique_graph_name).get(); - - // Collect all nodes' outputs and nodes' name - for (int i = 0; i < graph.MaxNodeIndex(); ++i) { - auto node = graph.GetNode(i); - if (node == nullptr) { - continue; - } - - for (const auto& output : node->OutputDefs()) { - context->output_args.insert(output->Name()); - } - } - - // Go thru all node's inputs - for (int i = 0; i < graph.MaxNodeIndex(); ++i) { - auto node = graph.GetNode(i); - if (node == nullptr) { - continue; - } - - for (const auto& input : node->InputDefs()) { - if (context->output_args.find(input->Name()) != context->output_args.end()) { - continue; - } - // This input arg is not the output of another node so must come from either a graph input or an initializer. - context->inputs_and_initializers[input->Name()] = input; - } - } -} - -// Set outer scope values for subgraphs and add thoes values as top-level graph's inputs if needed. -void NvExecutionProvider::SetGraphOuterScopeValuesAndInputs(Graph& graph_build, - const Graph& graph) const { - // Iterate all the nodes and recurse into inner most subgraph first for both newly built graph and original graph - for (int i = 0; i < graph_build.MaxNodeIndex(); ++i) { - auto graph_build_node = graph_build.GetNode(i); - if (graph_build_node == nullptr) { - continue; - } - - auto graph_build_map = graph_build_node->GetAttributeNameToMutableSubgraphMap(); - std::unordered_map> subgraph_map; - const Node* graph_node = nullptr; - - // Find corresponding original graph node's subgraphs - for (int j = 0; j < graph.MaxNodeIndex(); ++j) { - if (graph.GetNode(j) && graph.GetNode(j)->Name() == graph_build_node->Name()) { - graph_node = graph.GetNode(j); - subgraph_map = graph_node->GetAttributeNameToSubgraphMap(); - break; - } - } - - for (auto& entry : graph_build_map) { - auto attr_name = entry.first; - Graph* subgraph_build = entry.second; - if (subgraph_map.find(attr_name) != subgraph_map.end()) { - // recurse into subgraph - const Graph* subgraph = subgraph_map.at(attr_name); - SetGraphOuterScopeValuesAndInputs(*subgraph_build, *subgraph); - } - } - } - - // Start from the inner most subgraph first and check whether its outer scope values are existed in the - // newly built graph. If not, we need to add those outer scope values as explicit inputs to the top-level - // of newly built graph. - if (graph_build.ParentNode()) { - auto top_level_graph = &graph_build; - while (top_level_graph->MutableParentGraph()) { - top_level_graph = top_level_graph->MutableParentGraph(); - } - std::string unique_graph_name = GetUniqueGraphName(*top_level_graph); - if (subgraph_context_map_.find(unique_graph_name) == subgraph_context_map_.end()) { - LOGS_DEFAULT(ERROR) << "[NvTensorRTRTX EP] Can't find top-level graph context. \ - Please check BuildSubGraphContext() has built the graph context correctly."; - return; - } - - SubGraphContext* context = subgraph_context_map_.at(unique_graph_name).get(); - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Subgraph name is " << graph_build.Name(); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Its parent node is " << graph.ParentNode()->Name(); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Its parent node's implicit inputs:"; - - // Iterate all the implicit inputs to set outer scope value for the newly built subgraph - for (const auto& input : graph.ParentNode()->ImplicitInputDefs()) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] \t" << input->Name(); - - // The node arg in parent node's implicit inputs could be used for parent node's other subgraph, for example - // "If" op has two subgraphs. So we need to make sure that the node arg is used in current subgraph only. - // (GetNodeArg searches for specific node arg in all node args in the graph) - if (graph_build.GetNodeArg(input->Name())) { - graph_build.AddOuterScopeNodeArg(input->Name()); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] \t" << input->Name() << " is used in this subgraph"; - - if (context && - (context->manually_added_graph_inputs.find(input->Name()) != context->manually_added_graph_inputs.end())) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] \t" << input->Name() << " is already been added as an explicit input to graph"; - continue; - } - - // Handle the case where this outer scope value is not existed in any outer scope levels of the - // newly built graph (the newly built graph is the subgraph of the original graph). Need to add - // the outer scope value as an explicit input to the top-level of newly built graph. - if (!IsOuterScopeValue(graph_build, input->Name())) { - const auto& name = input->Name(); - auto graph_inputs_including_initializers = top_level_graph->GetInputsIncludingInitializers(); - auto added_graph_input = std::find_if(graph_inputs_including_initializers.begin(), - graph_inputs_including_initializers.end(), - [&name](const NodeArg* entry) { return entry->Name() == name; }); - - if (added_graph_input == graph_inputs_including_initializers.end()) { - if (context) { - auto type_proto = ONNX_NAMESPACE::TypeProto::Create(); - type_proto->copy_from(input->TypeAsProto()); - auto& n_input = top_level_graph->GetOrCreateNodeArg(name, type_proto.get()); - context->manually_added_graph_inputs[n_input.Name()] = &n_input; - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] \t" << n_input.Name() << " is added as an explicit input into the newly built graph"; - } - } - } - } - } - } -} - -// If ORT TRT manually sets graph input in NvExecutionProvider::SetGraphOuterScopeValuesAndInputs(), -// we have to manully set all the graph inputs in order to pass Graph::Resolve() -void NvExecutionProvider::SetAllGraphInputs(Graph& graph) const { - // If ORT TRT doesn't manully set graph input in NvExecutionProvider::SetGraphOuterScopeValuesAndInputs(), - // Graph::Resolve() will help set graph inputs in Graph::SetGraphInputsOutputs(), so no need to set graph inputs here. - std::string unique_graph_name = GetUniqueGraphName(graph); - if (subgraph_context_map_.find(unique_graph_name) == subgraph_context_map_.end() || - subgraph_context_map_[unique_graph_name].get()->manually_added_graph_inputs.size() == 0) { - return; - } - - SubGraphContext* context = subgraph_context_map_[unique_graph_name].get(); - std::vector graph_inputs_including_initializers; - std::unordered_set graph_inputs_including_initializers_set; - - for (const auto& entry : context->inputs_and_initializers) { - graph_inputs_including_initializers.push_back(entry.second); - graph_inputs_including_initializers_set.insert(entry.first); - } - - for (const auto& entry : context->manually_added_graph_inputs) { - if (graph_inputs_including_initializers_set.find(entry.first) == graph_inputs_including_initializers_set.end()) { - graph_inputs_including_initializers.push_back(entry.second); - graph_inputs_including_initializers_set.insert(entry.first); - } - } - - for (const auto& node_arg : graph.GetInputsIncludingInitializers()) { - if (graph_inputs_including_initializers_set.find(node_arg->Name()) == graph_inputs_including_initializers_set.end()) { - graph_inputs_including_initializers.push_back(node_arg); - graph_inputs_including_initializers_set.insert(node_arg->Name()); - } - } - - graph.SetInputs(graph_inputs_including_initializers); -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc deleted file mode 100644 index afb9d8107392c..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ /dev/null @@ -1,128 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_options.h" - -#include "core/session/onnxruntime_session_options_config_keys.h" -#include "core/common/make_string.h" -#include "core/common/parse_string.h" -#include "core/framework/provider_options_utils.h" -#include "core/providers/cuda/cuda_common.h" - -namespace onnxruntime { -NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const ProviderOptions& options, - const ConfigOptions& session_options) { - NvExecutionProviderInfo info{}; - void* user_compute_stream = nullptr; - void* user_aux_stream_array = nullptr; - void* onnx_bytestream = nullptr; - void* external_data_bytestream = nullptr; - ORT_THROW_IF_ERROR( - ProviderOptionsParser{} - .AddValueParser( - nv::provider_option_names::kDeviceId, - [&info](const std::string& value_str) -> Status { - ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, info.device_id)); - int num_devices{}; - CUDA_RETURN_IF_ERROR(cudaGetDeviceCount(&num_devices)); - ORT_RETURN_IF_NOT( - 0 <= info.device_id && info.device_id < num_devices, - "Invalid device ID: ", info.device_id, - ", must be between 0 (inclusive) and ", num_devices, " (exclusive)."); - return Status::OK(); - }) - .AddAssignmentToReference(nv::provider_option_names::kHasUserComputeStream, info.has_user_compute_stream) - .AddValueParser( - nv::provider_option_names::kUserComputeStream, - [&user_compute_stream](const std::string& value_str) -> Status { - size_t address; - ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, address)); - user_compute_stream = reinterpret_cast(address); - return Status::OK(); - }) - .AddValueParser( - nv::provider_option_names::kUserAuxStreamArray, - [&user_aux_stream_array](const std::string& value_str) -> Status { - size_t address; - ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, address)); - user_aux_stream_array = reinterpret_cast(address); - return Status::OK(); - }) - .AddAssignmentToReference(nv::provider_option_names::kMaxWorkspaceSize, info.max_workspace_size) - .AddAssignmentToReference(nv::provider_option_names::kMaxSharedMemSize, info.max_shared_mem_size) - .AddAssignmentToReference(nv::provider_option_names::kLengthAuxStreamArray, info.auxiliary_streams) - .AddAssignmentToReference(nv::provider_option_names::kDumpSubgraphs, info.dump_subgraphs) - .AddAssignmentToReference(nv::provider_option_names::kDetailedBuildLog, info.detailed_build_log) - .AddAssignmentToReference(nv::provider_option_names::kProfilesMinShapes, info.profile_min_shapes) - .AddAssignmentToReference(nv::provider_option_names::kProfilesMaxShapes, info.profile_max_shapes) - .AddAssignmentToReference(nv::provider_option_names::kProfilesOptShapes, info.profile_opt_shapes) - .AddAssignmentToReference(nv::provider_option_names::kCudaGraphEnable, info.cuda_graph_enable) - .AddAssignmentToReference(nv::provider_option_names::kUseExternalDataInitializer, info.use_external_data_initializer) - .AddAssignmentToReference(nv::provider_option_names::kMultiProfileEnable, info.multi_profile_enable) - .AddAssignmentToReference(nv::provider_option_names::kRuntimeCacheFile, info.runtime_cache_path) - .Parse(options)); // add new provider option here. - - info.user_compute_stream = user_compute_stream; - info.has_user_compute_stream = (user_compute_stream != nullptr); - info.user_aux_stream_array = user_aux_stream_array; - info.onnx_bytestream = onnx_bytestream; - info.external_data_bytestream = external_data_bytestream; - - // EP context settings - // when EP context is enabled, default is to embed the engine in the context model - // weight stripped engine is always enabled when EP context is enabled - - const auto ep_context_enable = session_options.GetConfigOrDefault(kOrtSessionOptionEpContextEnable, "0"); - if (ep_context_enable == "0") { - info.dump_ep_context_model = false; - } else if (ep_context_enable == "1") { - info.dump_ep_context_model = true; - // We want to reenable weightless engines as soon constant initializers are supported as inputs - info.weight_stripped_engine_enable = false; - } else { - ORT_THROW("Invalid ", kOrtSessionOptionEpContextEnable, " must 0 or 1"); - } - info.ep_context_file_path = session_options.GetConfigOrDefault(kOrtSessionOptionEpContextFilePath, ""); - - info.compile_only_mode = (session_options.GetConfigOrDefault(kOrtSessionOptionCompileOnly, "0") == "1"); - - // If embed mode is not specified, default to 1 if dump_ep_context_model is true, otherwise 0 - auto embed_mode = std::stoi(session_options.GetConfigOrDefault(kOrtSessionOptionEpContextEmbedMode, "-1")); - if (embed_mode == -1) { - if (info.dump_ep_context_model) - embed_mode = 1; - else - embed_mode = 0; - } - - if (0 <= embed_mode || embed_mode < 2) { - info.ep_context_embed_mode = embed_mode; - } else { - ORT_THROW("Invalid ", kOrtSessionOptionEpContextEmbedMode, " must 0 or 1"); - } - - return info; -} - -ProviderOptions NvExecutionProviderInfo::ToProviderOptions(const NvExecutionProviderInfo& info) { - const ProviderOptions options{ - {nv::provider_option_names::kDeviceId, MakeStringWithClassicLocale(info.device_id)}, - {nv::provider_option_names::kHasUserComputeStream, MakeStringWithClassicLocale(info.has_user_compute_stream)}, - {nv::provider_option_names::kUserComputeStream, MakeStringWithClassicLocale(reinterpret_cast(info.user_compute_stream))}, - {nv::provider_option_names::kUserAuxStreamArray, MakeStringWithClassicLocale(reinterpret_cast(info.user_aux_stream_array))}, - {nv::provider_option_names::kMaxWorkspaceSize, MakeStringWithClassicLocale(info.max_workspace_size)}, - {nv::provider_option_names::kMaxSharedMemSize, MakeStringWithClassicLocale(info.max_shared_mem_size)}, - {nv::provider_option_names::kLengthAuxStreamArray, MakeStringWithClassicLocale(info.auxiliary_streams)}, - {nv::provider_option_names::kDumpSubgraphs, MakeStringWithClassicLocale(info.dump_subgraphs)}, - {nv::provider_option_names::kDetailedBuildLog, MakeStringWithClassicLocale(info.detailed_build_log)}, - {nv::provider_option_names::kProfilesMinShapes, MakeStringWithClassicLocale(info.profile_min_shapes)}, - {nv::provider_option_names::kProfilesMaxShapes, MakeStringWithClassicLocale(info.profile_max_shapes)}, - {nv::provider_option_names::kProfilesOptShapes, MakeStringWithClassicLocale(info.profile_opt_shapes)}, - {nv::provider_option_names::kCudaGraphEnable, MakeStringWithClassicLocale(info.cuda_graph_enable)}, - {nv::provider_option_names::kUseExternalDataInitializer, MakeStringWithClassicLocale(info.use_external_data_initializer)}, - {nv::provider_option_names::kRuntimeCacheFile, MakeStringWithClassicLocale(info.runtime_cache_path)}}; - return options; -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h deleted file mode 100644 index 779d3f6b67c1a..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include - -#include "core/framework/ortdevice.h" -#include "core/framework/provider_options.h" -#include "core/framework/framework_provider_common.h" -#include "core/framework/library_handles.h" -#include "core/session/onnxruntime_c_api.h" -#include "core/providers/shared_library/provider_api.h" - -#define TRT_DEFAULT_OPTIMIZER_LEVEL 3 - -namespace onnxruntime { -// Information needed to construct trt execution providers. -struct NvExecutionProviderInfo { - int device_id{0}; - bool has_user_compute_stream{false}; - void* user_compute_stream{nullptr}; - void* user_aux_stream_array{nullptr}; - int max_partition_iterations{1000}; - int min_subgraph_size{1}; - size_t max_workspace_size{0}; - size_t max_shared_mem_size{0}; - bool dump_subgraphs{false}; - std::string engine_cache_path{""}; - bool weight_stripped_engine_enable{false}; - std::string onnx_model_folder_path{""}; - const void* onnx_bytestream{nullptr}; - size_t onnx_bytestream_size{0}; - bool use_external_data_initializer{false}; - const void* external_data_bytestream{nullptr}; - size_t external_data_bytestream_size{0}; - bool engine_decryption_enable{false}; - std::string engine_decryption_lib_path{""}; - bool force_sequential_engine_build{false}; - std::string runtime_cache_path{""}; - bool detailed_build_log{false}; - bool sparsity_enable{false}; - int auxiliary_streams{-1}; - std::string extra_plugin_lib_paths{""}; - std::string profile_min_shapes{""}; - std::string profile_max_shapes{""}; - std::string profile_opt_shapes{""}; - bool cuda_graph_enable{true}; - bool multi_profile_enable{false}; - bool dump_ep_context_model{false}; - std::string ep_context_file_path{""}; - int ep_context_embed_mode{0}; - // Set internally by OrtCompileAPI::CompileModel(). When true, the EP skips GPU - // deserialization and execution context creation since the session will not run inference. - bool compile_only_mode{false}; - std::string engine_cache_prefix{""}; - std::string op_types_to_exclude{""}; - - static NvExecutionProviderInfo FromProviderOptions(const ProviderOptions& options, - const ConfigOptions& session_options); - static ProviderOptions ToProviderOptions(const NvExecutionProviderInfo& info); - std::vector custom_op_domain_list; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h deleted file mode 100644 index c57322e46bfe8..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h +++ /dev/null @@ -1,301 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include "flatbuffers/idl.h" -#include "nv_includes.h" -#include "core/providers/cuda/cuda_pch.h" -#include "core/common/path_string.h" -#include "core/framework/murmurhash3.h" -#include "core/providers/cuda/shared_inc/cuda_call.h" - -namespace fs = std::filesystem; - -namespace onnxruntime { - -/* - * Get number of profile setting. - * - * profile_min_shapes/profile_max_shapes/profile_opt_shapes may contain multiple profile settings. - * Note: TRT EP currently only supports one profile setting. - * - * { - * tensor_a: [[dim_0_value_0, dim_1_value_1, dim_2_value_2]], - * tensor_b: [[dim_0_value_3, dim_1_value_4, dim_2_value_5]] - * } - * - */ -static int GetNumProfiles(std::unordered_map>>& profile_shapes) { - int num_profile = 0; - for (auto it = profile_shapes.begin(); it != profile_shapes.end(); it++) { - num_profile = static_cast(it->second.size()); - if (num_profile > 0) { - break; - } - } - return num_profile; -} - -/* - * Get cache by name - * - */ -static std::string GetCachePath(const std::string& root, const std::string& name) { - if (root.empty()) { - return name; - } else { - fs::path path = root; - path.append(name); - return path.string(); - } -} - -/* - * Get compute capability - * - */ -static std::string GetComputeCapability(const cudaDeviceProp& prop) { - const std::string compute_capability = std::to_string(prop.major * 10 + prop.minor); - return compute_capability; -} - -/** - * - * Helper class to generate engine id via model name/model content/env metadata - * - * - * The TensorRT Execution Provider is used in multiple sessions and the underlying infrastructure caches - * compiled kernels, so the name must be unique and deterministic across models and sessions. - * - */ -static HashValue TRTGenerateId(const GraphViewer& graph_viewer, std::string trt_version, std::string cuda_version) { - HashValue model_hash = 0; - - // find the top level graph - const Graph* cur_graph = &graph_viewer.GetGraph(); - while (cur_graph->IsSubgraph()) { - cur_graph = cur_graph->ParentGraph(); - } - - const Graph& main_graph = *cur_graph; - uint32_t hash[4] = {0, 0, 0, 0}; - - auto hash_str = [&hash](const std::string& str) { - MurmurHash3::x86_128(str.data(), gsl::narrow_cast(str.size()), hash[0], &hash); - }; - - // Use the model's file name instead of the entire path to avoid cache regeneration if path changes - if (main_graph.ModelPath().has_filename()) { - std::string model_name = PathToUTF8String(main_graph.ModelPath().filename()); - - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Model name is " << model_name; - // Ensure enough characters are hashed in case model names are too short - const size_t model_name_length = model_name.size(); - constexpr size_t hash_string_length = 500; - std::string repeat_model_name = model_name; - for (size_t i = model_name_length; i > 0 && i < hash_string_length; i += model_name_length) { - repeat_model_name += model_name; - } - hash_str(repeat_model_name); - } else { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Model path is empty"; - } - - // fingerprint current graph by hashing graph inputs - for (const auto* node_arg : graph_viewer.GetInputsIncludingInitializers()) { - hash_str(node_arg->Name()); - } - - // hashing output of each node - const int number_of_ort_nodes = graph_viewer.NumberOfNodes(); - std::vector nodes_vector(number_of_ort_nodes); - std::iota(std::begin(nodes_vector), std::end(nodes_vector), 0); - const std::vector& node_index = graph_viewer.GetNodesInTopologicalOrder(); - for (const auto& index : nodes_vector) { - const auto& node = graph_viewer.GetNode(node_index[index]); - for (const auto* node_arg : node->OutputDefs()) { - if (node_arg->Exists()) { - hash_str(node_arg->Name()); - } - } - } - -#ifdef __linux__ - hash_str("LINUX"); -#elif defined(_WIN32) - hash_str("WINDOWS"); -#endif - -#ifdef ORT_VERSION - hash_str(ORT_VERSION); -#endif - -#ifdef CUDA_VERSION - hash_str(cuda_version); -#endif - -#if defined(NV_TENSORRT_MAJOR) && defined(NV_TENSORRT_MINOR) - hash_str(trt_version); -#endif - - model_hash = hash[0] | (uint64_t(hash[1]) << 32); - - // return the current unique id - return model_hash; -} - -static bool ValidateProfileShapes(std::unordered_map>>& profile_min_shapes, - std::unordered_map>>& profile_max_shapes, - std::unordered_map>>& profile_opt_shapes) { - if (profile_min_shapes.empty() && profile_max_shapes.empty() && profile_opt_shapes.empty()) { - return true; - } - - if ((profile_min_shapes.size() != profile_max_shapes.size()) && - (profile_min_shapes.size() != profile_opt_shapes.size()) && - (profile_max_shapes.size() != profile_opt_shapes.size())) { - return false; - } - - std::unordered_map>>::iterator it; - for (it = profile_min_shapes.begin(); it != profile_min_shapes.end(); it++) { - auto input_name = it->first; - auto num_profile = it->second.size(); - - // input_name must also be in max/opt profile - if ((profile_max_shapes.find(input_name) == profile_max_shapes.end()) || - (profile_opt_shapes.find(input_name) == profile_opt_shapes.end())) { - return false; - } - - // number of profiles should be the same - if ((num_profile != profile_max_shapes[input_name].size()) || - (num_profile != profile_opt_shapes[input_name].size())) { - return false; - } - } - - return true; -} - -/* - * Make input-name and shape as a pair. - * This helper function is being used by ParseProfileShapes(). - * - * For example: - * The input string is "input_id:32x1", - * after the string is being parsed, the pair object is returned as below. - * pair("input_id", [32, 1]) - * - * Return true if string can be successfully parsed or false if string has wrong format. - */ -static bool MakeInputNameShapePair(std::string pair_string, std::pair>& pair) { - if (pair_string.empty()) { - return true; - } - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] " << pair_string; - - std::stringstream input_string_stream(pair_string); - char first_delim = ':'; - char second_delim = 'x'; - std::string input_name; - std::string shape; - std::getline(input_string_stream, input_name, first_delim); - std::getline(input_string_stream, shape, first_delim); - - std::vector shapes; - std::stringstream shape_string_stream(shape); - std::string value; - while (std::getline(shape_string_stream, value, second_delim)) { - shapes.push_back(std::stoi(value)); - } - - // wrong input string - if (input_name.empty() || shapes.empty()) { - return false; - } - - pair.first = input_name; - pair.second = shapes; - - return true; -} - -/* - * Parse explicit profile min/max/opt shapes from Nv EP provider options. - * - * For example: - * The provider option is --trt_profile_min_shapes="input_id:32x1,attention_mask:32x1,input_id:32x41,attention_mask:32x41", - * after string is being parsed, the profile shapes has two profiles and is being represented as below. - * {"input_id": [[32, 1], [32, 41]], "attention_mask": [[32, 1], [32, 41]]} - * - * Return true if string can be successfully parsed or false if string has wrong format. - */ -static bool ParseProfileShapes(std::string profile_shapes_string, std::unordered_map>>& profile_shapes) { - if (profile_shapes_string.empty()) { - return true; - } - - std::stringstream input_string_stream(profile_shapes_string); - char delim = ','; - std::string input_name_with_shape; // input_name:shape, ex: "input_id:32x1" - while (std::getline(input_string_stream, input_name_with_shape, delim)) { - std::pair> pair; - if (!MakeInputNameShapePair(input_name_with_shape, pair)) { - return false; - } - - std::string input_name = pair.first; - if (profile_shapes.find(input_name) == profile_shapes.end()) { - std::vector> profile_shape_vector; - profile_shapes[input_name] = profile_shape_vector; - } - profile_shapes[input_name].push_back(pair.second); - - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] " << input_name; - std::string shape_string = ""; - for (auto v : pair.second) { - shape_string += std::to_string(v); - shape_string += ", "; - } - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] " << shape_string; - } - - return true; -} - -/* - * Checks if there is a an element with value `-1` in nvinfer1::Dims - */ -static bool checkTrtDimIsDynamic(nvinfer1::Dims dims) { - for (int j = 0, end = dims.nbDims; j < end; ++j) { - if (dims.d[j] == -1) { - return true; - } - } - return false; -} - -/* - * Checks if an nvinfer1::ITensor signales a dynamic shape, - * either due to dynamic shapes or due to it being a shape tensor - */ -static bool checkTrtTensorIsDynamic(nvinfer1::ITensor* tensor) { - if (tensor->isShapeTensor()) { - return true; - } else { - // Execution tensor - return checkTrtDimIsDynamic(tensor->getDimensions()); - } -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_file_utils.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_file_utils.h deleted file mode 100644 index 159aba0507ffb..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_file_utils.h +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once -#include -#include -#include -#include -#include -#include "core/providers/shared_library/provider_api.h" - -namespace onnxruntime { -namespace file_utils { - -inline std::vector ReadFile(const std::string& path) { - if (!std::filesystem::exists(path)) { - LOGS_DEFAULT(INFO) << "TensorRT RTX could not find the file and will create a new one " << path << std::endl; - return {}; - } - std::ifstream file(path, std::ios::in | std::ios::binary); - if (!file) { - ORT_THROW("Failed to open file: " + path); - } - file.seekg(0, std::ios::end); - std::streamsize size = file.tellg(); - file.seekg(0, std::ios::beg); - std::vector buffer(size); - if (size > 0 && !file.read(buffer.data(), size)) { - ORT_THROW("Failed to read file: " + path); - } - return buffer; -} - -inline void WriteFile(const std::string& path, const void* data, size_t size) { - if (std::filesystem::exists(path)) { - std::ofstream file(path, std::ios::out | std::ios::binary | std::ios::trunc); - if (!file) { - ORT_THROW("Failed to open file for writing: " + path); - } - file.write(static_cast(data), size); - } else { - LOGS_DEFAULT(INFO) << "TensorRT RTX a new file cache was written to " << path << std::endl; - // Create new file - std::ofstream file(path, std::ios::out | std::ios::binary); - if (!file) { - ORT_THROW("Failed to create file: " + path); - } - file.write(static_cast(data), size); - } -} - -inline void WriteFile(const std::string& path, const std::vector& data) { WriteFile(path, data.data(), data.size()); } - -} // namespace file_utils -} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_includes.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_includes.h deleted file mode 100644 index a4e3777008560..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_includes.h +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#pragma once - -// File to include the required TRT headers with workarounds for warnings we can't fix or not fixed yet. -#if defined(_MSC_VER) -#pragma warning(push) -#pragma warning(disable : 4100) // Ignore warning C4100: unreferenced formal parameter -#pragma warning(disable : 4996) // Ignore warning C4996: 'nvinfer1::IPluginV2' was declared deprecated -#endif - -#include -#include -#include -#include - -#if defined(_MSC_VER) -#pragma warning(pop) -#endif diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_platform_utils.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_platform_utils.h deleted file mode 100644 index f3298a8449157..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_platform_utils.h +++ /dev/null @@ -1,67 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#pragma once - -#include -#include "core/common/path_string.h" - -#ifdef _WIN32 -#include -#else -#include -#endif - -// The filename extension for a shared library is different per platform -#ifdef _WIN32 -#define LIBRARY_PREFIX -#define LIBRARY_EXTENSION ORT_TSTR(".dll") -#elif defined(__APPLE__) -#define LIBRARY_PREFIX "lib" -#define LIBRARY_EXTENSION ".dylib" -#else -#define LIBRARY_PREFIX "lib" -#define LIBRARY_EXTENSION ".so" -#endif - -namespace onnxruntime { -inline PathString GetEPLibraryDirectory() { -#ifdef _WIN32 - HMODULE hModule = NULL; - // Get handle to the DLL executing this code - if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | - GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, - reinterpret_cast(GetEPLibraryDirectory), - &hModule)) { - return PathString(); - } - - wchar_t buffer[MAX_PATH]; - DWORD len = GetModuleFileNameW(hModule, buffer, MAX_PATH); - if (len == 0 || len >= MAX_PATH) { - return PathString(); - } - - std::wstring path(buffer); - size_t lastSlash = path.find_last_of(L"\\/"); - if (lastSlash != std::wstring::npos) { - return PathString(path.substr(0, lastSlash + 1)); - } - return PathString(); -#else - // Linux and other Unix-like platforms - Dl_info dl_info; - - if (dladdr((void*)&GetEPLibraryDirectory, &dl_info) == 0 || dl_info.dli_fname == nullptr) { - return PathString(); - } - - std::string so_path(dl_info.dli_fname); - size_t last_slash = so_path.find_last_of('/'); - if (last_slash != std::string::npos) { - return PathString(so_path.substr(0, last_slash + 1)); - } - return PathString(); -#endif -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc deleted file mode 100644 index af979ce89d7d8..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc +++ /dev/null @@ -1,2001 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "core/providers/shared_library/provider_api.h" -#include "core/session/onnxruntime_c_api.h" -#include "core/framework/provider_options.h" -#include "core/framework/plugin_ep_stream.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_options.h" -#include "core/providers/nv_tensorrt_rtx/nv_execution_provider_custom_ops.h" -#include "core/providers/cuda/cuda_stream_handle.h" - -// D3D12 headers for graphics interop on Windows -#if defined(_WIN32) && USE_DX_INTEROP -#include -#endif -// NVML for driver version checking (optional; see HAVE_NVML in CMake) -#if defined(HAVE_NVML) -#include -#endif - -#include "onnx_ctx_model_helper.h" -#include "nv_provider_factory.h" -#include "nv_execution_provider.h" -#include "nv_provider_factory_creator.h" -#include "nv_data_transfer.h" -#include "nv_allocator.h" -#include "nv_scoped_context.h" - -using namespace onnxruntime; - -// External declarations -namespace onnxruntime { -extern TensorrtLogger& GetTensorrtLogger(bool verbose_log); -} - -namespace onnxruntime { - -void InitializeRegistry(); -void DeleteRegistry(); - -struct ProviderInfo_Nv_Impl final : ProviderInfo_Nv { - OrtStatus* GetCurrentGpuDeviceId(_In_ int* device_id) override { - auto cuda_err = cudaGetDevice(device_id); - if (cuda_err != cudaSuccess) { - return CreateStatus(ORT_FAIL, "Failed to get device id."); - } - return nullptr; - } - - OrtStatus* GetTensorRTCustomOpDomainList(std::vector& domain_list, const std::string extra_plugin_lib_paths) override { - common::Status status = CreateTensorRTCustomOpDomainList(domain_list, extra_plugin_lib_paths); - if (!status.IsOK()) { - return CreateStatus(ORT_FAIL, "[NvTensorRTRTX EP] Can't create custom ops for TRT plugins."); - } - return nullptr; - } - - OrtStatus* ReleaseCustomOpDomainList(std::vector& domain_list) override { - ReleaseTensorRTCustomOpDomainList(domain_list); - return nullptr; - } -} g_info; - -struct NvProviderFactory : IExecutionProviderFactory { - NvProviderFactory(const NvExecutionProviderInfo& info) : info_{info} {} - ~NvProviderFactory() override {} - - std::unique_ptr CreateProvider() override; - std::unique_ptr CreateProvider(const OrtSessionOptions& session_options, - const OrtLogger& session_logger) override; - - private: - NvExecutionProviderInfo info_; -}; - -std::unique_ptr NvProviderFactory::CreateProvider() { - return std::make_unique(info_); -} - -std::unique_ptr NvProviderFactory::CreateProvider(const OrtSessionOptions& session_options, const OrtLogger& session_logger) { - const ConfigOptions& config_options = session_options.GetConfigOptions(); - const std::unordered_map& config_options_map = config_options.GetConfigOptionsMap(); - - // The implementation of the SessionOptionsAppendExecutionProvider C API function automatically adds EP options to - // the session option configurations with the key prefix "ep..". - // We extract those EP options to create a new "provider options" key/value map. - std::string lowercase_ep_name = kNvTensorRTRTXExecutionProvider; - std::transform(lowercase_ep_name.begin(), lowercase_ep_name.end(), lowercase_ep_name.begin(), [](unsigned char c) { - return static_cast(std::tolower(c)); - }); - - ProviderOptions provider_options; - std::string key_prefix = "ep."; - key_prefix += lowercase_ep_name; - key_prefix += "."; - - for (const auto& [key, value] : config_options_map) { - if (key.rfind(key_prefix, 0) == 0) { - provider_options[key.substr(key_prefix.size())] = value; - } - } - NvExecutionProviderInfo info = onnxruntime::NvExecutionProviderInfo::FromProviderOptions(provider_options, config_options); - - auto ep = std::make_unique(info); - ep->SetLogger(reinterpret_cast(&session_logger)); - return ep; -} - -struct Nv_Provider : Provider { - void* GetInfo() override { return &g_info; } - std::shared_ptr CreateExecutionProviderFactory(int device_id) override { - NvExecutionProviderInfo info; - info.device_id = device_id; - - return std::make_shared(info); - } - - std::shared_ptr CreateExecutionProviderFactory(const void* param) override { - if (param == nullptr) { - LOGS_DEFAULT(ERROR) << "[NvTensorRTRTX EP] Passed NULL options to CreateExecutionProviderFactory()"; - return nullptr; - } - - std::array pointers_array = *reinterpret_cast*>(param); - const ProviderOptions* provider_options = reinterpret_cast(pointers_array[0]); - const ConfigOptions* config_options = reinterpret_cast(pointers_array[1]); - - if (provider_options == nullptr) { - LOGS_DEFAULT(ERROR) << "[NvTensorRTRTX EP] Passed NULL ProviderOptions to CreateExecutionProviderFactory()"; - return nullptr; - } - - NvExecutionProviderInfo info = onnxruntime::NvExecutionProviderInfo::FromProviderOptions(*provider_options, *config_options); - return std::make_shared(info); - } - - Status CreateIExecutionProvider(const OrtHardwareDevice* const* /*devices*/, - const OrtKeyValuePairs* const* /*ep_metadata*/, - size_t /*num_devices*/, - ProviderOptions& provider_options, - const OrtSessionOptions& session_options, - const OrtLogger& logger, - std::unique_ptr& ep) override { - const ConfigOptions* config_options = &session_options.GetConfigOptions(); - - std::array configs_array = {&provider_options, config_options}; - const void* arg = reinterpret_cast(&configs_array); - auto ep_factory = CreateExecutionProviderFactory(arg); - ep = ep_factory->CreateProvider(session_options, logger); - - return Status::OK(); - } - - void Initialize() override { - InitializeRegistry(); - } - - void Shutdown() override { - DeleteRegistry(); - // ensure that there is no longer an active context which will be destroyed when unloading cudart - CUcontext cu_context = 0; - CU_CALL_THROW(cuCtxGetCurrent(&cu_context)); - if (cu_context) { - CU_CALL_THROW(cuCtxPopCurrent(&cu_context)); - } - } - -} g_provider; - -} // namespace onnxruntime - -extern "C" { - -ORT_API(onnxruntime::Provider*, GetProvider) { - return &onnxruntime::g_provider; -} -} - -// -// Plug-in EP infrastructure -// - -#include "core/session/abi_devices.h" -#include "onnxruntime_config.h" // for ORT_VERSION - -struct ErrorHelper { - static const OrtApi* ort_api; - - static OrtStatus* ToOrtStatus(const Status& status) { - if (status.IsOK()) { - return nullptr; // no error - } - - return ort_api->CreateStatus(static_cast(status.Code()), - status.ErrorMessage().c_str()); - } -}; - -const OrtApi* ErrorHelper::ort_api = nullptr; - -#define RETURN_IF_ERROR(fn) \ - do { \ - OrtStatus* _status = (fn); \ - if (_status != nullptr) { \ - return _status; \ - } \ - } while (0) - -#define RETURN_IF_STATUS_NOTOK(fn) \ - do { \ - Status _status = (fn); \ - if (!_status.IsOK()) { \ - return ErrorHelper::ToOrtStatus(_status); \ - } \ - } while (0) - -#define CUDA_RETURN_IF_ERROR(expr) RETURN_IF_STATUS_NOTOK(CUDA_CALL(expr)) - -struct NvTrtRtxOrtAllocator : OrtAllocator { - NvTrtRtxOrtAllocator(const OrtMemoryInfo* mem_info, const OrtApi& api) : memory_info_{mem_info} { - version = ORT_API_VERSION; - Alloc = AllocImpl; - AllocOnStream = AllocOnStreamImpl; - Free = FreeImpl; - Info = InfoImpl; - Reserve = AllocImpl; // no special behavior for Reserve so use AllocImpl - GetStats = nullptr; // GetStatsImpl. The CUDA allocators don't have stats currently so we can skip. - Shrink = nullptr; - - const OrtEpApi& ep_api = *api.GetEpApi(); - const OrtMemoryDevice* mem_device = ep_api.MemoryInfo_GetMemoryDevice(mem_info); - uint32_t device_id = ep_api.MemoryDevice_GetDeviceId(mem_device); - const char* name = nullptr; - auto* status = api.MemoryInfoGetName(mem_info, &name); - static_cast(status); // GetName never fails - - if (ep_api.MemoryDevice_GetMemoryType(mem_device) == OrtDeviceMemoryType_HOST_ACCESSIBLE) { - allocator_ = std::make_unique(device_id, name); - } else { - allocator_ = std::make_unique(device_id, name); - } - } - - static void* ORT_API_CALL AllocImpl(struct OrtAllocator* this_, size_t size) { - auto& impl = *static_cast(this_); - return impl.allocator_->Alloc(size); - } - - static void* ORT_API_CALL AllocOnStreamImpl(struct OrtAllocator* this_, size_t size, OrtSyncStream* stream) { - auto& impl = *static_cast(this_); - return impl.allocator_->AllocOnStream(size, stream); - } - - static void ORT_API_CALL FreeImpl(struct OrtAllocator* this_, void* p) { - auto& impl = *static_cast(this_); - impl.allocator_->Free(p); - } - - static const struct OrtMemoryInfo* ORT_API_CALL InfoImpl(const struct OrtAllocator* this_) { - const NvTrtRtxOrtAllocator& impl = *static_cast(this_); - return impl.memory_info_; - } - - private: - const OrtMemoryInfo* memory_info_; - std::unique_ptr allocator_; -}; - -struct NvTrtRtxDataTransferImpl : OrtDataTransferImpl { - NvTrtRtxDataTransferImpl(const OrtApi& ort_api_in) - : ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()} { - ort_version_supported = ORT_API_VERSION; - CanCopy = CanCopyImpl; - CopyTensors = CopyTensorsImpl; - Release = ReleaseImpl; - } - - static bool CanCopyImpl(const OrtDataTransferImpl* this_ptr, - const OrtMemoryDevice* src_memory_device, - const OrtMemoryDevice* dst_memory_device) noexcept { - const auto& impl = *static_cast(this_ptr); - - // logic copied from GPUDataTransfer::CanCopy - OrtMemoryInfoDeviceType src_type = impl.ep_api.MemoryDevice_GetDeviceType(src_memory_device); - OrtMemoryInfoDeviceType dst_type = impl.ep_api.MemoryDevice_GetDeviceType(dst_memory_device); - auto src_vendor_id = impl.ep_api.MemoryDevice_GetVendorId(src_memory_device); - auto dst_vendor_id = impl.ep_api.MemoryDevice_GetVendorId(dst_memory_device); - - if ((src_type == OrtDevice::GPU && src_vendor_id != OrtDevice::VendorIds::NVIDIA) || - (dst_type == OrtDevice::GPU && dst_vendor_id != OrtDevice::VendorIds::NVIDIA)) { - return false; - } - - // copy must be GPU to GPU or between GPU and CPU - return (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_GPU) || - (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_CPU) || - (src_type == OrtMemoryInfoDeviceType_CPU && dst_type == OrtMemoryInfoDeviceType_GPU); - } - - static OrtStatus* CopyTensorsImpl(OrtDataTransferImpl* this_ptr, - const OrtValue** src_tensors, - OrtValue** dst_tensors, - OrtSyncStream** streams, - size_t num_tensors) noexcept { - auto& impl = *static_cast(this_ptr); - bool need_stream_sync = false; - - for (size_t idx = 0; idx < num_tensors; ++idx) { - const OrtValue* src_tensor = src_tensors[idx]; - OrtValue* dst_tensor = dst_tensors[idx]; - OrtSyncStream* stream = streams ? streams[idx] : nullptr; - - const OrtMemoryDevice* src_device = impl.ep_api.Value_GetMemoryDevice(src_tensor); - const OrtMemoryDevice* dst_device = impl.ep_api.Value_GetMemoryDevice(dst_tensor); - - size_t bytes; - RETURN_IF_ERROR(impl.ort_api.GetTensorSizeInBytes(src_tensor, &bytes)); - - const void* src_data = nullptr; - void* dst_data = nullptr; - RETURN_IF_ERROR(impl.ort_api.GetTensorData(src_tensor, &src_data)); - RETURN_IF_ERROR(impl.ort_api.GetTensorMutableData(dst_tensor, &dst_data)); - - OrtMemoryInfoDeviceType src_type = impl.ep_api.MemoryDevice_GetDeviceType(src_device); - OrtMemoryInfoDeviceType dst_type = impl.ep_api.MemoryDevice_GetDeviceType(dst_device); - OrtDeviceMemoryType src_mem_type = impl.ep_api.MemoryDevice_GetMemoryType(src_device); - OrtDeviceMemoryType dst_mem_type = impl.ep_api.MemoryDevice_GetMemoryType(dst_device); - - const bool src_is_gpu_default = src_type == OrtMemoryInfoDeviceType_GPU && - src_mem_type == OrtDeviceMemoryType_DEFAULT; - const bool dst_is_gpu_default = dst_type == OrtMemoryInfoDeviceType_GPU && - dst_mem_type == OrtDeviceMemoryType_DEFAULT; - - cudaStream_t cuda_stream = nullptr; - if (stream) { - cuda_stream = static_cast(impl.ort_api.SyncStream_GetHandle(stream)); - } - - if (dst_is_gpu_default) { - if (src_is_gpu_default) { - // Copy only if the two addresses are different. - if (dst_data != src_data) { - if (cuda_stream) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_data, src_data, bytes, cudaMemcpyDeviceToDevice, cuda_stream)); - - } else { - CUDA_RETURN_IF_ERROR(cudaMemcpy(dst_data, src_data, bytes, cudaMemcpyDeviceToDevice)); - - // For device memory to device memory copy, no host-side synchronization is performed by cudaMemcpy. - // see https://docs.nvidia.com/cuda/cuda-runtime-api/api-sync-behavior.html - need_stream_sync = true; - } - } - } else { - // copy from pinned or non-pinned CPU memory to GPU - if (cuda_stream) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_data, src_data, bytes, cudaMemcpyHostToDevice, cuda_stream)); - } else { - CUDA_RETURN_IF_ERROR(cudaMemcpy(dst_data, src_data, bytes, cudaMemcpyHostToDevice)); - - if (src_mem_type != OrtDeviceMemoryType_HOST_ACCESSIBLE) { - // For cudaMemcpy from pageable host memory to device memory, DMA to final destination may not - // have completed. - // see https://docs.nvidia.com/cuda/cuda-runtime-api/api-sync-behavior.html - need_stream_sync = true; - } - } - } - } else if (src_is_gpu_default) { - // copying from GPU to CPU memory, this is blocking - - if (cuda_stream) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dst_data, src_data, bytes, cudaMemcpyDeviceToHost, cuda_stream)); - - } else { - CUDA_RETURN_IF_ERROR(cudaMemcpy(dst_data, src_data, bytes, cudaMemcpyDeviceToHost)); - } - } else { - // copying between CPU accessible memory - - if (dst_data != src_data) { - if (cuda_stream) { - if (src_mem_type == OrtDeviceMemoryType_HOST_ACCESSIBLE) { - // sync the stream first to make sure the data arrived - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream)); - } - } - - memcpy(dst_data, src_data, bytes); - } - } - } - - if (need_stream_sync) { - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(nullptr)); - } - - return nullptr; - } - - static void ReleaseImpl(OrtDataTransferImpl* /*this_ptr*/) noexcept { - // no-op as we have a single shared instance in OrtEpFactory which is returned from CreateDataTransferImpl, and is - // owned by and freed by the factory. - } - - const OrtApi& ort_api; - const OrtEpApi& ep_api; -}; - -struct NvTrtRtxSyncNotificationImpl : OrtSyncNotificationImpl { - static OrtStatus* Create(cudaStream_t stream, const OrtApi& ort_api, - std::unique_ptr& notification) { - notification.reset(new NvTrtRtxSyncNotificationImpl(stream, ort_api)); // can't use make_unique with private ctor - CUDA_RETURN_IF_ERROR(cudaEventCreateWithFlags(¬ification->event_, cudaEventDisableTiming)); - - return nullptr; - } - - static void ReleaseImpl(_In_ OrtSyncNotificationImpl* this_ptr) noexcept { - delete static_cast(this_ptr); - } - - static OrtStatus* ActivateImpl(_In_ OrtSyncNotificationImpl* this_ptr) noexcept { - auto& impl = *static_cast(this_ptr); - CUDA_RETURN_IF_ERROR(cudaEventRecord(impl.event_, impl.stream_)); - - return nullptr; - } - - static OrtStatus* WaitOnDeviceImpl(_In_ OrtSyncNotificationImpl* this_ptr, - _In_ OrtSyncStream* consumer_stream) noexcept { - auto& impl = *static_cast(this_ptr); - - // setup the consumer stream to wait on our event. - void* consumer_handle = impl.ort_api.SyncStream_GetHandle(consumer_stream); - CUDA_RETURN_IF_ERROR(cudaStreamWaitEvent(static_cast(consumer_handle), impl.event_)); - - return nullptr; - } - - static OrtStatus* WaitOnHostImpl(_In_ OrtSyncNotificationImpl* this_ptr) noexcept { - auto& impl = *static_cast(this_ptr); - CUDA_RETURN_IF_ERROR(cudaEventSynchronize(impl.event_)); - - return nullptr; - } - - ~NvTrtRtxSyncNotificationImpl() { - cudaEventDestroy(event_); - } - - private: - NvTrtRtxSyncNotificationImpl(cudaStream_t stream, const OrtApi& ort_api_in) - : stream_{stream}, ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()} { - ort_version_supported = ORT_API_VERSION; - Activate = ActivateImpl; - WaitOnDevice = WaitOnDeviceImpl; - WaitOnHost = WaitOnHostImpl; - Release = ReleaseImpl; - } - - cudaStream_t stream_; - cudaEvent_t event_; - - const OrtApi& ort_api; - const OrtEpApi& ep_api; -}; - -struct NvTrtRtxSyncStreamImpl : OrtSyncStreamImpl { - NvTrtRtxSyncStreamImpl(cudaStream_t&& stream, - const OrtDevice& device, - AllocatorPtr cpu_allocator, - bool release_cpu_buffer_on_cuda_stream, - const OrtApi& ort_api_in) - : stream_{ - stream, device, cpu_allocator, release_cpu_buffer_on_cuda_stream, /*own*/ true, - /*external_cudnn_handle*/ nullptr, - /*external_cublas_handle*/ nullptr, - // ep_info is used by GetResource which seems to be a somewhat ugly way to make arbitrary info that is - // unrelated to the stream available to a custom op. - // avoiding adding GetResource to OrtSyncStreamImpl as we should have a cleaner setup for custom ops, - // so this argument value isn't used and doesn't matter. - /*ep_info*/ CUDAExecutionProviderInfo{}}, - ort_api{ort_api_in} { - ort_version_supported = ORT_API_VERSION; - GetHandle = GetHandleImpl; - CreateNotification = CreateNotificationImpl; - Flush = FlushImpl; - OnSessionRunEnd = OnSessionRunEndImpl; - Release = ReleaseImpl; - } - - static void ReleaseImpl(_In_ OrtSyncStreamImpl* this_ptr) noexcept { - delete static_cast(this_ptr); - } - - static void* GetHandleImpl(_In_ OrtSyncStreamImpl* this_ptr) noexcept { - auto& impl = *static_cast(this_ptr); - return impl.stream_.GetHandle(); - } - - static OrtStatus* CreateNotificationImpl(_In_ OrtSyncStreamImpl* this_ptr, - _Outptr_ OrtSyncNotificationImpl** notification_impl) noexcept { - auto& impl = *static_cast(this_ptr); - *notification_impl = nullptr; - - std::unique_ptr notification; - cudaStream_t cuda_stream = static_cast(impl.stream_.GetHandle()); - - RETURN_IF_ERROR(NvTrtRtxSyncNotificationImpl::Create(cuda_stream, impl.ort_api, notification)); - *notification_impl = notification.release(); - - return nullptr; - } - - static OrtStatus* FlushImpl(_In_ OrtSyncStreamImpl* this_ptr) noexcept { - auto& impl = *static_cast(this_ptr); - impl.stream_.Flush(); - - return nullptr; - } - - static OrtStatus* OnSessionRunEndImpl(_In_ OrtSyncStreamImpl* this_ptr) noexcept { - auto& impl = *static_cast(this_ptr); - RETURN_IF_STATUS_NOTOK(impl.stream_.CleanUpOnRunEnd()); - - return nullptr; - } - - private: - // this is a little onion-ish as CudaStream is a onnxruntime::Stream and this is an OrtSyncStreamImpl that will be - // used via plugin_ep::Stream, which is also an onnxruntime::Stream. in a 'real' plugin EP implementation - // CudaStream would go away and the logic it has would be implemented directly here. - CudaStream stream_; - const OrtApi& ort_api; -}; - -// External Resource Import Implementation (D3D12 to CUDA) -/** - * @brief Derived handle for imported external memory from D3D12 to CUDA. - * - * Derives from OrtExternalMemoryHandle (base struct) and adds CUDA-specific fields. - * This struct holds the CUDA external memory object and the mapped device pointer - * that can be used for zero-copy tensor creation. - */ -struct NvTrtRtxExternalMemoryHandle : OrtExternalMemoryHandle { - CUexternalMemory ext_memory; ///< CUDA external memory object - CUdeviceptr mapped_ptr; ///< Mapped device pointer for tensor access - bool is_dedicated; ///< Whether the D3D12 resource is a dedicated allocation - - NvTrtRtxExternalMemoryHandle(const OrtExternalMemoryDescriptor& descriptor_in) - : ext_memory(nullptr), mapped_ptr(0), is_dedicated(true) { - // Initialize base struct fields - version = ORT_API_VERSION; - descriptor = descriptor_in; - ep_device = nullptr; - Release = ReleaseCallback; - } - - static void ORT_API_CALL ReleaseCallback(_In_ OrtExternalMemoryHandle* handle) noexcept { - if (handle == nullptr) return; - auto derived = std::unique_ptr( - static_cast(handle)); - // Destroy the external memory object (also releases mapped buffer) - if (derived->ext_memory != nullptr) { - cuDestroyExternalMemory(derived->ext_memory); - } - } -}; - -/** - * @brief Derived handle for imported external semaphore from D3D12 fence to CUDA. - * - * Derives from OrtExternalSemaphoreHandle (base struct) and adds CUDA-specific fields. - * D3D12 timeline fences are imported as CUDA external semaphores, enabling - * GPU-GPU synchronization between D3D12 and CUDA streams. - */ -struct NvTrtRtxExternalSemaphoreHandle : OrtExternalSemaphoreHandle { - CUexternalSemaphore ext_semaphore; ///< CUDA external semaphore object - - NvTrtRtxExternalSemaphoreHandle(const OrtExternalSemaphoreDescriptor& descriptor_in) - : ext_semaphore(nullptr) { - // Initialize base struct fields - version = ORT_API_VERSION; - descriptor = descriptor_in; - ep_device = nullptr; - Release = ReleaseCallback; - } - - static void ORT_API_CALL ReleaseCallback(_In_ OrtExternalSemaphoreHandle* handle) noexcept { - if (handle == nullptr) return; - auto derived = std::unique_ptr( - static_cast(handle)); - // Destroy the external semaphore object - if (derived->ext_semaphore != nullptr) { - cuDestroyExternalSemaphore(derived->ext_semaphore); - } - } -}; - -/** - * @brief Implementation of OrtExternalResourceImporterImpl for NvTensorRtRtx EP. - * - * This struct implements the external resource importer interface using CUDA Driver APIs - * to import D3D12 shared resources and timeline fences for zero-copy import. - * - * Supported handle types: - * - ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE → CU_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE - * - ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_HEAP → CU_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_HEAP - * - ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE → CU_EXTERNAL_SEMAPHORE_HANDLE_TYPE_D3D12_FENCE - */ -struct NvTrtRtxExternalResourceImporterImpl : OrtExternalResourceImporterImpl { - NvTrtRtxExternalResourceImporterImpl(const OrtEpDevice* ep_device, const OrtApi& ort_api_in) - : ep_device_{ep_device}, ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()} { - ort_version_supported = ORT_API_VERSION; - - // Memory operations - CanImportMemory = CanImportMemoryImpl; - ImportMemory = ImportMemoryImpl; - ReleaseMemory = ReleaseMemoryImpl; - CreateTensorFromMemory = CreateTensorFromMemoryImpl; - - // Semaphore operations - CanImportSemaphore = CanImportSemaphoreImpl; - ImportSemaphore = ImportSemaphoreImpl; - ReleaseSemaphore = ReleaseSemaphoreImpl; - WaitSemaphore = WaitSemaphoreImpl; - SignalSemaphore = SignalSemaphoreImpl; - - // Release - Release = ReleaseImpl; - } - - static bool ORT_API_CALL CanImportMemoryImpl( - _In_ const OrtExternalResourceImporterImpl* this_ptr, - _In_ OrtExternalMemoryHandleType handle_type) noexcept { - (void)this_ptr; - // CUDA supports both D3D12 resource and heap handles -#if defined(_WIN32) - return handle_type == ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE || - handle_type == ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_HEAP || - handle_type == ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_WIN32; -#elif __linux__ - return handle_type == ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_OPAQUE_FD; -#endif - } - - static OrtStatus* ORT_API_CALL ImportMemoryImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ const OrtExternalMemoryDescriptor* desc, - _Outptr_ OrtExternalMemoryHandle** out_handle) noexcept { - auto& impl = *static_cast(this_ptr); - - if (desc == nullptr || out_handle == nullptr) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "Invalid arguments to ImportMemory"); - } - - // Validate descriptor version - check minimum supported version for forward compatibility - if (desc->version < ORT_API_VERSION) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "OrtExternalMemoryDescriptor version too old"); - } - - *out_handle = nullptr; - - // Validate handle type - if (!CanImportMemoryImpl(this_ptr, desc->handle_type)) { - return impl.ort_api.CreateStatus(ORT_NOT_IMPLEMENTED, - "Unsupported external memory handle type for CUDA import"); - } - - // Validate offset does not exceed allocation size - if (desc->offset_bytes > desc->size_bytes) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "offset_bytes exceeds size_bytes in OrtExternalMemoryDescriptor"); - } - - // Set CUDA device for this EP. The imported external memory handle is associated with - // the device where it was imported and remains valid regardless of subsequent cudaSetDevice - // calls. Multi-GPU scenarios with different sessions/EPs work correctly because each - // importer is bound to its EP's device via ep_device_->device_memory_info. - ScopedContext ctx(impl.DeviceId()); - - // Map ORT handle type to CUDA handle type - CUexternalMemoryHandleType cu_handle_type; - bool is_dedicated = true; - switch (desc->handle_type) { - case ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE: - cu_handle_type = CU_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE; - is_dedicated = true; // D3D12 committed resources are dedicated - break; - case ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_HEAP: - cu_handle_type = CU_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_HEAP; - is_dedicated = false; // D3D12 heaps are not dedicated - break; - case ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_WIN32: - cu_handle_type = CU_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32; - is_dedicated = false; // API header currently, documents that this handle currently is non-dedicated - break; - case ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_OPAQUE_FD: - cu_handle_type = CU_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD; - is_dedicated = false; // API header currently, documents that this handle currently is non-dedicated - break; - default: - // Should not reach here - CanImportMemory already validated handle type - return impl.ort_api.CreateStatus(ORT_EP_FAIL, "Unexpected external memory handle type"); - } - - // Setup external memory handle descriptor - CUDA_EXTERNAL_MEMORY_HANDLE_DESC ext_mem_desc = {}; - ext_mem_desc.type = cu_handle_type; -#if defined(_WIN32) - ext_mem_desc.handle.win32.handle = desc->native_handle; -#else - ext_mem_desc.handle.fd = static_cast(reinterpret_cast((desc->native_handle))); -#endif - ext_mem_desc.size = desc->size_bytes; - ext_mem_desc.flags = is_dedicated ? CUDA_EXTERNAL_MEMORY_DEDICATED : 0; - - // Import the external memory - CUexternalMemory ext_memory = nullptr; - CUresult cu_result = cuImportExternalMemory(&ext_memory, &ext_mem_desc); - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "cuImportExternalMemory failed: "; - error_msg += error_str ? error_str : "unknown error"; - return impl.ort_api.CreateStatus(ORT_EP_FAIL, error_msg.c_str()); - } - - // Map the external memory to get a device pointer - CUDA_EXTERNAL_MEMORY_BUFFER_DESC buffer_desc = {}; - buffer_desc.offset = desc->offset_bytes; - buffer_desc.size = desc->size_bytes - desc->offset_bytes; - buffer_desc.flags = 0; - - CUdeviceptr mapped_ptr = 0; - cu_result = cuExternalMemoryGetMappedBuffer(&mapped_ptr, ext_memory, &buffer_desc); - if (cu_result != CUDA_SUCCESS) { - cuDestroyExternalMemory(ext_memory); - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "cuExternalMemoryGetMappedBuffer failed: "; - error_msg += error_str ? error_str : "unknown error"; - return impl.ort_api.CreateStatus(ORT_EP_FAIL, error_msg.c_str()); - } - - // Create and return the derived handle (cast to base pointer) - OrtExternalMemoryDescriptor descriptor = {}; // make a copy for the handle - descriptor.version = ORT_API_VERSION; - descriptor.handle_type = desc->handle_type; - descriptor.size_bytes = desc->size_bytes; - descriptor.offset_bytes = desc->offset_bytes; - auto handle = std::make_unique(descriptor); - handle->ep_device = impl.ep_device_; - handle->ext_memory = ext_memory; - handle->mapped_ptr = mapped_ptr; - handle->is_dedicated = is_dedicated; - - *out_handle = handle.release(); - return nullptr; - } - - static void ORT_API_CALL ReleaseMemoryImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ OrtExternalMemoryHandle* handle) noexcept { - (void)this_ptr; - - if (handle == nullptr) { - return; - } - - // The handle has a Release callback that does the actual cleanup - // This method is called from OrtExternalResourceImporterImpl::ReleaseMemory - // The Release callback in the handle will call the static ReleaseCallback - auto mem_handle = std::unique_ptr( - static_cast(handle)); - - // Destroy the external memory object (also releases mapped buffer) - if (mem_handle->ext_memory != nullptr) { - cuDestroyExternalMemory(mem_handle->ext_memory); - } - } - - static OrtStatus* ORT_API_CALL CreateTensorFromMemoryImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ const OrtExternalMemoryHandle* mem_handle, - _In_ const OrtExternalTensorDescriptor* tensor_desc, - _Outptr_ OrtValue** out_tensor) noexcept { - auto& impl = *static_cast(this_ptr); - - if (mem_handle == nullptr || tensor_desc == nullptr || out_tensor == nullptr) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "Invalid arguments to CreateTensorFromMemory"); - } - - // Validate descriptor version - check minimum supported version for forward compatibility - if (tensor_desc->version < ORT_API_VERSION) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "OrtExternalTensorDescriptor version too old"); - } - - *out_tensor = nullptr; - - auto* handle = static_cast(mem_handle); - - // Validate tensor offset does not exceed available buffer size - size_t available_size = handle->descriptor.size_bytes - handle->descriptor.offset_bytes; - if (tensor_desc->offset_bytes > available_size) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "tensor offset_bytes exceeds available imported memory size"); - } - - // Calculate the data pointer with tensor offset - void* data_ptr = reinterpret_cast(handle->mapped_ptr + tensor_desc->offset_bytes); - - // Get memory info from the EP device (the importer is associated with the OrtEpDevice) - const OrtMemoryInfo* memory_info = impl.ep_device_->device_memory_info; - - // Create tensor that references the imported memory. The tensor does not own the memory - - // the user manages the lifetime of both the OrtValue and OrtExternalMemoryHandle. - // The user must keep the handle alive while the tensor is in use. - // No deleter is needed since this is for inference inputs/outputs where the user controls lifetime. - OrtStatus* status = impl.ort_api.CreateTensorWithDataAsOrtValue( - memory_info, - data_ptr, - available_size - tensor_desc->offset_bytes, - tensor_desc->shape, - tensor_desc->rank, - tensor_desc->element_type, - out_tensor); - - return status; - } - - static bool ORT_API_CALL CanImportSemaphoreImpl( - _In_ const OrtExternalResourceImporterImpl* this_ptr, - _In_ OrtExternalSemaphoreType type) noexcept { - (void)this_ptr; -#if defined(_WIN32) - return type == ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE || type == ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_WIN32; -#else - return type == ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_OPAQUE_FD; -#endif - } - - static OrtStatus* ORT_API_CALL ImportSemaphoreImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ const OrtExternalSemaphoreDescriptor* desc, - _Outptr_ OrtExternalSemaphoreHandle** out_handle) noexcept { - auto& impl = *static_cast(this_ptr); - - if (desc == nullptr || out_handle == nullptr) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "Invalid arguments to ImportSemaphore"); - } - - // Validate descriptor version - check minimum supported version for forward compatibility - if (desc->version < ORT_API_VERSION) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "OrtExternalSemaphoreDescriptor version too old"); - } - - *out_handle = nullptr; - - // Validate semaphore type - if (!CanImportSemaphoreImpl(this_ptr, desc->type)) { - return impl.ort_api.CreateStatus(ORT_NOT_IMPLEMENTED, - "Unsupported external semaphore type for CUDA import"); - } - - ScopedContext ctx(impl.DeviceId()); - - // Setup external semaphore handle descriptor for D3D12 fence - CUDA_EXTERNAL_SEMAPHORE_HANDLE_DESC ext_sem_desc = {}; - switch (desc->type) { - case ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE: - ext_sem_desc.type = CU_EXTERNAL_SEMAPHORE_HANDLE_TYPE_D3D12_FENCE; - break; - case ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_WIN32: - ext_sem_desc.type = CU_EXTERNAL_SEMAPHORE_HANDLE_TYPE_TIMELINE_SEMAPHORE_WIN32; - break; - case ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_OPAQUE_FD: - ext_sem_desc.type = CU_EXTERNAL_SEMAPHORE_HANDLE_TYPE_TIMELINE_SEMAPHORE_FD; - break; - default: - // Should not reach here - ImportSemaphoreImpl already validated handle type - return impl.ort_api.CreateStatus(ORT_EP_FAIL, "Unexpected external memory handle type"); - } -#if defined(_WIN32) - ext_sem_desc.handle.win32.handle = desc->native_handle; -#else - ext_sem_desc.handle.fd = static_cast(reinterpret_cast(desc->native_handle)); -#endif - ext_sem_desc.flags = 0; - - // Import the external semaphore - CUexternalSemaphore ext_semaphore = nullptr; - CUresult cu_result = cuImportExternalSemaphore(&ext_semaphore, &ext_sem_desc); - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "cuImportExternalSemaphore failed: "; - error_msg += error_str ? error_str : "unknown error"; - return impl.ort_api.CreateStatus(ORT_EP_FAIL, error_msg.c_str()); - } - - // Create and return the derived handle (cast to base pointer) - auto handle = std::make_unique(*desc); - - // Populate base struct fields - handle->ep_device = impl.ep_device_; - // Populate derived fields - handle->ext_semaphore = ext_semaphore; - - *out_handle = handle.release(); // Return base pointer - return nullptr; - } - - static void ORT_API_CALL ReleaseSemaphoreImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ OrtExternalSemaphoreHandle* handle) noexcept { - (void)this_ptr; - - if (handle == nullptr) { - return; - } - - auto sem_handle = std::unique_ptr( - static_cast(handle)); - - if (sem_handle->ext_semaphore != nullptr) { - cuDestroyExternalSemaphore(sem_handle->ext_semaphore); - } - } - - static OrtStatus* ORT_API_CALL WaitSemaphoreImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ OrtExternalSemaphoreHandle* handle, - _In_ OrtSyncStream* stream, - _In_ uint64_t value) noexcept { - auto& impl = *static_cast(this_ptr); - - if (handle == nullptr || stream == nullptr) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "Invalid arguments to WaitSemaphore"); - } - - auto* sem_handle = static_cast(handle); - - // Get the CUDA stream from OrtSyncStream - cudaStream_t cuda_stream = static_cast(impl.ort_api.SyncStream_GetHandle(stream)); - - // Setup wait parameters for D3D12 fence (timeline semaphore) - CUDA_EXTERNAL_SEMAPHORE_WAIT_PARAMS wait_params = {}; - wait_params.params.fence.value = value; - wait_params.flags = 0; - - // Wait on the external semaphore asynchronously - CUresult cu_result = cuWaitExternalSemaphoresAsync( - &sem_handle->ext_semaphore, - &wait_params, - 1, // numExtSems - cuda_stream); - - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "cuWaitExternalSemaphoresAsync failed: "; - error_msg += error_str ? error_str : "unknown error"; - return impl.ort_api.CreateStatus(ORT_EP_FAIL, error_msg.c_str()); - } - - return nullptr; - } - - static OrtStatus* ORT_API_CALL SignalSemaphoreImpl( - _In_ OrtExternalResourceImporterImpl* this_ptr, - _In_ OrtExternalSemaphoreHandle* handle, - _In_ OrtSyncStream* stream, - _In_ uint64_t value) noexcept { - auto& impl = *static_cast(this_ptr); - - if (handle == nullptr || stream == nullptr) { - return impl.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "Invalid arguments to SignalSemaphore"); - } - - auto* sem_handle = static_cast(handle); - - // Get the CUDA stream from OrtSyncStream - cudaStream_t cuda_stream = static_cast(impl.ort_api.SyncStream_GetHandle(stream)); - - // Setup signal parameters for D3D12 fence / VK timeline semaphore - CUDA_EXTERNAL_SEMAPHORE_SIGNAL_PARAMS signal_params = {}; - signal_params.params.fence.value = value; - signal_params.flags = 0; - - // Signal the external semaphore asynchronously - CUresult cu_result = cuSignalExternalSemaphoresAsync( - &sem_handle->ext_semaphore, - &signal_params, - 1, // numExtSems - cuda_stream); - - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "cuSignalExternalSemaphoresAsync failed: "; - error_msg += error_str ? error_str : "unknown error"; - return impl.ort_api.CreateStatus(ORT_EP_FAIL, error_msg.c_str()); - } - - return nullptr; - } - - static void ORT_API_CALL ReleaseImpl(_In_ OrtExternalResourceImporterImpl* this_ptr) noexcept { - delete static_cast(this_ptr); - } - - /// @brief Get the CUDA device ID from the EP device's memory info. - int DeviceId() const { - return ep_device_->device_memory_info->device.Id(); - } - - private: - const OrtEpDevice* ep_device_; - const OrtApi& ort_api; - const OrtEpApi& ep_api; -}; - -// OrtEpApi infrastructure to be able to use the NvTensorRTRTX EP as an OrtEpFactory for auto EP selection. -struct NvTensorRtRtxEpFactory : OrtEpFactory { - using MemoryInfoUniquePtr = std::unique_ptr>; - - NvTensorRtRtxEpFactory(const OrtApi& ort_api_in, - const OrtLogger& default_logger_in) : OrtEpFactory{}, - ort_api{ort_api_in}, - ep_api{*ort_api_in.GetEpApi()}, - default_logger{default_logger_in}, - data_transfer_impl{ort_api_in} { - GetName = GetNameImpl; - GetVendor = GetVendorImpl; - GetVendorId = GetVendorIdImpl; - GetVersion = GetVersionImpl; - GetSupportedDevices = GetSupportedDevicesImpl; - CreateEp = CreateEpImpl; - ReleaseEp = ReleaseEpImpl; - - CreateAllocator = CreateAllocatorImpl; - ReleaseAllocator = ReleaseAllocatorImpl; - - CreateDataTransfer = CreateDataTransferImpl; - - IsStreamAware = IsStreamAwareImpl; - CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl; - ValidateCompiledModelCompatibilityInfo = ValidateCompiledModelCompatibilityInfoImpl; - - // Graphics interop support (added in ORT 1.25) - InitGraphicsInterop = InitGraphicsInteropImpl; - DeinitGraphicsInterop = DeinitGraphicsInteropImpl; - - CreateExternalResourceImporterForDevice = CreateExternalResourceImporterForDeviceImpl; - GetHardwareDeviceIncompatibilityDetails = GetHardwareDeviceIncompatibilityDetailsImpl; - ort_version_supported = ORT_API_VERSION; // Set to the ORT version we were compiled with. - } - - // Returns the name for the EP. Each unique factory configuration must have a unique name. - // Ex: a factory that supports NPU should have a different than a factory that supports GPU. - static const char* GetNameImpl(const OrtEpFactory* this_ptr) noexcept { - const auto* factory = static_cast(this_ptr); - return factory->ep_name.c_str(); - } - - static const char* GetVendorImpl(const OrtEpFactory* this_ptr) noexcept { - const auto* factory = static_cast(this_ptr); - return factory->vendor.c_str(); - } - - static uint32_t GetVendorIdImpl(const OrtEpFactory* this_ptr) noexcept { - const auto* factory = static_cast(this_ptr); - return factory->vendor_id; - } - - static const char* ORT_API_CALL GetVersionImpl(const OrtEpFactory* /*this_ptr*/) noexcept { - return ORT_VERSION; - } - - /** - * @brief Checks if a given OrtHardwareDevice is a supported NVIDIA GPU. - * - * This function verifies if the provided hardware device corresponds to a physical - * NVIDIA GPU that meets the minimum compute capability requirements for this execution provider. - * - * The check is performed by: - * 1. Extracting the LUID (Locally Unique Identifier) from the device's metadata. - * 2. Converting the string LUID to a 64-bit integer. - * 3. Iterating through all available CUDA devices on the system. - * 4. For each CUDA device, constructing its 64-bit LUID from its properties. - * 5. Comparing the LUIDs. If a match is found, it checks if the device's - * compute capability is at least 8.0 (Ampere) or newer. - * - * @param device The OrtHardwareDevice to check. - * @param major Optional output parameter for the major compute capability version. - * @param minor Optional output parameter for the minor compute capability version. - * @return True if the device is a supported NVIDIA GPU, false otherwise. - */ - bool IsOrtHardwareDeviceSupported(const OrtHardwareDevice& device, int* major = nullptr, int* minor = nullptr) { -#if _WIN32 - const auto& metadata_entries = device.metadata.Entries(); - const auto it = metadata_entries.find("LUID"); - if (it == metadata_entries.end()) { - return false; - } - - uint64_t target_luid; - try { - target_luid = std::stoull(it->second); - } catch (const std::exception&) { - return false; - } - - int device_count = 0; - if (cudaGetDeviceCount(&device_count) != cudaSuccess) { - return false; - } - - for (int i = 0; i < device_count; ++i) { - cudaDeviceProp prop; - if (cudaGetDeviceProperties(&prop, i) != cudaSuccess) { - continue; - } - - // The LUID is an 8-byte value, valid on Windows when luidDeviceNodeMask is non-zero. - // We reconstruct the 64-bit integer representation from the raw bytes. - if (prop.luidDeviceNodeMask == 0) { - continue; - } - - // Copy LUID into uint64_t to avoid unaligned read and strict aliasing (luid is a byte array). - static_assert(sizeof(prop.luid) == sizeof(uint64_t), "cudaDeviceProp::luid should be 8 bytes"); - uint64_t current_luid; - memcpy(¤t_luid, prop.luid, sizeof(current_luid)); - - if (current_luid == target_luid) { - // Set output parameters if provided - if (major != nullptr) { - *major = prop.major; - } - if (minor != nullptr) { - *minor = prop.minor; - } - // Ampere architecture or newer is required. - return prop.major >= 8; - } - } - - return false; -#else - const auto& metadata_entries = device.metadata.Entries(); - const auto it = metadata_entries.find("pci_bus_id"); - if (it == metadata_entries.end()) { - return false; - } - auto& target_id = it->second; - int cuda_device_idx = 0; - if (cudaDeviceGetByPCIBusId(&cuda_device_idx, target_id.c_str()) != cudaSuccess) { - return false; - } - - cudaDeviceProp prop; - if (cudaGetDeviceProperties(&prop, cuda_device_idx) != cudaSuccess) { - return false; - } - // Set output parameters if provided - if (major != nullptr) { - *major = prop.major; - } - if (minor != nullptr) { - *minor = prop.minor; - } - // Ampere architecture or newer is required. - return prop.major >= 8; -#endif - } - - /** - * @brief Parses NVML driver version string and compares with minimum required version. - * - * NVML returns driver versions as "major.minor" on Windows and "major.minor.patch" on Linux. - * Only major and minor are needed for the minimum-version checks below. - * - * @param driver_version_str NVML driver version string (e.g., "581.80" or "555.42.06") - * @param min_version_str Minimum required version string (e.g., "570.00") - * @return true if driver_version >= min_version, false otherwise - */ - static bool CompareNVMLDriverVersion(std::string_view driver_version_str, std::string_view min_version_str) { - auto parseVersion = [](std::string_view version_str, int& major, int& minor) -> bool { - size_t dot_pos = version_str.find('.'); - if (dot_pos == std::string::npos || dot_pos == version_str.length() - 1 || dot_pos == 0) { - return false; - } - - auto parsePart = [](std::string_view part, int& value) -> bool { - if (part.empty()) { - return false; - } - - int parsed_value = 0; - const char* const begin = part.data(); - const char* const end = begin + part.size(); - const auto result = std::from_chars(begin, end, parsed_value); - if (result.ec != std::errc{} || result.ptr != end || parsed_value < 0) { - return false; - } - - value = parsed_value; - return true; - }; - - const size_t minor_start = dot_pos + 1; - const size_t second_dot_pos = version_str.find('.', minor_start); - const size_t minor_length = second_dot_pos == std::string::npos - ? version_str.length() - minor_start - : second_dot_pos - minor_start; - if (!parsePart(version_str.substr(0, dot_pos), major) || - !parsePart(version_str.substr(minor_start, minor_length), minor)) { - return false; - } - - // Linux NVML versions include a patch component. Validate it when present, but it is not part of - // the minimum-version comparison. - if (second_dot_pos != std::string::npos) { - int patch = 0; - return parsePart(version_str.substr(second_dot_pos + 1), patch); - } - - return true; - }; - - int driver_major = 0, driver_minor = 0; - int min_major = 0, min_minor = 0; - - if (!parseVersion(driver_version_str, driver_major, driver_minor) || - !parseVersion(min_version_str, min_major, min_minor)) { - // If parsing fails, conservatively treat as incompatible (fail safe) - return false; - } - - // Compare versions numerically - if (driver_major > min_major) return true; - if (driver_major < min_major) return false; - return driver_minor >= min_minor; - } - - /** - * @brief Checks for hardware device incompatibility reasons with NvTensorRTRTX EP. - * - * This function is called by ORT's GetHardwareDeviceEpIncompatibilityDetails() API - * to provide diagnostic information about why a device may be incompatible with - * this execution provider. - * - * Currently checks: - * - Compute capability: Requires Ampere (8.0) or newer GPU architecture - * - Driver version: Uses NVML to check NVIDIA graphics driver version - * - CC 8.x-11.x devices: Requires R555 or newer - * - Blackwell (CC 12.x): Requires R570 or newer - * - * @param hw The hardware device to check for incompatibility. - * @param details Pre-allocated incompatibility details object initialized by ORT. - * @return nullptr on success (compatible or details set), OrtStatus on error. - */ - OrtStatus* GetHardwareDeviceIncompatibilityDetailsInternal( - const OrtHardwareDevice* hw, - OrtDeviceEpIncompatibilityDetails* details) { - if (hw == nullptr || details == nullptr) { - return ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] Invalid arguments: hw or details is null"); - } - - // Check if the device is a GPU from NVIDIA vendor - OrtHardwareDeviceType device_type = ort_api.HardwareDevice_Type(hw); - uint32_t hardware_vendor_id = ort_api.HardwareDevice_VendorId(hw); - - if (device_type != OrtHardwareDeviceType::OrtHardwareDeviceType_GPU || - hardware_vendor_id != vendor_id) { - // Not a NVIDIA GPU - device type/vendor incompatible - uint32_t reasons = OrtDeviceEpIncompatibility_DEVICE_INCOMPATIBLE; - return ep_api.DeviceEpIncompatibilityDetails_SetDetails( - details, - reasons, - 0, // error_code - "NvTensorRTRTX EP only supports NVIDIA GPU devices"); - } - - // Check compute capability and get major/minor for driver version check - int compute_capability_major = 0; - int compute_capability_minor = 0; - if (!IsOrtHardwareDeviceSupported(*hw, &compute_capability_major, &compute_capability_minor)) { - uint32_t reasons = OrtDeviceEpIncompatibility_DEVICE_INCOMPATIBLE; - // If both are still 0, the CUDA device lookup failed (device not matched by LUID/PCI bus ID) - if (compute_capability_major == 0 && compute_capability_minor == 0) { - return ep_api.DeviceEpIncompatibilityDetails_SetDetails( - details, - reasons, - 0, // error_code - "NvTensorRTRTX EP could not determine the compute capability of this NVIDIA GPU. " - "Ensure the CUDA runtime is installed and the device is accessible."); - } - // Device architecture not supported - compute capability too low - std::string cc_string = std::to_string(compute_capability_major) + "." + std::to_string(compute_capability_minor); - std::string msg = "NvTensorRTRTX EP does not support GPU with Compute Capability " + cc_string + - ". Requires Ampere (8.0+) or newer GPU architecture."; - return ep_api.DeviceEpIncompatibilityDetails_SetDetails( - details, - reasons, - 0, // error_code - msg.c_str()); - } - -#if defined(HAVE_NVML) - // Determine minimum driver version based on GPU architecture - const char* min_driver_version = nullptr; - if (compute_capability_major >= 12) { - // Blackwell architecture (CC 12.x) requires the R570 branch or newer. - min_driver_version = "570.00"; - } else if (compute_capability_major >= 8) { - // CC 8.x-11.x devices use the R555 branch minimum. -#if _WIN32 - min_driver_version = "555.85"; -#else - min_driver_version = "555.42"; -#endif - } else { - // Should not reach here (already checked compute capability above) -#if _WIN32 - min_driver_version = "555.85"; -#else - min_driver_version = "555.42"; -#endif - } - - // Initialize NVML and get driver version - nvmlReturn_t nvml_result = nvmlInit_v2(); - if (nvml_result != NVML_SUCCESS) { - uint32_t reasons = OrtDeviceEpIncompatibility_DRIVER_INCOMPATIBLE; - std::string msg = "Failed to initialize NVML (" + std::string(nvmlErrorString(nvml_result)) + - "). NVML may be unavailable due to missing or incompatible NVIDIA driver, insufficient permissions, or runtime/container restrictions."; - return ep_api.DeviceEpIncompatibilityDetails_SetDetails( - details, - reasons, - 0, // error_code - msg.c_str()); - } - - char driver_version_str[NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE] = {0}; - nvml_result = nvmlSystemGetDriverVersion(driver_version_str, sizeof(driver_version_str)); - - // Shutdown NVML before returning - nvmlShutdown(); - - if (nvml_result != NVML_SUCCESS) { - uint32_t reasons = OrtDeviceEpIncompatibility_DRIVER_INCOMPATIBLE; - std::string msg = "Failed to query NVIDIA driver version: " + std::string(nvmlErrorString(nvml_result)); - return ep_api.DeviceEpIncompatibilityDetails_SetDetails( - details, - reasons, - 0, // error_code - msg.c_str()); - } - - // Compare driver version with minimum required - if (!CompareNVMLDriverVersion(driver_version_str, min_driver_version)) { - uint32_t reasons = OrtDeviceEpIncompatibility_DRIVER_INCOMPATIBLE; - std::string msg = "NVIDIA driver version " + std::string(driver_version_str) + - " is too old. Minimum required: " + min_driver_version + " or higher"; - - return ep_api.DeviceEpIncompatibilityDetails_SetDetails( - details, - reasons, - 0, // error_code (could store parsed version if needed) - msg.c_str()); - } -#endif // HAVE_NVML - - // Device is compatible - details are already initialized with default values by ORT - return nullptr; - } - - static OrtStatus* ORT_API_CALL GetHardwareDeviceIncompatibilityDetailsImpl( - OrtEpFactory* this_ptr, - const OrtHardwareDevice* hw, - OrtDeviceEpIncompatibilityDetails* details) noexcept { - auto* factory = static_cast(this_ptr); - try { - return factory->GetHardwareDeviceIncompatibilityDetailsInternal(hw, details); - } catch (const std::exception&) { - return factory->ort_api.CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] Exception in GetHardwareDeviceIncompatibilityDetails"); - } catch (...) { - return factory->ort_api.CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] Unknown exception in GetHardwareDeviceIncompatibilityDetails"); - } - } - - // Creates and returns OrtEpDevice instances for all OrtHardwareDevices that this factory supports. - // An EP created with this factory is expected to be able to execute a model with *all* supported - // hardware devices at once. A single instance of NvTensorRtRtx EP is not currently setup to partition a model among - // multiple different NvTensorRtRtx backends at once (e.g, npu, cpu, gpu), so this factory instance is set to only - // support one backend: gpu. To support a different backend, like npu, create a different factory instance - // that only supports NPU. - static OrtStatus* GetSupportedDevicesImpl(OrtEpFactory* this_ptr, - const OrtHardwareDevice* const* devices, - size_t num_devices, - OrtEpDevice** ep_devices, - size_t max_ep_devices, - size_t* p_num_ep_devices) noexcept { - size_t& num_ep_devices = *p_num_ep_devices; - auto* factory = static_cast(this_ptr); - - int num_cuda_devices = 0; - cudaGetDeviceCount(&num_cuda_devices); - RETURN_IF_ERROR(factory->CreateMemoryInfoForDevices(num_cuda_devices)); - - int16_t device_id = 0; - for (size_t i = 0; i < num_devices && num_ep_devices < max_ep_devices; ++i) { - const OrtHardwareDevice& device = *devices[i]; - - if (factory->ort_api.HardwareDevice_Type(&device) == OrtHardwareDeviceType::OrtHardwareDeviceType_GPU && - factory->ort_api.HardwareDevice_VendorId(&device) == factory->vendor_id && - factory->IsOrtHardwareDeviceSupported(device)) { - OrtKeyValuePairs* ep_options = nullptr; - OrtKeyValuePairs* ep_metadata = nullptr; - factory->ort_api.CreateKeyValuePairs(&ep_options); - factory->ort_api.CreateKeyValuePairs(&ep_metadata); - factory->ort_api.AddKeyValuePair(ep_options, "device_id", std::to_string(device_id).c_str()); - - RETURN_IF_ERROR(factory->ort_api.GetEpApi()->CreateEpDevice(factory, &device, ep_metadata, ep_options, - &ep_devices[num_ep_devices])); - - factory->ort_api.ReleaseKeyValuePairs(ep_options); - factory->ort_api.ReleaseKeyValuePairs(ep_metadata); - - const OrtMemoryInfo* gpu_mem_info = factory->gpu_memory_infos[device_id].get(); - const OrtMemoryInfo* host_accessible_mem_info = factory->host_accessible_memory_infos[device_id].get(); - - RETURN_IF_ERROR(factory->ep_api.EpDevice_AddAllocatorInfo(ep_devices[num_ep_devices], gpu_mem_info)); - RETURN_IF_ERROR(factory->ep_api.EpDevice_AddAllocatorInfo(ep_devices[num_ep_devices], host_accessible_mem_info)); - num_ep_devices++; - device_id++; - } - } - return nullptr; - } - - static OrtStatus* CreateEpImpl(OrtEpFactory* /*this_ptr*/, - _In_reads_(num_devices) const OrtHardwareDevice* const* /*devices*/, - _In_reads_(num_devices) const OrtKeyValuePairs* const* /*ep_metadata*/, - _In_ size_t /*num_devices*/, - _In_ const OrtSessionOptions* /*session_options*/, - _In_ const OrtLogger* /*logger*/, - _Out_ OrtEp** /*ep*/) noexcept { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, "[NvTensorRTRTX EP] EP factory does not support this method."); - } - - static void ReleaseEpImpl(OrtEpFactory* /*this_ptr*/, OrtEp* /*ep*/) noexcept { - // no-op as we never create an EP here. - } - - static OrtStatus* ORT_API_CALL CreateAllocatorImpl(OrtEpFactory* this_ptr, - const OrtMemoryInfo* memory_info, - const OrtKeyValuePairs* /*allocator_options*/, - OrtAllocator** allocator) noexcept { - auto& factory = *static_cast(this_ptr); - auto allocator_ = std::make_unique(memory_info, factory.ort_api); - *allocator = allocator_.release(); - return nullptr; - } - - static void ORT_API_CALL ReleaseAllocatorImpl(OrtEpFactory* /*this*/, OrtAllocator* allocator) noexcept { - delete static_cast(allocator); - } - - static OrtStatus* ORT_API_CALL CreateDataTransferImpl(OrtEpFactory* this_ptr, - OrtDataTransferImpl** data_transfer) noexcept { - auto& factory = *static_cast(this_ptr); - *data_transfer = &factory.data_transfer_impl; - return nullptr; - } - - static bool ORT_API_CALL IsStreamAwareImpl(const OrtEpFactory* /*this_ptr*/) noexcept { - return true; - } - - static OrtStatus* ORT_API_CALL CreateSyncStreamForDeviceImpl(OrtEpFactory* this_ptr, - const OrtMemoryDevice* memory_device, - const OrtKeyValuePairs* /*stream_options*/, - OrtSyncStreamImpl** ort_stream) noexcept { - auto& factory = *static_cast(this_ptr); - - auto device_id = factory.ep_api.MemoryDevice_GetDeviceId(memory_device); - cudaStream_t stream = nullptr; - CUcontext cig_context = factory.GetCigContext(device_id); - - try { - // Push context so stream is created on it; pop on scope exit to restore thread's prior context. - if (cig_context != nullptr) { - ScopedContext scoped(cig_context); - CUDA_RETURN_IF_ERROR(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); - } else { - ScopedContext scoped(device_id); - CUDA_RETURN_IF_ERROR(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); - } - } catch (const std::exception& e) { - return onnxruntime::CreateStatus(ORT_FAIL, e.what()); - } catch (...) { - return onnxruntime::CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] Failed to set context current for stream creation."); - } - - const OrtDevice* ort_device = static_cast(memory_device); - - auto impl = std::make_unique(std::move(stream), *ort_device, nullptr, - /*release_cpu_buffer_on_cuda_stream*/ true, - factory.ort_api); - *ort_stream = impl.release(); - return nullptr; - } - - /** @brief Create an external resource importer for D3D12 to CUDA import. - * - * This enables zero-copy import of D3D12 shared resources and timeline fences. - * The implementation uses CUDA Driver APIs (cuImportExternalMemory, cuImportExternalSemaphore). - * - * @param this_ptr The OrtEpFactory instance. - * @param ep_device The OrtEpDevice to create the importer for. - * @param out_importer Output parameter set to the created OrtExternalResourceImporterImpl. - * @return nullptr on success, OrtStatus with error on failure. - */ - static OrtStatus* ORT_API_CALL CreateExternalResourceImporterForDeviceImpl( - OrtEpFactory* this_ptr, - const OrtEpDevice* ep_device, - OrtExternalResourceImporterImpl** out_importer) noexcept { - auto& factory = *static_cast(this_ptr); - - if (out_importer == nullptr) { - return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "out_importer cannot be nullptr"); - } - - *out_importer = nullptr; - - // Create the external resource importer - auto importer = std::make_unique(ep_device, factory.ort_api); - *out_importer = importer.release(); - - return nullptr; - } - - /** - * This function is called by the public C API GetModelCompatibilityForEpDevices. - * It uses TensorRT RTX runtime directly to call runtime->getEngineValidity() to check the 64-byte engine header. - * - * @param this_ptr Factory instance pointer - * @param devices Hardware devices (not used, validation is done against current system) - * @param num_devices Number of devices - * @param compatibility_info Hex-encoded 64-byte TensorRT RTX engine header (128 hex characters) - * @param model_compatibility Output parameter for compatibility status - * @return OrtStatus* nullptr on success, error status on failure - */ - static OrtStatus* ORT_API_CALL ValidateCompiledModelCompatibilityInfoImpl( - OrtEpFactory* this_ptr, - const OrtHardwareDevice* const* devices, - size_t num_devices, - const char* compatibility_info, - OrtCompiledModelCompatibility* model_compatibility) noexcept { - auto& factory = *static_cast(this_ptr); - - // Validate input parameters - if (compatibility_info == nullptr || model_compatibility == nullptr) { - return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] Invalid arguments: compatibility_info or model_compatibility is null"); - } - - // Device parameters not used for header validation - ORT_UNUSED_PARAMETER(devices); - ORT_UNUSED_PARAMETER(num_devices); - - try { - // If no compatibility info provided, validation not applicable - if (compatibility_info[0] == '\0') { - *model_compatibility = OrtCompiledModelCompatibility_EP_NOT_APPLICABLE; - return nullptr; - } - - // Decode hex string to binary - std::vector engine_header; - try { - engine_header = HexStringToBinary(std::string(compatibility_info)); - } catch (const std::exception& ex) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Failed to decode engine header: " << ex.what(); - *model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - return nullptr; - } - - // Validate header size (keep in sync with TensorRT engine header size) - if (engine_header.size() != kTensorRTEngineHeaderSize) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Invalid header size: " << engine_header.size() - << " bytes (expected " << kTensorRTEngineHeaderSize << ")"; - *model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - return nullptr; - } - - // Create TensorRT runtime for validation - static std::mutex runtime_creation_mutex; - std::unique_ptr runtime; - { - std::lock_guard lock(runtime_creation_mutex); - TensorrtLogger& trt_logger = GetTensorrtLogger(false); - runtime.reset(nvinfer1::createInferRuntime(trt_logger)); - } - - if (!runtime) { - LOGS_DEFAULT(ERROR) << "[NvTensorRTRTX EP] Failed to create TensorRT runtime"; - return factory.ort_api.CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] Failed to create TensorRT runtime"); - } - - // Use TensorRT's getEngineValidity to check compatibility - uint64_t diagnostics = 0; - nvinfer1::EngineValidity validity = runtime->getEngineValidity( - engine_header.data(), - engine_header.size(), - &diagnostics); - - // Map TensorRT validity to ORT compatibility status - switch (validity) { - case nvinfer1::EngineValidity::kVALID: - *model_compatibility = OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL; - break; - - case nvinfer1::EngineValidity::kSUBOPTIMAL: - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine compatible but recompilation recommended " - << "(diagnostics: 0x" << std::hex << diagnostics << std::dec << ")"; - *model_compatibility = OrtCompiledModelCompatibility_EP_SUPPORTED_PREFER_RECOMPILATION; - break; - - case nvinfer1::EngineValidity::kINVALID: - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine incompatible with this system " - << "(diagnostics: 0x" << std::hex << diagnostics << std::dec << ")"; - *model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - break; - - default: - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Unknown validity status: " - << static_cast(validity); - *model_compatibility = OrtCompiledModelCompatibility_EP_UNSUPPORTED; - break; - } - - return nullptr; - - } catch (const std::exception& ex) { - std::string error_msg = std::string("[NvTensorRTRTX EP] Exception during validation: ") + ex.what(); - LOGS_DEFAULT(ERROR) << error_msg; - return factory.ort_api.CreateStatus(ORT_FAIL, error_msg.c_str()); - } catch (...) { - LOGS_DEFAULT(ERROR) << "[NvTensorRTRTX EP] Unknown exception during validation"; - return factory.ort_api.CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] Unknown exception during validation"); - } - } - - // Initialize graphics interop for a device. Creates a CIG (CUDA Interop Graphics) context - // bound to the provided graphics command queue/device. - // This follows Scott's suggestion to pass in everything required so we don't end up with multiple - // init function signatures, and the factory stores the queue to utilize in stream creation. - static OrtStatus* ORT_API_CALL InitGraphicsInteropImpl(OrtEpFactory* this_ptr, - const OrtEpDevice* ep_device, - const OrtGraphicsInteropConfig* config) noexcept { - if (ep_device == nullptr) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] InitGraphicsInterop: ep_device is null"); - } - if (config == nullptr) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] InitGraphicsInterop: config is null"); - } - if (config->version != ORT_API_VERSION) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] InitGraphicsInterop: config version does not match ORT_API_VERSION"); - } - - auto& factory = *static_cast(this_ptr); - - // Extract device_id from OrtEpDevice options - const OrtKeyValuePairs* ep_options = factory.ort_api.EpDevice_EpOptions(ep_device); - const char* device_id_str = factory.ort_api.GetKeyValue(ep_options, "device_id"); - if (device_id_str == nullptr) { - return onnxruntime::CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] InitGraphicsInterop: device_id not found in ep_device options"); - } - char* parse_end = nullptr; - errno = 0; - long device_id_long = strtol(device_id_str, &parse_end, 10); - if (parse_end == device_id_str || *parse_end != '\0' || errno == ERANGE || - device_id_long < static_cast(INT32_MIN) || device_id_long > static_cast(INT32_MAX)) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] InitGraphicsInterop: invalid device_id in ep_device options"); - } - int32_t device_id = static_cast(device_id_long); - - // Validate graphics API - if (config->graphics_api == ORT_GRAPHICS_API_NONE) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] InitGraphicsInterop: graphics_api cannot be NONE"); - } - - // Initialize CUDA driver API - CUresult cu_result = cuInit(0); - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "[NvTensorRTRTX EP] Failed to initialize CUDA driver API: "; - error_msg += error_str ? error_str : "unknown error"; - return onnxruntime::CreateStatus(ORT_FAIL, error_msg.c_str()); - } - - // Get CUDA device properties to retrieve LUID - cudaDeviceProp cuda_prop; - cudaError_t cuda_err = cudaGetDeviceProperties(&cuda_prop, device_id); - if (cuda_err != cudaSuccess) { - std::string error_msg = "[NvTensorRTRTX EP] Failed to get CUDA device properties: "; - error_msg += cudaGetErrorString(cuda_err); - return onnxruntime::CreateStatus(ORT_FAIL, error_msg.c_str()); - } - - // Create CIG context based on graphics API type - CUcontext cig_context = nullptr; - - if (config->graphics_api == ORT_GRAPHICS_API_D3D12) { -#if defined(_WIN32) && USE_DX_INTEROP - // command_queue is optional (performance optimization for GPU-side sync). Without it, skip CIG context. - if (config->command_queue == nullptr) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] InitGraphicsInterop: command_queue not passed; skipping graphics interop context (streams will use default context)."; - return nullptr; // Success; Interop API still works, streams will use default context - } - - ID3D12CommandQueue* d3d12_queue = reinterpret_cast(config->command_queue); - - // Get device from command queue (no separate device param; avoids mismatch) - ID3D12Device* d3d12_device = nullptr; - HRESULT hr = d3d12_queue->GetDevice(IID_PPV_ARGS(&d3d12_device)); - if (FAILED(hr) || d3d12_device == nullptr) { - return onnxruntime::CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] InitGraphicsInterop: failed to get D3D12 device from command queue"); - } - - // Get LUID from CUDA device (OrtEpDevice identifies the GPU via device_id) - if (cuda_prop.luidDeviceNodeMask == 0) { - d3d12_device->Release(); - return onnxruntime::CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] CUDA device does not have a valid LUID"); - } - uint64_t cuda_luid; - memcpy(&cuda_luid, cuda_prop.luid, sizeof(cuda_luid)); - - // Ensure graphics device from command queue and OrtEpDevice GPU are logically the same - LUID d3d12_luid = d3d12_device->GetAdapterLuid(); - uint64_t d3d12_luid_64 = (static_cast(d3d12_luid.HighPart) << 32) | d3d12_luid.LowPart; - d3d12_device->Release(); - - if (d3d12_luid_64 != cuda_luid) { - return onnxruntime::CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] D3D12 device LUID does not match CUDA device LUID"); - } - - // Create CIG context bound to D3D12 command queue - - CUctxCigParam cig_param = {CIG_DATA_TYPE_D3D12_COMMAND_QUEUE, d3d12_queue}; - CUctxCreateParams ctx_params = {nullptr, 0, &cig_param}; - - cu_result = cuCtxCreate_v4(&cig_context, &ctx_params, 0, device_id); - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "[NvTensorRTRTX EP] Failed to create CIG context for D3D12: "; - error_msg += error_str ? error_str : "unknown error"; - return onnxruntime::CreateStatus(ORT_FAIL, error_msg.c_str()); - } -#else - return onnxruntime::CreateStatus(ORT_NOT_IMPLEMENTED, - "[NvTensorRTRTX EP] D3D12 CIG context creation not supported on this platform"); -#endif - } else if (config->graphics_api == ORT_GRAPHICS_API_VULKAN) { - int cig_supported{false}; - if (cudaSuccess != cudaDeviceGetAttribute(&cig_supported, cudaDevAttrVulkanCigSupported, device_id)) { - return onnxruntime::CreateStatus(ORT_EP_FAIL, - "[NvTensorRTRTX EP] Could not determine CiG support for CUDA device"); - } - if (!cig_supported) { - LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] InitGraphicsInterop: CiG for Vulkan is not supported on the given device. Will use the default CUDA context"; - return nullptr; - } - const char* nv_blob_ptr_str = factory.ort_api.GetKeyValue(config->additional_options, onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data); - if (!nv_blob_ptr_str) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] InitGraphicsInterop: Can't enable CUDA in Graphics (CiG) for Vulkan without onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data"; - return nullptr; - } - uint64_t nv_blob_ptr = std::stoull(nv_blob_ptr_str); - if (nv_blob_ptr == 0) { - return onnxruntime::CreateStatus(ORT_EP_FAIL, - "[NvTensorRTRTX EP] Could not parse provided values for onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data or onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data_len"); - } - - CUctxCigParam cig_params{}; - cig_params.sharedDataType = CIG_DATA_TYPE_NV_BLOB; - cig_params.sharedData = reinterpret_cast(nv_blob_ptr); - CUctxCreateParams params{}; - params.cigParams = &cig_params; - - cu_result = cuCtxCreate_v4(&cig_context, ¶ms, 0, device_id); - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "[NvTensorRTRTX EP] Failed to create CIG context for Vulkan: "; - error_msg += error_str ? error_str : "unknown error"; - return onnxruntime::CreateStatus(ORT_FAIL, error_msg.c_str()); - } - } else { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] Unsupported graphics API for CIG context"); - } - - // Store the CIG context for this device - { - std::lock_guard lock(factory.cig_contexts_mutex_); - factory.cig_contexts_[device_id] = cig_context; - } - - return nullptr; - } - - // Deinitialize graphics interop for a device. Destroys the CIG context (cuCtxDestroy) and removes - // it from the map. Callers must release dependent resources (OrtSyncStream, semaphore, importer) - // before calling Deinit so the context is not in use. - static OrtStatus* ORT_API_CALL DeinitGraphicsInteropImpl(OrtEpFactory* this_ptr, - const OrtEpDevice* ep_device) noexcept { - if (ep_device == nullptr) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] DeinitGraphicsInterop: ep_device is null"); - } - - auto& factory = *static_cast(this_ptr); - - // Extract device_id from OrtEpDevice options - const OrtKeyValuePairs* ep_options = factory.ort_api.EpDevice_EpOptions(ep_device); - const char* device_id_str = factory.ort_api.GetKeyValue(ep_options, "device_id"); - if (device_id_str == nullptr) { - return onnxruntime::CreateStatus(ORT_FAIL, - "[NvTensorRTRTX EP] DeinitGraphicsInterop: device_id not found in ep_device options"); - } - char* parse_end = nullptr; - errno = 0; - long device_id_long = strtol(device_id_str, &parse_end, 10); - if (parse_end == device_id_str || *parse_end != '\0' || errno == ERANGE || - device_id_long < static_cast(INT32_MIN) || device_id_long > static_cast(INT32_MAX)) { - return onnxruntime::CreateStatus(ORT_INVALID_ARGUMENT, - "[NvTensorRTRTX EP] DeinitGraphicsInterop: invalid device_id in ep_device options"); - } - int32_t device_id = static_cast(device_id_long); - - CUcontext cig_context = nullptr; - { - std::lock_guard lock(factory.cig_contexts_mutex_); - auto it = factory.cig_contexts_.find(device_id); - if (it != factory.cig_contexts_.end()) { - cig_context = it->second; - factory.cig_contexts_.erase(it); - } - } - - // Destroy the CIG context. Callers must release dependent resources (OrtSyncStream, - // external semaphore/importer, etc.) before Deinit so the context is not in use. - if (cig_context != nullptr) { - CUresult cu_result = cuCtxDestroy(cig_context); - if (cu_result != CUDA_SUCCESS) { - const char* error_str = nullptr; - cuGetErrorString(cu_result, &error_str); - std::string error_msg = "[NvTensorRTRTX EP] DeinitGraphicsInterop: cuCtxDestroy failed: "; - error_msg += error_str ? error_str : "unknown error"; - return onnxruntime::CreateStatus(ORT_FAIL, error_msg.c_str()); - } - } - - return nullptr; - } - - // Get the CIG CUDA context for a device (for use in stream creation) - CUcontext GetCigContext(int32_t device_id) const { - std::lock_guard lock(cig_contexts_mutex_); - auto it = cig_contexts_.find(device_id); - if (it != cig_contexts_.end()) { - return it->second; - } - return nullptr; - } - - OrtStatus* CreateMemoryInfoForDevices(int num_devices) { - gpu_memory_infos.reserve(num_devices); - host_accessible_memory_infos.reserve(num_devices); - - for (int device_id = 0; device_id < num_devices; ++device_id) { - OrtMemoryInfo* mem_info = nullptr; - RETURN_IF_ERROR(ort_api.CreateMemoryInfo_V2("NvTensorRTRTX", OrtMemoryInfoDeviceType_GPU, - /*vendor*/ OrtDevice::VendorIds::NVIDIA, - /* device_id */ device_id, - OrtDeviceMemoryType_DEFAULT, - /*alignment*/ 0, - OrtAllocatorType::OrtDeviceAllocator, - &mem_info)); - gpu_memory_infos.emplace_back(MemoryInfoUniquePtr(mem_info, ort_api.ReleaseMemoryInfo)); - - mem_info = nullptr; - RETURN_IF_ERROR(ort_api.CreateMemoryInfo_V2("NvTensorRTRTX host accessible", OrtMemoryInfoDeviceType_GPU, - /*vendor*/ OrtDevice::VendorIds::NVIDIA, - /* device_id */ device_id, - OrtDeviceMemoryType_HOST_ACCESSIBLE, - /*alignment*/ 0, - OrtAllocatorType::OrtDeviceAllocator, - &mem_info)); - host_accessible_memory_infos.emplace_back(MemoryInfoUniquePtr(mem_info, ort_api.ReleaseMemoryInfo)); - } - return nullptr; - } - - private: - const OrtApi& ort_api; - const OrtEpApi& ep_api; - const OrtLogger& default_logger; - const std::string ep_name{kNvTensorRTRTXExecutionProvider}; - const std::string vendor{"NVIDIA"}; - - // NVIDIA vendor ID. Refer to the ACPI ID registry (search NVIDIA): https://uefi.org/ACPI_ID_List - const uint32_t vendor_id{0x10de}; - - std::vector gpu_memory_infos; - std::vector host_accessible_memory_infos; - - // we use a shared instance for the OrtDataTransferImpl instead of creating a new one on every call to - NvTrtRtxDataTransferImpl data_transfer_impl; - - // CIG contexts per device (keyed by device_id). - // Threading: InitGraphicsInteropImpl, DeinitGraphicsInteropImpl, and CreateSyncStreamForDeviceImpl - // (via GetCigContext) may be called from different threads. All access is guarded by cig_contexts_mutex_. - mutable std::mutex cig_contexts_mutex_; - std::unordered_map cig_contexts_; - - NvTensorRtRtxEpFactory(const NvTensorRtRtxEpFactory&) = delete; - NvTensorRtRtxEpFactory& operator=(const NvTensorRtRtxEpFactory&) = delete; - - NvTensorRtRtxEpFactory(NvTensorRtRtxEpFactory&&) = default; - NvTensorRtRtxEpFactory& operator=(NvTensorRtRtxEpFactory&&) = default; -}; - -extern "C" { -// -// Public symbols -// -OrtStatus* CreateEpFactories(const char* /*registration_name*/, const OrtApiBase* ort_api_base, - const OrtLogger* default_logger, - OrtEpFactory** factories, size_t max_factories, size_t* num_factories) { - const OrtApi* ort_api = ort_api_base->GetApi(ORT_API_VERSION); - - // Factory could use registration_name or define its own EP name. - auto factory = std::make_unique(*ort_api, *default_logger); - - if (max_factories < 1) { - return ort_api->CreateStatus(ORT_INVALID_ARGUMENT, - "Not enough space to return EP factory. Need at least one."); - } - - factories[0] = factory.release(); - *num_factories = 1; - - return nullptr; -} - -OrtStatus* ReleaseEpFactory(OrtEpFactory* factory) { - delete static_cast(factory); - return nullptr; -} -} diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.h deleted file mode 100644 index 5672c5dda632e..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.h +++ /dev/null @@ -1,17 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "onnxruntime_c_api.h" -#include "core/framework/provider_options.h" - -namespace onnxruntime { -struct ProviderInfo_Nv { - virtual OrtStatus* GetCurrentGpuDeviceId(_In_ int* device_id) = 0; - virtual OrtStatus* GetTensorRTCustomOpDomainList(std::vector& domain_list, const std::string extra_plugin_lib_paths) = 0; - virtual OrtStatus* ReleaseCustomOpDomainList(std::vector& domain_list) = 0; - - protected: - ~ProviderInfo_Nv() = default; // Can only be destroyed through a subclass instance -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory_creator.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory_creator.h deleted file mode 100644 index 6b2e516211257..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory_creator.h +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include - -#include "core/framework/provider_options.h" -#include "core/providers/providers.h" - -namespace onnxruntime { -struct SessionOptions; - -// defined in provider_bridge_ort.cc -struct NvProviderFactoryCreator { - static std::shared_ptr Create(int device_id); - static std::shared_ptr Create(const ProviderOptions& provider_options_map, - const SessionOptions* session_options); -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_scoped_context.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_scoped_context.h deleted file mode 100644 index 8a16533b01c7b..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_scoped_context.h +++ /dev/null @@ -1,42 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "nv_includes.h" -#include "core/providers/cuda/cuda_pch.h" -#include "core/providers/cuda/shared_inc/cuda_call.h" - -namespace onnxruntime { -struct ScopedContext { - explicit ScopedContext(int device_id) : pushed_(true) { - CUcontext cu_context = 0; - CU_CALL_THROW(cuCtxGetCurrent(&cu_context)); - if (!cu_context) { - // cuCtxGetCurrent succeeded but returned nullptr, which indicates that no CUDA context - // is currently set for this thread. This implicates that there is not user created context. - // We use runtime API to initialize a context for the specified device. - CUDA_CALL_THROW(cudaSetDevice(device_id)); - CU_CALL_THROW(cuCtxGetCurrent(&cu_context)); - } - CU_CALL_THROW(cuCtxPushCurrent(cu_context)); - } - - /** \brief Push an existing context (e.g. CIG context); pop on destruction. */ - explicit ScopedContext(CUcontext ctx) : pushed_(ctx != nullptr) { - if (ctx != nullptr) { - CU_CALL_THROW(cuCtxPushCurrent(ctx)); - } - } - - ScopedContext(const ScopedContext&) = delete; - - ~ScopedContext() { - if (pushed_) { - cuCtxPopCurrent(nullptr); - } - } - - private: - bool pushed_ = true; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc deleted file mode 100644 index 9311c598ce1b5..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ /dev/null @@ -1,341 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include - -#include "onnx_ctx_model_helper.h" -#include "core/common/path_utils.h" -#include "core/providers/cuda/shared_inc/cuda_call.h" -#include "core/framework/execution_provider.h" -#include "nv_execution_provider.h" - -namespace onnxruntime { -extern TensorrtLogger& GetTensorrtLogger(bool verbose_log); - -/* - * Convert binary data to hex string - */ -std::string BinaryToHexString(const void* data, size_t size) { - static const char hex_chars[] = "0123456789abcdef"; - const uint8_t* bytes = static_cast(data); - std::string result; - result.reserve(size * 2); - - for (size_t i = 0; i < size; ++i) { - result.push_back(hex_chars[(bytes[i] >> 4) & 0xF]); - result.push_back(hex_chars[bytes[i] & 0xF]); - } - return result; -} - -/* - * Convert hex string back to binary - */ -std::vector HexStringToBinary(const std::string& hex) { - if (hex.size() % 2 != 0) { - ORT_THROW("Hex string must have even length"); - } - - std::vector result; - result.reserve(hex.size() / 2); - - for (size_t i = 0; i < hex.size(); i += 2) { - uint8_t byte = 0; - - // High nibble - char c = hex[i]; - byte |= (c >= '0' && c <= '9') ? static_cast((c - '0') << 4) : (c >= 'a' && c <= 'f') ? static_cast((c - 'a' + 10) << 4) - : (c >= 'A' && c <= 'F') ? static_cast((c - 'A' + 10) << 4) - : 0; - - // Low nibble - c = hex[i + 1]; - byte |= (c >= '0' && c <= '9') ? static_cast(c - '0') : (c >= 'a' && c <= 'f') ? static_cast(c - 'a' + 10) - : (c >= 'A' && c <= 'F') ? static_cast(c - 'A' + 10) - : 0; - - result.push_back(byte); - } - return result; -} - -/* - * Check whether the graph has the EP context contrib op. - * The op can contain the precompiled engine info for TRT EP to directly load the engine. - * - * Note: Please see more details about "EPContext" contrib op in contrib_defs.cc - */ -bool GraphHasCtxNode(const GraphViewer& graph_viewer, size_t& node_idx) { - for (int i = 0; i < graph_viewer.MaxNodeIndex(); ++i) { - auto node = graph_viewer.GetNode(i); - if (node != nullptr && node->OpType() == EPCONTEXT_OP) { - // Only match EPContext nodes that belong to this EP. - // If the source attribute is present and doesn't match, skip the node. - const auto& attrs = node->GetAttributes(); - if (attrs.count(SOURCE) > 0 && - attrs.at(SOURCE).s() != kNvTensorRTRTXExecutionProvider) { - continue; - } - node_idx = i; - return true; - } - } - return false; -} - -const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer) { - // find the top level graph - const Graph* cur_graph = &graph_viewer.GetGraph(); - while (cur_graph->IsSubgraph()) { - cur_graph = cur_graph->ParentGraph(); - } - - const Graph& main_graph = *cur_graph; - return main_graph.ModelPath(); -} - -/* - * Update ep_cache_context attribute of the EP context node with the given engine binary data - */ -void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, - char* engine_data, - size_t size) { - ONNX_NAMESPACE::GraphProto* graph_proto = model_proto->mutable_graph(); - ONNX_NAMESPACE::NodeProto* node_proto = graph_proto->mutable_node(0); - - for (int i = 0; i < node_proto->attribute_size(); ++i) { - ONNX_NAMESPACE::AttributeProto* attribute_proto = node_proto->mutable_attribute(i); - if (attribute_proto->name() == EP_CACHE_CONTEXT) { - std::string engine_data_str = ""; - if (size > 0) { - engine_data_str.assign(engine_data, size); - } - attribute_proto->set_s(engine_data_str); - } - } -} - -/* - * Create EP context node where engine information is embedded - */ -Status CreateCtxNode(const GraphViewer& graph_viewer, - Graph& graph_build, - const std::string engine_cache_path, - char* engine_data, - size_t size, - const int64_t embed_mode, - const std::string compute_capability, - const std::string onnx_model_path, - const std::string& ep_context_node_name, - int32_t trt_version) { - // Get graph inputs and outputs - std::vector inputs, outputs; - for (auto input : graph_viewer.GetInputs()) { - auto& n_input = graph_build.GetOrCreateNodeArg(input->Name(), input->TypeAsProto()); - inputs.push_back(&n_input); - } - - for (auto output : graph_viewer.GetOutputs()) { - auto& n_output = graph_build.GetOrCreateNodeArg(output->Name(), output->TypeAsProto()); - outputs.push_back(&n_output); - } - - // Create EP context node attributes - auto attr_embed_mode = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_main_context = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_ep_cache_context = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_sdk_version = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_hw_architecture = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_onnx_filename = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_partition_name = ONNX_NAMESPACE::AttributeProto::Create(); - auto attr_source = ONNX_NAMESPACE::AttributeProto::Create(); - std::string engine_data_str = ""; - attr_main_context->set_name(MAIN_CONTEXT); - attr_main_context->set_type(onnx::AttributeProto_AttributeType_INT); - attr_main_context->set_i(0); // we do not support a main context node but each has it's own engine payload - attr_embed_mode->set_name(EMBED_MODE); - attr_embed_mode->set_type(onnx::AttributeProto_AttributeType_INT); - attr_embed_mode->set_i(embed_mode); - attr_ep_cache_context->set_name(EP_CACHE_CONTEXT); - attr_ep_cache_context->set_type(onnx::AttributeProto_AttributeType_STRING); - if (embed_mode) { - if (size > 0) { - engine_data_str.assign(engine_data, size); - } - attr_ep_cache_context->set_s(engine_data_str); - } else { - std::string engine_cache_filename = PathToUTF8String(std::filesystem::path(engine_cache_path).filename().native()); - attr_ep_cache_context->set_s(engine_cache_filename); - std::fstream engine_cache_file(engine_cache_path, std::ios::binary | std::ios::out); - if (engine_cache_file.is_open()) { - engine_cache_file.write(engine_data, size); - engine_cache_file.close(); - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "NvTensorRTRTX EP could not write cache to ", engine_cache_path); - } - } - - attr_hw_architecture->set_name(COMPUTE_CAPABILITY); - attr_hw_architecture->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_hw_architecture->set_s(compute_capability); - - attr_partition_name->set_name(PARTITION_NAME); - attr_partition_name->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_partition_name->set_s(ep_context_node_name); // includes hash of the subgraph that was built - - attr_onnx_filename->set_name(ONNX_MODEL_FILENAME); - attr_onnx_filename->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_onnx_filename->set_s(PathToUTF8String(std::filesystem::path(onnx_model_path).filename().native())); - - attr_sdk_version->set_name(SDK_VERSION); - attr_sdk_version->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_sdk_version->set_s(std::to_string(trt_version)); - - attr_source->set_name(SOURCE); - attr_source->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_source->set_s(kNvTensorRTRTXExecutionProvider); - - auto node_attributes = ONNX_NAMESPACE::NodeAttributes::Create(); - constexpr int num_attributes = 8; - node_attributes->reserve(num_attributes); - node_attributes->emplace(MAIN_CONTEXT, *attr_main_context); - node_attributes->emplace(EMBED_MODE, *attr_embed_mode); - node_attributes->emplace(EP_CACHE_CONTEXT, *attr_ep_cache_context); - node_attributes->emplace(COMPUTE_CAPABILITY, *attr_hw_architecture); - node_attributes->emplace(PARTITION_NAME, *attr_partition_name); - node_attributes->emplace(ONNX_MODEL_FILENAME, *attr_onnx_filename); - node_attributes->emplace(SDK_VERSION, *attr_sdk_version); - node_attributes->emplace(SOURCE, *attr_source); - - // Create EP context node - graph_build.AddNode(ep_context_node_name, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); - ORT_ENFORCE(graph_build.Resolve().IsOK()); - return Status::OK(); -} - -Status TensorRTCacheModelHandler::GetEpContextFromGraph(const Node& node) { - auto& attrs = node.GetAttributes(); - - const int64_t embed_mode = attrs.at(EMBED_MODE).i(); - - if (embed_mode) { - // Get engine from byte stream. - const std::string& context_binary = attrs.at(EP_CACHE_CONTEXT).s(); - *(trt_engine_) = std::unique_ptr(trt_runtime_->deserializeCudaEngine(const_cast(context_binary.c_str()), - static_cast(context_binary.length()))); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Read engine as binary data from \"ep_cache_context\" attribute of ep context node and deserialized it"; - if (!(*trt_engine_)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP could not deserialize engine from binary data"); - } - - if (weight_stripped_engine_refit_) { - const std::string onnx_model_filename = attrs.at(ONNX_MODEL_FILENAME).s(); - auto status = NvExecutionProvider::RefitEngine(onnx_model_filename, - onnx_model_folder_path_, - onnx_model_bytestream_, - onnx_model_bytestream_size_, - onnx_external_data_bytestream_, - onnx_external_data_bytestream_size_, - (*trt_engine_).get(), - detailed_build_log_); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - } else { - // Get engine from cache file. - std::string cache_path = attrs.at(EP_CACHE_CONTEXT).s(); - - // Validate that the cache path does not escape the model directory. - // Rejects absolute paths, ".." traversal, and symlink-based escapes. - std::filesystem::path ctx_model_dir(path_utils::GetDirOrParentPath(ep_context_model_path_)); - ORT_RETURN_IF_ERROR(utils::ValidateExternalDataPathFromDir(ctx_model_dir, std::filesystem::path(cache_path))); - - // The engine cache and context model (current model) should be in the same directory - auto engine_cache_path = ctx_model_dir.append(cache_path); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] GetEpContextFromGraph engine_cache_path: " + engine_cache_path.string(); - - if (!std::filesystem::exists(engine_cache_path)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP can't find engine cache: " + engine_cache_path.string() + - ". Please make sure engine cache is in the same directory or sub-directory of context model."); - } - - size_t file_length = 0; - auto path_str = ToPathString(engine_cache_path.string()); - - Env::MappedMemoryPtr engine_buf; - const auto& env = GetDefaultEnv(); - ORT_RETURN_IF_ERROR(env.GetFileLength(path_str.c_str(), file_length)); - if (!file_length) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP could not read engine from cache: " + engine_cache_path.string()); - } - ORT_RETURN_IF_ERROR(env.MapFileIntoMemory(path_str.c_str(), 0, file_length, engine_buf)); - - *(trt_engine_) = std::unique_ptr(trt_runtime_->deserializeCudaEngine(engine_buf.get(), file_length)); - if (!(*trt_engine_)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP could not deserialize engine from cache: " + engine_cache_path.string()); - } - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] DeSerialized " + engine_cache_path.string(); - - if (weight_stripped_engine_refit_) { - const std::string onnx_model_filename = attrs.at(ONNX_MODEL_FILENAME).s(); - std::string weight_stripped_engine_cache = engine_cache_path.string(); - auto status = NvExecutionProvider::RefitEngine(onnx_model_filename, - onnx_model_folder_path_, - onnx_model_bytestream_, - onnx_model_bytestream_size_, - onnx_external_data_bytestream_, - onnx_external_data_bytestream_size_, - (*trt_engine_).get(), - detailed_build_log_); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); - } - } - } - return Status::OK(); -} - -/* - * The sanity check for EP context contrib op. - */ -bool TensorRTCacheModelHandler::ValidateEPCtxNode(const Node& node) { - auto& attrs = node.GetAttributes(); - - // Show the warning if compute capability is not matched - if (attrs.count(COMPUTE_CAPABILITY) > 0) { - std::string model_compute_capability = attrs.at(COMPUTE_CAPABILITY).s(); - // Verify if engine was compiled with ampere+ hardware compatibility enabled - if (model_compute_capability == "80+") { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine is compatible to all Ampere+ GPU (except Jetson)"; - if (std::stoi(compute_capability_) < 80) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] However, this GPU doesn't match. The compute capability of the GPU: " << compute_capability_; - } - } else if (model_compute_capability != compute_capability_) { - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Engine was compiled for a different compatibility level and might not work or perform suboptimal"; - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] The compute capability of the engine: " << model_compute_capability; - LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] The compute capability of the GPU: " << compute_capability_; - } - } - - // "embed_mode" attr and "ep_cache_context" attr should be present - assert(attrs.count(EMBED_MODE) > 0); - assert(attrs.count(EP_CACHE_CONTEXT) > 0); - - const int64_t embed_mode = attrs.at(EMBED_MODE).i(); - if (embed_mode == 1) { - // engine binary data - // LOGS_DEFAULT(WARNING) << EPCONTEXT_WARNING; - } - - return true; -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h deleted file mode 100644 index 120f36433d0a8..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ /dev/null @@ -1,93 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include -#include -#include -#include - -#include "core/providers/nv_tensorrt_rtx/nv_includes.h" -#include "core/providers/shared_library/provider_api.h" - -namespace onnxruntime { - -static const std::string EPCONTEXT_OP = "EPContext"; -static const std::string MAIN_CONTEXT = "main_context"; -static const std::string EMBED_MODE = "embed_mode"; -static const std::string EP_CACHE_CONTEXT = "ep_cache_context"; -static const std::string COMPUTE_CAPABILITY = "hardware_architecture"; -static const std::string ONNX_MODEL_FILENAME = "onnx_model_filename"; -static const std::string PARTITION_NAME = "partition_name"; -static const std::string SDK_VERSION = "ep_sdk_version"; -static const std::string SOURCE = "source"; -static const std::string EPCONTEXT_OP_DOMAIN = "com.microsoft"; - -// TensorRT does not currently expose a header size define; keep in sync with TRT engine serialization header size. -constexpr size_t kTensorRTEngineHeaderSize = 64; -// Helper functions for engine header validation -std::string BinaryToHexString(const void* data, size_t size); -std::vector HexStringToBinary(const std::string& hex); -bool GraphHasCtxNode(const GraphViewer& graph_viewer, size_t& node_idx); -const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer); -Status CreateCtxNode(const GraphViewer& graph_viewer, - Graph& graph_build, - const std::string engine_cache_path, - char* engine_data, - size_t size, - const int64_t embed_mode, - const std::string compute_capability, - const std::string onnx_model_path, - const std::string& ep_context_node_name, - int trt_version); -void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, - char* engine_data, - size_t size); - -class TensorRTCacheModelHandler { - public: - TensorRTCacheModelHandler(std::unique_ptr* trt_engine, - nvinfer1::IRuntime* trt_runtime, - std::string ep_context_model_path, - std::string compute_capability, - bool weight_stripped_engine_refit, - std::string onnx_model_folder_path, - const void* onnx_model_bytestream, - size_t onnx_model_bytestream_size, - const void* onnx_external_data_bytestream, - size_t onnx_external_data_bytestream_size, - bool detailed_build_log) - : trt_engine_(trt_engine), - trt_runtime_(trt_runtime), - ep_context_model_path_(ep_context_model_path), - compute_capability_(compute_capability), - weight_stripped_engine_refit_(weight_stripped_engine_refit), - onnx_model_folder_path_(onnx_model_folder_path), - onnx_model_bytestream_(onnx_model_bytestream), - onnx_model_bytestream_size_(onnx_model_bytestream_size), - onnx_external_data_bytestream_(onnx_external_data_bytestream), - onnx_external_data_bytestream_size_(onnx_external_data_bytestream_size), - detailed_build_log_(detailed_build_log) { - } - ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TensorRTCacheModelHandler); - - bool ValidateEPCtxNode(const Node& node); - - Status GetEpContextFromGraph(const Node& node); - - private: - std::unique_ptr* trt_engine_; - nvinfer1::IRuntime* trt_runtime_; - std::string ep_context_model_path_; // If using context model, it implies context model and engine cache is in the same directory - std::string compute_capability_; - bool weight_stripped_engine_refit_; - std::string onnx_model_folder_path_; - const void* onnx_model_bytestream_; - size_t onnx_model_bytestream_size_; - const void* onnx_external_data_bytestream_; - size_t onnx_external_data_bytestream_size_; - bool detailed_build_log_; -}; // TRTCacheModelHandler -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/symbols.def b/onnxruntime/core/providers/nv_tensorrt_rtx/symbols.def deleted file mode 100644 index 3afed01da1966..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/symbols.def +++ /dev/null @@ -1,4 +0,0 @@ -EXPORTS - GetProvider - CreateEpFactories - ReleaseEpFactory diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/symbols.txt b/onnxruntime/core/providers/nv_tensorrt_rtx/symbols.txt deleted file mode 100644 index e69de29bb2d1d..0000000000000 diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/version_script.lds b/onnxruntime/core/providers/nv_tensorrt_rtx/version_script.lds deleted file mode 100644 index 251e39e089275..0000000000000 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/version_script.lds +++ /dev/null @@ -1,11 +0,0 @@ -#_init and _fini should be local -VERS_1.0 { - global: - GetProvider; - CreateEpFactories; - ReleaseEpFactory; - - # Hide everything else. - local: - *; -}; diff --git a/onnxruntime/core/providers/provider_factory_creators.h b/onnxruntime/core/providers/provider_factory_creators.h index 987385ace2b2d..f7c8dc1e9b8f9 100644 --- a/onnxruntime/core/providers/provider_factory_creators.h +++ b/onnxruntime/core/providers/provider_factory_creators.h @@ -70,10 +70,6 @@ #include "core/providers/tensorrt/tensorrt_provider_factory_creator.h" #endif -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) -#include "core/providers/nv_tensorrt_rtx/nv_provider_factory_creator.h" -#endif - #if defined(USE_VITISAI) || defined(USE_VITISAI_PROVIDER_INTERFACE) #include "core/providers/vitisai/vitisai_provider_factory_creator.h" #endif diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index a13938922bc54..64e9fa737acdf 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/onnxruntime/core/providers/shared_library/provider_api.h @@ -313,7 +313,6 @@ constexpr const char* kDnnlExecutionProvider = "DnnlExecutionProvider"; constexpr const char* kOpenVINOExecutionProvider = "OpenVINOExecutionProvider"; constexpr const char* kVitisAIExecutionProvider = "VitisAIExecutionProvider"; constexpr const char* kTensorrtExecutionProvider = "TensorrtExecutionProvider"; -constexpr const char* kNvTensorRTRTXExecutionProvider = "NvTensorRTRTXExecutionProvider"; constexpr const char* kMIGraphXExecutionProvider = "MIGraphXExecutionProvider"; constexpr const char* kQnnExecutionProvider = "QNNExecutionProvider"; constexpr const char* kCpuExecutionProvider = "CPUExecutionProvider"; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 520103e687241..c1b239ef96a6b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc @@ -337,7 +337,7 @@ WebGpuContextConfig ParseWebGpuContextConfig( // Compile-only mode (skip Dawn adapter/device creation so graph transformation can run device-free) // is derived from the session config kOrtSessionOptionCompileOnly, which the Compile API sets - // automatically -- same signal other EPs use (e.g. NV TensorRT RTX). Not a WebGPU-specific option. + // automatically. This is not a WebGPU-specific option. config.compile_only = config_options.GetConfigOrDefault(kOrtSessionOptionCompileOnly, "0") == "1"; std::string max_storage_buffer_binding_size_str; diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index e9015b3e40b09..8e46524d89831 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -3024,7 +3024,7 @@ common::Status InferenceSession::Initialize() { return Status::OK(); }; - // Enable DQ->MatMulNBits fusion if NvTensorRTRTX EP is registered. + // Enable DQ->MatMulNBits fusion if the NvTensorRTRTX plugin EP is registered. if (execution_providers_.Get(onnxruntime::kNvTensorRTRTXExecutionProvider) != nullptr) { if (session_options_.config_options.GetConfigOrDefault( kOrtSessionOptionsEnableDQMatMulNBitsFusion, "") == "") { diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 88d15124e09a1..3a9a6f4c7e4ec 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -104,7 +104,6 @@ using EtwRegistrationManager_EtwInternalCallback = EtwRegistrationManager::EtwIn #include "core/providers/migraphx/migraphx_provider_factory_creator.h" #include "core/providers/openvino/openvino_provider_factory_creator.h" #include "core/providers/tensorrt/tensorrt_provider_factory_creator.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_factory_creator.h" #include "core/providers/vitisai/vitisai_provider_factory_creator.h" #include "core/providers/qnn/qnn_provider_factory_creator.h" @@ -118,16 +117,13 @@ using EtwRegistrationManager_EtwInternalCallback = EtwRegistrationManager::EtwIn #include "core/providers/cuda/cuda_provider_options.h" #include "core/providers/cann/cann_provider_options.h" #include "core/providers/dnnl/dnnl_provider_options.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_factory.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_options.h" #if defined(_WIN32) && !defined(NDEBUG) && defined(ONNXRUNTIME_ENABLE_MEMLEAK_CHECK) #include "core/platform/windows/debug_alloc.h" #endif -#if !defined(ORT_MINIMAL_BUILD) && \ - (defined(USE_TENSORRT) || defined(USE_TENSORRT_PROVIDER_INTERFACE) || \ - defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE)) +#if !defined(ORT_MINIMAL_BUILD) && \ + (defined(USE_TENSORRT) || defined(USE_TENSORRT_PROVIDER_INTERFACE)) #include "core/session/onnxruntime_session_options_config_keys.h" #endif @@ -160,8 +156,6 @@ ProviderInfo_Dnnl& GetProviderInfo_Dnnl(); ProviderHostCPU& GetProviderHostCPU(); ProviderInfo_MIGraphX* TryGetProviderInfo_MIGraphX(); ProviderInfo_MIGraphX& GetProviderInfo_MIGraphX(); -ProviderInfo_Nv* TryGetProviderInfo_Nv(); -ProviderInfo_Nv& GetProviderInfo_Nv(); ProviderInfo_OpenVINO* TryGetProviderInfo_OpenVINO(); ONNX_NAMESPACE::OpSchema CreateSchema(const std::string& domain, const std::vector& ops); @@ -2130,7 +2124,6 @@ static ProviderLibrary s_library_tensorrt(LIBRARY_PREFIX ORT_TSTR("onnxruntime_p false #endif ); -static ProviderLibrary s_library_nv(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_nv_tensorrt_rtx") LIBRARY_EXTENSION); static ProviderLibrary s_library_migraphx(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_migraphx") LIBRARY_EXTENSION); // QNN EP can be built either as a static library or a shared library. Can safely define s_library_qnn even if static. @@ -2147,7 +2140,6 @@ void UnloadSharedProviders() { s_library_shared.Unload(); s_library_migraphx.Unload(); s_library_qnn.Unload(); - s_library_nv.Unload(); } // Used by test code @@ -2296,28 +2288,6 @@ std::shared_ptr TensorrtProviderFactoryCreator::Creat ORT_CATCH_LOG_RETURN_NULLPTR; } -std::shared_ptr NvProviderFactoryCreator::Create(int device_id) { - ORT_TRY { - return s_library_nv.Get().CreateExecutionProviderFactory(device_id); - } - ORT_CATCH_LOG_RETURN_NULLPTR; -} - -std::shared_ptr NvProviderFactoryCreator::Create( - const ProviderOptions& provider_options, const SessionOptions* session_options) { - ORT_TRY { - const ConfigOptions* config_options = nullptr; - if (session_options != nullptr) { - config_options = &session_options->config_options; - } - - std::array configs_array = {&provider_options, config_options}; - const void* arg = reinterpret_cast(&configs_array); - return s_library_nv.Get().CreateExecutionProviderFactory(arg); - } - ORT_CATCH_LOG_RETURN_NULLPTR; -} - std::shared_ptr MIGraphXProviderFactoryCreator::Create(const ProviderOptions& provider_options) { return s_library_migraphx.Get().CreateExecutionProviderFactory(&provider_options); } @@ -2437,20 +2407,6 @@ ProviderInfo_TensorRT& GetProviderInfo_TensorRT() { ORT_THROW("TensorRT Provider not available, can't get interface for it"); } -ProviderInfo_Nv* TryGetProviderInfo_Nv() { - ORT_TRY { - return reinterpret_cast(s_library_nv.Get().GetInfo()); - } - ORT_CATCH_LOG_RETURN_NULLPTR; -} - -ProviderInfo_Nv& GetProviderInfo_Nv() { - if (auto* info = TryGetProviderInfo_Nv()) - return *info; - - ORT_THROW("NV Provider not available, can't get interface for it"); -} - ProviderInfo_CUDA* TryGetProviderInfo_CUDA() { ORT_TRY { return reinterpret_cast(s_library_cuda.Get().GetInfo()); diff --git a/onnxruntime/core/session/provider_registration.cc b/onnxruntime/core/session/provider_registration.cc index 9fef7a611075e..2055894e8aa11 100644 --- a/onnxruntime/core/session/provider_registration.cc +++ b/onnxruntime/core/session/provider_registration.cc @@ -27,9 +27,6 @@ #include "core/providers/dml/dml_provider_factory_creator.h" #endif -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) -#include "core/providers/nv_tensorrt_rtx/nv_provider_options.h" -#endif using namespace onnxruntime; namespace onnxruntime { @@ -101,7 +98,6 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, JS, VitisAI, CoreML, - NvTensorRtRtx, // TensorRt EP for RTX GPUs. MIGraphX, CPU }; @@ -112,7 +108,7 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, const char* canonical_name = nullptr; }; - static std::array supported_eps = { + static std::array supported_eps = { EpToAppend{EpID::DML, "DML", kDmlExecutionProvider}, EpToAppend{EpID::QNN, "QNN", kQnnExecutionProvider}, EpToAppend{EpID::OpenVINO, "OpenVINO", kOpenVINOExecutionProvider}, @@ -124,7 +120,6 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, EpToAppend{EpID::JS, "JS", kJsExecutionProvider}, EpToAppend{EpID::VitisAI, "VitisAI", kVitisAIExecutionProvider}, EpToAppend{EpID::CoreML, "CoreML", kCoreMLExecutionProvider}, - EpToAppend{EpID::NvTensorRtRtx, "NvTensorRtRtx", kNvTensorRTRTXExecutionProvider}, EpToAppend{EpID::MIGraphX, "MIGraphX", kMIGraphXExecutionProvider}, EpToAppend{EpID::CPU, "CPU", kCpuExecutionProvider}}; @@ -184,6 +179,20 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, str_builder.str().c_str()); }; +#if !defined(ORT_MINIMAL_BUILD) + // The provider-bridge NvTensorRtRtx EP is deprecated. Report the plugin EP flow with the deprecation error. + // Minimal builds never included this EP, so they keep the generic unknown-provider error to save binary size. + if (strcmp(provider_name, "NvTensorRtRtx") == 0 || strcmp(provider_name, kNvTensorRTRTXExecutionProvider) == 0) { + return OrtApis::CreateStatus( + ORT_INVALID_ARGUMENT, + (std::string("The '") + provider_name + + "' execution provider is no longer built into ONNX Runtime. Use the NVIDIA TensorRT RTX plugin EP " + "instead: register its library with RegisterExecutionProviderLibrary, select an OrtEpDevice from " + "GetEpDevices, and append it with SessionOptionsAppendExecutionProvider_V2.") + .c_str()); + } +#endif // !defined(ORT_MINIMAL_BUILD) + auto ep_to_append_iter = std::find_if(supported_eps.begin(), supported_eps.end(), [&provider_name](const EpToAppend& elem) -> bool { return (strcmp(provider_name, elem.short_name) == 0) || @@ -317,19 +326,6 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, options->provider_factories.push_back(CoreMLProviderFactoryCreator::Create(provider_options)); #else status = create_not_supported_status(); -#endif - break; - } - case EpID::NvTensorRtRtx: { -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) - auto factory = onnxruntime::NvProviderFactoryCreator::Create(provider_options, &(options->value)); - if (factory) { - options->provider_factories.push_back(factory); - } else { - status = create_failed_to_load_provider_status(); - } -#else - status = create_not_supported_status(); #endif break; } diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index e3597b6cce158..4ef834862c0ec 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -700,58 +700,22 @@ def __init__( def _create_inference_session(self, providers, provider_options, disabled_optimizers=None): available_providers = C.get_available_providers() - # Validate that TensorrtExecutionProvider and NvTensorRTRTXExecutionProvider are not both specified - if providers: - has_tensorrt = any( - provider == "TensorrtExecutionProvider" - or (isinstance(provider, tuple) and provider[0] == "TensorrtExecutionProvider") - for provider in providers - ) - has_tensorrt_rtx = any( - provider == "NvTensorRTRTXExecutionProvider" - or (isinstance(provider, tuple) and provider[0] == "NvTensorRTRTXExecutionProvider") - for provider in providers + provider_names = [provider[0] if isinstance(provider, tuple) else provider for provider in providers or []] + has_tensorrt = "TensorrtExecutionProvider" in provider_names + # NvTensorRTRTXExecutionProvider is only available as a plugin EP registered by the application. + has_tensorrt_rtx = "NvTensorRTRTXExecutionProvider" in provider_names + if has_tensorrt and has_tensorrt_rtx: + raise ValueError( + "Cannot enable both 'TensorrtExecutionProvider' and 'NvTensorRTRTXExecutionProvider' " + "in the same session." ) - if has_tensorrt and has_tensorrt_rtx: - raise ValueError( - "Cannot enable both 'TensorrtExecutionProvider' and 'NvTensorRTRTXExecutionProvider' " - "in the same session." - ) - # Tensorrt and TensorRT RTX can fall back to CUDA if it's explicitly assigned. All others fall back to CPU. - if "NvTensorRTRTXExecutionProvider" in available_providers: - if ( - providers - and any( - provider == "CUDAExecutionProvider" - or (isinstance(provider, tuple) and provider[0] == "CUDAExecutionProvider") - for provider in providers - ) - and any( - provider == "NvTensorRTRTXExecutionProvider" - or (isinstance(provider, tuple) and provider[0] == "NvTensorRTRTXExecutionProvider") - for provider in providers - ) - ): - self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] - else: - self._fallback_providers = ["CPUExecutionProvider"] - elif "TensorrtExecutionProvider" in available_providers: - if ( - providers - and any( - provider == "CUDAExecutionProvider" - or (isinstance(provider, tuple) and provider[0] == "CUDAExecutionProvider") - for provider in providers - ) - and any( - provider == "TensorrtExecutionProvider" - or (isinstance(provider, tuple) and provider[0] == "TensorrtExecutionProvider") - for provider in providers - ) - ): - self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] - else: - self._fallback_providers = ["CPUExecutionProvider"] + + # TensorRT and TensorRT RTX can fall back to CUDA if it's explicitly assigned. All others fall back to CPU. + if "CUDAExecutionProvider" in provider_names and any( + name in provider_names and name in available_providers + for name in ("TensorrtExecutionProvider", "NvTensorRTRTXExecutionProvider") + ): + self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] else: self._fallback_providers = ["CPUExecutionProvider"] @@ -848,15 +812,6 @@ def _register_ep_custom_ops(self, session_options, providers, provider_options, ): C.register_tensorrt_plugins_as_custom_ops(session_options, providers[i][1]) - if providers[i] in available_providers and providers[i] == "NvTensorRTRTXExecutionProvider": - C.register_nv_tensorrt_rtx_plugins_as_custom_ops(session_options, provider_options[i]) - elif ( - isinstance(providers[i], tuple) - and providers[i][0] in available_providers - and providers[i][0] == "NvTensorrtRTXExecutionProvider" - ): - C.register_nv_tensorrt_rtx_plugins_as_custom_ops(session_options, providers[i][1]) - def make_get_initializer_location_func_wrapper( get_initializer_location_func: GetInitializerLocationFunc, diff --git a/onnxruntime/python/onnxruntime_pybind_schema.cc b/onnxruntime/python/onnxruntime_pybind_schema.cc index 1b5b888110e27..2a72e36987be5 100644 --- a/onnxruntime/python/onnxruntime_pybind_schema.cc +++ b/onnxruntime/python/onnxruntime_pybind_schema.cc @@ -46,9 +46,6 @@ void addGlobalSchemaFunctions(pybind11::module& m) { #ifdef USE_TENSORRT onnxruntime::TensorrtProviderFactoryCreator::Create(0), #endif -#ifdef USE_NV - onnxruntime::NvProviderFactoryCreator::Create(0), -#endif #ifdef USE_MIGRAPHX onnxruntime::MIGraphXProviderFactoryCreator::Create(0), #endif diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 1211899ceb103..641ad98c1afea 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -731,38 +731,6 @@ void RegisterTensorRTPluginsAsCustomOps(PySessionOptions& so, const ProviderOpti } #endif -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) -void RegisterNvTensorRTRtxPluginsAsCustomOps(PySessionOptions& so, const ProviderOptions& options) { - if (auto* nv_tensorrt_rtx_provider_info = TryGetProviderInfo_Nv()) { - auto is_already_in_domains = [&](std::string& domain_name, std::vector& domains) { - for (auto ptr : domains) { - if (domain_name == ptr->domain_) { - return true; - } - } - return false; - }; - - std::string extra_plugin_lib_paths = ""; - const auto it = options.find("extra_plugin_lib_paths"); - if (it != options.end()) { - extra_plugin_lib_paths = it->second; - } - std::vector custom_op_domains; - nv_tensorrt_rtx_provider_info->GetTensorRTCustomOpDomainList(custom_op_domains, extra_plugin_lib_paths); - for (auto ptr : custom_op_domains) { - if (!is_already_in_domains(ptr->domain_, so.custom_op_domains_)) { - so.custom_op_domains_.push_back(ptr); - } else { - LOGS_DEFAULT(WARNING) << "The custom op domain name " << ptr->domain_ << " is already in session option."; - } - } - } else { - ORT_THROW("Please install TensorRT libraries as mentioned in the GPU requirements page, make sure they're in the PATH or LD_LIBRARY_PATH, and that your GPU is supported."); - } -} -#endif - #if !defined(ORT_MINIMAL_BUILD) // Find a registered plugin EP device matching the given EP name and optional device_id from provider options. // Returns nullptr if no matching device is found. @@ -1236,28 +1204,6 @@ static std::shared_ptr CreateExecutionProviderFactory << "TensorRT-ExecutionProvider.html#requirements to ensure all dependencies are met."; #endif - } else if (type == kNvTensorRTRTXExecutionProvider) { -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) - if (Env::Default().GetEnvironmentVar("ORT_NV_TENSORRT_RTX_UNAVAILABLE").empty()) { - auto it = provider_options_map.find(type); - if (it != provider_options_map.end()) { - ProviderOptions info = it->second; - if (std::shared_ptr nv_tensorrt_rtx_provider_factory = onnxruntime::NvProviderFactoryCreator::Create( - info, &session_options)) { - return nv_tensorrt_rtx_provider_factory; - } - } else { - if (std::shared_ptr nv_tensorrt_rtx_provider_factory = onnxruntime::NvProviderFactoryCreator::Create(cuda_device_id)) { - return nv_tensorrt_rtx_provider_factory; - } - } - } - LOGS_DEFAULT(WARNING) << "Failed to create " - << type - << ". Please reference " - << "https://onnxruntime.ai/docs/execution-providers/" - << "TensorRT-ExecutionProvider.html#requirements to ensure all dependencies are met."; -#endif } else if (type == kMIGraphXExecutionProvider) { #if defined(USE_MIGRAPHX) || defined(USE_MIGRAPHX_PROVIDER_INTERFACE) auto it = provider_options_map.find(type); @@ -2053,12 +1999,6 @@ Returns a dictionary with: "Register TensorRT plugins as custom ops."); #endif -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) - m.def( - "register_nv_tensorrt_rtx_plugins_as_custom_ops", [](PySessionOptions& so, const ProviderOptions& options) { RegisterNvTensorRTRtxPluginsAsCustomOps(so, options); }, - "Register NV TensorRT RTX plugins as custom ops."); -#endif - #ifdef ENABLE_ATEN m.def("register_aten_op_executor", [](const std::string& is_tensor_argument_address_str, const std::string& aten_op_executor_address_str) -> void { @@ -2195,7 +2135,7 @@ void addObjectMethods(py::module& m, ExecutionProviderRegistrationFn ep_registra type = OrtDevice::GPU; vendor = OrtDevice::VendorIds::MICROSOFT; } else if (type == OrtDevice::GPU) { -#if USE_CUDA || USE_NV || USE_NV_PROVIDER_INTERFACE || USE_CUDA_PROVIDER_INTERFACE +#if USE_CUDA || USE_CUDA_PROVIDER_INTERFACE vendor = OrtDevice::VendorIds::NVIDIA; #elif USE_MIGRAPHX vendor = OrtDevice::VendorIds::AMD; diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 71607fefbfdb3..012034408d36d 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -33,7 +33,6 @@ struct OrtStatus { #include "core/providers/providers.h" #include "core/providers/provider_factory_creators.h" #include "core/providers/tensorrt/tensorrt_provider_options.h" -#include "core/providers/nv_tensorrt_rtx/nv_provider_options.h" #if defined(USE_CUDA) #define BACKEND_PROC "GPU" @@ -120,9 +119,6 @@ struct OrtStatus { #if defined(USE_TENSORRT) || defined(USE_TENSORRT_PROVIDER_INTERFACE) #include "core/providers/tensorrt/tensorrt_provider_factory.h" #endif -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) -#include "core/providers/nv_tensorrt_rtx/nv_provider_factory.h" -#endif #if defined(USE_MIGRAPHX) || defined(USE_MIGRAPHX_PROVIDER_INTERFACE) #include "core/providers/migraphx/migraphx_provider_factory.h" #include "core/providers/migraphx/migraphx_execution_provider_info.h" @@ -171,13 +167,6 @@ ProviderInfo_TensorRT& GetProviderInfo_TensorRT(); } // namespace onnxruntime #endif -#if defined(USE_NV) || defined(USE_NV_PROVIDER_INTERFACE) -namespace onnxruntime { -ProviderInfo_Nv* TryGetProviderInfo_Nv(); -ProviderInfo_Nv& GetProviderInfo_Nv(); -} // namespace onnxruntime -#endif - #ifdef USE_CANN namespace onnxruntime { ProviderInfo_CANN* TryGetProviderInfo_CANN(); diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_custom_op.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep_custom_op.h index c37038a727067..63a97e14e9821 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_custom_op.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_custom_op.h @@ -24,7 +24,7 @@ // - In Compile(), the EP executes its compiled bits to perform inference for // the fused custom node. // -// Note: Approach #2 is suitable for plugin TRT RTX EP to support TRT plugins. +// Note: Approach #2 is suitable for plugin EPs that support custom operator libraries. struct CustomMulKernel : MulKernel { CustomMulKernel(const OrtApi& ort_api, diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc index 2cf22da327a2a..ca82b70c59a75 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc @@ -89,6 +89,9 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL CreateExternalResourceImporterForDevice = CreateExternalResourceImporterForDeviceImpl; + InitGraphicsInterop = InitGraphicsInteropImpl; + DeinitGraphicsInterop = DeinitGraphicsInteropImpl; + GetNumCustomOpDomains = GetNumCustomOpDomainsImpl; GetCustomOpDomains = GetCustomOpDomainsImpl; ValidateCompiledModelCompatibilityInfo = ValidateCompiledModelCompatibilityInfoImpl; @@ -611,6 +614,44 @@ OrtStatus* ORT_API_CALL ExampleEpFactory::CreateExternalResourceImporterForDevic return nullptr; } +OrtStatus* ORT_API_CALL ExampleEpFactory::InitGraphicsInteropImpl(OrtEpFactory* this_ptr, + const OrtEpDevice* /*ep_device*/, + const OrtGraphicsInteropConfig* config) noexcept { + auto& factory = *static_cast(this_ptr); + + // OrtGraphicsInteropConfig was added in ORT API version 25. + if (config->version < 25) { + return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, + "ExampleEP: OrtGraphicsInteropConfig version must be at least 25"); + } + + // The example EP has no real graphics device; it only validates the requested API and tracks state. + if (config->graphics_api != ORT_GRAPHICS_API_D3D12 && config->graphics_api != ORT_GRAPHICS_API_VULKAN) { + return factory.ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "ExampleEP: unsupported graphics API"); + } + + std::lock_guard lock{factory.mutex_}; + if (factory.graphics_interop_initialized_) { + return factory.ort_api.CreateStatus(ORT_FAIL, "ExampleEP: graphics interop is already initialized"); + } + + factory.graphics_interop_initialized_ = true; + return nullptr; +} + +OrtStatus* ORT_API_CALL ExampleEpFactory::DeinitGraphicsInteropImpl(OrtEpFactory* this_ptr, + const OrtEpDevice* /*ep_device*/) noexcept { + auto& factory = *static_cast(this_ptr); + + std::lock_guard lock{factory.mutex_}; + if (!factory.graphics_interop_initialized_) { + return factory.ort_api.CreateStatus(ORT_FAIL, "ExampleEP: graphics interop is not initialized"); + } + + factory.graphics_interop_initialized_ = false; + return nullptr; +} + OrtStatus* ORT_API_CALL ExampleEpFactory::ValidateCompiledModelCompatibilityInfoImpl( OrtEpFactory* this_ptr, const OrtHardwareDevice* const* devices, diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h index 20103a7f45594..404958f753a76 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.h @@ -91,6 +91,13 @@ class ExampleEpFactory : public OrtEpFactory, public ApiPtrs { const OrtEpDevice* ep_device, OrtExternalResourceImporterImpl** out_importer) noexcept; + static OrtStatus* ORT_API_CALL InitGraphicsInteropImpl(OrtEpFactory* this_ptr, + const OrtEpDevice* ep_device, + const OrtGraphicsInteropConfig* config) noexcept; + + static OrtStatus* ORT_API_CALL DeinitGraphicsInteropImpl(OrtEpFactory* this_ptr, + const OrtEpDevice* ep_device) noexcept; + static OrtStatus* ORT_API_CALL GetHardwareDeviceIncompatibilityDetailsImpl( OrtEpFactory* this_ptr, const OrtHardwareDevice* hw, @@ -133,7 +140,10 @@ class ExampleEpFactory : public OrtEpFactory, public ApiPtrs { bool arena_allocator_using_default_settings_{true}; std::unique_ptr arena_allocator_; // shared device allocator that uses an arena uint32_t num_arena_users_{0}; - std::mutex mutex_; // mutex to protect arena_allocator_ and num_arena_users_ + std::mutex mutex_; // mutex to protect arena_allocator_, num_arena_users_ and graphics_interop_initialized_ + + // The example EP has a single device, so one flag tracks its graphics interop state. + bool graphics_interop_initialized_{false}; std::unique_ptr data_transfer_impl_; // data transfer implementation for this factory diff --git a/onnxruntime/test/autoep/test_graphics_interop.cc b/onnxruntime/test/autoep/test_graphics_interop.cc new file mode 100644 index 0000000000000..2192c7b8a0f3a --- /dev/null +++ b/onnxruntime/test/autoep/test_graphics_interop.cc @@ -0,0 +1,116 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Tests for InitGraphicsInteropForEpDevice / DeinitGraphicsInteropForEpDevice using the example plugin EPs. +// The example_plugin_ep factory implements the graphics interop callbacks without a real graphics device; the +// kernel registry example EP does not implement them. + +#include + +#include + +#include "core/session/onnxruntime_cxx_api.h" + +#include "test/autoep/test_autoep_utils.h" + +extern std::unique_ptr ort_env; + +namespace onnxruntime { +namespace test { + +#if !defined(ORT_MINIMAL_BUILD) + +namespace { + +OrtGraphicsInteropConfig MakeConfig(OrtGraphicsApi graphics_api) { + OrtGraphicsInteropConfig config = {}; + config.version = ORT_API_VERSION; + config.graphics_api = graphics_api; + config.command_queue = nullptr; + config.additional_options = nullptr; + return config; +} + +// Checks the error code without calling GetErrorCode on an OK (null) status. +void ExpectErrorCode(OrtStatus* ort_status, OrtErrorCode expected) { + Ort::Status status(ort_status); + ASSERT_FALSE(status.IsOK()) << "Expected error code " << expected; + EXPECT_EQ(status.GetErrorCode(), expected) << status.GetErrorMessage(); +} + +} // namespace + +TEST(GraphicsInteropTest, InitAndDeinitAreForwardedToEpFactory) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + const OrtEpDevice* ep_device = example_ep.get(); + + const OrtInteropApi& interop_api = Ort::GetInteropApi(); + const OrtGraphicsInteropConfig config = MakeConfig(ORT_GRAPHICS_API_D3D12); + + // The example EP rejects deinit before init, so this error shows the call reached the factory. + { + Ort::Status status(interop_api.DeinitGraphicsInteropForEpDevice(ep_device)); + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.GetErrorCode(), ORT_FAIL); + EXPECT_NE(status.GetErrorMessage().find("not initialized"), std::string::npos) << status.GetErrorMessage(); + } + + { + Ort::Status status(interop_api.InitGraphicsInteropForEpDevice(ep_device, &config)); + ASSERT_TRUE(status.IsOK()) << status.GetErrorMessage(); + } + + { + Ort::Status status(interop_api.InitGraphicsInteropForEpDevice(ep_device, &config)); + ASSERT_FALSE(status.IsOK()); + EXPECT_NE(status.GetErrorMessage().find("already initialized"), std::string::npos) << status.GetErrorMessage(); + } + + { + Ort::Status status(interop_api.DeinitGraphicsInteropForEpDevice(ep_device)); + ASSERT_TRUE(status.IsOK()) << status.GetErrorMessage(); + } +} + +TEST(GraphicsInteropTest, EpFactoryErrorIsReturned) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + + const OrtGraphicsInteropConfig config = MakeConfig(ORT_GRAPHICS_API_NONE); + Ort::Status status(Ort::GetInteropApi().InitGraphicsInteropForEpDevice(example_ep.get(), &config)); + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.GetErrorCode(), ORT_INVALID_ARGUMENT); + EXPECT_NE(status.GetErrorMessage().find("unsupported graphics API"), std::string::npos) + << status.GetErrorMessage(); +} + +TEST(GraphicsInteropTest, NotImplementedWhenEpFactoryDoesNotSupportIt) { + RegisteredEpDeviceUniquePtr kernel_registry_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_kernel_registry_info, + kernel_registry_ep)); + const OrtEpDevice* ep_device = kernel_registry_ep.get(); + + const OrtInteropApi& interop_api = Ort::GetInteropApi(); + const OrtGraphicsInteropConfig config = MakeConfig(ORT_GRAPHICS_API_D3D12); + + ExpectErrorCode(interop_api.InitGraphicsInteropForEpDevice(ep_device, &config), ORT_NOT_IMPLEMENTED); + ExpectErrorCode(interop_api.DeinitGraphicsInteropForEpDevice(ep_device), ORT_NOT_IMPLEMENTED); +} + +TEST(GraphicsInteropTest, NullArgumentsAreRejected) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + + const OrtInteropApi& interop_api = Ort::GetInteropApi(); + const OrtGraphicsInteropConfig config = MakeConfig(ORT_GRAPHICS_API_D3D12); + + ExpectErrorCode(interop_api.InitGraphicsInteropForEpDevice(nullptr, &config), ORT_INVALID_ARGUMENT); + ExpectErrorCode(interop_api.InitGraphicsInteropForEpDevice(example_ep.get(), nullptr), ORT_INVALID_ARGUMENT); + ExpectErrorCode(interop_api.DeinitGraphicsInteropForEpDevice(nullptr), ORT_INVALID_ARGUMENT); +} + +#endif // !defined(ORT_MINIMAL_BUILD) + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/common/cuda_op_test_utils.cc b/onnxruntime/test/common/cuda_op_test_utils.cc index fbd9b0a33c7c0..ebbde17adc28d 100644 --- a/onnxruntime/test/common/cuda_op_test_utils.cc +++ b/onnxruntime/test/common/cuda_op_test_utils.cc @@ -3,7 +3,7 @@ #include -#if defined(USE_CUDA) || defined(USE_NV) +#if defined(USE_CUDA) #include "cuda_runtime_api.h" #endif @@ -15,7 +15,7 @@ int GetCudaArchitecture() { // Usually, we test on a single GPU or multiple GPUs of same architecture, so it's fine to cache the result. static int cuda_arch = -1; -#if defined(USE_CUDA) || defined(USE_NV) +#if defined(USE_CUDA) if (cuda_arch == -1) { int current_device_id = 0; cudaGetDevice(¤t_device_id); diff --git a/onnxruntime/test/perftest/command_args_parser.cc b/onnxruntime/test/perftest/command_args_parser.cc index b4997d3615be5..6c99daec8648c 100644 --- a/onnxruntime/test/perftest/command_args_parser.cc +++ b/onnxruntime/test/perftest/command_args_parser.cc @@ -57,7 +57,9 @@ ABSL_FLAG(std::string, F, "", "[Usage]: -f \"dimension_denotation1:override_value1\" -f \"dimension_denotation2:override_value2\" ... or" " -f \"dimension_denotation1:override_value1 dimension_denotation2 : override_value2... \". Override value must > 0."); ABSL_FLAG(std::string, m, "duration", "Specifies the test mode. Value could be 'duration' or 'times'."); -ABSL_FLAG(std::string, e, "cpu", "Specifies the provider 'cpu','cuda','dnnl','tensorrt', 'nvtensorrtrtx', 'openvino', 'dml', 'acl', 'nnapi', 'coreml', 'qnn', 'snpe', 'migraphx', 'xnnpack', 'vitisai' or 'webgpu'."); +ABSL_FLAG(std::string, e, "cpu", + "Specifies the provider 'cpu','cuda','dnnl','tensorrt', 'openvino', 'dml', 'acl', 'nnapi', " + "'coreml', 'qnn', 'snpe', 'migraphx', 'xnnpack', 'vitisai' or 'webgpu'."); ABSL_FLAG(size_t, r, DefaultPerformanceTestConfig().run_config.repeated_times, "Specifies the repeated times if running in 'times' test mode."); ABSL_FLAG(size_t, t, DefaultPerformanceTestConfig().run_config.duration_in_seconds, "Specifies the seconds to run for 'duration' mode."); ABSL_FLAG(std::string, p, "", "Specifies the profile name to enable profiling and dump the profile data to the file."); @@ -191,7 +193,9 @@ ABSL_FLAG(uint32_t, hold_ms_after_session_creation, DefaultPerformanceTestConfig "When used with -n, keeps the process alive for the specified number of milliseconds after session creation.\n" "Prints 'SESSION_READY' to stdout before sleeping. Useful for multi-process memory measurements."); ABSL_FLAG(bool, l, DefaultPerformanceTestConfig().model_info.load_via_path, "Provides file as binary in memory by using fopen before session creation."); -ABSL_FLAG(bool, g, DefaultPerformanceTestConfig().run_config.enable_cuda_io_binding, "[TensorRT RTX | TensorRT | CUDA] Enables tensor input and output bindings on CUDA before session run."); +ABSL_FLAG(bool, g, DefaultPerformanceTestConfig().run_config.enable_cuda_io_binding, + "[Plugin EP | TensorRT | CUDA] Enables tensor input and output bindings on the device before session run. " + "For plugin EPs, binds to the selected EP device's allocator."); ABSL_FLAG(bool, X, DefaultPerformanceTestConfig().run_config.use_extensions, "Registers custom ops from onnxruntime-extensions."); ABSL_FLAG(std::string, plugin_ep_libs, "", "Specifies a list of plugin execution provider (EP) registration names and their corresponding shared libraries to register.\n" @@ -370,8 +374,6 @@ bool CommandLineParser::ParseArguments(PerformanceTestConfig& test_config, int a test_config.machine_config.provider_type_name = onnxruntime::kVitisAIExecutionProvider; } else if (ep == "webgpu") { test_config.machine_config.provider_type_name = onnxruntime::kWebGpuExecutionProvider; - } else if (ep == "nvtensorrtrtx") { - test_config.machine_config.provider_type_name = onnxruntime::kNvTensorRTRTXExecutionProvider; } else { return false; } diff --git a/onnxruntime/test/perftest/common_utils.cc b/onnxruntime/test/perftest/common_utils.cc index d09e872daefa2..05fc3dedf414a 100644 --- a/onnxruntime/test/perftest/common_utils.cc +++ b/onnxruntime/test/perftest/common_utils.cc @@ -236,25 +236,6 @@ std::optional GetPluginEpAllocator(Ort::Env& env, return std::nullopt; } -bool UsesNvidiaDevice(Ort::Env& env, const PerformanceTestConfig& test_config) { - const auto ep_devices = env.GetEpDevices(); - const auto& plugin_eps = test_config.machine_config.plugin_provider_type_list; - - const bool uses_nvidia_device = std::any_of(ep_devices.begin(), ep_devices.end(), - [&plugin_eps](const Ort::ConstEpDevice& ep_device) { - const std::string ep_name(ep_device.EpName()); - if (std::find(plugin_eps.begin(), plugin_eps.end(), ep_name) == plugin_eps.end()) { - return false; - } - - const auto hardware_device = ep_device.Device(); - return hardware_device.Type() == OrtDevice::GPU && - hardware_device.VendorId() == OrtDevice::VendorIds::NVIDIA; - }); - - return uses_nvidia_device; -} - } // namespace utils } // namespace perftest } // namespace onnxruntime diff --git a/onnxruntime/test/perftest/ort_test_session.cc b/onnxruntime/test/perftest/ort_test_session.cc index cad70ae0e7eb7..078ff062351c7 100644 --- a/onnxruntime/test/perftest/ort_test_session.cc +++ b/onnxruntime/test/perftest/ort_test_session.cc @@ -60,19 +60,25 @@ RunTiming OnnxRuntimeTestSession::Run() { } RunTiming timing; - if (CUDA == device_memory_name_) { + if (CUDA == device_memory_name_ || use_plugin_ep_io_binding_) { Ort::IoBinding io_binding(session_); auto mem_info = allocator_.GetInfo(); for (size_t i = 0; i < input_names_.size(); ++i) { io_binding.BindInput(input_names_[i], input[i]); } - for (auto& name : output_names_) { - io_binding.BindOutput(name.c_str(), mem_info); + // String outputs require host-accessible memory, so let the session allocate them on the CPU. + const auto cpu_mem_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + for (size_t i = 0; i < output_names_.size(); ++i) { + io_binding.BindOutput(output_names_[i].c_str(), is_output_string_[i] ? cpu_mem_info : mem_info); } - // Use async execution and rely on IO binding's SynchronizeOutputs to synchronize. - run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); + // Plugin EPs may synchronize at run end (OnRunEnd / stream flush) without implementing OrtEp::Sync, in which case + // IoBinding::SynchronizeOutputs() does not wait. Only defer synchronization for the built-in CUDA/TensorRT path. + if (!use_plugin_ep_io_binding_) { + // Use async execution and rely on IO binding's SynchronizeOutputs to synchronize. + run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); + } // do not time IO binding creation start = std::chrono::high_resolution_clock::now(); @@ -118,20 +124,28 @@ OnnxRuntimeTestSession::OnnxRuntimeTestSession(Ort::Env& env, std::random_device std::vector selected_ep_devices = perftest::utils::AppendPluginExecutionProviders(env, session_options, performance_test_config); - if (performance_test_config.run_config.enable_cuda_io_binding && - perftest::utils::UsesNvidiaDevice(env, performance_test_config) && - device_memory_name_.empty()) { - device_memory_name_ = CUDA; + // Pick an allocator from the plugin EP devices unless the user explicitly requested to force the CPU allocator. + if (performance_test_config.plugin_ep_force_cpu_allocator) { + fprintf(stdout, "[Plugin EP] Forcing CPU allocator (--plugin_ep_force_cpu_allocator was specified).\n"); + } else if (auto plugin_ep_allocator = perftest::utils::GetPluginEpAllocator(env, selected_ep_devices)) { + allocator_ = plugin_ep_allocator->allocator; + plugin_ep_allocator_selection_ = std::move(plugin_ep_allocator); } - // Pick an allocator from the plugin EP devices unless IO binding already set one, - // or the user explicitly requested to force the CPU allocator. - if (device_memory_name_.empty()) { - if (performance_test_config.plugin_ep_force_cpu_allocator) { - fprintf(stdout, "[Plugin EP] Forcing CPU allocator (--plugin_ep_force_cpu_allocator was specified).\n"); - } else if (auto plugin_ep_allocator = perftest::utils::GetPluginEpAllocator(env, selected_ep_devices)) { - allocator_ = plugin_ep_allocator->allocator; - plugin_ep_allocator_selection_ = std::move(plugin_ep_allocator); + // IO binding binds outputs to allocator_'s memory info, which comes from the selected EP device (including its + // device id), so it is only used when that allocator is device memory. + if (performance_test_config.run_config.enable_cuda_io_binding) { + use_plugin_ep_io_binding_ = IsAllocatorDeviceOnly(); + if (!use_plugin_ep_io_binding_) { + if (performance_test_config.plugin_ep_force_cpu_allocator) { + fprintf(stdout, + "[Plugin EP] -g is ignored because --plugin_ep_force_cpu_allocator was specified. " + "IO binding is disabled.\n"); + } else { + fprintf(stdout, + "[Plugin EP] -g requires a single selected EP device with a device allocator. " + "IO binding is disabled.\n"); + } } } } @@ -261,32 +275,6 @@ OnnxRuntimeTestSession::OnnxRuntimeTestSession(Ort::Env& env, std::random_device session_options.AppendExecutionProvider_CUDA(cuda_options); #else ORT_THROW("TensorRT is not supported in this build\n"); -#endif - } else if (provider_name_ == onnxruntime::kNvTensorRTRTXExecutionProvider) { -#ifdef USE_NV -#ifdef _MSC_VER - std::string opt_string = ToUTF8String(performance_test_config.run_config.ep_runtime_config_string); -#else - std::string opt_string = performance_test_config.run_config.ep_runtime_config_string; -#endif - ParseSessionConfigs(opt_string, provider_options); - if (!provider_options.empty()) { - std::cout << "Setting NV TensorRT RTX provider options to:\n"; - for (const auto& provider_option : provider_options) { - std::cout << "\t" << provider_option.first << ":" << provider_option.second << "\n"; - } - } - if (performance_test_config.run_config.enable_cuda_io_binding) { - device_memory_name_ = CUDA; - if (cudaStreamCreate(&stream_) != cudaError_t::cudaSuccess) { - ORT_THROW("Unable to create CUDA stream for IOBinding"); - } - auto stream_str = std::to_string(reinterpret_cast(stream_)); - provider_options["user_compute_stream"] = stream_str; - } - session_options.AppendExecutionProvider("NvTensorRtRtx", provider_options); -#else - ORT_THROW("NV TensorRT RTX is not supported in this build\n"); #endif } else if (provider_name_ == onnxruntime::kQnnExecutionProvider) { #ifdef USE_QNN @@ -1032,7 +1020,10 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)"); // String tensors require host-accessible memory; skip pre-allocation for them when allocator_ // is device-only and let the session allocate them itself. bool is_string_output = tensor_info.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; - if (is_dynamic || (device_memory_name_.empty() && !plugin_ep_allocator_selection_.has_value()) || + is_output_string_.push_back(is_string_output); + // With plugin EP IO binding, Run() binds fresh outputs and never uses outputs_. + if (is_dynamic || use_plugin_ep_io_binding_ || + (device_memory_name_.empty() && !plugin_ep_allocator_selection_.has_value()) || (is_string_output && IsAllocatorDeviceOnly())) { outputs_.emplace_back(Ort::Value(nullptr)); } else { @@ -1061,7 +1052,7 @@ bool OnnxRuntimeTestSession::IsAllocatorDeviceOnly() const { if (plugin_ep_allocator_selection_.has_value()) { return !plugin_ep_allocator_selection_->is_host_accessible; } -#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) +#if defined(USE_CUDA) || defined(USE_TENSORRT) if (device_memory_name_ == CUDA) { return true; } @@ -1173,7 +1164,7 @@ static void InitializeTensorWithSeed(int32_t seed, Ort::Value& tensor) { void OnnxRuntimeTestSession::CreateAndStoreGeneratedInput(size_t test_data_id, size_t input_idx, const std::vector& dims, ONNXTensorElementDataType element_type, int32_t seed) { -#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) +#if defined(USE_CUDA) || defined(USE_TENSORRT) if (device_memory_name_ == CUDA) { Ort::AllocatorWithDefaultOptions default_allocator; Ort::Value default_tensor = Ort::Value::CreateTensor(default_allocator, dims.data(), diff --git a/onnxruntime/test/perftest/ort_test_session.h b/onnxruntime/test/perftest/ort_test_session.h index aa1f6d265d275..8b319815edafa 100644 --- a/onnxruntime/test/perftest/ort_test_session.h +++ b/onnxruntime/test/perftest/ort_test_session.h @@ -10,7 +10,7 @@ #include "test_session.h" #include "utils.h" -#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) +#if defined(USE_CUDA) || defined(USE_TENSORRT) #include #endif @@ -80,10 +80,14 @@ class OnnxRuntimeTestSession : public TestSession { // Per-output dynamic-shape flag, aligned with outputs_. Lets Run() reset only the outputs that // actually have dynamic shapes instead of discarding every pre-allocated buffer. std::vector is_output_dynamic_; + // Per-output string-type flag, used to bind string outputs to host memory under IO binding. + std::vector is_output_string_; std::atomic round_robin_counter_{0}; bool use_round_robin_{false}; -#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) - cudaStream_t stream_; // Device stream if required by IO bindings + // True when the selected plugin EP device's allocator is used for IO binding (-g). + bool use_plugin_ep_io_binding_{false}; +#if defined(USE_CUDA) || defined(USE_TENSORRT) + cudaStream_t stream_{nullptr}; // Device stream if required by IO bindings #endif Ort::ArenaCfg cuda_mempool_arena_cfg_{nullptr}; }; diff --git a/onnxruntime/test/perftest/utils.h b/onnxruntime/test/perftest/utils.h index b9a52aeb86688..b27687de57d2e 100644 --- a/onnxruntime/test/perftest/utils.h +++ b/onnxruntime/test/perftest/utils.h @@ -40,8 +40,6 @@ std::vector AppendPluginExecutionProviders(Ort::Env& env, Ort::SessionOptions& session_options, const PerformanceTestConfig& test_config); -bool UsesNvidiaDevice(Ort::Env& env, const PerformanceTestConfig& test_config); - struct PluginEpAllocatorSelection { Ort::UnownedAllocator allocator; bool is_host_accessible = false; diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_ort_interop_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_ort_interop_test.cc deleted file mode 100644 index 25b1567939026..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_ort_interop_test.cc +++ /dev/null @@ -1,504 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -#include "core/graph/onnx_protobuf.h" -#include "core/graph/model.h" -#include "core/session/inference_session.h" -#include "test/providers/provider_test_utils.h" -#include "test/unittest_util/framework_test_utils.h" -#include "test/util/include/scoped_env_vars.h" -#include "test/util/include/asserts.h" -#include "test/common/trt_op_test_utils.h" -#include "test/common/random_generator.h" -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" -#include "test/unittest_util/conversion.h" - -#include -#include -#include - -#if defined(USE_DX_INTEROP) && USE_DX_INTEROP && defined(_WIN32) -#include -#include -#include -using Microsoft::WRL::ComPtr; -#endif - -using namespace ONNX_NAMESPACE; -using namespace ::onnxruntime::logging; - -extern std::unique_ptr ort_env; - -namespace onnxruntime { -namespace test { - -#if defined(USE_DX_INTEROP) && USE_DX_INTEROP && defined(_WIN32) -namespace { - -// Load d3d12.dll at runtime and get D3D12CreateDevice so the test exe does not need to link d3d12.lib. -// Caller must keep .module loaded for the duration of D3D12 API use (do not FreeLibrary). -struct D3D12CreateDeviceLoadResult { - HMODULE module = nullptr; - typedef HRESULT(WINAPI* PFN_D3D12CreateDevice)(IUnknown* pAdapter, D3D_FEATURE_LEVEL MinimumFeatureLevel, - REFIID riid, void** ppDevice); - PFN_D3D12CreateDevice pfn = nullptr; -}; -inline D3D12CreateDeviceLoadResult LoadD3D12CreateDevice() { - D3D12CreateDeviceLoadResult r; - r.module = LoadLibraryW(L"d3d12.dll"); - if (r.module) { - r.pfn = reinterpret_cast( - GetProcAddress(r.module, "D3D12CreateDevice")); - } - return r; -} - -void CreateD3D12Buffer(ID3D12Device* pDevice, size_t size, ID3D12Resource** ppResource, - D3D12_RESOURCE_STATES initState) { - D3D12_RESOURCE_DESC bufferDesc = {}; - bufferDesc.MipLevels = 1; - bufferDesc.Format = DXGI_FORMAT_UNKNOWN; - bufferDesc.Width = size; - bufferDesc.Height = 1; - bufferDesc.Flags = D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS; - bufferDesc.DepthOrArraySize = 1; - bufferDesc.SampleDesc.Count = 1; - bufferDesc.SampleDesc.Quality = 0; - bufferDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - bufferDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - - D3D12_HEAP_PROPERTIES heapProps = {}; - heapProps.Type = D3D12_HEAP_TYPE_DEFAULT; - heapProps.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN; - heapProps.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN; - heapProps.CreationNodeMask = 1; - heapProps.VisibleNodeMask = 1; - - HRESULT hr = pDevice->CreateCommittedResource( - &heapProps, D3D12_HEAP_FLAG_NONE, &bufferDesc, initState, nullptr, IID_PPV_ARGS(ppResource)); - if (FAILED(hr)) { - GTEST_FAIL() << "Failed creating D3D12 resource, HRESULT: 0x" << std::hex << hr; - } -} - -void CreateUploadBuffer(ID3D12Device* pDevice, size_t size, ID3D12Resource** ppResource) { - D3D12_HEAP_PROPERTIES heapProps = {}; - heapProps.Type = D3D12_HEAP_TYPE_UPLOAD; - heapProps.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN; - heapProps.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN; - heapProps.CreationNodeMask = 1; - heapProps.VisibleNodeMask = 1; - - D3D12_RESOURCE_DESC bufferDesc = {}; - bufferDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - bufferDesc.Alignment = 0; - bufferDesc.Width = size; - bufferDesc.Height = 1; - bufferDesc.DepthOrArraySize = 1; - bufferDesc.MipLevels = 1; - bufferDesc.Format = DXGI_FORMAT_UNKNOWN; - bufferDesc.SampleDesc.Count = 1; - bufferDesc.SampleDesc.Quality = 0; - bufferDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - bufferDesc.Flags = D3D12_RESOURCE_FLAG_NONE; - - HRESULT hr = pDevice->CreateCommittedResource( - &heapProps, D3D12_HEAP_FLAG_NONE, &bufferDesc, - D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, IID_PPV_ARGS(ppResource)); - if (FAILED(hr)) { - GTEST_FAIL() << "Failed creating D3D12 upload resource, HRESULT: 0x" << std::hex << hr; - } -} - -void CreateReadBackBuffer(ID3D12Device* pDevice, size_t size, ID3D12Resource** ppResource) { - D3D12_HEAP_PROPERTIES heapProps = {}; - heapProps.Type = D3D12_HEAP_TYPE_READBACK; - heapProps.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN; - heapProps.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN; - heapProps.CreationNodeMask = 1; - heapProps.VisibleNodeMask = 1; - - D3D12_RESOURCE_DESC bufferDesc = {}; - bufferDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - bufferDesc.Alignment = 0; - bufferDesc.Width = size; - bufferDesc.Height = 1; - bufferDesc.DepthOrArraySize = 1; - bufferDesc.MipLevels = 1; - bufferDesc.Format = DXGI_FORMAT_UNKNOWN; - bufferDesc.SampleDesc.Count = 1; - bufferDesc.SampleDesc.Quality = 0; - bufferDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - bufferDesc.Flags = D3D12_RESOURCE_FLAG_NONE; - - HRESULT hr = pDevice->CreateCommittedResource( - &heapProps, D3D12_HEAP_FLAG_NONE, &bufferDesc, - D3D12_RESOURCE_STATE_COPY_DEST, nullptr, IID_PPV_ARGS(ppResource)); - if (FAILED(hr)) { - GTEST_FAIL() << "Failed creating D3D12 readback resource, HRESULT: 0x" << std::hex << hr; - } -} - -void FlushAndWait(ID3D12Device* pDevice, ID3D12CommandQueue* pQueue) { - HANDLE hEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr); - ComPtr pFence; - pDevice->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&pFence)); - pQueue->Signal(pFence.Get(), 1); - pFence->SetEventOnCompletion(1, hEvent); - WaitForSingleObject(hEvent, INFINITE); - CloseHandle(hEvent); -} - -} // namespace - -// Test InitGraphicsInteropForEpDevice with command_queue = nullptr (graceful exit). -// Only built when D3D12 interop is enabled; otherwise Init would return ORT_NOT_IMPLEMENTED. -TEST(NvExecutionProviderTest, GraphicsInteropInitWithoutCommandQueue) { - ASSERT_NE(ort_env.get(), nullptr); - - RegisteredEpDeviceUniquePtr nv_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_ep); - const OrtEpDevice* ep_device = nv_ep.get(); - ASSERT_NE(ep_device, nullptr); - - const OrtInteropApi& interop_api = Ort::GetInteropApi(); - OrtGraphicsInteropConfig config = {}; - config.version = ORT_API_VERSION; - config.graphics_api = ORT_GRAPHICS_API_D3D12; - config.command_queue = nullptr; - config.additional_options = nullptr; - - { - Ort::Status init_status(interop_api.InitGraphicsInteropForEpDevice(ep_device, &config)); - ASSERT_TRUE(init_status.IsOK()) << init_status.GetErrorMessage(); - } - { - Ort::Status deinit_status(interop_api.DeinitGraphicsInteropForEpDevice(ep_device)); - ASSERT_TRUE(deinit_status.IsOK()) << deinit_status.GetErrorMessage(); - } -} - -// Test full D3D12 graphics interop: create D3D12 device and command queue, -// init graphics interop, create sync stream, release, deinit. -TEST(NvExecutionProviderTest, GraphicsInteropD3D12InitStreamDeinit) { - ASSERT_NE(ort_env.get(), nullptr); - - D3D12CreateDeviceLoadResult d3d12 = LoadD3D12CreateDevice(); - if (!d3d12.pfn) { - GTEST_SKIP() << "d3d12.dll or D3D12CreateDevice not available"; - } - ComPtr pDevice; - ComPtr pCommandQueue; - HRESULT hr = d3d12.pfn(nullptr, D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&pDevice)); - if (FAILED(hr)) { - GTEST_SKIP() << "D3D12 device creation failed, HRESULT: 0x" << std::hex << hr; - } - (void)d3d12.module; // keep d3d12.dll loaded for the duration of the test - - D3D12_COMMAND_QUEUE_DESC queueDesc = {}; - queueDesc.Type = D3D12_COMMAND_LIST_TYPE_COMPUTE; - queueDesc.Flags = D3D12_COMMAND_QUEUE_FLAG_NONE; - hr = pDevice->CreateCommandQueue(&queueDesc, IID_PPV_ARGS(&pCommandQueue)); - if (FAILED(hr)) { - GTEST_SKIP() << "D3D12 command queue creation failed, HRESULT: 0x" << std::hex << hr; - } - - RegisteredEpDeviceUniquePtr nv_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_ep); - const OrtEpDevice* ep_device = nv_ep.get(); - ASSERT_NE(ep_device, nullptr); - - const OrtInteropApi& interop_api = Ort::GetInteropApi(); - OrtGraphicsInteropConfig config = {}; - config.version = ORT_API_VERSION; - config.graphics_api = ORT_GRAPHICS_API_D3D12; - config.command_queue = pCommandQueue.Get(); - config.additional_options = nullptr; - - { - Ort::Status init_status(interop_api.InitGraphicsInteropForEpDevice(ep_device, &config)); - if (!init_status.IsOK()) { - GTEST_SKIP() << "InitGraphicsInteropForEpDevice failed (e.g. DX interop not built): " - << init_status.GetErrorMessage(); - } - } - - const OrtApi& ort_api = Ort::GetApi(); - OrtSyncStream* stream = nullptr; - { - Ort::Status stream_status(ort_api.CreateSyncStreamForEpDevice(ep_device, nullptr, &stream)); - ASSERT_TRUE(stream_status.IsOK()) << stream_status.GetErrorMessage(); - } - ASSERT_NE(stream, nullptr); - ort_api.ReleaseSyncStream(stream); - - { - Ort::Status deinit_status(interop_api.DeinitGraphicsInteropForEpDevice(ep_device)); - ASSERT_TRUE(deinit_status.IsOK()) << deinit_status.GetErrorMessage(); - } -} - -// Full D3D12 interop + inference: mirrors SimpleDXInterop_cig_only Main.cpp with random input data. -// Uses ORT Interop API for semaphore (ImportSemaphore, WaitSemaphore, SignalSemaphore) and runs inference on a Relu model. -TEST(NvExecutionProviderTest, GraphicsInteropD3D12FullInference) { - ASSERT_NE(ort_env.get(), nullptr); - - constexpr int image_dim = 64; // Smaller than 1080 for faster unit test - const size_t tensor_num_elements = 3 * image_dim * image_dim; - const size_t tensor_byte_size = tensor_num_elements * sizeof(uint16_t); - - // Create simple Relu model (1, 3, image_dim, image_dim) FLOAT16 - PathString model_name = ORT_TSTR("nv_interop_inference_test.onnx"); - clearFileIfExists(model_name); - { - onnxruntime::Model model("interop_test", false, DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - ONNX_NAMESPACE::TypeProto tensor_type; - tensor_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(image_dim); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(image_dim); - - auto& input_arg = graph.GetOrCreateNodeArg("input", &tensor_type); - auto& output_arg = graph.GetOrCreateNodeArg("output", &tensor_type); - graph.AddNode("relu_node", "Relu", "Relu operation", {&input_arg}, {&output_arg}); - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_STATUS_OK(onnxruntime::Model::Save(model, model_name)); - } - - // Random input data in valid float16 range (not raw uint16 bit patterns) - std::vector cpuInputHalf(tensor_num_elements); - std::vector cpuOutputHalf(tensor_num_elements); - { - RandomValueGenerator random{}; - std::vector shape{3, image_dim, image_dim}; - std::vector input_fp16 = random.Uniform(shape, -2.f, 2.f); - memcpy(cpuInputHalf.data(), input_fp16.data(), tensor_byte_size); - } - - D3D12CreateDeviceLoadResult d3d12 = LoadD3D12CreateDevice(); - if (!d3d12.pfn) { - GTEST_SKIP() << "d3d12.dll or D3D12CreateDevice not available"; - } - ComPtr pDevice; - ComPtr pCommandQueue; - HRESULT hr = d3d12.pfn(nullptr, D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&pDevice)); - if (FAILED(hr)) { - GTEST_SKIP() << "D3D12 device creation failed, HRESULT: 0x" << std::hex << hr; - } - (void)d3d12.module; // keep d3d12.dll loaded for the duration of the test - D3D12_COMMAND_QUEUE_DESC queueDesc = {}; - queueDesc.Type = D3D12_COMMAND_LIST_TYPE_COMPUTE; - queueDesc.Flags = D3D12_COMMAND_QUEUE_FLAG_NONE; - hr = pDevice->CreateCommandQueue(&queueDesc, IID_PPV_ARGS(&pCommandQueue)); - if (FAILED(hr)) { - GTEST_SKIP() << "D3D12 command queue creation failed, HRESULT: 0x" << std::hex << hr; - } - - RegisteredEpDeviceUniquePtr nv_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_ep); - const OrtEpDevice* trt_ep_device = nv_ep.get(); - ASSERT_NE(trt_ep_device, nullptr); - - const OrtApi& ort_api = Ort::GetApi(); - const OrtInteropApi& interop_api = Ort::GetInteropApi(); - - OrtGraphicsInteropConfig graphics_config = {}; - graphics_config.version = ORT_API_VERSION; - graphics_config.graphics_api = ORT_GRAPHICS_API_D3D12; - graphics_config.command_queue = pCommandQueue.Get(); - graphics_config.additional_options = nullptr; - { - Ort::Status init_status(interop_api.InitGraphicsInteropForEpDevice(trt_ep_device, &graphics_config)); - if (!init_status.IsOK()) { - GTEST_SKIP() << "InitGraphicsInteropForEpDevice failed: " << init_status.GetErrorMessage(); - } - } - - // Optional: D3D12 fence/semaphore for GPU sync between D3D12 and ORT. Only call - // CreateExternalResourceImporterForDevice when the EP implements it (e.g. not NV TensorRT RTX). - // Otherwise run inference with CPU sync (FlushAndWait) instead. - OrtExternalResourceImporter* importer = nullptr; - OrtExternalSemaphoreHandle* ort_sem_handle = nullptr; - ComPtr pFence; - HANDLE sharedFenceHandle = nullptr; - const bool ep_supports_external_importer = - (strcmp(ort_api.EpDevice_EpName(trt_ep_device), Utils::nv_tensorrt_rtx_ep_info.registration_name.c_str()) != 0); - if (ep_supports_external_importer) { - Ort::Status s(interop_api.CreateExternalResourceImporterForDevice(trt_ep_device, &importer)); - if (s.IsOK() && importer != nullptr) { - enum FenceState { - FENCE_UPLOAD_DONE = 1, - FENCE_KERNEL_DONE = 2 - }; - pDevice->CreateFence(0, D3D12_FENCE_FLAG_SHARED, IID_PPV_ARGS(&pFence)); - pDevice->CreateSharedHandle(pFence.Get(), nullptr, GENERIC_ALL, nullptr, &sharedFenceHandle); - ASSERT_NE(sharedFenceHandle, nullptr); - OrtExternalSemaphoreDescriptor sem_desc = {}; - sem_desc.version = ORT_API_VERSION; - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE; - sem_desc.native_handle = sharedFenceHandle; - Ort::Status import_s(interop_api.ImportSemaphore(importer, &sem_desc, &ort_sem_handle)); - ASSERT_TRUE(import_s.IsOK()) << import_s.GetErrorMessage(); - ASSERT_NE(ort_sem_handle, nullptr); - } - } - - OrtSyncStream* stream = nullptr; - { - Ort::Status s(ort_api.CreateSyncStreamForEpDevice(trt_ep_device, nullptr, &stream)); - ASSERT_TRUE(s.IsOK()) << s.GetErrorMessage(); - } - ASSERT_NE(stream, nullptr); - - // IHV-agnostic memory info for binding D3D12 buffers - OrtMemoryInfo* memory_info_agnostic = nullptr; - const OrtHardwareDevice* hw_device = ort_api.EpDevice_Device(trt_ep_device); - UINT vID = ort_api.HardwareDevice_VendorId(hw_device); - { - Ort::Status mem_status(ort_api.CreateMemoryInfo_V2( - "Device_Agnostic", OrtMemoryInfoDeviceType_GPU, vID, 0, - OrtDeviceMemoryType_DEFAULT, 0, OrtArenaAllocator, &memory_info_agnostic)); - ASSERT_TRUE(mem_status.IsOK()) << mem_status.GetErrorMessage(); - } - ASSERT_NE(memory_info_agnostic, nullptr); - - // Session options: user compute stream for non-CPU EPs - Ort::SessionOptions sessionOptions; - sessionOptions.SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL); - sessionOptions.DisableMemPattern(); - sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); - sessionOptions.AddConfigEntry("session.disable_cpu_ep_fallback", "1"); - sessionOptions.SetEpSelectionPolicy(OrtExecutionProviderDevicePolicy_PREFER_GPU); - - const OrtEpDevice* const* ep_devices = nullptr; - size_t num_ep_devices = 0; - { - Ort::Status get_devices_status(ort_api.GetEpDevices(*ort_env, &ep_devices, &num_ep_devices)); - ASSERT_TRUE(get_devices_status.IsOK()) << get_devices_status.GetErrorMessage(); - } - char stream_address[32]; - sprintf_s(stream_address, "%llu", static_cast(reinterpret_cast(ort_api.SyncStream_GetHandle(stream)))); - const char* option_keys[] = {"user_compute_stream", "has_user_compute_stream"}; - const char* option_values[] = {stream_address, "1"}; - for (size_t i = 0; i < num_ep_devices; i++) { - if (strcmp(ort_api.EpDevice_EpName(ep_devices[i]), "CPUExecutionProvider") != 0) { - Ort::Status append_ep_status(ort_api.SessionOptionsAppendExecutionProvider_V2( - sessionOptions, *ort_env, &ep_devices[i], 1, option_keys, option_values, 2)); - ASSERT_TRUE(append_ep_status.IsOK()) << append_ep_status.GetErrorMessage(); - } - } - - // D3D12 resources - ComPtr pInput; - ComPtr pOutput; - ComPtr pUploadRes; - ComPtr pDownloadRes; - CreateD3D12Buffer(pDevice.Get(), tensor_byte_size, pInput.GetAddressOf(), D3D12_RESOURCE_STATE_COPY_DEST); - CreateD3D12Buffer(pDevice.Get(), tensor_byte_size, pOutput.GetAddressOf(), D3D12_RESOURCE_STATE_COPY_SOURCE); - CreateUploadBuffer(pDevice.Get(), tensor_byte_size, pUploadRes.GetAddressOf()); - CreateReadBackBuffer(pDevice.Get(), tensor_byte_size, pDownloadRes.GetAddressOf()); - - ComPtr pAllocatorCopy; - pDevice->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_COMPUTE, IID_PPV_ARGS(&pAllocatorCopy)); - ComPtr pUploadCommandList; - pDevice->CreateCommandList(1, D3D12_COMMAND_LIST_TYPE_COMPUTE, pAllocatorCopy.Get(), nullptr, - IID_PPV_ARGS(&pUploadCommandList)); - pUploadCommandList->CopyResource(pInput.Get(), pUploadRes.Get()); - pUploadCommandList->Close(); - ComPtr pDownloadCommandList; - pDevice->CreateCommandList(1, D3D12_COMMAND_LIST_TYPE_COMPUTE, pAllocatorCopy.Get(), nullptr, - IID_PPV_ARGS(&pDownloadCommandList)); - pDownloadCommandList->CopyResource(pDownloadRes.Get(), pOutput.Get()); - pDownloadCommandList->Close(); - - // Session and inference in inner scope so TensorRT engine is destroyed before we destroy the CIG context in Deinit. - { - Ort::Session session(*ort_env, model_name.c_str(), sessionOptions); - Ort::IoBinding ioBinding(session); - Ort::AllocatorWithDefaultOptions allocator; - Ort::AllocatedStringPtr input_name = session.GetInputNameAllocated(0, allocator); - Ort::AllocatedStringPtr output_name = session.GetOutputNameAllocated(0, allocator); - int64_t input_dim[] = {1, 3, image_dim, image_dim}; - int64_t output_dim[] = {1, 3, image_dim, image_dim}; - - // Upload random input to D3D12 upload buffer - void* pData = nullptr; - pUploadRes->Map(0, nullptr, &pData); - memcpy(pData, cpuInputHalf.data(), tensor_byte_size); - pUploadRes->Unmap(0, nullptr); - - Ort::Value input_tensor(Ort::Value::CreateTensor( - memory_info_agnostic, reinterpret_cast(pInput->GetGPUVirtualAddress()), - tensor_byte_size, input_dim, 4, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16)); - Ort::Value output_tensor(Ort::Value::CreateTensor( - memory_info_agnostic, reinterpret_cast(pOutput->GetGPUVirtualAddress()), - tensor_byte_size, output_dim, 4, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16)); - ioBinding.BindInput(input_name.get(), input_tensor); - ioBinding.BindOutput(output_name.get(), output_tensor); - ioBinding.SynchronizeInputs(); - - // Run 2 iterations: upload -> sync -> inference -> download - const bool use_semaphore_sync = (importer != nullptr && ort_sem_handle != nullptr); - enum FenceState { - FENCE_UPLOAD_DONE = 1, - FENCE_KERNEL_DONE = 2 - }; - Ort::RunOptions run_options; - run_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); - for (int iter = 0; iter < 2; iter++) { - ID3D12CommandList* upload_list = pUploadCommandList.Get(); - pCommandQueue->ExecuteCommandLists(1, &upload_list); - if (use_semaphore_sync) { - pCommandQueue->Signal(pFence.Get(), FENCE_UPLOAD_DONE); - Ort::Status s(interop_api.WaitSemaphore(importer, ort_sem_handle, stream, FENCE_UPLOAD_DONE)); - ASSERT_TRUE(s.IsOK()) << s.GetErrorMessage(); - } else { - FlushAndWait(pDevice.Get(), pCommandQueue.Get()); - } - session.Run(run_options, ioBinding); - if (use_semaphore_sync) { - Ort::Status s(interop_api.SignalSemaphore(importer, ort_sem_handle, stream, FENCE_KERNEL_DONE)); - ASSERT_TRUE(s.IsOK()) << s.GetErrorMessage(); - pCommandQueue->Wait(pFence.Get(), FENCE_KERNEL_DONE); - } - ID3D12CommandList* download_list = pDownloadCommandList.Get(); - pCommandQueue->ExecuteCommandLists(1, &download_list); - FlushAndWait(pDevice.Get(), pCommandQueue.Get()); - } - - // Read back output and validate Relu: interpret as float16, compare to max(0, input) - void* pOutputData = nullptr; - pDownloadRes->Map(0, nullptr, &pOutputData); - memcpy(cpuOutputHalf.data(), pOutputData, tensor_byte_size); - pDownloadRes->Unmap(0, nullptr); - std::vector input_float(tensor_num_elements); - std::vector output_float(tensor_num_elements); - ConvertMLFloat16ToFloat(reinterpret_cast(cpuInputHalf.data()), input_float.data(), tensor_num_elements); - ConvertMLFloat16ToFloat(reinterpret_cast(cpuOutputHalf.data()), output_float.data(), tensor_num_elements); - const float tol = 1e-2f; // float16 limited precision - for (size_t i = 0; i < tensor_num_elements; i++) { - float expected = std::max(0.f, input_float[i]); - ASSERT_NEAR(output_float[i], expected, tol) << "Relu mismatch at index " << i; - } - } - - // Cleanup: release resources that use the interop context before DeinitGraphicsInteropForEpDevice - ort_api.ReleaseSyncStream(stream); - if (importer != nullptr) { - interop_api.ReleaseExternalSemaphoreHandle(ort_sem_handle); - interop_api.ReleaseExternalResourceImporter(importer); - CloseHandle(sharedFenceHandle); - } - ort_api.ReleaseMemoryInfo(memory_info_agnostic); - { - Ort::Status deinit_status(interop_api.DeinitGraphicsInteropForEpDevice(trt_ep_device)); - ASSERT_TRUE(deinit_status.IsOK()) << deinit_status.GetErrorMessage(); - } -} - -#endif // USE_DX_INTEROP && _WIN32 - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc deleted file mode 100644 index 2e34fd58a2628..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc +++ /dev/null @@ -1,608 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#include "core/graph/onnx_protobuf.h" -#include "core/session/inference_session.h" -#include "test/providers/provider_test_utils.h" -#include "test/unittest_util/framework_test_utils.h" -#include "test/util/include/default_providers.h" - -#include "test/util/include/scoped_env_vars.h" -#include "test/common/trt_op_test_utils.h" -#include "test/common/random_generator.h" -#include "test/common/cuda_op_test_utils.h" -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" - -#include -#include - -using namespace std; -using namespace ONNX_NAMESPACE; -using namespace ::onnxruntime::logging; -extern std::unique_ptr ort_env; -namespace onnxruntime { - -namespace test { - -// Helper function to check if GPU is Blackwell (SM 12.0+) or above -// Returns true if requirement is met -// Returns false if CUDA is unavailable or GPU is below SM 12.0 -static bool IsBlackwellOrAbove() { - constexpr int kBlackwellMinCapability = 1200; // SM 12.0 = 12 * 100 + 0 * 10 - int cuda_arch = GetCudaArchitecture(); - - // Check if CUDA is available - if (cuda_arch == -1) { - return false; - } - - return cuda_arch >= kBlackwellMinCapability; -} - -TEST(NvExecutionProviderTest, ContextEmbedAndReload) { - PathString model_name = ORT_TSTR("nv_execution_provider_test.onnx"); - PathString model_name_ctx = ORT_TSTR("nv_execution_provider_test_ctx.onnx"); - auto model_name_ctx_str = PathToUTF8(model_name_ctx); - clearFileIfExists(model_name_ctx); - std::string graph_name = "test"; - std::vector dims = {1, 3, 2}; - - CreateBaseModel(model_name, graph_name, dims); - // AOT time - { - auto start = std::chrono::high_resolution_clock::now(); - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name.c_str(), so); - auto stop = std::chrono::high_resolution_clock::now(); - std::cout << "Session creation AOT: " << std::chrono::duration_cast((stop - start)).count() << " ms" << std::endl; - - auto io_binding = generate_io_binding(session_object); - session_object.Run(run_options, io_binding); - } - - // JIT time - { - auto start = std::chrono::high_resolution_clock::now(); - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name_ctx.c_str(), so); - auto stop = std::chrono::high_resolution_clock::now(); - std::cout << "Session creation JIT: " << std::chrono::duration_cast((stop - start)).count() << " ms" << std::endl; - - auto io_binding = generate_io_binding(session_object); - session_object.Run(run_options, io_binding); - } -} - -TEST(NvExecutionProviderTest, ContextEmbedAndReloadDynamic) { - PathString model_name = ORT_TSTR("nv_execution_provider_dyn_test.onnx"); - PathString model_name_ctx = ORT_TSTR("nv_execution_provider_dyn_test_ctx.onnx"); - auto model_name_ctx_str = PathToUTF8(model_name_ctx); - clearFileIfExists(model_name_ctx); - std::string graph_name = "test"; - std::vector dims = {1, -1, -1}; - - CreateBaseModel(model_name, graph_name, dims); - - // AOT time - { - auto start = std::chrono::high_resolution_clock::now(); - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name.c_str(), so); - auto stop = std::chrono::high_resolution_clock::now(); - std::cout << "Session creation AOT: " << std::chrono::duration_cast((stop - start)).count() << " ms" << std::endl; - - auto io_binding = generate_io_binding(session_object); - session_object.Run(run_options, io_binding); - } - - // JIT time - { - auto start = std::chrono::high_resolution_clock::now(); - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name_ctx.c_str(), so); - auto stop = std::chrono::high_resolution_clock::now(); - std::cout << "Session creation JIT: " << std::chrono::duration_cast((stop - start)).count() << " ms" << std::endl; - - std::map> shape_overwrites; - shape_overwrites["X"] = {1, 5, 5}; - shape_overwrites["Y"] = {1, 5, 1}; - auto io_binding = generate_io_binding(session_object, shape_overwrites); - session_object.Run(run_options, io_binding); - } -} - -TEST(NvExecutionProviderTest, ContextEmbedAndReloadDataDynamic) { - PathString model_name = ORT_TSTR("nv_execution_provider_data_dyn_test.onnx"); - PathString model_name_ctx = ORT_TSTR("nv_execution_provider_data_dyn_test_ctx.onnx"); - auto model_name_ctx_str = PathToUTF8(model_name_ctx); - clearFileIfExists(model_name_ctx); - std::string graph_name = "test"; - std::vector dims = {1, -1, -1}; - - CreateBaseModel(model_name, graph_name, dims); - - // AOT time - { - auto start = std::chrono::high_resolution_clock::now(); - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name.c_str(), so); - auto stop = std::chrono::high_resolution_clock::now(); - std::cout << "Session creation AOT: " << std::chrono::duration_cast((stop - start)).count() << " ms" << std::endl; - - auto io_binding = generate_io_binding(session_object); - session_object.Run(run_options, io_binding); - } - - // JIT time - { - auto start = std::chrono::high_resolution_clock::now(); - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name_ctx.c_str(), so); - auto stop = std::chrono::high_resolution_clock::now(); - std::cout << "Session creation JIT: " << std::chrono::duration_cast((stop - start)).count() << " ms" << std::endl; - - std::map> shape_overwrites; - shape_overwrites["X"] = {1, 5, 5}; - shape_overwrites["Y"] = {1, 5, 5}; - auto io_binding = generate_io_binding(session_object, shape_overwrites); - session_object.Run(run_options, io_binding); - } -} - -std::string getTypeAsName(ONNX_NAMESPACE::TensorProto_DataType dtype) { - switch (dtype) { - case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: - return "fp64"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: - return "fp32"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: - return "fp16"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: - return "bf16"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: - return "int64"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: - return "int32"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: - return "int8"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: - return "uint8"; - case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4: - return "int4"; - default: - return "Unkwon type"; - } -} - -class TypeTests : public ::testing::TestWithParam { - public: -}; - -TEST_P(TypeTests, IOTypes) { - const std::string dtype_name = getTypeAsName(GetParam()); - ASSERT_FALSE(dtype_name.empty()); - const std::string model_name_str = "nv_execution_provider_" + dtype_name + ".onnx"; - const PathString model_name = ToPathString(model_name_str); - const std::string graph_name = "test" + dtype_name; - const std::vector dims = {1, 5, 10}; - - CreateBaseModel(model_name, graph_name, dims, false, GetParam()); - - // AOT time - { - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(*ort_env, model_name.c_str(), so); - - auto io_binding = generate_io_binding(session_object); - session_object.Run(run_options, io_binding); - } -} - -TEST(NvExecutionProviderTest, TestSessionOutputs) { - /* - * Model #1: - * - * "input" ---> TopK --- - * |---> "scores" - * |--- Less ---> "Less_output_0" - * |--- Div ---> "Div_output_0" - * |--- Mod ---> "labels" - */ - { - Ort::SessionOptions session_options; - session_options.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - // topk_and_multiple_graph_outputs.onnx contains input with dynamic dimension N - // setting override to avoid shape mismatch - session_options.AddFreeDimensionOverrideByName("N", 300); - - auto model_path = ORT_TSTR("testdata/topk_and_multiple_graph_outputs.onnx"); - Ort::Session session(*ort_env, model_path, session_options); - - size_t output_count = session.GetOutputCount(); - ASSERT_TRUE(output_count == 4); - } - - /* - * Model #2: - * - * "X" ---> Dropout ---> MatMul ---> "Y" - * ^ | - * | | - * "W" ------ ----> Can't be graph's output - * - */ - { - Ort::SessionOptions session_options; - session_options.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - - auto model_path = ORT_TSTR("testdata/node_output_not_used.onnx"); - Ort::Session session(*ort_env, model_path, session_options); - - size_t output_count = session.GetOutputCount(); - ASSERT_TRUE(output_count == 1); - } -} - -INSTANTIATE_TEST_SUITE_P(NvExecutionProviderTest, TypeTests, - ::testing::Values(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, - ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, - ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16, - ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, - ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 - // disabled low precision integer types since a specific quantize/dequantize model is required - // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8, - // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8, - // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4 - ), - [](const testing::TestParamInfo& info) { return getTypeAsName(info.param); }); - -static bool SessionHasEp(Ort::Session& session, const char* ep_name) { - // Access the underlying InferenceSession. - const OrtSession* ort_session = session; - const InferenceSession* s = reinterpret_cast(ort_session); - bool has_ep = false; - - for (const auto& provider : s->GetRegisteredProviderTypes()) { - if (provider == ep_name) { - has_ep = true; - break; - } - } - return has_ep; -} - -// Tests autoEP feature to automatically select an EP that supports the GPU. -TEST(NvExecutionProviderTest, AutoEp_PreferGpu) { - PathString model_name = ORT_TSTR("nv_execution_provider_auto_ep.onnx"); - std::string graph_name = "test"; - - std::vector dims = {1, 3, 2}; - - CreateBaseModel(model_name, graph_name, dims); - - { -#if _WIN32 - ort_env->RegisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider, ORT_TSTR("onnxruntime_providers_nv_tensorrt_rtx.dll")); -#else - ort_env->RegisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider, ORT_TSTR("libonnxruntime_providers_nv_tensorrt_rtx.so")); -#endif - - Ort::SessionOptions so; - so.SetEpSelectionPolicy(OrtExecutionProviderDevicePolicy_PREFER_GPU); - Ort::Session session_object(*ort_env, model_name.c_str(), so); - EXPECT_TRUE(SessionHasEp(session_object, kNvTensorRTRTXExecutionProvider)); - } - - ort_env->UnregisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider); -} - -TEST(NvExecutionProviderTest, GetSharedAllocator) { - const OrtApi& c_api = Ort::GetApi(); - RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); - - const auto* ep_memory_info = c_api.EpDevice_MemoryInfo(nv_tensorrt_rtx_ep.get(), OrtDeviceMemoryType_DEFAULT); - - // validate there is a shared allocator - OrtAllocator* allocator = nullptr; - ASSERT_ORTSTATUS_OK(c_api.GetSharedAllocator(*ort_env, ep_memory_info, &allocator)); - ASSERT_NE(allocator, nullptr); - - const auto* ep_host_accessible_memory_info = c_api.EpDevice_MemoryInfo(nv_tensorrt_rtx_ep.get(), OrtDeviceMemoryType_HOST_ACCESSIBLE); - OrtAllocator* host_accessible_allocator = nullptr; - ASSERT_ORTSTATUS_OK(c_api.GetSharedAllocator(*ort_env, ep_host_accessible_memory_info, &host_accessible_allocator)); - ASSERT_NE(host_accessible_allocator, nullptr); -} - -TEST(NvExecutionProviderTest, LoadUnloadPluginLibrary) { - const std::filesystem::path& library_path = Utils::nv_tensorrt_rtx_ep_info.library_path; - const std::string& registration_name = Utils::nv_tensorrt_rtx_ep_info.registration_name; - - const OrtApi* c_api = &Ort::GetApi(); - // this should load the library and create OrtEpDevice - ASSERT_ORTSTATUS_OK(Ort::GetApi().RegisterExecutionProviderLibrary(*ort_env, registration_name.c_str(), - library_path.c_str())); - - const OrtEpDevice* const* ep_devices = nullptr; - size_t num_devices = 0; - - ASSERT_ORTSTATUS_OK(Ort::GetApi().GetEpDevices(*ort_env, &ep_devices, &num_devices)); - // should be at least one device for the example EP - auto num_test_ep_devices = std::count_if(ep_devices, ep_devices + num_devices, - [®istration_name, &c_api](const OrtEpDevice* device) { - // the example uses the registration name for the EP name - // but that is not a requirement and the two can differ. - return c_api->EpDevice_EpName(device) == registration_name; - }); - ASSERT_GE(num_test_ep_devices, 1) << "Expected at least one OrtEpDevice to have been created by the test library."; - - // and this should unload it - ASSERT_ORTSTATUS_OK(Ort::GetApi().UnregisterExecutionProviderLibrary(*ort_env, - registration_name.c_str())); -} - -TEST(NvExecutionProviderTest, LoadUnloadPluginLibraryCxxApi) { - const std::filesystem::path& library_path = Utils::nv_tensorrt_rtx_ep_info.library_path; - const std::string& registration_name = Utils::nv_tensorrt_rtx_ep_info.registration_name; - const OrtApi* c_api = &Ort::GetApi(); - // this should load the library and create OrtEpDevice - ort_env->RegisterExecutionProviderLibrary(registration_name.c_str(), library_path.c_str()); - - std::vector ep_devices = ort_env->GetEpDevices(); - - auto test_ep_device = std::find_if(ep_devices.begin(), ep_devices.end(), - [®istration_name, &c_api](const Ort::ConstEpDevice& device) { - return device.EpName() == registration_name; - }); - ASSERT_NE(test_ep_device, ep_devices.end()) << "Expected an OrtEpDevice to have been created by the test library."; - - // test all the C++ getters. expected values are from \onnxruntime\test\autoep\library\example_plugin_ep.cc - ASSERT_STREQ(test_ep_device->EpVendor(), "NVIDIA"); - - auto metadata = test_ep_device->EpMetadata(); - ASSERT_STREQ(metadata.GetValue(kOrtEpDevice_EpMetadataKey_Version), ORT_VERSION); - - // the GPU device info will vary by machine so check for the lowest common denominator values - Ort::ConstHardwareDevice device = test_ep_device->Device(); - ASSERT_EQ(device.Type(), OrtHardwareDeviceType_GPU); - ASSERT_GE(device.VendorId(), 0); - ASSERT_GE(device.DeviceId(), 0); - ASSERT_NE(device.Vendor(), nullptr); - Ort::ConstKeyValuePairs device_metadata = device.Metadata(); - std::unordered_map metadata_entries = device_metadata.GetKeyValuePairs(); - ASSERT_GT(metadata_entries.size(), 0); // should have at least SPDRP_HARDWAREID on Windows - - // and this should unload it without throwing - ort_env->UnregisterExecutionProviderLibrary(registration_name.c_str()); -} - -TEST(NvExecutionProviderTest, DataTransfer) { - const OrtApi& c_api = Ort::GetApi(); - RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); - const OrtEpDevice* ep_device = nv_tensorrt_rtx_ep.get(); - - const OrtMemoryInfo* device_memory_info = c_api.EpDevice_MemoryInfo(ep_device, OrtDeviceMemoryType_DEFAULT); - - // create a tensor using the default CPU allocator - Ort::AllocatorWithDefaultOptions cpu_allocator; - std::vector shape{2, 3, 4}; // shape doesn't matter - const size_t num_elements = 2 * 3 * 4; - - RandomValueGenerator random{}; - std::vector input_data = random.Gaussian(shape, 0.0f, 2.f); - Ort::Value cpu_tensor = Ort::Value::CreateTensor(cpu_allocator.GetInfo(), - input_data.data(), input_data.size(), - shape.data(), shape.size()); - - // create an on-device Tensor using the NV TensorRT RTX EP GPU allocator. - - OrtAllocator* allocator = nullptr; - ASSERT_ORTSTATUS_OK(c_api.GetSharedAllocator(*ort_env, device_memory_info, &allocator)); - ASSERT_NE(allocator, nullptr); - Ort::Value device_tensor = Ort::Value::CreateTensor(allocator, shape.data(), shape.size()); - - std::vector src_tensor_ptrs{cpu_tensor}; - std::vector dst_tensor_ptrs{device_tensor}; - - ASSERT_ORTSTATUS_OK(c_api.CopyTensors(*ort_env, src_tensor_ptrs.data(), dst_tensor_ptrs.data(), nullptr, - src_tensor_ptrs.size())); - - // Copy data back from device_tensor to a new CPU tensor and verify the contents - - // Create a new CPU tensor to receive the data - Ort::Value cpu_tensor_copy = Ort::Value::CreateTensor(cpu_allocator, shape.data(), shape.size()); - - std::vector src_tensor_ptrs_back{device_tensor}; - std::vector dst_tensor_ptrs_back{cpu_tensor_copy}; - - ASSERT_ORTSTATUS_OK(c_api.CopyTensors(*ort_env, src_tensor_ptrs_back.data(), dst_tensor_ptrs_back.data(), nullptr, - src_tensor_ptrs_back.size())); - - const float* src_data = nullptr; - ASSERT_ORTSTATUS_OK(c_api.GetTensorData(cpu_tensor, reinterpret_cast(&src_data))); - - const float* cpu_copy_data = nullptr; - ASSERT_ORTSTATUS_OK(c_api.GetTensorData(cpu_tensor_copy, reinterpret_cast(&cpu_copy_data))); - - ASSERT_NE(src_data, cpu_copy_data) << "Should have copied between two different memory locations"; - - size_t bytes; - ASSERT_ORTSTATUS_OK(c_api.GetTensorSizeInBytes(cpu_tensor, &bytes)); - ASSERT_EQ(bytes, num_elements * sizeof(float)); - - auto src_span = gsl::make_span(src_data, num_elements); - auto cpu_copy_span = gsl::make_span(cpu_copy_data, num_elements); - - EXPECT_THAT(cpu_copy_span, ::testing::ContainerEq(src_span)); - - // must release this before we unload the EP and the allocator is deleted - device_tensor = Ort::Value(); -} - -TEST(NvExecutionProviderTest, FP8CustomOpModel) { - if (!IsBlackwellOrAbove()) { - GTEST_SKIP() << "Test requires SM 12.0+ GPU (Blackwell+)"; - } - - PathString model_name = ORT_TSTR("nv_execution_provider_fp8_quantize_dequantize_test.onnx"); - clearFileIfExists(model_name); - std::string graph_name = "nv_execution_provider_fp8_quantize_dequantize_graph"; - - // Create a model with TRT_FP8QuantizeLinear -> TRT_FP8DequantizeLinear (FP16 -> FP8 -> FP16, per-tensor quantization) - CreateFP8CustomOpModel(model_name, graph_name); - - // Verify the model file was created - ASSERT_TRUE(std::filesystem::exists(model_name)); - - // Create session and register execution provider explicitly - // This ensures custom ops are registered before model is loaded - SessionOptions so; - so.session_logid = "NvExecutionProviderTest.FP8CustomOpModel"; - InferenceSession session_object{so, GetEnvironment()}; - - // Register TRTRTX EP - this will register custom ops - std::unique_ptr execution_provider = DefaultNvTensorRTRTXExecutionProvider(); - ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(std::move(execution_provider))); - - // Load and initialize model - ASSERT_STATUS_OK(session_object.Load(model_name)); - ASSERT_STATUS_OK(session_object.Initialize()); - - // Create input data (FP16, shape [4, 64]) - std::vector input_data(4 * 64); - for (size_t i = 0; i < input_data.size(); ++i) { - input_data[i] = MLFloat16(static_cast(i % 100) / 100.0f); - } - - // Create input tensor - std::vector input_shape = {4, 64}; - OrtValue input_tensor; - Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(input_shape), - input_data.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), - input_tensor); - - // Prepare feeds - NameMLValMap feeds; - feeds.insert(std::make_pair("X", input_tensor)); - - // Prepare outputs - std::vector output_names; - output_names.push_back("Y"); - - // Run inference - std::vector fetches; - RunOptions run_options; - ASSERT_STATUS_OK(session_object.Run(run_options, feeds, output_names, &fetches)); - - // Verify output tensor is valid - ASSERT_EQ(fetches.size(), 1u); - ASSERT_TRUE(fetches[0].IsTensor()); - - const auto& output_tensor = fetches[0].Get(); - auto output_shape = output_tensor.Shape(); - ASSERT_EQ(output_shape.NumDimensions(), 2u); - ASSERT_EQ(output_shape[0], 4); - ASSERT_EQ(output_shape[1], 64); - - // Verify output is FLOAT16 - ASSERT_TRUE(output_tensor.IsDataType()); - - LOGS_DEFAULT(INFO) << "[NvExecutionProviderTest] TRT FP8 custom ops model run completed successfully"; -} - -TEST(NvExecutionProviderTest, FP4CustomOpModel) { - if (!IsBlackwellOrAbove()) { - GTEST_SKIP() << "Test requires SM 12.0+ GPU (Blackwell+)"; - } - - PathString model_name = ORT_TSTR("nv_execution_provider_fp4_dynamic_quantize_test.onnx"); - clearFileIfExists(model_name); - std::string graph_name = "nv_execution_provider_fp4_dynamic_quantize_graph"; - - // Create a model with TRT_FP4DynamicQuantize node - CreateFP4CustomOpModel(model_name, graph_name); - - // Verify the model file was created - ASSERT_TRUE(std::filesystem::exists(model_name)); - - // Create session and register execution provider explicitly - // This ensures custom ops are registered before model is loaded - SessionOptions so; - so.session_logid = "NvExecutionProviderTest.FP4CustomOpModel"; - InferenceSession session_object{so, GetEnvironment()}; - - // Register TRTRTX EP - this will register custom ops - std::unique_ptr execution_provider = DefaultNvTensorRTRTXExecutionProvider(); - ASSERT_STATUS_OK(session_object.RegisterExecutionProvider(std::move(execution_provider))); - - // Load and initialize model - ASSERT_STATUS_OK(session_object.Load(model_name)); - ASSERT_STATUS_OK(session_object.Initialize()); - - // Create input data (FP16, shape [64, 64]) - std::vector input_data(64 * 64); - for (size_t i = 0; i < input_data.size(); ++i) { - input_data[i] = MLFloat16(static_cast(i % 100) / 100.0f); - } - - // Create input tensor - std::vector input_shape = {64, 64}; - OrtValue input_tensor; - Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(input_shape), - input_data.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), - input_tensor); - - // Prepare feeds - NameMLValMap feeds; - feeds.insert(std::make_pair("X", input_tensor)); - - // Prepare outputs - std::vector output_names; - output_names.push_back("X_dequantized"); - - // Run inference - std::vector fetches; - RunOptions run_options; - ASSERT_STATUS_OK(session_object.Run(run_options, feeds, output_names, &fetches)); - - // Verify output tensor is valid - ASSERT_EQ(fetches.size(), 1u); - ASSERT_TRUE(fetches[0].IsTensor()); - - const auto& output_tensor = fetches[0].Get(); - auto output_shape = output_tensor.Shape(); - ASSERT_EQ(output_shape.NumDimensions(), 2u); - ASSERT_EQ(output_shape[0], 64); - ASSERT_EQ(output_shape[1], 64); - - // Verify output is FLOAT16 - ASSERT_TRUE(output_tensor.IsDataType()); - - LOGS_DEFAULT(INFO) << "[NvExecutionProviderTest] TRT FP4 dynamic quantize model run completed successfully"; -} - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc deleted file mode 100644 index 5a5cb71f73e45..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc +++ /dev/null @@ -1,361 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#include "core/common/path_utils.h" -#include "core/graph/onnx_protobuf.h" -#include "test/unittest_util/framework_test_utils.h" -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" - -#include -#include - -extern std::unique_ptr ort_env; - -namespace onnxruntime { - -namespace test { - -RegisteredEpDeviceUniquePtr AppendTrtEtxEP(Ort::SessionOptions& session_options, std::unordered_map& option_map) { - RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; - /// Since this test runs after other tests that use registration interface this test has to use it as well - /// windows as otherwise the kernel registry inside the EP will not be populated. The legacy APis ony call the initialize once. - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); - auto ep_devices = ort_env->GetEpDevices(); - Ort::ConstEpDevice selected_device; - for (auto& device : ep_devices) { - if (!std::strcmp(device.EpName(), kNvTensorRTRTXExecutionProvider)) { - selected_device = device; - } - } - session_options.AppendExecutionProvider_V2(*ort_env, {selected_device}, option_map); - return nv_tensorrt_rtx_ep; -} - -std::vector readBinaryFile(const PathString& filename) { - std::ifstream file(filename, std::ios::binary); - if (!file.is_open()) { - throw std::runtime_error("Could not open file: " + PathToUTF8String(filename)); - } - - file.seekg(0, std::ios::end); - std::streamsize filesize = file.tellg(); - file.seekg(0, std::ios::beg); - - std::vector buffer(filesize); - if (!file.read(reinterpret_cast(buffer.data()), filesize)) { - throw std::runtime_error("Could not read file: " + PathToUTF8String(filename)); - } - - return buffer; -} - -struct CompileParam { - bool embed_mode; - bool bytestream_io; - bool external_initialzier_for_parser = false; - const std::string to_string() const { - return "embed_mode_" + std::to_string(embed_mode) + "_bytestream_io_" + std::to_string(bytestream_io) + "_ext_init_" + std::to_string(external_initialzier_for_parser); - ; - } -}; -class CompileApiTest - : public testing::TestWithParam { - public: - const CompileParam& GetCompileParam() const { - return GetParam(); - } -}; - -void SmallModelTest(CompileParam test_param, bool fully_supported_model) { - std::string test_name = test_param.to_string(); - if (!fully_supported_model) - test_name += "_fast_gelu"; - PathString model_name = path_utils::MakePathString("nv_execution_provider_compile_" + test_name + ".onnx"); - PathString model_name_ctx = path_utils::MakePathString("nv_execution_provider_compile_" + test_name + "_ctx.onnx"); - clearFileIfExists(model_name_ctx); - std::string graph_name = "test"; - std::vector dims = {1, 3, 2}; - - CreateBaseModel(model_name, graph_name, dims, !fully_supported_model); - - Ort::SessionOptions session_options; - std::unordered_map option_map{ - {onnxruntime::nv::provider_option_names::kUseExternalDataInitializer, std::to_string(test_param.external_initialzier_for_parser)}}; - auto ep = AppendTrtEtxEP(session_options, option_map); - - Ort::ModelCompilationOptions model_compile_options(*ort_env, session_options); - model_compile_options.SetEpContextEmbedMode(test_param.embed_mode); - - void* output_context = nullptr; - size_t output_context_size = 0; - std::vector input_onnx; - if (test_param.bytestream_io) { - input_onnx = readBinaryFile(model_name); - model_compile_options.SetInputModelFromBuffer(input_onnx.data(), input_onnx.size()); - model_compile_options.SetOutputModelBuffer(Ort::AllocatorWithDefaultOptions(), &output_context, &output_context_size); - } else { - model_compile_options.SetInputModelPath(model_name.c_str()); - model_compile_options.SetOutputModelPath(model_name_ctx.c_str()); - } - // AOT time - ASSERT_TRUE(Ort::CompileModel(*ort_env, model_compile_options).IsOK()); - - // JIT time - Ort::Session session_object{nullptr}; - if (test_param.bytestream_io) { - session_object = Ort::Session(*ort_env, output_context, output_context_size, session_options); - } else { - session_object = Ort::Session(*ort_env, model_name_ctx.c_str(), session_options); - } - auto io_binding = generate_io_binding(session_object); - Ort::RunOptions run_options; - session_object.Run(run_options, io_binding); -} - -TEST_P(CompileApiTest, SmallModel) { - const auto& test_param = GetCompileParam(); - SmallModelTest(test_param, true); -} - -TEST_P(CompileApiTest, SmallSplitModel) { - const auto& test_param = GetCompileParam(); - SmallModelTest(test_param, false); -} - -TEST_P(CompileApiTest, LargeModel) { - const auto& test_param = GetCompileParam(); - // with embed mode == 1 the resulting file will be over the 2GB proto limit - if (test_param.embed_mode == 1) { - GTEST_SKIP(); - } - std::string test_name = test_param.to_string(); - PathString model_name = path_utils::MakePathString("nv_execution_provider_compile_large_" + test_name + ".onnx"); - PathString external_data_name = path_utils::MakePathString("nv_execution_provider_compile_large_" + test_name + ".onnx_data"); - PathString model_name_ctx = path_utils::MakePathString("nv_execution_provider_compile_large_" + test_name + "_ctx.onnx"); - PathString model_name_ctx_data = path_utils::MakePathString("nv_execution_provider_compile_large_" + test_name + "_ctx.onnx_data"); - clearFileIfExists(model_name_ctx); - clearFileIfExists(model_name_ctx_data); - // This accelerates test iterations if the large model was already generated - if (!std::filesystem::exists(model_name) || !std::filesystem::exists(external_data_name)) { - CreateLargeLLMModel(model_name, external_data_name); - } - - Ort::SessionOptions session_options; - std::unordered_map option_map{ - {onnxruntime::nv::provider_option_names::kUseExternalDataInitializer, - std::to_string(test_param.bytestream_io || test_param.external_initialzier_for_parser)}}; - auto ep = AppendTrtEtxEP(session_options, option_map); - - Ort::ModelCompilationOptions model_compile_options(*ort_env, session_options); - model_compile_options.SetEpContextEmbedMode(test_param.embed_mode); - - void* output_context = nullptr; - size_t output_context_size = 0; - std::vector input_onnx, input_data; - std::vector file_names; - std::vector file_buffers; - std::vector lengths; - if (test_param.bytestream_io) { - input_onnx = readBinaryFile(model_name); - input_data = readBinaryFile(external_data_name); - file_names = {external_data_name}; - file_buffers = {input_data.data()}; - lengths = {input_data.size()}; - session_options.AddExternalInitializersFromFilesInMemory(file_names, file_buffers, lengths); - - model_compile_options.SetInputModelFromBuffer(input_onnx.data(), input_onnx.size()); - model_compile_options.SetOutputModelBuffer(Ort::AllocatorWithDefaultOptions(), &output_context, &output_context_size); - } else { - model_compile_options.SetInputModelPath(model_name.c_str()); - model_compile_options.SetOutputModelPath(model_name_ctx.c_str()); - model_compile_options.SetOutputModelExternalInitializersFile(model_name_ctx_data.c_str(), 1024); - } - - // AOT time - ASSERT_TRUE(Ort::CompileModel(*ort_env, model_compile_options).IsOK()); - - // JIT time - std::unique_ptr session; - if (test_param.bytestream_io) { - session = std::make_unique(*ort_env, output_context, output_context_size, session_options); - } else { - session = std::make_unique(*ort_env, model_name_ctx.c_str(), session_options); - } - - auto io_binding = generate_io_binding(*session); - Ort::RunOptions run_options; - session->Run(run_options, io_binding); -} - -INSTANTIATE_TEST_SUITE_P( - NvExecutionProviderTest, CompileApiTest, - ::testing::Values( - CompileParam{true, false}, - CompileParam{false, false}, - CompileParam{true, true}, - CompileParam{false, true}, - // test with external initializers for parser - CompileParam{true, true, true}, - CompileParam{true, false, true}), - [](const testing::TestParamInfo& info) { - return info.param.to_string(); - }); - -/* - * Helper to create a synthetic EPContext ONNX model with a specific "source" attribute. - * Uses raw ONNX protobuf to bypass schema validation (EPContext is a contrib op). - */ -void CreateSyntheticEPContextModel(const PathString& model_path, - const std::string& source_attr, - bool include_source_attr = true) { - ONNX_NAMESPACE::ModelProto model; - model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); - auto* opset = model.add_opset_import(); - opset->set_domain(""); - opset->set_version(11); - auto* ms_opset = model.add_opset_import(); - ms_opset->set_domain("com.microsoft"); - ms_opset->set_version(1); - - auto* graph = model.mutable_graph(); - graph->set_name("EPContextSourceTest"); - - // Input - auto* input = graph->add_input(); - input->set_name("input"); - auto* input_type = input->mutable_type()->mutable_tensor_type(); - input_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - input_type->mutable_shape()->add_dim()->set_dim_value(1); - input_type->mutable_shape()->add_dim()->set_dim_value(3); - - // Output - auto* output = graph->add_output(); - output->set_name("output"); - auto* output_type = output->mutable_type()->mutable_tensor_type(); - output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - output_type->mutable_shape()->add_dim()->set_dim_value(1); - output_type->mutable_shape()->add_dim()->set_dim_value(3); - - // EPContext node - auto* node = graph->add_node(); - node->set_op_type("EPContext"); - node->set_domain("com.microsoft"); - node->set_name("ep_context_node"); - node->add_input("input"); - node->add_output("output"); - - // embed_mode attribute - auto* attr_embed = node->add_attribute(); - attr_embed->set_name("embed_mode"); - attr_embed->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); - attr_embed->set_i(1); - - // ep_cache_context attribute (dummy data) - auto* attr_cache = node->add_attribute(); - attr_cache->set_name("ep_cache_context"); - attr_cache->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); - attr_cache->set_s("dummy_context_data"); - - // source attribute (conditionally added) - if (include_source_attr) { - auto* attr_source = node->add_attribute(); - attr_source->set_name("source"); - attr_source->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); - attr_source->set_s(source_attr); - } - - // Save to file - std::ofstream ofs(model_path, std::ios::binary); - ASSERT_TRUE(ofs.is_open()); - ASSERT_TRUE(model.SerializeToOstream(&ofs)); -} - -/* - * Test: NvTensorRTRTX EP should NOT claim an EPContext node whose "source" - * attribute belongs to a different EP (e.g., OpenVINO). - * - * Expected: Session initialization fails because no EP claims the node. - */ -TEST(NvExecutionProviderTest, EPContextNode_ForeignSourceSkipped) { - PathString model_path = path_utils::MakePathString("ep_context_foreign_source_nv.onnx"); - CreateSyntheticEPContextModel(model_path, "OpenVINOExecutionProvider"); - - Ort::SessionOptions session_options; - std::unordered_map option_map; - auto ep = AppendTrtEtxEP(session_options, option_map); - - // Loading a model with a foreign-source EPContext node should fail during - // session creation because the NvTensorRTRTX EP correctly skips the node - // and no other EP can handle it. - try { - Ort::Session session(*ort_env, model_path.c_str(), session_options); - FAIL() << "Expected session creation to fail for EPContext node with foreign source"; - } catch (const Ort::Exception& e) { - std::string error_msg = e.what(); - EXPECT_TRUE(error_msg.find("EPContext") != std::string::npos) - << "Error should mention EPContext. Actual: " << error_msg; - } - - // Clean up - std::filesystem::remove(model_path); -} - -/* - * Test: NvTensorRTRTX EP should NOT claim an EPContext node whose "source" - * attribute is set to the classic TensorRT EP name. - */ -TEST(NvExecutionProviderTest, EPContextNode_ClassicTrtSourceSkipped) { - PathString model_path = path_utils::MakePathString("ep_context_classic_trt_source_nv.onnx"); - CreateSyntheticEPContextModel(model_path, "TensorrtExecutionProvider"); - - Ort::SessionOptions session_options; - std::unordered_map option_map; - auto ep = AppendTrtEtxEP(session_options, option_map); - - try { - Ort::Session session(*ort_env, model_path.c_str(), session_options); - FAIL() << "Expected session creation to fail for EPContext node with classic TRT source"; - } catch (const Ort::Exception& e) { - std::string error_msg = e.what(); - EXPECT_TRUE(error_msg.find("EPContext") != std::string::npos) - << "Error should mention EPContext. Actual: " << error_msg; - } - - // Clean up - std::filesystem::remove(model_path); -} - -/* - * Test: NvTensorRTRTX EP should still claim an EPContext node that has NO - * "source" attribute (backward compatibility with legacy context models). - * - * Expected: The EP claims the node. It may fail later during engine - * deserialization (since context data is synthetic), but the error must NOT - * be "is not compatible with any execution provider", which would indicate - * the node was not claimed at all. - */ -TEST(NvExecutionProviderTest, EPContextNode_NoSourceAttribute_BackwardCompat) { - PathString model_path = path_utils::MakePathString("ep_context_no_source_nv.onnx"); - CreateSyntheticEPContextModel(model_path, "", /*include_source_attr=*/false); - - Ort::SessionOptions session_options; - std::unordered_map option_map; - auto ep = AppendTrtEtxEP(session_options, option_map); - - try { - Ort::Session session(*ort_env, model_path.c_str(), session_options); - // If session creation succeeds, backward compatibility is working. - } catch (const Ort::Exception& e) { - std::string error_msg = e.what(); - // The node should have been claimed by the EP. Any failure should be - // EP-internal (e.g., bad engine data), NOT the "not compatible" error - // that indicates no EP claimed the node. - EXPECT_TRUE(error_msg.find("is not compatible with any execution provider") == std::string::npos) - << "Legacy EPContext node without source should still be claimed by EP. Error: " << error_msg; - } - - // Clean up - std::filesystem::remove(model_path); -} - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_external_resource_importer_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_external_resource_importer_test.cc deleted file mode 100644 index 232477c96b4f4..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_external_resource_importer_test.cc +++ /dev/null @@ -1,1191 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -// This test validates the D3D12 ↔ CUDA external resource import functionality -// for the NvTensorRtRtx execution provider. -// -// Test Coverage: -// 1. External Resource Importer creation and destruction -// 2. Memory import capability check (D3D12 Resource & Heap) -// 3. Semaphore import capability check (D3D12 Fence) -// 4. D3D12 shared resource import to CUDA -// 5. Tensor creation from imported external memory -// 6. D3D12 timeline fence import for GPU synchronization -// 7. Wait/Signal semaphore operations -// 8. Full inference pipeline with zero-copy external memory -// 9. Full inference pipeline with zero-copy external memory and a CIG context - -#include "core/graph/onnx_protobuf.h" -#include "core/session/inference_session.h" -#include "core/session/onnxruntime_cxx_api.h" -#include "test/providers/provider_test_utils.h" -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" -#include "test/unittest_util/framework_test_utils.h" -#include "test/util/include/scoped_env_vars.h" -#include "test/common/random_generator.h" - -#include -#include -#include -#include -#include -#include - -#if defined(_WIN32) -#include -#include -#include -using Microsoft::WRL::ComPtr; -#endif - -// Include CUDA headers for pointer attribute verification -#include -#include -#include - -using namespace std; -using namespace ONNX_NAMESPACE; -using namespace ::onnxruntime::logging; - -extern std::unique_ptr ort_env; - -namespace onnxruntime { -namespace test { - -#if defined(_WIN32) - -// Dynamic CUDA driver function loader -class CudaDriverLoader { - private: - HMODULE cuda_driver_dll_ = nullptr; - - // CUDA Driver API function pointers - using cuCtxCreate_v4_t = CUresult (*)(CUcontext*, CUctxCreateParams*, unsigned int, CUdevice); - using cuCtxDestroy_t = CUresult (*)(CUcontext); - using cuCtxGetCurrent_t = CUresult (*)(CUcontext*); - using cuCtxSetCurrent_t = CUresult (*)(CUcontext); - - public: - cuCtxCreate_v4_t cuCtxCreate_v4_fn = nullptr; - cuCtxDestroy_t cuCtxDestroy_fn = nullptr; - cuCtxSetCurrent_t cuCtxSetCurrent_fn = nullptr; - cuCtxGetCurrent_t cuCtxGetCurrent_fn = nullptr; - - CudaDriverLoader() { - // Load CUDA driver library dynamically - cuda_driver_dll_ = LoadLibraryA("nvcuda.dll"); - if (cuda_driver_dll_) { - cuCtxCreate_v4_fn = reinterpret_cast( - GetProcAddress(cuda_driver_dll_, "cuCtxCreate_v4")); - cuCtxDestroy_fn = reinterpret_cast( - GetProcAddress(cuda_driver_dll_, "cuCtxDestroy")); - cuCtxSetCurrent_fn = reinterpret_cast( - GetProcAddress(cuda_driver_dll_, "cuCtxSetCurrent")); - cuCtxGetCurrent_fn = reinterpret_cast( - GetProcAddress(cuda_driver_dll_, "cuCtxGetCurrent")); - } - } - - ~CudaDriverLoader() { - if (cuda_driver_dll_) { - FreeLibrary(cuda_driver_dll_); - } - } - - bool IsLoaded() const { - return cuda_driver_dll_ != nullptr && - cuCtxCreate_v4_fn != nullptr && - cuCtxSetCurrent_fn != nullptr && - cuCtxDestroy_fn != nullptr && - cuCtxGetCurrent_fn != nullptr; - } -}; - -// Helper functions for D3D12 resource creation -class D3D12ResourceHelper { - public: - static void CreateSharedBuffer(ID3D12Device* device, - size_t size, - ID3D12Resource** out_resource, - D3D12_RESOURCE_STATES initial_state = D3D12_RESOURCE_STATE_COMMON) { - D3D12_RESOURCE_DESC desc = {}; - desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - desc.Alignment = 0; - desc.Width = size; - desc.Height = 1; - desc.DepthOrArraySize = 1; - desc.MipLevels = 1; - desc.Format = DXGI_FORMAT_UNKNOWN; - desc.SampleDesc.Count = 1; - desc.SampleDesc.Quality = 0; - desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - desc.Flags = D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS; - - D3D12_HEAP_PROPERTIES heap_props = {}; - heap_props.Type = D3D12_HEAP_TYPE_DEFAULT; - heap_props.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN; - heap_props.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN; - heap_props.CreationNodeMask = 1; - heap_props.VisibleNodeMask = 1; - - // Create with SHARED heap flag for cross-API import - HRESULT hr = device->CreateCommittedResource( - &heap_props, - D3D12_HEAP_FLAG_SHARED, - &desc, - initial_state, - nullptr, - IID_PPV_ARGS(out_resource)); - - if (FAILED(hr)) { - GTEST_FAIL() << "Failed to create shared D3D12 buffer, HRESULT: 0x" << std::hex << hr; - } - } - - static void CreateUploadBuffer(ID3D12Device* device, size_t size, ID3D12Resource** out_resource) { - D3D12_RESOURCE_DESC desc = {}; - desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - desc.Alignment = 0; - desc.Width = size; - desc.Height = 1; - desc.DepthOrArraySize = 1; - desc.MipLevels = 1; - desc.Format = DXGI_FORMAT_UNKNOWN; - desc.SampleDesc.Count = 1; - desc.SampleDesc.Quality = 0; - desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - desc.Flags = D3D12_RESOURCE_FLAG_NONE; - - D3D12_HEAP_PROPERTIES heap_props = {}; - heap_props.Type = D3D12_HEAP_TYPE_UPLOAD; - - HRESULT hr = device->CreateCommittedResource( - &heap_props, - D3D12_HEAP_FLAG_NONE, - &desc, - D3D12_RESOURCE_STATE_GENERIC_READ, - nullptr, - IID_PPV_ARGS(out_resource)); - - if (FAILED(hr)) { - GTEST_FAIL() << "Failed to create upload buffer, HRESULT: 0x" << std::hex << hr; - } - } - - static void CreateReadbackBuffer(ID3D12Device* device, size_t size, ID3D12Resource** out_resource) { - D3D12_RESOURCE_DESC desc = {}; - desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; - desc.Alignment = 0; - desc.Width = size; - desc.Height = 1; - desc.DepthOrArraySize = 1; - desc.MipLevels = 1; - desc.Format = DXGI_FORMAT_UNKNOWN; - desc.SampleDesc.Count = 1; - desc.SampleDesc.Quality = 0; - desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - desc.Flags = D3D12_RESOURCE_FLAG_NONE; - - D3D12_HEAP_PROPERTIES heap_props = {}; - heap_props.Type = D3D12_HEAP_TYPE_READBACK; - - HRESULT hr = device->CreateCommittedResource( - &heap_props, - D3D12_HEAP_FLAG_NONE, - &desc, - D3D12_RESOURCE_STATE_COPY_DEST, - nullptr, - IID_PPV_ARGS(out_resource)); - - if (FAILED(hr)) { - GTEST_FAIL() << "Failed to create readback buffer, HRESULT: 0x" << std::hex << hr; - } - } - - static void FlushAndWait(ID3D12Device* device, ID3D12CommandQueue* queue) { - ComPtr fence; - HRESULT hr = device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence)); - if (FAILED(hr)) { - GTEST_FAIL() << "Failed to create fence for flush, HRESULT: 0x" << std::hex << hr; - } - - HANDLE event = CreateEvent(nullptr, FALSE, FALSE, nullptr); - queue->Signal(fence.Get(), 1); - fence->SetEventOnCompletion(1, event); - WaitForSingleObject(event, INFINITE); - CloseHandle(event); - } -}; - -// Test Fixture -class NvExecutionProviderExternalResourceImporterTest : public testing::Test { - protected: - void SetUp() override { - // Get the ORT API - ort_api_ = &Ort::GetApi(); - ort_interop_api_ = &Ort::GetInteropApi(); - - // Try to create D3D12 device - HRESULT hr = D3D12CreateDevice(nullptr, D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&d3d12_device_)); - if (FAILED(hr)) { - d3d12_available_ = false; - return; - } - d3d12_available_ = true; - - // Create command queue - D3D12_COMMAND_QUEUE_DESC queue_desc = {}; - queue_desc.Type = D3D12_COMMAND_LIST_TYPE_COMPUTE; - queue_desc.Flags = D3D12_COMMAND_QUEUE_FLAG_NONE; - hr = d3d12_device_->CreateCommandQueue(&queue_desc, IID_PPV_ARGS(&command_queue_)); - if (FAILED(hr)) { - d3d12_available_ = false; - return; - } - - // Create command allocator and list - hr = d3d12_device_->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_COMPUTE, - IID_PPV_ARGS(&command_allocator_)); - if (FAILED(hr)) { - d3d12_available_ = false; - return; - } - - hr = d3d12_device_->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_COMPUTE, - command_allocator_.Get(), nullptr, - IID_PPV_ARGS(&command_list_)); - if (FAILED(hr)) { - d3d12_available_ = false; - return; - } - command_list_->Close(); - // Register NvTensorRtRtx EP - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, registered_ep_); - if (registered_ep_.get() == nullptr) { - ep_available_ = false; - return; - } - ep_available_ = true; - ep_device_ = registered_ep_.get(); - } - - void TearDown() override { - // Release resources - command_list_.Reset(); - command_allocator_.Reset(); - command_queue_.Reset(); - d3d12_device_.Reset(); - } - - bool IsD3D12Available() const { return d3d12_available_; } - bool IsEPAvailable() const { return ep_available_; } - - ComPtr d3d12_device_; - ComPtr command_queue_; - ComPtr command_allocator_; - ComPtr command_list_; - const OrtApi* ort_api_ = nullptr; - const OrtInteropApi* ort_interop_api_ = nullptr; - RegisteredEpDeviceUniquePtr registered_ep_; // RAII - auto-unregisters EP - const OrtEpDevice* ep_device_ = nullptr; - bool d3d12_available_ = false; - bool ep_available_ = false; -}; - -// Test: External Resource Importer Creation -TEST_F(NvExecutionProviderExternalResourceImporterTest, CreateExternalResourceImporter) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - // Create external resource importer - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - - if (status != nullptr) { - std::string error = ort_api_->GetErrorMessage(status); - ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "CreateExternalResourceImporterForDevice not supported: " << error; - } - - if (importer == nullptr) { - // EP doesn't support external resource import yet - GTEST_SKIP() << "External resource import not yet implemented by this EP"; - } - - // Release the importer - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Memory Import Capability Check -TEST_F(NvExecutionProviderExternalResourceImporterTest, CanImportMemoryCapabilities) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "External resource import not supported"; - } - - // Check D3D12 Resource support - bool can_import_resource = false; - status = ort_interop_api_->CanImportMemory( - importer, ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE, &can_import_resource); - ASSERT_EQ(status, nullptr) << "CanImportMemory for D3D12_RESOURCE should succeed"; - EXPECT_TRUE(can_import_resource) << "Should support D3D12 Resource import"; - - // Check D3D12 Heap support - bool can_import_heap = false; - status = ort_interop_api_->CanImportMemory( - importer, ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_HEAP, &can_import_heap); - ASSERT_EQ(status, nullptr) << "CanImportMemory for D3D12_HEAP should succeed"; - EXPECT_TRUE(can_import_heap) << "Should support D3D12 Heap import"; - - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Semaphore Import Capability Check -TEST_F(NvExecutionProviderExternalResourceImporterTest, CanImportSemaphoreCapabilities) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "External resource import not supported"; - } - - // Check D3D12 Fence support - bool can_import_fence = false; - status = ort_interop_api_->CanImportSemaphore( - importer, ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE, &can_import_fence); - ASSERT_EQ(status, nullptr) << "CanImportSemaphore for D3D12_FENCE should succeed"; - EXPECT_TRUE(can_import_fence) << "Should support D3D12 Fence import"; - - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Import D3D12 Shared Resource -TEST_F(NvExecutionProviderExternalResourceImporterTest, ImportD3D12SharedResource) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "External resource import not supported"; - } - - // Create a shared D3D12 buffer - const size_t buffer_size = 1024 * sizeof(float); - ComPtr d3d12_buffer; - D3D12ResourceHelper::CreateSharedBuffer(d3d12_device_.Get(), buffer_size, &d3d12_buffer); - - // Create shared handle - HANDLE shared_handle = nullptr; - HRESULT hr = d3d12_device_->CreateSharedHandle(d3d12_buffer.Get(), nullptr, GENERIC_ALL, nullptr, &shared_handle); - ASSERT_TRUE(SUCCEEDED(hr)) << "Failed to create shared handle"; - - // Import the memory - OrtExternalMemoryDescriptor mem_desc = {}; - mem_desc.version = ORT_API_VERSION; - mem_desc.handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE; - mem_desc.native_handle = shared_handle; - mem_desc.size_bytes = buffer_size; - mem_desc.offset_bytes = 0; - - OrtExternalMemoryHandle* mem_handle = nullptr; - status = ort_interop_api_->ImportMemory(importer, &mem_desc, &mem_handle); - ASSERT_EQ(status, nullptr) << "ImportMemory should succeed"; - ASSERT_NE(mem_handle, nullptr) << "Memory handle should not be null"; - - // Release memory handle - ort_interop_api_->ReleaseExternalMemoryHandle(mem_handle); - - // Close shared handle - CloseHandle(shared_handle); - - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Create Tensor from Imported Memory -TEST_F(NvExecutionProviderExternalResourceImporterTest, CreateTensorFromImportedMemory) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "External resource import not supported"; - } - - // Create tensor shape: [1, 3, 32, 32] (batch, channels, height, width) - const int64_t batch = 1, channels = 3, height = 32, width = 32; - const int64_t shape[] = {batch, channels, height, width}; - const size_t num_elements = batch * channels * height * width; - const size_t buffer_size = num_elements * sizeof(float); - - // Create shared D3D12 buffer - ComPtr d3d12_buffer; - D3D12ResourceHelper::CreateSharedBuffer(d3d12_device_.Get(), buffer_size, &d3d12_buffer); - - // Create shared handle - HANDLE shared_handle = nullptr; - HRESULT hr = d3d12_device_->CreateSharedHandle(d3d12_buffer.Get(), nullptr, GENERIC_ALL, nullptr, &shared_handle); - ASSERT_TRUE(SUCCEEDED(hr)); - - // Import the memory - OrtExternalMemoryDescriptor mem_desc = {}; - mem_desc.version = ORT_API_VERSION; - mem_desc.handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE; - mem_desc.native_handle = shared_handle; - mem_desc.size_bytes = buffer_size; - mem_desc.offset_bytes = 0; - - OrtExternalMemoryHandle* mem_handle = nullptr; - status = ort_interop_api_->ImportMemory(importer, &mem_desc, &mem_handle); - ASSERT_EQ(status, nullptr); - - // Create tensor from imported memory - OrtExternalTensorDescriptor tensor_desc = {}; - tensor_desc.version = ORT_API_VERSION; - tensor_desc.element_type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; - tensor_desc.shape = shape; - tensor_desc.rank = 4; - tensor_desc.offset_bytes = 0; - - OrtValue* tensor = nullptr; - status = ort_interop_api_->CreateTensorFromMemory(importer, mem_handle, &tensor_desc, &tensor); - ASSERT_EQ(status, nullptr) << "CreateTensorFromMemory should succeed"; - ASSERT_NE(tensor, nullptr) << "Tensor should not be null"; - - // Verify tensor properties - OrtTensorTypeAndShapeInfo* type_info = nullptr; - status = ort_api_->GetTensorTypeAndShape(tensor, &type_info); - ASSERT_EQ(status, nullptr); - - size_t rank = 0; - ort_api_->GetDimensionsCount(type_info, &rank); - EXPECT_EQ(rank, 4u); - - std::vector actual_shape(rank); - ort_api_->GetDimensions(type_info, actual_shape.data(), rank); - EXPECT_EQ(actual_shape[0], batch); - EXPECT_EQ(actual_shape[1], channels); - EXPECT_EQ(actual_shape[2], height); - EXPECT_EQ(actual_shape[3], width); - - ONNXTensorElementDataType elem_type; - ort_api_->GetTensorElementType(type_info, &elem_type); - EXPECT_EQ(elem_type, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); - - ort_api_->ReleaseTensorTypeAndShapeInfo(type_info); - - // Get the tensor's data pointer and verify it's CUDA device memory - // This proves the D3D12 to CUDA memory import actually happened - void* tensor_data = nullptr; - status = ort_api_->GetTensorMutableData(tensor, &tensor_data); - ASSERT_EQ(status, nullptr) << "GetTensorMutableData should succeed"; - ASSERT_NE(tensor_data, nullptr) << "Tensor data pointer should not be null"; - - // Use cudaPointerGetAttributes to verify this is CUDA device memory - cudaPointerAttributes attrs; - cudaError_t cuda_err = cudaPointerGetAttributes(&attrs, tensor_data); - ASSERT_EQ(cuda_err, cudaSuccess) << "cudaPointerGetAttributes failed: " << cudaGetErrorString(cuda_err); - EXPECT_EQ(attrs.type, cudaMemoryTypeDevice) - << "Memory should be CUDA device memory, but got type " << attrs.type - << " (cudaMemoryTypeDevice=" << cudaMemoryTypeDevice << ")"; - EXPECT_NE(attrs.device, -1) << "Device should be valid"; - - // Cleanup - ort_api_->ReleaseValue(tensor); - ort_interop_api_->ReleaseExternalMemoryHandle(mem_handle); - CloseHandle(shared_handle); - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Import D3D12 Timeline Fence -TEST_F(NvExecutionProviderExternalResourceImporterTest, ImportD3D12Fence) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "External resource import not supported"; - } - - // Create a D3D12 fence with SHARED flag for cross-API import - ComPtr d3d12_fence; - HRESULT hr = d3d12_device_->CreateFence(0, D3D12_FENCE_FLAG_SHARED, IID_PPV_ARGS(&d3d12_fence)); - ASSERT_TRUE(SUCCEEDED(hr)) << "Failed to create D3D12 fence"; - - // Create shared handle - HANDLE shared_handle = nullptr; - hr = d3d12_device_->CreateSharedHandle(d3d12_fence.Get(), nullptr, GENERIC_ALL, nullptr, &shared_handle); - ASSERT_TRUE(SUCCEEDED(hr)) << "Failed to create shared fence handle"; - - // Import the semaphore - OrtExternalSemaphoreDescriptor sem_desc = {}; - sem_desc.version = ORT_API_VERSION; - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE; - sem_desc.native_handle = shared_handle; - - OrtExternalSemaphoreHandle* sem_handle = nullptr; - status = ort_interop_api_->ImportSemaphore(importer, &sem_desc, &sem_handle); - ASSERT_EQ(status, nullptr) << "ImportSemaphore should succeed"; - ASSERT_NE(sem_handle, nullptr) << "Semaphore handle should not be null"; - - // Release semaphore handle - ort_interop_api_->ReleaseExternalSemaphoreHandle(sem_handle); - - // Close shared handle - CloseHandle(shared_handle); - - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Wait and Signal Semaphore -TEST_F(NvExecutionProviderExternalResourceImporterTest, WaitAndSignalSemaphore) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - GTEST_SKIP() << "External resource import not supported"; - } - - // Create a CUDA stream via ORT - OrtSyncStream* ort_stream = nullptr; - status = ort_api_->CreateSyncStreamForEpDevice(ep_device_, nullptr, &ort_stream); - ASSERT_EQ(status, nullptr) << "CreateSyncStreamForEpDevice should succeed"; - - // Create a D3D12 fence - ComPtr d3d12_fence; - HRESULT hr = d3d12_device_->CreateFence(0, D3D12_FENCE_FLAG_SHARED, IID_PPV_ARGS(&d3d12_fence)); - ASSERT_TRUE(SUCCEEDED(hr)); - - HANDLE shared_handle = nullptr; - hr = d3d12_device_->CreateSharedHandle(d3d12_fence.Get(), nullptr, GENERIC_ALL, nullptr, &shared_handle); - ASSERT_TRUE(SUCCEEDED(hr)); - - // Import semaphore - OrtExternalSemaphoreDescriptor sem_desc = {}; - sem_desc.version = ORT_API_VERSION; - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE; - sem_desc.native_handle = shared_handle; - - OrtExternalSemaphoreHandle* sem_handle = nullptr; - status = ort_interop_api_->ImportSemaphore(importer, &sem_desc, &sem_handle); - ASSERT_EQ(status, nullptr); - - // Signal the fence from D3D12 side - uint64_t signal_value = 1; - command_queue_->Signal(d3d12_fence.Get(), signal_value); - - // Wait on the fence from CUDA side - status = ort_interop_api_->WaitSemaphore(importer, sem_handle, ort_stream, signal_value); - ASSERT_EQ(status, nullptr) << "WaitSemaphore should succeed"; - - // Signal from CUDA side - uint64_t cuda_signal_value = 2; - status = ort_interop_api_->SignalSemaphore(importer, sem_handle, ort_stream, cuda_signal_value); - ASSERT_EQ(status, nullptr) << "SignalSemaphore should succeed"; - - // Synchronize by getting the native stream handle and calling cudaStreamSynchronize - // (In real code, the signal enqueued on the stream will complete when the stream is flushed) - void* stream_handle = ort_api_->SyncStream_GetHandle(ort_stream); - ASSERT_NE(stream_handle, nullptr); - // Note: In production code, you'd call cudaStreamSynchronize((cudaStream_t)stream_handle) - // For this test, we just verify the handle is valid - - // Verify D3D12 can see the signaled value - HANDLE wait_event = CreateEvent(nullptr, FALSE, FALSE, nullptr); - d3d12_fence->SetEventOnCompletion(cuda_signal_value, wait_event); - DWORD wait_result = WaitForSingleObject(wait_event, 5000); // 5 second timeout - CloseHandle(wait_event); - EXPECT_EQ(wait_result, WAIT_OBJECT_0) << "D3D12 should see the fence signaled by CUDA"; - - // Cleanup - ort_interop_api_->ReleaseExternalSemaphoreHandle(sem_handle); - CloseHandle(shared_handle); - ort_api_->ReleaseSyncStream(ort_stream); - ort_interop_api_->ReleaseExternalResourceImporter(importer); -} - -// Test: Full Inference with External Memory (E2E) -// This test validates the complete D3D12 to CUDA interop pipeline: -// 1. Create D3D12 shared resources and fences -// 2. Import them into CUDA via OrtExternalResourceImporter -// 3. Create ORT tensors from imported memory -// 4. Run inference with proper synchronization -// 5. Verify output correctness -TEST_F(NvExecutionProviderExternalResourceImporterTest, FullInferenceWithExternalMemory) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - - // Create a simple ReLU model using shared utility pattern - PathString model_path = ORT_TSTR("external_mem_relu_test.onnx"); - { - onnxruntime::Model model("relu_test", false, DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - - ONNX_NAMESPACE::TypeProto tensor_type; - tensor_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(64); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(64); - - auto& input_arg = graph.GetOrCreateNodeArg("X", &tensor_type); - auto& output_arg = graph.GetOrCreateNodeArg("Y", &tensor_type); - graph.AddNode("relu", "Relu", "ReLU operation", {&input_arg}, {&output_arg}); - - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_STATUS_OK(onnxruntime::Model::Save(model, model_path)); - } - - const int64_t batch = 1, channels = 3, dim = 64; - const int64_t shape[] = {batch, channels, dim, dim}; - const size_t num_elements = batch * channels * dim * dim; - const size_t buffer_size = num_elements * sizeof(float); - - // Create external resource importer - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - clearFileIfExists(model_path); - GTEST_SKIP() << "External resource import not supported"; - } - - // Create CUDA stream via ORT - OrtSyncStream* ort_stream = nullptr; - status = ort_api_->CreateSyncStreamForEpDevice(ep_device_, nullptr, &ort_stream); - ASSERT_EQ(status, nullptr); - - // Create shared D3D12 buffers for input and output - ComPtr input_buffer, output_buffer; - D3D12ResourceHelper::CreateSharedBuffer(d3d12_device_.Get(), buffer_size, &input_buffer); - D3D12ResourceHelper::CreateSharedBuffer(d3d12_device_.Get(), buffer_size, &output_buffer); - - // Create shared handles for cross-API import - HANDLE input_handle = nullptr, output_handle = nullptr; - d3d12_device_->CreateSharedHandle(input_buffer.Get(), nullptr, GENERIC_ALL, nullptr, &input_handle); - d3d12_device_->CreateSharedHandle(output_buffer.Get(), nullptr, GENERIC_ALL, nullptr, &output_handle); - - // Import memory into CUDA - OrtExternalMemoryDescriptor input_mem_desc = {}; - input_mem_desc.version = ORT_API_VERSION; - input_mem_desc.handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE; - input_mem_desc.native_handle = input_handle; - input_mem_desc.size_bytes = buffer_size; - input_mem_desc.offset_bytes = 0; - - OrtExternalMemoryDescriptor output_mem_desc = input_mem_desc; - output_mem_desc.native_handle = output_handle; - - OrtExternalMemoryHandle *input_mem = nullptr, *output_mem = nullptr; - status = ort_interop_api_->ImportMemory(importer, &input_mem_desc, &input_mem); - ASSERT_EQ(status, nullptr) << "ImportMemory for input should succeed (proves cuImportExternalMemory called)"; - status = ort_interop_api_->ImportMemory(importer, &output_mem_desc, &output_mem); - ASSERT_EQ(status, nullptr) << "ImportMemory for output should succeed"; - - // Create ORT tensors from imported memory - OrtExternalTensorDescriptor tensor_desc = {}; - tensor_desc.version = ORT_API_VERSION; - tensor_desc.element_type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; - tensor_desc.shape = shape; - tensor_desc.rank = 4; - tensor_desc.offset_bytes = 0; - - OrtValue *input_tensor = nullptr, *output_tensor = nullptr; - status = ort_interop_api_->CreateTensorFromMemory(importer, input_mem, &tensor_desc, &input_tensor); - ASSERT_EQ(status, nullptr); - status = ort_interop_api_->CreateTensorFromMemory(importer, output_mem, &tensor_desc, &output_tensor); - ASSERT_EQ(status, nullptr); - - // Verify the tensor data pointers are CUDA device memory - void* input_data_ptr = nullptr; - void* output_data_ptr = nullptr; - status = ort_api_->GetTensorMutableData(input_tensor, &input_data_ptr); - ASSERT_EQ(status, nullptr); - status = ort_api_->GetTensorMutableData(output_tensor, &output_data_ptr); - ASSERT_EQ(status, nullptr); - - cudaPointerAttributes input_attrs, output_attrs; - ASSERT_EQ(cudaPointerGetAttributes(&input_attrs, input_data_ptr), cudaSuccess); - ASSERT_EQ(cudaPointerGetAttributes(&output_attrs, output_data_ptr), cudaSuccess); - EXPECT_EQ(input_attrs.type, cudaMemoryTypeDevice) << "Input tensor must be CUDA device memory"; - EXPECT_EQ(output_attrs.type, cudaMemoryTypeDevice) << "Output tensor must be CUDA device memory"; - - // Create D3D12 fence for bidirectional synchronization - ComPtr sync_fence; - d3d12_device_->CreateFence(0, D3D12_FENCE_FLAG_SHARED, IID_PPV_ARGS(&sync_fence)); - HANDLE fence_handle = nullptr; - d3d12_device_->CreateSharedHandle(sync_fence.Get(), nullptr, GENERIC_ALL, nullptr, &fence_handle); - - OrtExternalSemaphoreDescriptor sem_desc = {}; - sem_desc.version = ORT_API_VERSION; - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE; - sem_desc.native_handle = fence_handle; - - OrtExternalSemaphoreHandle* sem_handle = nullptr; - status = ort_interop_api_->ImportSemaphore(importer, &sem_desc, &sem_handle); - ASSERT_EQ(status, nullptr) << "ImportSemaphore should succeed"; - - // Setup test data via D3D12 upload buffer - ComPtr upload_buffer; - D3D12ResourceHelper::CreateUploadBuffer(d3d12_device_.Get(), buffer_size, &upload_buffer); - - // Generate test data: alternating positive and negative values for ReLU verification - std::vector test_data(num_elements); - for (size_t i = 0; i < num_elements; ++i) { - test_data[i] = (i % 2 == 0) ? static_cast(i + 1) : -static_cast(i + 1); - } - - void* upload_ptr = nullptr; - upload_buffer->Map(0, nullptr, &upload_ptr); - memcpy(upload_ptr, test_data.data(), buffer_size); - upload_buffer->Unmap(0, nullptr); - - // Copy upload buffer to input buffer via D3D12 - command_allocator_->Reset(); - command_list_->Reset(command_allocator_.Get(), nullptr); - command_list_->CopyBufferRegion(input_buffer.Get(), 0, upload_buffer.Get(), 0, buffer_size); - command_list_->Close(); - - ID3D12CommandList* cmd_lists[] = {command_list_.Get()}; - command_queue_->ExecuteCommandLists(1, cmd_lists); - - // Signal fence after D3D12 upload completes - uint64_t upload_complete_value = 1; - command_queue_->Signal(sync_fence.Get(), upload_complete_value); - - // Make CUDA wait for D3D12 upload to complete - status = ort_interop_api_->WaitSemaphore(importer, sem_handle, ort_stream, upload_complete_value); - ASSERT_EQ(status, nullptr) << "WaitSemaphore should succeed"; - - // Setup ORT session with user_compute_stream - Ort::SessionOptions session_options; - session_options.SetExecutionMode(ORT_SEQUENTIAL); - session_options.DisableMemPattern(); - session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL); - - // Add the NvTensorRtRtx EP with user stream - status = ort_api_->SessionOptionsAppendExecutionProvider_V2( - session_options, *ort_env, &ep_device_, 1, - nullptr, nullptr, 0); - ASSERT_EQ(status, nullptr); - - // Create session - Ort::Session session(*ort_env, model_path.c_str(), session_options); - - // Create IoBinding and bind external tensors - Ort::IoBinding io_binding(session); - Ort::AllocatorWithDefaultOptions allocator; - - Ort::AllocatedStringPtr input_name = session.GetInputNameAllocated(0, allocator); - Ort::AllocatedStringPtr output_name = session.GetOutputNameAllocated(0, allocator); - Ort::Value input(input_tensor); - Ort::Value output(output_tensor); - io_binding.BindInput(input_name.get(), input); - io_binding.BindOutput(output_name.get(), output); - io_binding.SynchronizeInputs(); - // Run inference. ORT submits all work to the stream before returning, so we signal the async semaphore below. - Ort::RunOptions run_options; - run_options.SetSyncStream(ort_stream); - run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); - session.Run(run_options, io_binding); - - // Signal from CUDA that inference is complete - uint64_t inference_complete_value = 2; - status = ort_interop_api_->SignalSemaphore(importer, sem_handle, ort_stream, inference_complete_value); - ASSERT_EQ(status, nullptr) << "SignalSemaphore should succeed"; - - // Wait on D3D12 for CUDA inference to complete - command_queue_->Wait(sync_fence.Get(), inference_complete_value); - - // Copy output to readback buffer - ComPtr readback_buffer; - D3D12ResourceHelper::CreateReadbackBuffer(d3d12_device_.Get(), buffer_size, &readback_buffer); - - command_allocator_->Reset(); - command_list_->Reset(command_allocator_.Get(), nullptr); - command_list_->CopyBufferRegion(readback_buffer.Get(), 0, output_buffer.Get(), 0, buffer_size); - command_list_->Close(); - - command_queue_->ExecuteCommandLists(1, cmd_lists); - D3D12ResourceHelper::FlushAndWait(d3d12_device_.Get(), command_queue_.Get()); - - // Read back and verify ReLU output: max(0, x) - std::vector output_data(num_elements); - void* readback_ptr = nullptr; - readback_buffer->Map(0, nullptr, &readback_ptr); - memcpy(output_data.data(), readback_ptr, buffer_size); - readback_buffer->Unmap(0, nullptr); - - // Verify ReLU correctness - for (size_t i = 0; i < num_elements; ++i) { - float expected = std::max(0.0f, test_data[i]); - EXPECT_FLOAT_EQ(output_data[i], expected) - << "Mismatch at index " << i << ": input=" << test_data[i] - << ", expected=" << expected << ", got=" << output_data[i]; - } - - // Note: io_binding takes ownership of input_tensor and output_tensor, so don't release them manually - - // Cleanup - ort_interop_api_->ReleaseExternalSemaphoreHandle(sem_handle); - ort_interop_api_->ReleaseExternalMemoryHandle(output_mem); - ort_interop_api_->ReleaseExternalMemoryHandle(input_mem); - CloseHandle(fence_handle); - CloseHandle(output_handle); - CloseHandle(input_handle); - ort_api_->ReleaseSyncStream(ort_stream); - ort_interop_api_->ReleaseExternalResourceImporter(importer); - - clearFileIfExists(model_path); -} - -// Test: Full Inference with External Memory (E2E) -// This test validates the complete D3D12 to CUDA interop pipeline: -// 1. Create D3D12 shared resources and fences -// 2. Import them into CUDA via OrtExternalResourceImporter -// 3. Create ORT tensors from imported memory -// 4. Run inference with proper synchronization -// 5. Verify output correctness -TEST_F(NvExecutionProviderExternalResourceImporterTest, FullInferenceWithExternalMemoryCIG) { - if (!IsD3D12Available()) { - GTEST_SKIP() << "D3D12 not available"; - } - - // Push CIG context - CUctxCigParam ctxCigParams; - ctxCigParams.sharedDataType = CIG_DATA_TYPE_D3D12_COMMAND_QUEUE; - ctxCigParams.sharedData = command_queue_.Get(); - CUctxCreateParams ctxParams = {nullptr, 0, &ctxCigParams}; - CUcontext cig_context; - LUID d3d12_luid_struct = d3d12_device_->GetAdapterLuid(); - uint64_t d3d12_luid = (static_cast(d3d12_luid_struct.HighPart) << 32) | d3d12_luid_struct.LowPart; - int cuda_device_count = 0; - int cuda_device_id = 0; - bool found = false; - ASSERT_EQ(cudaGetDeviceCount(&cuda_device_count), cudaSuccess); - for (; cuda_device_id < cuda_device_count; ++cuda_device_id) { - cudaDeviceProp prop; - ASSERT_EQ(cudaGetDeviceProperties(&prop, cuda_device_id), cudaSuccess); - uint64_t cuda_luid; - std::memcpy(&cuda_luid, prop.luid, sizeof(char) * 8); - if (cuda_luid == d3d12_luid) { - found = true; - break; - } - } - ASSERT_TRUE(found); - CudaDriverLoader cuda_driver_loader; - ASSERT_TRUE(cuda_driver_loader.IsLoaded()); - ASSERT_EQ(cuda_driver_loader.cuCtxCreate_v4_fn(&cig_context, &ctxParams, 0, cuda_device_id), CUDA_SUCCESS); - ASSERT_EQ(cuda_driver_loader.cuCtxSetCurrent_fn(cig_context), CUDA_SUCCESS); - - if (!IsEPAvailable()) { - GTEST_SKIP() << "NvTensorRtRtx EP not available"; - } - { - // Create a simple ReLU model using shared utility pattern - PathString model_path = ORT_TSTR("external_mem_relu_test.onnx"); - { - onnxruntime::Model model("relu_test", false, DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - - ONNX_NAMESPACE::TypeProto tensor_type; - tensor_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(64); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(64); - - auto& input_arg = graph.GetOrCreateNodeArg("X", &tensor_type); - auto& output_arg = graph.GetOrCreateNodeArg("Y", &tensor_type); - graph.AddNode("relu", "Relu", "ReLU operation", {&input_arg}, {&output_arg}); - - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_STATUS_OK(onnxruntime::Model::Save(model, model_path)); - } - - const int64_t batch = 1, channels = 3, dim = 64; - const int64_t shape[] = {batch, channels, dim, dim}; - const size_t num_elements = batch * channels * dim * dim; - const size_t buffer_size = num_elements * sizeof(float); - - // Create external resource importer - OrtExternalResourceImporter* importer = nullptr; - OrtStatus* status = ort_interop_api_->CreateExternalResourceImporterForDevice(ep_device_, &importer); - if (status != nullptr || importer == nullptr) { - if (status != nullptr) ort_api_->ReleaseStatus(status); - clearFileIfExists(model_path); - GTEST_SKIP() << "External resource import not supported"; - } - - // Create CUDA stream via ORT - OrtSyncStream* ort_stream = nullptr; - status = ort_api_->CreateSyncStreamForEpDevice(ep_device_, nullptr, &ort_stream); - ASSERT_EQ(status, nullptr); - - // Create shared D3D12 buffers for input and output - ComPtr input_buffer, output_buffer; - D3D12ResourceHelper::CreateSharedBuffer(d3d12_device_.Get(), buffer_size, &input_buffer); - D3D12ResourceHelper::CreateSharedBuffer(d3d12_device_.Get(), buffer_size, &output_buffer); - - // Create shared handles for cross-API import - HANDLE input_handle = nullptr, output_handle = nullptr; - d3d12_device_->CreateSharedHandle(input_buffer.Get(), nullptr, GENERIC_ALL, nullptr, &input_handle); - d3d12_device_->CreateSharedHandle(output_buffer.Get(), nullptr, GENERIC_ALL, nullptr, &output_handle); - - // Import memory into CUDA - OrtExternalMemoryDescriptor input_mem_desc = {}; - input_mem_desc.version = ORT_API_VERSION; - input_mem_desc.handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_D3D12_RESOURCE; - input_mem_desc.native_handle = input_handle; - input_mem_desc.size_bytes = buffer_size; - input_mem_desc.offset_bytes = 0; - - OrtExternalMemoryDescriptor output_mem_desc = input_mem_desc; - output_mem_desc.native_handle = output_handle; - - OrtExternalMemoryHandle *input_mem = nullptr, *output_mem = nullptr; - status = ort_interop_api_->ImportMemory(importer, &input_mem_desc, &input_mem); - ASSERT_EQ(status, nullptr) << "ImportMemory for input should succeed (proves cuImportExternalMemory called)"; - status = ort_interop_api_->ImportMemory(importer, &output_mem_desc, &output_mem); - ASSERT_EQ(status, nullptr) << "ImportMemory for output should succeed"; - - // Create ORT tensors from imported memory - OrtExternalTensorDescriptor tensor_desc = {}; - tensor_desc.version = ORT_API_VERSION; - tensor_desc.element_type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; - tensor_desc.shape = shape; - tensor_desc.rank = 4; - tensor_desc.offset_bytes = 0; - - OrtValue *input_tensor = nullptr, *output_tensor = nullptr; - status = ort_interop_api_->CreateTensorFromMemory(importer, input_mem, &tensor_desc, &input_tensor); - ASSERT_EQ(status, nullptr); - status = ort_interop_api_->CreateTensorFromMemory(importer, output_mem, &tensor_desc, &output_tensor); - ASSERT_EQ(status, nullptr); - - // Verify the tensor data pointers are CUDA device memory - void* input_data_ptr = nullptr; - void* output_data_ptr = nullptr; - status = ort_api_->GetTensorMutableData(input_tensor, &input_data_ptr); - ASSERT_EQ(status, nullptr); - status = ort_api_->GetTensorMutableData(output_tensor, &output_data_ptr); - ASSERT_EQ(status, nullptr); - - cudaPointerAttributes input_attrs, output_attrs; - ASSERT_EQ(cudaPointerGetAttributes(&input_attrs, input_data_ptr), cudaSuccess); - ASSERT_EQ(cudaPointerGetAttributes(&output_attrs, output_data_ptr), cudaSuccess); - EXPECT_EQ(input_attrs.type, cudaMemoryTypeDevice) << "Input tensor must be CUDA device memory"; - EXPECT_EQ(output_attrs.type, cudaMemoryTypeDevice) << "Output tensor must be CUDA device memory"; - - // Create D3D12 fence for bidirectional synchronization - ComPtr sync_fence; - d3d12_device_->CreateFence(0, D3D12_FENCE_FLAG_SHARED, IID_PPV_ARGS(&sync_fence)); - HANDLE fence_handle = nullptr; - d3d12_device_->CreateSharedHandle(sync_fence.Get(), nullptr, GENERIC_ALL, nullptr, &fence_handle); - - OrtExternalSemaphoreDescriptor sem_desc = {}; - sem_desc.version = ORT_API_VERSION; - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_D3D12_FENCE; - sem_desc.native_handle = fence_handle; - - OrtExternalSemaphoreHandle* sem_handle = nullptr; - status = ort_interop_api_->ImportSemaphore(importer, &sem_desc, &sem_handle); - ASSERT_EQ(status, nullptr) << "ImportSemaphore should succeed"; - - // Setup test data via D3D12 upload buffer - ComPtr upload_buffer; - D3D12ResourceHelper::CreateUploadBuffer(d3d12_device_.Get(), buffer_size, &upload_buffer); - - // Generate test data: alternating positive and negative values for ReLU verification - std::vector test_data(num_elements); - for (size_t i = 0; i < num_elements; ++i) { - test_data[i] = (i % 2 == 0) ? static_cast(i + 1) : -static_cast(i + 1); - } - - void* upload_ptr = nullptr; - upload_buffer->Map(0, nullptr, &upload_ptr); - memcpy(upload_ptr, test_data.data(), buffer_size); - upload_buffer->Unmap(0, nullptr); - - // Copy upload buffer to input buffer via D3D12 - command_allocator_->Reset(); - command_list_->Reset(command_allocator_.Get(), nullptr); - command_list_->CopyBufferRegion(input_buffer.Get(), 0, upload_buffer.Get(), 0, buffer_size); - command_list_->Close(); - - ID3D12CommandList* cmd_lists[] = {command_list_.Get()}; - command_queue_->ExecuteCommandLists(1, cmd_lists); - - // Signal fence after D3D12 upload completes - uint64_t upload_complete_value = 1; - command_queue_->Signal(sync_fence.Get(), upload_complete_value); - - // Make CUDA wait for D3D12 upload to complete - status = ort_interop_api_->WaitSemaphore(importer, sem_handle, ort_stream, upload_complete_value); - ASSERT_EQ(status, nullptr) << "WaitSemaphore should succeed"; - - // Setup ORT session with user_compute_stream - Ort::SessionOptions session_options; - session_options.SetExecutionMode(ORT_SEQUENTIAL); - session_options.DisableMemPattern(); - session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL); - - // Configure to use our CUDA stream - char stream_address[32]; - size_t stream_addr_val = reinterpret_cast(ort_api_->SyncStream_GetHandle(ort_stream)); - sprintf_s(stream_address, "%llu", static_cast(stream_addr_val)); - const char* option_keys[] = { - // TODO we should no longer require to set the compute stream at this point but there are too many cudaSetDevice calls from allocators and stream handling (NVBUG 5822116) - onnxruntime::nv::provider_option_names::kUserComputeStream, - onnxruntime::nv::provider_option_names::kHasUserComputeStream, - onnxruntime::nv::provider_option_names::kMaxSharedMemSize, - // TRT will create itss own context to create streams if we do not manually provide aux streams - onnxruntime::nv::provider_option_names::kLengthAuxStreamArray, - onnxruntime::nv::provider_option_names::kUserAuxStreamArray, - onnxruntime::nv::provider_option_names::kCudaGraphEnable, - }; - char aux_stream_address[32]; - size_t aux_streams[] = {stream_addr_val}; - sprintf_s(aux_stream_address, "%llu", reinterpret_cast(aux_streams)); - std::string max_shared_mem_size = std::to_string(1024 * 28); // 28 KiB - const char* option_values[] = { - stream_address, - "1", - max_shared_mem_size.c_str(), - "1", - aux_stream_address, - "0"}; - - // Add the NvTensorRtRtx EP with user stream - status = ort_api_->SessionOptionsAppendExecutionProvider_V2( - session_options, *ort_env, &ep_device_, 1, option_keys, option_values, 6); - ASSERT_EQ(status, nullptr); - - // Create session - Ort::Session session(*ort_env, model_path.c_str(), session_options); - - // Create IoBinding and bind external tensors - Ort::IoBinding io_binding(session); - Ort::AllocatorWithDefaultOptions allocator; - - Ort::AllocatedStringPtr input_name = session.GetInputNameAllocated(0, allocator); - Ort::AllocatedStringPtr output_name = session.GetOutputNameAllocated(0, allocator); - - Ort::Value input(input_tensor); - Ort::Value output(output_tensor); - io_binding.BindInput(input_name.get(), input); - io_binding.BindOutput(output_name.get(), output); - io_binding.SynchronizeInputs(); - // Run inference. ORT submits all work to the stream before returning, so we signal the async semaphore below. - Ort::RunOptions run_options; - run_options.SetSyncStream(ort_stream); - run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); - session.Run(run_options, io_binding); - - // Signal from CUDA that inference is complete - uint64_t inference_complete_value = 2; - status = ort_interop_api_->SignalSemaphore(importer, sem_handle, ort_stream, inference_complete_value); - ASSERT_EQ(status, nullptr) << "SignalSemaphore should succeed"; - - // Wait on D3D12 for CUDA inference to complete - command_queue_->Wait(sync_fence.Get(), inference_complete_value); - - // Copy output to readback buffer - ComPtr readback_buffer; - D3D12ResourceHelper::CreateReadbackBuffer(d3d12_device_.Get(), buffer_size, &readback_buffer); - - command_allocator_->Reset(); - command_list_->Reset(command_allocator_.Get(), nullptr); - command_list_->CopyBufferRegion(readback_buffer.Get(), 0, output_buffer.Get(), 0, buffer_size); - command_list_->Close(); - - command_queue_->ExecuteCommandLists(1, cmd_lists); - D3D12ResourceHelper::FlushAndWait(d3d12_device_.Get(), command_queue_.Get()); - - // Read back and verify ReLU output: max(0, x) - std::vector output_data(num_elements); - void* readback_ptr = nullptr; - readback_buffer->Map(0, nullptr, &readback_ptr); - memcpy(output_data.data(), readback_ptr, buffer_size); - readback_buffer->Unmap(0, nullptr); - - // Verify ReLU correctness - for (size_t i = 0; i < num_elements; ++i) { - float expected = std::max(0.0f, test_data[i]); - EXPECT_FLOAT_EQ(output_data[i], expected) - << "Mismatch at index " << i << ": input=" << test_data[i] - << ", expected=" << expected << ", got=" << output_data[i]; - } - CUcontext cu_context; - ASSERT_EQ(cuda_driver_loader.cuCtxGetCurrent_fn(&cu_context), CUDA_SUCCESS); - ASSERT_EQ(cu_context, cig_context); - // Note: io_binding takes ownership of input_tensor and output_tensor, so don't release them manually - - // Cleanup - ort_interop_api_->ReleaseExternalSemaphoreHandle(sem_handle); - ort_interop_api_->ReleaseExternalMemoryHandle(output_mem); - ort_interop_api_->ReleaseExternalMemoryHandle(input_mem); - CloseHandle(fence_handle); - CloseHandle(output_handle); - CloseHandle(input_handle); - ort_api_->ReleaseSyncStream(ort_stream); - ort_interop_api_->ReleaseExternalResourceImporter(importer); - clearFileIfExists(model_path); - } - // all associated objects with the context must be destroyed before destroying the CIG context which ORT inherited - ASSERT_EQ(cuda_driver_loader.cuCtxDestroy_fn(cig_context), CUDA_SUCCESS); -} - -#endif // _WIN32 - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_options_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_options_test.cc deleted file mode 100644 index 8ba03ad19b755..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_options_test.cc +++ /dev/null @@ -1,81 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#include "core/graph/onnx_protobuf.h" -#include "core/session/inference_session.h" -#include "test/providers/provider_test_utils.h" - -#include "test/util/include/scoped_env_vars.h" -#include "test/common/trt_op_test_utils.h" -#include "test/common/random_generator.h" -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" - -#include -#include - -using namespace std; -using namespace ONNX_NAMESPACE; -using namespace ::onnxruntime::logging; -extern std::unique_ptr ort_env; -namespace onnxruntime { - -namespace test { -size_t countFilesInDirectory(const std::string& dir_path) { - return std::distance(std::filesystem::directory_iterator(dir_path), std::filesystem::directory_iterator{}); -} - -TEST(NvExecutionProviderTest, RuntimeCaching) { - PathString model_name = ORT_TSTR("nv_execution_provider_runtime_caching.onnx"); - PathString model_name_ctx = ORT_TSTR("nv_execution_provider_runtime_caching_ctx.onnx"); - auto model_name_ctx_str = PathToUTF8(model_name_ctx); - clearFileIfExists(model_name_ctx); - std::string graph_name = "test"; - std::vector dims = {1, 3, 2}; - std::string runtime_cache_name = "./runtime_cache/"; - if (std::filesystem::exists(runtime_cache_name)) { - std::filesystem::remove_all(runtime_cache_name); - } - CreateBaseModel(model_name, graph_name, dims); - // AOT time - { - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); - so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {{"nv_runtime_cache_path", runtime_cache_name.c_str()}}); - Ort::Session session_object(*ort_env, model_name.c_str(), so); - - auto io_binding = generate_io_binding(session_object); - session_object.Run(run_options, io_binding); - } - // the cache will be dumped to disk upon session destruction - ASSERT_TRUE(std::filesystem::exists(runtime_cache_name.c_str())); - ASSERT_TRUE(1 == countFilesInDirectory(runtime_cache_name)); - - // use existing cache - { - Ort::SessionOptions so; - Ort::RunOptions run_options; - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {{"nv_runtime_cache_path", runtime_cache_name.c_str()}}); - Ort::Session session_object(*ort_env, model_name_ctx.c_str(), so); - } - ASSERT_TRUE(1 == countFilesInDirectory(runtime_cache_name)); - - // create new cache - { - Ort::SessionOptions so; - Ort::RunOptions run_options; - std::string new_cache_name = "/tmp/runtime_cache_new/"; - if (std::filesystem::exists(new_cache_name)) { - std::filesystem::remove_all(new_cache_name); - } - so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {{"nv_runtime_cache_path", new_cache_name.c_str()}}); - { - Ort::Session session_object(*ort_env, model_name_ctx.c_str(), so); - } - // the cache will be dumped to disk upon session destruction - ASSERT_TRUE(std::filesystem::exists(new_cache_name.c_str())); - ASSERT_TRUE(1 == countFilesInDirectory(new_cache_name)); - } -} -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_vulkan_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_vulkan_test.cc deleted file mode 100644 index 918e137025807..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_vulkan_test.cc +++ /dev/null @@ -1,895 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include "core/session/onnxruntime_cxx_api.h" -#include "core/graph/model.h" -#include -#include -#include -#include -#if defined(_WIN32) -#define VK_USE_PLATFORM_WIN32_KHR -#endif -#define VULKAN_HPP_DISPATCH_LOADER_DYNAMIC 1 -#define VULKAN_HPP_NO_EXCEPTIONS -#define VULKAN_HPP_NO_SMART_HANDLE -#define VULKAN_HPP_NO_CONSTRUCTORS -#include - -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" -#include "test/providers/provider_test_utils.h" - -#if __linux__ -#include -#include -#define HMODULE void* -#define CUDALIB "libcuda.so.1" -#define DLOPENOPTIONS , RTLD_LAZY -#else -#define STRINGIFY(a) _STRINGIFY(a) -#define _STRINGIFY(a) #a -#include -#define dlopen LoadLibraryA -#define dlclose FreeLibrary -#define dlsym GetProcAddress -#define RTLD_LAZY -#define CUDALIB "nvcuda.dll" -#define DLOPENOPTIONS -#endif - -// Small util for Go-like defer -#define DEFER(resource, x) \ - std::shared_ptr resource##_finalizer(nullptr, [&](...) { x; }) - -extern std::unique_ptr ort_env; - -namespace onnxruntime::test { - -namespace { -// Dynamic CUDA driver function loader -class CudaDriverLoader { - private: - HMODULE cuda_driver_dll_ = nullptr; - - // CUDA Driver API function pointers - using cuCtxCreate_v4_t = CUresult (*)(CUcontext*, CUctxCreateParams*, unsigned int, CUdevice); - using cuCtxDestroy_t = CUresult (*)(CUcontext); - using cuCtxGetCurrent_t = CUresult (*)(CUcontext*); - using cuCtxSetCurrent_t = CUresult (*)(CUcontext); - using cuDeviceGetAttribute_t = CUresult (*)(int*, CUdevice_attribute attrib, CUdevice dev); - - public: - cuCtxCreate_v4_t cuCtxCreate_v4_fn = nullptr; - cuCtxDestroy_t cuCtxDestroy_fn = nullptr; - cuCtxSetCurrent_t cuCtxSetCurrent_fn = nullptr; - cuCtxGetCurrent_t cuCtxGetCurrent_fn = nullptr; - cuDeviceGetAttribute_t cuDeviceGetAttribute_fn = nullptr; - - CudaDriverLoader() { - cuda_driver_dll_ = dlopen(CUDALIB DLOPENOPTIONS); - if (cuda_driver_dll_) { - cuCtxCreate_v4_fn = reinterpret_cast( - dlsym(cuda_driver_dll_, "cuCtxCreate_v4")); - cuCtxDestroy_fn = reinterpret_cast( - dlsym(cuda_driver_dll_, "cuCtxDestroy")); - cuCtxSetCurrent_fn = reinterpret_cast( - dlsym(cuda_driver_dll_, "cuCtxSetCurrent")); - cuCtxGetCurrent_fn = reinterpret_cast( - dlsym(cuda_driver_dll_, "cuCtxGetCurrent")); - cuDeviceGetAttribute_fn = reinterpret_cast( - dlsym(cuda_driver_dll_, "cuDeviceGetAttribute")); - } - } - - ~CudaDriverLoader() { - if (cuda_driver_dll_) { - dlclose(cuda_driver_dll_); - } - } - - bool IsLoaded() const { - return cuda_driver_dll_ != nullptr && - cuCtxCreate_v4_fn != nullptr && - cuCtxSetCurrent_fn != nullptr && - cuCtxDestroy_fn != nullptr && - cuCtxGetCurrent_fn != nullptr && - cuDeviceGetAttribute_fn != nullptr; - } -}; - -struct NvDevice { - VkPhysicalDevice phys_dev{}; - VkPhysicalDeviceProperties props{}; - VkPhysicalDeviceVulkan11Properties id_props{}; - std::vector ep_device_candidates; - bool has_cig_extension{}; -}; - -struct ExportableTimelineSemaphore { - VkSemaphore vk_handle{}; - void* native_handle{}; - OrtExternalSemaphoreHandle* ort_handle{}; -}; - -struct ExportableBuffer { - VkBuffer buffer{}; - VkBufferView view{}; - VkDeviceMemory memory{}; - void* native_handle{}; - OrtExternalMemoryHandle* ort_handle{}; -}; - -struct TestParameters { - bool force_cig_if_supported{}; - bool allow_cig{}; - bool use_dmabuf{}; - bool use_init_graphics_interop_call{}; -}; - -struct VkResources { - vk::detail::DispatchLoaderDynamic loader; - VkInstance instance{}; - VkDevice device{}; - VkPhysicalDevice phys_device{}; - std::vector buffers; - std::vector semaphores; - std::vector nv_devices; - VkQueue queue{}; - VkExternalComputeQueueNV ex_compute_queue{}; - OrtExternalResourceImporter* importer{}; - VkCommandPool cmd_pool{}; - VkCommandBuffer upload_cmd_buf{}; - VkCommandBuffer download_cmd_buf{}; - std::optional ep_device; - - ~VkResources() { - auto& interop = Ort::GetInteropApi(); - // Release ORT external handles first (they reference Vulkan objects) - for (auto& sem : semaphores) { - if (sem.ort_handle) { - interop.ReleaseExternalSemaphoreHandle(sem.ort_handle); - sem.ort_handle = nullptr; - } - } - for (auto& buf : buffers) { - if (buf.ort_handle) { - interop.ReleaseExternalMemoryHandle(buf.ort_handle); - buf.ort_handle = nullptr; - } - } - if (importer) { - interop.ReleaseExternalResourceImporter(importer); - importer = nullptr; - } - if (ep_device) { - // We deinit even if we never inited for graphics - EXPECT_TRUE(Ort::Status(interop.DeinitGraphicsInteropForEpDevice(*ep_device)).IsOK()); - } - - // Destroy Vulkan objects in correct order: buffer views -> buffers -> memory, then semaphores, queue, device, instance - if (device != VK_NULL_HANDLE && loader.vkDeviceWaitIdle) { - loader.vkDeviceWaitIdle(device); - for (auto& buf : buffers) { - if (buf.view != VK_NULL_HANDLE) { - loader.vkDestroyBufferView(device, buf.view, nullptr); - buf.view = VK_NULL_HANDLE; - } - if (buf.buffer != VK_NULL_HANDLE) { - loader.vkDestroyBuffer(device, buf.buffer, nullptr); - buf.buffer = VK_NULL_HANDLE; - } - if (buf.memory != VK_NULL_HANDLE) { - loader.vkFreeMemory(device, buf.memory, nullptr); - buf.memory = VK_NULL_HANDLE; - } - } - for (auto& sem : semaphores) { - if (sem.vk_handle != VK_NULL_HANDLE) { - loader.vkDestroySemaphore(device, sem.vk_handle, nullptr); - sem.vk_handle = VK_NULL_HANDLE; - } - } - if (ex_compute_queue != VK_NULL_HANDLE) { - loader.vkDestroyExternalComputeQueueNV(device, ex_compute_queue, nullptr); - ex_compute_queue = VK_NULL_HANDLE; - } - if (upload_cmd_buf != VK_NULL_HANDLE || download_cmd_buf != VK_NULL_HANDLE) { - std::vector to_free; - if (upload_cmd_buf != VK_NULL_HANDLE) to_free.push_back(upload_cmd_buf); - if (download_cmd_buf != VK_NULL_HANDLE) to_free.push_back(download_cmd_buf); - loader.vkFreeCommandBuffers(device, cmd_pool, static_cast(to_free.size()), to_free.data()); - upload_cmd_buf = VK_NULL_HANDLE; - download_cmd_buf = VK_NULL_HANDLE; - } - if (cmd_pool != VK_NULL_HANDLE) { - loader.vkDestroyCommandPool(device, cmd_pool, nullptr); - cmd_pool = VK_NULL_HANDLE; - } - loader.vkDestroyDevice(device, nullptr); - device = VK_NULL_HANDLE; - } - if (instance != VK_NULL_HANDLE && loader.vkDestroyInstance) { - loader.vkDestroyInstance(instance, nullptr); - instance = VK_NULL_HANDLE; - } - } -}; - -void init_vulkan_interop(VkResources& resources) { - resources.loader.init(); - - VkApplicationInfo app_info{}; - app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO; - app_info.applicationVersion = VK_MAKE_VERSION(0, 1, 0); - app_info.pApplicationName = "ORT"; - app_info.engineVersion = VK_MAKE_VERSION(0, 1, 0); - app_info.apiVersion = VK_API_VERSION_1_4; - app_info.pEngineName = "ORT"; - - std::vector instance_extensions{}; - VkInstanceCreateInfo instance_info{}; - instance_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO; - instance_info.pApplicationInfo = &app_info; - - auto instance_creation_result = resources.loader.vkCreateInstance(&instance_info, nullptr, &resources.instance); - if (instance_creation_result != VK_SUCCESS) { - GTEST_SKIP() << "Vulkan instance creation failed: skipping Vulkan interop test"; - } - resources.loader.init(vk::Instance(resources.instance)); - - std::vector physical_devices; - uint32_t count{}; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkEnumeratePhysicalDevices(resources.instance, &count, nullptr)); - physical_devices.resize(count); - EXPECT_EQ(VK_SUCCESS, resources.loader.vkEnumeratePhysicalDevices(resources.instance, &count, physical_devices.data())); - - auto ep_devices = ort_env->GetEpDevices(); -#if !defined(_WIN32) - std::regex pci_bus_id_pattern("([a-fA-F0-9]+):([a-fA-F0-9]+):([a-fA-F0-9]+)\\.([a-fA-F0-9]+)"); -#endif - - for (auto& p : physical_devices) { - NvDevice dev; - dev.phys_dev = p; - std::vector ext_props; - uint32_t num_extensions; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkEnumerateDeviceExtensionProperties(p, nullptr, &num_extensions, nullptr)); - ext_props.resize(num_extensions); - EXPECT_EQ(VK_SUCCESS, resources.loader.vkEnumerateDeviceExtensionProperties(p, nullptr, &num_extensions, ext_props.data())); - for (const auto& prop : ext_props) { - if (std::strcmp(prop.extensionName, VK_NV_EXTERNAL_COMPUTE_QUEUE_EXTENSION_NAME) == 0) { - dev.has_cig_extension = true; - break; - } - } - - VkPhysicalDevicePCIBusInfoPropertiesEXT pci_props{}; - pci_props.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PCI_BUS_INFO_PROPERTIES_EXT; - VkPhysicalDeviceVulkan11Properties id_props{}; - id_props.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_PROPERTIES; - id_props.pNext = &pci_props; - VkPhysicalDeviceProperties2 props{}; - props.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2; - props.pNext = &id_props; - resources.loader.vkGetPhysicalDeviceProperties2(p, &props); - if (props.properties.vendorID == 0x10DE) { - std::memcpy(&dev.props, &props, sizeof(props)); - std::memcpy(&dev.id_props, &id_props, sizeof(id_props)); - - for (const auto& d : ep_devices) { - if (d.Device().VendorId() == props.properties.vendorID && d.Device().DeviceId() == props.properties.deviceID) { -#if defined(_WIN32) - // verify the real device with LUID, but on Linux we only have UUID which we don't know for EpDevice - auto luid = d.Device().Metadata().GetValue("LUID"); - if (id_props.deviceLUIDValid && luid) { - LUID vk_luid; - std::memcpy(&vk_luid, dev.id_props.deviceLUID, sizeof(LUID)); - uint64_t ep_luid = std::stoull(luid); - uint64_t vk = (uint64_t(vk_luid.HighPart) << 32) | uint64_t(vk_luid.LowPart); - if (ep_luid != vk) { - continue; - } - } -#else - auto pci_bus_id = d.Device().Metadata().GetValue("pci_bus_id"); - if (pci_bus_id) { - std::cmatch matches; - if (std::regex_match(pci_bus_id, matches, pci_bus_id_pattern)) { - auto domain = std::stoull(matches[1].str(), nullptr, 16); - auto bus = std::stoull(matches[2].str(), nullptr, 16); - auto device = std::stoull(matches[3].str(), nullptr, 16); - auto function = std::stoull(matches[4].str(), nullptr, 16); - if (domain != pci_props.pciDomain || bus != pci_props.pciBus || device != pci_props.pciDevice || function != pci_props.pciFunction) { - continue; - } - } - } -#endif - dev.ep_device_candidates.push_back(d); - } - } - if (!dev.ep_device_candidates.empty()) { - resources.nv_devices.push_back(dev); - } - } - } - - if (resources.nv_devices.empty()) { - GTEST_SKIP() << "No Nv VK devices with EpDevices found"; - } - - // Try to run the test for first NV device - resources.phys_device = resources.nv_devices[0].phys_dev; - bool has_cig_extension = resources.nv_devices[0].has_cig_extension; - - VkPhysicalDeviceVulkan11Features vk11{}; - vk11.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; - VkPhysicalDeviceVulkan12Features vk12{}; - vk12.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES; - vk12.bufferDeviceAddress = true; - vk12.timelineSemaphore = true; - vk12.pNext = &vk11; - - VkExternalComputeQueueDeviceCreateInfoNV cig_create_info{}; - cig_create_info.sType = VK_STRUCTURE_TYPE_EXTERNAL_COMPUTE_QUEUE_DEVICE_CREATE_INFO_NV; - cig_create_info.reservedExternalQueues = 1; - cig_create_info.pNext = &vk12; - - float priority = 1.f; - VkDeviceQueueCreateInfo queue_info{}; - queue_info.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO; - queue_info.queueFamilyIndex = 0; - queue_info.queueCount = 1; - queue_info.pQueuePriorities = &priority; - - std::vector device_extensions{}; -#if defined(_WIN32) - device_extensions.push_back(VK_KHR_EXTERNAL_SEMAPHORE_WIN32_EXTENSION_NAME); - device_extensions.push_back(VK_KHR_EXTERNAL_MEMORY_WIN32_EXTENSION_NAME); -#else - device_extensions.push_back(VK_KHR_EXTERNAL_SEMAPHORE_FD_EXTENSION_NAME); - device_extensions.push_back(VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME); - // DMA_BUF is currently not tested, adding this to expand test for platforms where - // CUDA supports DMABUF import - device_extensions.push_back(VK_EXT_EXTERNAL_MEMORY_DMA_BUF_EXTENSION_NAME); -#endif - if (has_cig_extension) { - device_extensions.push_back(VK_NV_EXTERNAL_COMPUTE_QUEUE_EXTENSION_NAME); - } - device_extensions.push_back(VK_EXT_PCI_BUS_INFO_EXTENSION_NAME); - - VkDeviceCreateInfo device_info{}; - device_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO; - device_info.queueCreateInfoCount = 1; - device_info.pQueueCreateInfos = &queue_info; - device_info.enabledExtensionCount = uint32_t(device_extensions.size()); - device_info.ppEnabledExtensionNames = device_extensions.data(); - device_info.pNext = has_cig_extension ? (void*)&cig_create_info : &vk12; - auto device_creation_result = - resources.loader.vkCreateDevice(resources.phys_device, - &device_info, - nullptr, - &resources.device); - if (device_creation_result != VK_SUCCESS) { - GTEST_SKIP() << "Vulkan device creation failed: skipping Vulkan interop test"; - } - resources.loader.init(vk::Device(resources.device)); - - resources.loader.vkGetDeviceQueue(resources.device, 0, 0, &resources.queue); - // External compute queues are used in CiG mode - if (has_cig_extension) { - VkExternalComputeQueueCreateInfoNV ex_compute_queue_info{}; - ex_compute_queue_info.sType = VK_STRUCTURE_TYPE_EXTERNAL_COMPUTE_QUEUE_CREATE_INFO_NV; - ex_compute_queue_info.preferredQueue = resources.queue; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkCreateExternalComputeQueueNV(resources.device, &ex_compute_queue_info, nullptr, &resources.ex_compute_queue)); - } -} - -void create_timeline_semaphore(VkResources& resources, ExportableTimelineSemaphore& semaphore) { - // VK creation - VkSemaphoreCreateInfo sem_info{}; - sem_info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO; - VkSemaphoreTypeCreateInfo timeline_info{}; - timeline_info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_TYPE_CREATE_INFO; - timeline_info.initialValue = 0; - timeline_info.semaphoreType = VK_SEMAPHORE_TYPE_TIMELINE; - VkExportSemaphoreCreateInfoKHR export_info{}; - export_info.sType = VK_STRUCTURE_TYPE_EXPORT_SEMAPHORE_CREATE_INFO; -#if defined(_WIN32) - export_info.handleTypes = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT; -#else - export_info.handleTypes = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT; -#endif - sem_info.pNext = &timeline_info; - timeline_info.pNext = &export_info; - - EXPECT_EQ(VK_SUCCESS, resources.loader.vkCreateSemaphore(resources.device, &sem_info, nullptr, &semaphore.vk_handle)); - - // ORT logic - OrtExternalSemaphoreDescriptor sem_desc = {}; - sem_desc.version = ORT_API_VERSION; -#if defined(_WIN32) - VkSemaphoreGetWin32HandleInfoKHR info{}; - info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_GET_WIN32_HANDLE_INFO_KHR; - info.semaphore = semaphore.vk_handle; - info.handleType = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_WIN32_BIT; - HANDLE fd; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkGetSemaphoreWin32HandleKHR(resources.device, &info, &fd)); - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_WIN32; - semaphore.native_handle = (void*)(size_t)fd; -#else - VkSemaphoreGetFdInfoKHR info{}; - info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_GET_FD_INFO_KHR; - info.semaphore = semaphore.vk_handle; - info.handleType = VK_EXTERNAL_SEMAPHORE_HANDLE_TYPE_OPAQUE_FD_BIT; - int fd; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkGetSemaphoreFdKHR(resources.device, &info, &fd)); - sem_desc.type = ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_OPAQUE_FD; - semaphore.native_handle = (void*)(size_t)fd; -#endif - sem_desc.native_handle = semaphore.native_handle; - - EXPECT_TRUE(Ort::Status(Ort::GetInteropApi().ImportSemaphore(resources.importer, &sem_desc, &semaphore.ort_handle)).IsOK()); -} - -void allocate_buffer(VkResources& resources, - ExportableBuffer& export_buffer, - VkDeviceSize size, - VkMemoryPropertyFlags mem_prop_flags, - bool do_export, - [[maybe_unused]] bool use_dma) { - VkPhysicalDeviceMemoryProperties mem_props = {}; - resources.loader.vkGetPhysicalDeviceMemoryProperties(resources.phys_device, - &mem_props); -#if defined(_WIN32) - auto handle_type = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT; -#else - auto handle_type = use_dma ? VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT : VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; -#endif - uint32_t index = 0; - VkBufferCreateInfo buffer_info{}; - buffer_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; - buffer_info.queueFamilyIndexCount = 1; - buffer_info.pQueueFamilyIndices = &index; - buffer_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; - buffer_info.size = size; - buffer_info.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_2_STORAGE_TEXEL_BUFFER_BIT; - VkExternalMemoryBufferCreateInfo external_create_info{}; - external_create_info.sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_BUFFER_CREATE_INFO; - external_create_info.handleTypes = handle_type; - if (do_export) { - buffer_info.pNext = &external_create_info; - } - EXPECT_EQ(VK_SUCCESS, resources.loader.vkCreateBuffer(resources.device, &buffer_info, - nullptr, &export_buffer.buffer)); - VkMemoryRequirements mem_reqs = {}; - resources.loader.vkGetBufferMemoryRequirements(resources.device, - export_buffer.buffer, &mem_reqs); - int mem_idx = -1; - for (uint32_t i = 0; i < mem_props.memoryTypeCount; i++) { - if ((mem_reqs.memoryTypeBits & (1 << i)) && - (mem_props.memoryTypes[i].propertyFlags & mem_prop_flags) == - mem_prop_flags) { - mem_idx = i; - break; - } - } - EXPECT_NE(mem_idx, -1); - - VkMemoryAllocateInfo alloc_info{}; - alloc_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO; - alloc_info.memoryTypeIndex = mem_idx; - alloc_info.allocationSize = mem_reqs.size; - VkExportMemoryAllocateInfo export_info{}; - export_info.sType = VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO; - export_info.handleTypes = handle_type; - - if (do_export) { - alloc_info.pNext = &export_info; - } - EXPECT_EQ(VK_SUCCESS, resources.loader.vkAllocateMemory(resources.device, &alloc_info, - nullptr, &export_buffer.memory)); - EXPECT_EQ(VK_SUCCESS, resources.loader.vkBindBufferMemory(resources.device, - export_buffer.buffer, export_buffer.memory, 0)); - VkBufferViewCreateInfo buffer_view_info{}; - buffer_view_info.sType = VK_STRUCTURE_TYPE_BUFFER_VIEW_CREATE_INFO; - buffer_view_info.buffer = export_buffer.buffer; - buffer_view_info.format = VK_FORMAT_R32_SFLOAT; - buffer_view_info.offset = 0; - buffer_view_info.range = size; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkCreateBufferView(resources.device, &buffer_view_info, - nullptr, &export_buffer.view)); - if (!do_export) { - return; - } -#if defined(_WIN32) - VkMemoryGetWin32HandleInfoKHR info{}; - info.sType = VK_STRUCTURE_TYPE_MEMORY_GET_WIN32_HANDLE_INFO_KHR; - info.memory = export_buffer.memory; - info.handleType = handle_type; - HANDLE fd; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkGetMemoryWin32HandleKHR(resources.device, &info, &fd)); - export_buffer.native_handle = (void*)fd; -#else - VkMemoryGetFdInfoKHR info{}; - info.sType = VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR; - info.memory = export_buffer.memory; - info.handleType = handle_type; - int fd; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkGetMemoryFdKHR(resources.device, &info, &fd)); - export_buffer.native_handle = (void*)(size_t)fd; -#endif - - OrtExternalMemoryDescriptor mem_desc = {}; - mem_desc.version = ORT_API_VERSION; -#if defined(_WIN32) - mem_desc.handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_WIN32; -#else - EXPECT_FALSE(use_dma); // would need to use ORT_EXTERNAL_MEMORY_HANDLE_TYPE_DMABUF_FD - mem_desc.handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_OPAQUE_FD; -#endif - mem_desc.native_handle = export_buffer.native_handle; - mem_desc.size_bytes = size; - mem_desc.offset_bytes = 0; - - auto& interop_api = Ort::GetInteropApi(); - EXPECT_TRUE(Ort::Status(interop_api.ImportMemory(resources.importer, &mem_desc, &export_buffer.ort_handle)).IsOK()); -} - -Ort::Session configure_session(const PathString& model_path, Ort::SyncStream& ort_stream, const Ort::ConstEpDevice& ep_device, size_t* aux_streams_array) { - Ort::SessionOptions session_options; - session_options.SetExecutionMode(ORT_SEQUENTIAL); - session_options.DisableMemPattern(); - session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL); - Ort::KeyValuePairs ep_options; - ep_options.Add(onnxruntime::nv::provider_option_names::kUserComputeStream, std::to_string(size_t(ort_stream.GetHandle())).c_str()); - ep_options.Add(onnxruntime::nv::provider_option_names::kHasUserComputeStream, "1"); - ep_options.Add(onnxruntime::nv::provider_option_names::kMaxSharedMemSize, std::to_string(1024 * 28).c_str()); - ep_options.Add(onnxruntime::nv::provider_option_names::kUserAuxStreamArray, std::to_string((size_t)aux_streams_array).c_str()); - ep_options.Add(onnxruntime::nv::provider_option_names::kLengthAuxStreamArray, "1"); - ep_options.Add(onnxruntime::nv::provider_option_names::kCudaGraphEnable, "0"); - session_options.AppendExecutionProvider_V2(*ort_env, std::vector{ep_device}, ep_options); - - return Ort::Session(*ort_env, model_path.c_str(), session_options); -} -} // namespace - -void test_vulkan_interop(TestParameters& test_params) { - RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); - - VkResources resources; - init_vulkan_interop(resources); - - // Create a simple model: Y = X + 1 (Add with constant 1) - std::error_code ec{}; - auto model_path = ToPathString(std::filesystem::temp_directory_path(ec) / "external_mem_add_one_test.onnx"); - EXPECT_FALSE(ec); - clearFileIfExists(model_path); - { - onnxruntime::Model model("add_one_test", false, DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - - ONNX_NAMESPACE::TypeProto tensor_type; - tensor_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(64); - tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(64); - - auto& input_arg = graph.GetOrCreateNodeArg("X", &tensor_type); - auto& output_arg = graph.GetOrCreateNodeArg("Y", &tensor_type); - ONNX_NAMESPACE::TypeProto scalar_float; - scalar_float.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - scalar_float.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - auto& one_arg = graph.GetOrCreateNodeArg("one", &scalar_float); - ONNX_NAMESPACE::TensorProto one_proto; - one_proto.set_name("one"); - one_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - one_proto.add_dims(1); - one_proto.add_float_data(1.f); - graph.AddInitializedTensor(one_proto); - std::vector add_inputs = {&input_arg, &one_arg}; - graph.AddNode("add_one", "Add", "Y = X + 1", add_inputs, {&output_arg}); - - EXPECT_STATUS_OK(graph.Resolve()); - EXPECT_STATUS_OK(onnxruntime::Model::Save(model, model_path)); - } - DEFER(model_path, clearFileIfExists(model_path)); - - const int64_t batch = 1, channels = 3, dim = 64; - const int64_t shape[] = {batch, channels, dim, dim}; - const size_t num_elements = batch * channels * dim * dim; - const size_t buffer_size = num_elements * sizeof(float); - - if (resources.nv_devices.empty()) { - GTEST_SKIP() << "No NV devices found"; - } - const auto& ep_device = resources.nv_devices[0].ep_device_candidates[0]; - resources.ep_device = ep_device; - auto& interop_api = Ort::GetInteropApi(); - - // Create external resource importer - Ort::Status status(interop_api.CreateExternalResourceImporterForDevice(ep_device, &resources.importer)); - if (!status.IsOK() || resources.importer == nullptr) { - GTEST_SKIP() << "External resource import not supported"; - } - - // Check VK external memory (buffer) support - bool can_import_vk_buffer = false; -#if defined(_WIN32) - const auto vk_mem_handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_WIN32; -#else - EXPECT_FALSE(test_params.use_dmabuf); // would need to use a possible future ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_OPAQUE_FD - const auto vk_mem_handle_type = ORT_EXTERNAL_MEMORY_HANDLE_TYPE_VK_MEMORY_OPAQUE_FD; -#endif - status = Ort::Status(interop_api.CanImportMemory(resources.importer, vk_mem_handle_type, &can_import_vk_buffer)); - if (!status.IsOK() || !can_import_vk_buffer) { - GTEST_FAIL() << "VK external buffer import not supported"; - } - - // Check VK external semaphore (timeline) support - bool can_import_vk_semaphore = false; -#if defined(_WIN32) - const auto vk_sem_type = ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_WIN32; -#else - const auto vk_sem_type = ORT_EXTERNAL_SEMAPHORE_VK_TIMELINE_SEMAPHORE_OPAQUE_FD; -#endif - status = Ort::Status(interop_api.CanImportSemaphore(resources.importer, vk_sem_type, &can_import_vk_semaphore)); - if (!status.IsOK() || !can_import_vk_semaphore) { - GTEST_FAIL() << "VK external timeline semaphore import not supported"; - } - - bool has_cig_extension = resources.nv_devices[0].has_cig_extension; - std::vector external_compute_queue_data(64); - if (has_cig_extension) { - VkExternalComputeQueueDataParamsNV ex_compute_queue_params{}; - ex_compute_queue_params.sType = VK_STRUCTURE_TYPE_EXTERNAL_COMPUTE_QUEUE_DATA_PARAMS_NV; - ex_compute_queue_params.deviceIndex = 0; // device index within device group, so 0 for us - resources.loader.vkGetExternalComputeQueueDataNV(resources.ex_compute_queue, &ex_compute_queue_params, external_compute_queue_data.data()); - } - - int device_count; - CUdevice selected_device = -1; - EXPECT_EQ(cudaSuccess, cudaGetDeviceCount(&device_count)); - EXPECT_GT(device_count, 0); - for (int i = 0; i < device_count; i++) { - cudaDeviceProp props{}; - EXPECT_EQ(cudaSuccess, cudaGetDeviceProperties(&props, i)); - if (0 == std::memcmp(&props.uuid, resources.nv_devices[0].id_props.deviceUUID, sizeof(resources.nv_devices[0].id_props.deviceUUID))) { - selected_device = i; - break; - } - } - EXPECT_GT(selected_device, -1); - EXPECT_EQ(cudaSuccess, cudaSetDevice(selected_device)); - int cig_supported{false}; - EXPECT_EQ(cudaSuccess, cudaDeviceGetAttribute(&cig_supported, cudaDevAttrVulkanCigSupported, selected_device)); - CudaDriverLoader driver; - // int dma_supported{false}; - // EXPECT_EQ(CUDA_SUCCESS, driver.cuDeviceGetAttribute_fn(&dma_supported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, selected_device)); - - if (test_params.use_init_graphics_interop_call) { - Ort::KeyValuePairs kv; - kv.Add(onnxruntime::nv::provider_option_names::kExternalComputeQueueDataParamNV_data, std::to_string(reinterpret_cast(external_compute_queue_data.data())).c_str()); - OrtGraphicsInteropConfig interop_config{}; - interop_config.version = ORT_API_VERSION; - interop_config.graphics_api = OrtGraphicsApi::ORT_GRAPHICS_API_VULKAN; - interop_config.additional_options = kv.GetConst(); - EXPECT_TRUE(Ort::Status(interop_api.InitGraphicsInteropForEpDevice(ep_device, &interop_config)).IsOK()); - } else { - EXPECT_EQ(driver.IsLoaded(), true); - - if (test_params.force_cig_if_supported) { - if (!has_cig_extension) { - GTEST_SKIP() << "Skipping test because of missing VK_NV_external_compute_queue extension"; - } - if (!cig_supported) { - GTEST_SKIP() << "Skipping test because of CUDA device does not support \"Cuda in Graphics\" for Vulkan"; - } - } - - if (test_params.allow_cig && cig_supported && has_cig_extension) { - CUcontext ctx{}; - CUctxCigParam cig_params{}; - cig_params.sharedDataType = CIG_DATA_TYPE_NV_BLOB; - cig_params.sharedData = external_compute_queue_data.data(); - CUctxCreateParams params{}; - params.cigParams = &cig_params; - EXPECT_EQ(CUDA_SUCCESS, driver.cuCtxCreate_v4_fn(&ctx, ¶ms, 0, selected_device)); - EXPECT_EQ(CUDA_SUCCESS, driver.cuCtxSetCurrent_fn(ctx)); - } - } - - ExportableTimelineSemaphore input_ready{}; - create_timeline_semaphore(resources, input_ready); - resources.semaphores.push_back(input_ready); - ExportableTimelineSemaphore inference_done{}; - create_timeline_semaphore(resources, inference_done); - resources.semaphores.push_back(inference_done); - ExportableTimelineSemaphore download_done{}; - create_timeline_semaphore(resources, download_done); - resources.semaphores.push_back(download_done); - - ExportableBuffer upload_buffer{}; - allocate_buffer(resources, upload_buffer, buffer_size, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, false, test_params.use_dmabuf); - resources.buffers.push_back(upload_buffer); - ExportableBuffer input_buffer{}; - allocate_buffer(resources, input_buffer, buffer_size, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true, test_params.use_dmabuf); - resources.buffers.push_back(input_buffer); - ExportableBuffer output_buffer{}; - allocate_buffer(resources, output_buffer, buffer_size, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true, test_params.use_dmabuf); - resources.buffers.push_back(output_buffer); - - // Create command pool and command buffers for upload/download - VkCommandPoolCreateInfo pool_info{}; - pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO; - pool_info.queueFamilyIndex = 0; - pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkCreateCommandPool(resources.device, &pool_info, nullptr, &resources.cmd_pool)); - VkCommandBufferAllocateInfo alloc_info{}; - alloc_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - alloc_info.commandPool = resources.cmd_pool; - alloc_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - alloc_info.commandBufferCount = 2; - VkCommandBuffer cmd_bufs[2]; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkAllocateCommandBuffers(resources.device, &alloc_info, cmd_bufs)); - resources.upload_cmd_buf = cmd_bufs[0]; - resources.download_cmd_buf = cmd_bufs[1]; - - ExportableBuffer& upload_buf = resources.buffers[0]; - ExportableBuffer& input_buf = resources.buffers[1]; - ExportableBuffer& output_buf = resources.buffers[2]; - ExportableTimelineSemaphore& input_ready_sem = resources.semaphores[0]; - ExportableTimelineSemaphore& inference_done_sem = resources.semaphores[1]; - ExportableTimelineSemaphore& download_done_sem = resources.semaphores[2]; - - // Fill upload buffer with known values: 0.f, 1.f, 2.f, 3.f, ... - void* mapped = nullptr; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkMapMemory(resources.device, upload_buf.memory, 0, buffer_size, 0, &mapped)); - float* host_data = static_cast(mapped); - for (size_t i = 0; i < num_elements; ++i) { - host_data[i] = static_cast(i); - } - resources.loader.vkUnmapMemory(resources.device, upload_buf.memory); - - // Record upload command buffer: upload_buffer -> input_buffer (do not submit yet) - VkCommandBufferBeginInfo cmd_buf_info{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO, nullptr, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT, nullptr}; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkBeginCommandBuffer(resources.upload_cmd_buf, &cmd_buf_info)); - VkBufferCopy copy_region{}; - copy_region.size = buffer_size; - copy_region.srcOffset = 0; - copy_region.dstOffset = 0; - resources.loader.vkCmdCopyBuffer(resources.upload_cmd_buf, upload_buf.buffer, input_buf.buffer, 1, ©_region); - EXPECT_EQ(VK_SUCCESS, resources.loader.vkEndCommandBuffer(resources.upload_cmd_buf)); - - // Create ORT tensors from imported memory (input and output buffers) - OrtExternalTensorDescriptor tensor_desc = {}; - tensor_desc.version = ORT_API_VERSION; - tensor_desc.element_type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; - tensor_desc.shape = shape; - tensor_desc.rank = 4; - tensor_desc.offset_bytes = 0; - OrtValue* input_tensor = nullptr; - OrtValue* output_tensor = nullptr; - EXPECT_TRUE(Ort::Status(interop_api.CreateTensorFromMemory(resources.importer, input_buf.ort_handle, &tensor_desc, &input_tensor)).IsOK()); - EXPECT_TRUE(Ort::Status(interop_api.CreateTensorFromMemory(resources.importer, output_buf.ort_handle, &tensor_desc, &output_tensor)).IsOK()); - void* input_data_ptr = nullptr; - void* output_data_ptr = nullptr; - EXPECT_TRUE(Ort::Status(Ort::GetApi().GetTensorMutableData(input_tensor, &input_data_ptr)).IsOK()); - EXPECT_TRUE(Ort::Status(Ort::GetApi().GetTensorMutableData(output_tensor, &output_data_ptr)).IsOK()); - cudaPointerAttributes input_attrs, output_attrs; - ASSERT_EQ(cudaPointerGetAttributes(&input_attrs, input_data_ptr), cudaSuccess); - ASSERT_EQ(cudaPointerGetAttributes(&output_attrs, output_data_ptr), cudaSuccess); - EXPECT_EQ(input_attrs.type, cudaMemoryTypeDevice) << "Input tensor must be CUDA device memory"; - EXPECT_EQ(output_attrs.type, cudaMemoryTypeDevice) << "Output tensor must be CUDA device memory"; - - // Add the NvTensorRtRtx EP - // Configure to use our CUDA stream - auto ort_stream = ep_device.CreateSyncStream(); - size_t stream_addr_val = reinterpret_cast(ort_stream.GetHandle()); - size_t aux_streams[] = {stream_addr_val}; - auto session = configure_session(model_path, ort_stream, ep_device, aux_streams); - - Ort::IoBinding io_binding(session); - Ort::AllocatorWithDefaultOptions allocator; - Ort::AllocatedStringPtr input_name = session.GetInputNameAllocated(0, allocator); - Ort::AllocatedStringPtr output_name = session.GetOutputNameAllocated(0, allocator); - io_binding.BindInput(input_name.get(), Ort::Value(input_tensor)); - io_binding.BindOutput(output_name.get(), Ort::Value(output_tensor)); - io_binding.SynchronizeInputs(); - - // Semaphores: upload signals input_ready=1; ORT waits input_ready 1, runs, signals inference_done=2; download signals download_done=3 - const uint64_t input_ready_value = 1; - const uint64_t inference_done_value = 1; - const uint64_t download_done_value = 1; - - // Record download command buffer: output_buffer -> upload_buffer (do not submit yet) - auto begin_info = VkCommandBufferBeginInfo{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO, nullptr, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT, nullptr}; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkBeginCommandBuffer(resources.download_cmd_buf, &begin_info)); - resources.loader.vkCmdCopyBuffer(resources.download_cmd_buf, output_buf.buffer, upload_buf.buffer, 1, ©_region); - EXPECT_EQ(VK_SUCCESS, resources.loader.vkEndCommandBuffer(resources.download_cmd_buf)); - - // Submit upload first (signal input_ready = 1 when copy completes) - uint64_t signal_input_ready = input_ready_value; - VkTimelineSemaphoreSubmitInfo timeline_info = {}; - timeline_info.sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO; - timeline_info.signalSemaphoreValueCount = 1; - timeline_info.pSignalSemaphoreValues = &signal_input_ready; - VkSubmitInfo submit_upload = {}; - submit_upload.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_upload.pNext = &timeline_info; - submit_upload.commandBufferCount = 1; - submit_upload.pCommandBuffers = &resources.upload_cmd_buf; - submit_upload.signalSemaphoreCount = 1; - submit_upload.pSignalSemaphores = &input_ready_sem.vk_handle; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkQueueSubmit(resources.queue, 1, &submit_upload, VK_NULL_HANDLE)); - - EXPECT_TRUE(Ort::Status(interop_api.WaitSemaphore(resources.importer, input_ready_sem.ort_handle, ort_stream, input_ready_value)).IsOK()); - Ort::RunOptions run_options; - run_options.SetSyncStream(ort_stream); - run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); - session.Run(run_options, io_binding); - EXPECT_TRUE(Ort::Status(interop_api.SignalSemaphore(resources.importer, inference_done_sem.ort_handle, ort_stream, inference_done_value)).IsOK()); - - uint64_t wait_inference_done = inference_done_value; - uint64_t signal_download_done = download_done_value; - timeline_info.waitSemaphoreValueCount = 1; - timeline_info.pWaitSemaphoreValues = &wait_inference_done; - timeline_info.signalSemaphoreValueCount = 1; - timeline_info.pSignalSemaphoreValues = &signal_download_done; - VkPipelineStageFlags wait_dst_stage_mask = VkPipelineStageFlagBits::VK_PIPELINE_STAGE_TRANSFER_BIT; - VkSubmitInfo submit_download = {}; - submit_download.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit_download.pNext = &timeline_info; - submit_download.waitSemaphoreCount = 1; - submit_download.pWaitDstStageMask = &wait_dst_stage_mask; - submit_download.pWaitSemaphores = &inference_done_sem.vk_handle; - submit_download.commandBufferCount = 1; - submit_download.pCommandBuffers = &resources.download_cmd_buf; - submit_download.signalSemaphoreCount = 1; - submit_download.pSignalSemaphores = &download_done_sem.vk_handle; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkQueueSubmit(resources.queue, 1, &submit_download, VK_NULL_HANDLE)); - - VkSemaphoreWaitInfo wait_info = {}; - wait_info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_WAIT_INFO; - wait_info.semaphoreCount = 1; - wait_info.pSemaphores = &download_done_sem.vk_handle; - wait_info.pValues = &download_done_value; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkWaitSemaphores(resources.device, &wait_info, UINT64_MAX)); - - // Read back and verify: result should be input + 1 (0+1, 1+1, 2+1, 3+1, ...) - mapped = nullptr; - EXPECT_EQ(VK_SUCCESS, resources.loader.vkMapMemory(resources.device, upload_buf.memory, 0, buffer_size, 0, &mapped)); - host_data = static_cast(mapped); - for (size_t i = 0; i < num_elements; ++i) { - float expected = static_cast(i) + 1.f; - EXPECT_FLOAT_EQ(host_data[i], expected) << "index " << i; - } - resources.loader.vkUnmapMemory(resources.device, upload_buf.memory); -} - -TEST(NvExecutionProviderVulkanTest, VkCigDisabled) { - TestParameters params; - params.allow_cig = false; - test_vulkan_interop(params); -} - -TEST(NvExecutionProviderVulkanTest, VkInitGraphicsInterop) { - TestParameters params; - params.use_init_graphics_interop_call = true; - test_vulkan_interop(params); -} - -TEST(NvExecutionProviderVulkanTest, VkForceCig) { - TestParameters params; - params.allow_cig = true; - params.force_cig_if_supported = true; - test_vulkan_interop(params); -} - -} // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.cc deleted file mode 100644 index c86cf382d4653..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.cc +++ /dev/null @@ -1,824 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. - -// registration/selection is only supported on windows as there's no device discovery on other platforms - -#include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" - -#include -#include -#include - -#include "core/session/onnxruntime_cxx_api.h" -#include "test/util/include/api_asserts.h" -#include "core/graph/basic_types.h" -#include "core/graph/model.h" -#include "core/framework/tensorprotoutils.h" -#include "core/graph/onnx_protobuf.h" -#include "core/graph/model_saving_options.h" -#include "core/graph/schema_registry.h" -#include "test/util/include/scoped_env_vars.h" -#include "test/common/trt_op_test_utils.h" -#include "test/providers/provider_test_utils.h" -#include "test/unittest_util/framework_test_utils.h" - -namespace onnxruntime { -namespace test { - -Utils::NvTensorRtRtxEpInfo Utils::nv_tensorrt_rtx_ep_info; - -void Utils::GetEp(Ort::Env& env, const std::string& ep_name, const OrtEpDevice*& ep_device) { - const OrtApi& c_api = Ort::GetApi(); - const OrtEpDevice* const* ep_devices = nullptr; - size_t num_devices; - ASSERT_ORTSTATUS_OK(c_api.GetEpDevices(env, &ep_devices, &num_devices)); - - auto it = std::find_if(ep_devices, ep_devices + num_devices, - [&c_api, &ep_name](const OrtEpDevice* ep_device) { - // NV TensorRT RTX EP uses registration name as ep name - return c_api.EpDevice_EpName(ep_device) == ep_name; - }); - - if (it == ep_devices + num_devices) { - ep_device = nullptr; - } else { - ep_device = *it; - } -} - -void Utils::RegisterAndGetNvTensorRtRtxEp(Ort::Env& env, RegisteredEpDeviceUniquePtr& registered_ep) { - const OrtApi& c_api = Ort::GetApi(); - // this should load the library and create OrtEpDevice - ASSERT_ORTSTATUS_OK(c_api.RegisterExecutionProviderLibrary(env, - nv_tensorrt_rtx_ep_info.registration_name.c_str(), - nv_tensorrt_rtx_ep_info.library_path.c_str())); - const OrtEpDevice* nv_tensorrt_rtx_ep = nullptr; - GetEp(env, nv_tensorrt_rtx_ep_info.registration_name, nv_tensorrt_rtx_ep); - ASSERT_NE(nv_tensorrt_rtx_ep, nullptr); - - registered_ep = RegisteredEpDeviceUniquePtr(nv_tensorrt_rtx_ep, [&env, c_api](const OrtEpDevice* /*ep*/) { - c_api.UnregisterExecutionProviderLibrary(env, nv_tensorrt_rtx_ep_info.registration_name.c_str()); - }); -} - -void CreateBaseModel(const PathString& model_name, - std::string graph_name, - std::vector dims, - bool add_fast_gelu, - ONNX_NAMESPACE::TensorProto_DataType dtype, - const PathString& external_initializer_file) { - onnxruntime::Model model(graph_name, false, DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - std::vector inputs; - std::vector outputs; - - // FLOAT tensor - ONNX_NAMESPACE::TypeProto float_tensor; - float_tensor.mutable_tensor_type()->set_elem_type(dtype); - - for (auto dim : dims) { - float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - ONNX_NAMESPACE::TypeProto dyn_float_tensor; - dyn_float_tensor.mutable_tensor_type()->set_elem_type(dtype); - - auto& input_arg_1 = graph.GetOrCreateNodeArg("X", &float_tensor); - auto& input_arg_2 = graph.GetOrCreateNodeArg("Y", &float_tensor); - inputs.push_back(&input_arg_1); - inputs.push_back(&input_arg_2); - auto& output_arg = graph.GetOrCreateNodeArg("node_1_out_1", &float_tensor); - outputs.push_back(&output_arg); - graph.AddNode("node_1", "Add", "node 1.", inputs, outputs); - - auto& input_arg_3 = graph.GetOrCreateNodeArg("Z", &float_tensor); - inputs.clear(); - inputs.push_back(&output_arg); - inputs.push_back(&input_arg_3); - - auto& output_arg_2 = graph.GetOrCreateNodeArg("node_2_out_1", &float_tensor); - outputs.clear(); - outputs.push_back(&output_arg_2); - graph.AddNode("node_2", "Add", "node 2.", inputs, outputs); - - inputs.clear(); - inputs.push_back(&output_arg_2); - - if (add_fast_gelu) { - auto& output_arg_3 = graph.GetOrCreateNodeArg("node_3_out_1", &dyn_float_tensor); - outputs.clear(); - outputs.push_back(&output_arg_3); - - graph.AddNode("node_3", "FastGelu", "node 3.", inputs, outputs, - /* attributes */ nullptr, kMSDomain); - - inputs.clear(); - inputs.push_back(&output_arg_3); - } - - ONNX_NAMESPACE::TypeProto float_scalar; - float_scalar.mutable_tensor_type()->set_elem_type(dtype); - float_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - auto& input_scalar = graph.GetOrCreateNodeArg("S", &float_scalar); - inputs.push_back(&input_scalar); - - auto& output_arg_4 = graph.GetOrCreateNodeArg("O", &dyn_float_tensor); - - outputs.clear(); - outputs.push_back(&output_arg_4); - graph.AddNode("node_5", "Add", "node 5.", inputs, outputs); - - auto status = graph.Resolve(); - ASSERT_TRUE(status.IsOK()); - if (!external_initializer_file.empty()) { - ModelSavingOptions save_options(128); - status = Model::SaveWithExternalInitializers(model, model_name, external_initializer_file, save_options); - } else { - status = Model::Save(model, model_name); - } - ASSERT_TRUE(status.IsOK()); -} - -// Helper to create large initializers -ONNX_NAMESPACE::TensorProto CreateLargeWeight( - const std::string& name, - ONNX_NAMESPACE::TensorProto_DataType dtype, - const std::vector& shape, - float scale = 0.02f) { - ONNX_NAMESPACE::TensorProto tensor; - tensor.set_name(name); - tensor.set_data_type(dtype); - for (auto d : shape) tensor.add_dims(d); - // Here we fill with random floats, but for real data, use your trained weights. - size_t total_size = 1; - for (int64_t d : shape) total_size *= d; - std::random_device rd; - std::default_random_engine rng(rd()); - if (dtype == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { - std::vector data(total_size); - std::normal_distribution dist(0.0f, scale); - for (auto& v : data) v = dist(rng); - tensor.set_raw_data(data.data(), total_size * sizeof(float)); - } else if (dtype == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { - std::vector data(total_size); - std::normal_distribution dist(0.0f, scale); - for (auto& v : data) v = MLFloat16(dist(rng)); - tensor.set_raw_data(data.data(), total_size * sizeof(MLFloat16)); - } else { - throw std::runtime_error("Unsupported data type for large weight"); - } - return tensor; -} - -// Helper to add a GroupQueryAttention node -onnxruntime::NodeArg& AddGroupQueryAttention( - onnxruntime::Graph& graph, - onnxruntime::NodeArg& query, - onnxruntime::NodeArg& key, - onnxruntime::NodeArg& value, - int batch_size, - int head_dim, - int seq_len, - int num_heads, - int kv_num_heads, - float scale, - ONNX_NAMESPACE::TensorProto_DataType dtype, - const std::string& node_name) { - // KV cache - ONNX_NAMESPACE::TypeProto key_type; - key_type.mutable_tensor_type()->set_elem_type(dtype); - key_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(batch_size); - key_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(kv_num_heads); - key_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(seq_len); - key_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(head_dim); - auto& past_key = graph.GetOrCreateNodeArg(node_name + "_past_key", &key_type); - - ONNX_NAMESPACE::TypeProto value_type; - value_type.mutable_tensor_type()->set_elem_type(dtype); - value_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(batch_size); - value_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(kv_num_heads); - value_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(seq_len); - value_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(head_dim); - auto& past_value = graph.GetOrCreateNodeArg(node_name + "_past_value", &value_type); - - // Output - auto& output = graph.GetOrCreateNodeArg(node_name + "_output", nullptr); - - // Create required initializers for GroupQueryAttention - ONNX_NAMESPACE::TensorProto seqlens_k_tensor; - seqlens_k_tensor.set_name(node_name + "_seqlens_k"); - seqlens_k_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); - seqlens_k_tensor.add_dims(2); - seqlens_k_tensor.set_dims(0, batch_size); - seqlens_k_tensor.set_dims(0, 1); - seqlens_k_tensor.add_int32_data(seq_len - 1); // seqlens_k = total_sequence_length - 1 - graph.AddInitializedTensor(seqlens_k_tensor); - - ONNX_NAMESPACE::TensorProto total_seq_len_tensor; - total_seq_len_tensor.set_name(node_name + "_total_sequence_length"); - total_seq_len_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); - total_seq_len_tensor.add_int32_data(seq_len); - graph.AddInitializedTensor(total_seq_len_tensor); - - // Get the initializers that were created for this node - auto* seqlens_k = graph.GetNodeArg(node_name + "_seqlens_k"); - auto* total_sequence_length = graph.GetNodeArg(node_name + "_total_sequence_length"); - - auto& present_value = graph.GetOrCreateNodeArg(node_name + "_present_value", nullptr); - auto& present_key = graph.GetOrCreateNodeArg(node_name + "_present_key", nullptr); - - // Inputs - GroupQueryAttention requires at least 7 inputs (query, key, value, past_key, past_value, seqlens_k, total_sequence_length) - std::vector inputs = { - &query, // 0: query - &key, // 1: key - &value, // 2: value - &past_key, // 3: past_key (optional) - &past_value, // 4: past_value (optional) - seqlens_k, // 5: seqlens_k (required) - total_sequence_length, // 6: total_sequence_length (required) - // nullptr, // 7: cos_cache (optional) - // nullptr, // 8: sin_cache (optional) - // nullptr, // 9: position_ids (optional) - // nullptr, // 10: attention_bias (optional) - // nullptr // 11: head_sink (optional) - }; - - // Attributes - NodeAttributes attrs; - ONNX_NAMESPACE::AttributeProto attr_heads; - attr_heads.set_name("num_heads"); - attr_heads.set_type(onnx::AttributeProto_AttributeType_INT); - attr_heads.set_i(num_heads); - attrs["num_heads"] = attr_heads; - ONNX_NAMESPACE::AttributeProto attr_kv_num_heads; - attr_kv_num_heads.set_name("kv_num_heads"); - attr_kv_num_heads.set_type(onnx::AttributeProto_AttributeType_INT); - attr_kv_num_heads.set_i(kv_num_heads); - attrs["kv_num_heads"] = attr_kv_num_heads; - ONNX_NAMESPACE::AttributeProto attr_scale; - attr_scale.set_name("scale"); - attr_scale.set_type(onnx::AttributeProto_AttributeType_FLOAT); - attr_scale.set_f(scale); - attrs["scale"] = attr_scale; - - // Register node - graph.AddNode( - node_name, - "GroupQueryAttention", - "GroupQueryAttention Node", - inputs, - {&output, &present_key, &present_value}, - &attrs, - "com.microsoft"); - - return output; -} - -void CreateLargeLLMModel(const PathString& model_path, const PathString& external_data_path) { - // Model parameters (example: 24 layers, 4096 hidden dim, 32 attention heads, 8 kv heads => GQA) - int batch_size = 1; - int num_layers = 32; - int hidden_dim = 2048; - int q_num_heads = 8; - int kv_num_heads = 1; // GQA: q_num_heads > kv_num_heads, and divisible. - int seq_length = 128; // Short, for demonstration. - int vocab_size = 32000; - auto dtype = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; - - // Set up model/graph - onnxruntime::Model model("LLM_With_GQA", false, DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - - // Input - ONNX_NAMESPACE::TypeProto input_type; - input_type.mutable_tensor_type()->set_elem_type(dtype); - input_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(batch_size); - input_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(seq_length); - input_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(hidden_dim); - auto& input = graph.GetOrCreateNodeArg("input", &input_type); - - auto* current_arg = &input; - - // Repeated layers: [Attention + MLP] - for (int l = 0; l < num_layers; ++l) { - // KV cache - initialize with zeros for the first forward pass - int head_dim = hidden_dim / q_num_heads; - - // Split Q, K, V - auto& q_split = graph.GetOrCreateNodeArg("q_split_" + std::to_string(l), nullptr); - auto& k_split = graph.GetOrCreateNodeArg("k_split_" + std::to_string(l), nullptr); - auto& v_split = graph.GetOrCreateNodeArg("v_split_" + std::to_string(l), nullptr); - constexpr bool split = false; - if constexpr (split) { - // Attention weights (Q, K, V projections) - auto wqkv = CreateLargeWeight("wqkv_" + std::to_string(l), - dtype, {hidden_dim, hidden_dim * 3}); - graph.AddInitializedTensor(wqkv); - - // Q = input @ wq, K = input @ wk, V = input @ wv - auto& qkv_arg = graph.GetOrCreateNodeArg("qkv_" + std::to_string(l), nullptr); - graph.AddNode("QKV_Linear_" + std::to_string(l), "MatMul", "", {current_arg, graph.GetNodeArg(wqkv.name())}, {&qkv_arg}); - - NodeAttributes attrs_split; - ONNX_NAMESPACE::AttributeProto attr_split_axis; - attr_split_axis.set_name("axis"); - attr_split_axis.set_type(onnx::AttributeProto_AttributeType_INT); - attr_split_axis.set_i(-1); - attrs_split["axis"] = attr_split_axis; - ONNX_NAMESPACE::AttributeProto attr_split_num_outputs; - attr_split_num_outputs.set_name("num_outputs"); - attr_split_num_outputs.set_type(onnx::AttributeProto_AttributeType_INT); - attr_split_num_outputs.set_i(3); - attrs_split["num_outputs"] = attr_split_num_outputs; - graph.AddNode("Q_Split_" + std::to_string(l), "Split", "", {&qkv_arg}, {&q_split, &k_split, &v_split}, &attrs_split); - } else { - // Attention weights (Q, K, V projections) - auto wq = CreateLargeWeight("wq_" + std::to_string(l), - dtype, {hidden_dim, hidden_dim}); - graph.AddInitializedTensor(wq); - auto wk = CreateLargeWeight("wk_" + std::to_string(l), - dtype, {hidden_dim, head_dim * kv_num_heads}); - graph.AddInitializedTensor(wk); - auto wv = CreateLargeWeight("wv_" + std::to_string(l), - dtype, {hidden_dim, head_dim * kv_num_heads}); - graph.AddInitializedTensor(wv); - - // Q = input @ wq, K = input @ wk, V = input @ wv - graph.AddNode("Q_Linear_" + std::to_string(l), "MatMul", "", {current_arg, graph.GetNodeArg(wq.name())}, {&q_split}); - graph.AddNode("K_Linear_" + std::to_string(l), "MatMul", "", {current_arg, graph.GetNodeArg(wk.name())}, {&k_split}); - graph.AddNode("V_Linear_" + std::to_string(l), "MatMul", "", {current_arg, graph.GetNodeArg(wv.name())}, {&v_split}); - } - // Reshape Q, K, V - auto& q_reshaped = graph.GetOrCreateNodeArg("q_reshaped_" + std::to_string(l), nullptr); - auto& k_reshaped = graph.GetOrCreateNodeArg("k_reshaped_" + std::to_string(l), nullptr); - auto& v_reshaped = graph.GetOrCreateNodeArg("v_reshaped_" + std::to_string(l), nullptr); - - ONNX_NAMESPACE::TensorProto q_shape_tensor; - q_shape_tensor.set_name("q_shape_" + std::to_string(l)); - q_shape_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - q_shape_tensor.add_dims(3); - q_shape_tensor.add_int64_data(batch_size); - q_shape_tensor.add_int64_data(seq_length); - q_shape_tensor.add_int64_data(head_dim * q_num_heads); - graph.AddInitializedTensor(q_shape_tensor); - - ONNX_NAMESPACE::TensorProto k_shape_tensor; - k_shape_tensor.set_name("k_shape_" + std::to_string(l)); - k_shape_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - k_shape_tensor.add_dims(3); - k_shape_tensor.add_int64_data(batch_size); - k_shape_tensor.add_int64_data(seq_length); - k_shape_tensor.add_int64_data(head_dim * kv_num_heads); - graph.AddInitializedTensor(k_shape_tensor); - - ONNX_NAMESPACE::TensorProto v_shape_tensor; - v_shape_tensor.set_name("v_shape_" + std::to_string(l)); - v_shape_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - v_shape_tensor.add_dims(3); - v_shape_tensor.add_int64_data(batch_size); - v_shape_tensor.add_int64_data(seq_length); - v_shape_tensor.add_int64_data(head_dim * kv_num_heads); - graph.AddInitializedTensor(v_shape_tensor); - - graph.AddNode("Q_Reshape_" + std::to_string(l), "Reshape", "", {&q_split, graph.GetNodeArg(q_shape_tensor.name())}, {&q_reshaped}); - graph.AddNode("K_Reshape_" + std::to_string(l), "Reshape", "", {&k_split, graph.GetNodeArg(k_shape_tensor.name())}, {&k_reshaped}); - graph.AddNode("V_Reshape_" + std::to_string(l), "Reshape", "", {&v_split, graph.GetNodeArg(v_shape_tensor.name())}, {&v_reshaped}); - - // Replace standard attention with GQA - auto& attn_out = AddGroupQueryAttention( - graph, q_reshaped, k_reshaped, v_reshaped, - batch_size, head_dim, seq_length, q_num_heads, kv_num_heads, - 1.0f, dtype, - "GQA_" + std::to_string(l)); - - // Add an MLP block: (Linear + Activation + Linear) - auto w1 = CreateLargeWeight("mlp_w1_" + std::to_string(l), dtype, {hidden_dim, hidden_dim * 4}); - auto w2 = CreateLargeWeight("mlp_w2_" + std::to_string(l), dtype, {hidden_dim * 4, hidden_dim}); - graph.AddInitializedTensor(w1); - graph.AddInitializedTensor(w2); - - auto& mlp_hidden = graph.GetOrCreateNodeArg("mlp_hidden_" + std::to_string(l), nullptr); - graph.AddNode("MLP_1_" + std::to_string(l), "MatMul", "", {&attn_out, graph.GetNodeArg(w1.name())}, {&mlp_hidden}); - auto& relu_out = graph.GetOrCreateNodeArg("relu_" + std::to_string(l), nullptr); - graph.AddNode("Relu_" + std::to_string(l), "Relu", "", {&mlp_hidden}, {&relu_out}); - auto& mlp_out = graph.GetOrCreateNodeArg("mlp_out_" + std::to_string(l), nullptr); - graph.AddNode("MLP_2_" + std::to_string(l), "MatMul", "", {&relu_out, graph.GetNodeArg(w2.name())}, {&mlp_out}); - current_arg = &mlp_out; // For next layer. - } - - // Final projection to vocab - auto w_logits = CreateLargeWeight("w_logits", - dtype, {hidden_dim, vocab_size}); - graph.AddInitializedTensor(w_logits); - auto& output = graph.GetOrCreateNodeArg("logits", nullptr); - graph.AddNode("Output_Linear", "MatMul", "", {current_arg, graph.GetNodeArg(w_logits.name())}, {&output}); - - // Validate, Write as large model with external data - auto status = graph.Resolve(); - if (!status.IsOK()) throw std::runtime_error(status.ErrorMessage()); - - onnxruntime::ModelSavingOptions save_options(128); - status = onnxruntime::Model::SaveWithExternalInitializers( - model, model_path, external_data_path, save_options); - if (!status.IsOK()) throw std::runtime_error(status.ErrorMessage()); -} - -Ort::IoBinding generate_io_binding( - Ort::Session& session, - std::map> shape_overwrites, - OrtAllocator* allocator) { - Ort::IoBinding binding(session); - auto default_allocator = Ort::AllocatorWithDefaultOptions(); - if (allocator == nullptr) { - allocator = default_allocator; - } - const OrtMemoryInfo* info; - Ort::ThrowOnError(Ort::GetApi().AllocatorGetInfo(allocator, &info)); - Ort::MemoryInfo mem_info(info->name.c_str(), info->alloc_type, static_cast(info->device.Id()), info->mem_type); - - for (int input_idx = 0; input_idx < int(session.GetInputCount()); ++input_idx) { - auto input_name = session.GetInputNameAllocated(input_idx, Ort::AllocatorWithDefaultOptions()); - auto full_tensor_info = session.GetInputTypeInfo(input_idx); - auto tensor_info = full_tensor_info.GetTensorTypeAndShapeInfo(); - auto shape = tensor_info.GetShape(); - auto type = tensor_info.GetElementType(); - if (shape_overwrites.find(input_name.get()) == shape_overwrites.end()) { - for (auto& v : shape) { - if (v == -1) { - v = 1; - } - } - } else { - shape = shape_overwrites[input_name.get()]; - } - auto input_value = Ort::Value::CreateTensor(allocator, - shape.data(), - shape.size(), - type); - binding.BindInput(input_name.get(), input_value); - } - - for (int output_idx = 0; output_idx < int(session.GetOutputCount()); ++output_idx) { - auto output_name = session.GetOutputNameAllocated(output_idx, Ort::AllocatorWithDefaultOptions()); - binding.BindOutput(output_name.get(), mem_info); - } - return binding; -} - -#if !defined(DISABLE_FLOAT8_TYPES) -// Helper function to create operator schemas for TRT FP8 Q/DQ custom-ops -static std::vector CreateTRTFP8Schemas() { - std::vector schemas; - - // TRT_FP8QuantizeLinear schema - schemas.emplace_back(); - ONNX_NAMESPACE::OpSchema& fp8_quant_schema = schemas.back(); - fp8_quant_schema - .SetName("TRT_FP8QuantizeLinear") - .SetDomain("trt") - .SinceVersion(1) - .SetDoc("TensorRT FP8 Quantization - quantizes FP16 input to FP8") - .Input(0, "X", "Input tensor in FP16", "T1") - .Input(1, "scale", "Scale for quantization in FP16", "T1") - .Output(0, "Y", "Quantized output tensor in FP8", "T2") - .TypeConstraint("T1", {"tensor(float16)"}, "Input and scale must be float16") - .TypeConstraint("T2", {"tensor(float8e4m3fn)"}, "Output must be float8e4m3fn") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - // Output has same shape as input but FP8 type - auto input_type = ctx.getInputType(0); - if (input_type != nullptr && input_type->has_tensor_type()) { - auto output_type = ctx.getOutputType(0); - output_type->mutable_tensor_type()->set_elem_type( - ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN); - if (input_type->tensor_type().has_shape()) { - *output_type->mutable_tensor_type()->mutable_shape() = - input_type->tensor_type().shape(); - } - } - }); - - // TRT_FP8DequantizeLinear schema - schemas.emplace_back(); - ONNX_NAMESPACE::OpSchema& fp8_dequant_schema = schemas.back(); - fp8_dequant_schema - .SetName("TRT_FP8DequantizeLinear") - .SetDomain("trt") - .SinceVersion(1) - .SetDoc("TensorRT FP8 Dequantization - dequantizes FP8 input to FP16") - .Input(0, "X", "Quantized input tensor in FP8", "T1") - .Input(1, "scale", "Scale for dequantization in FP16", "T2") - .Output(0, "Y", "Dequantized output tensor in FP16", "T2") - .TypeConstraint("T1", {"tensor(float8e4m3fn)"}, "Input must be float8e4m3fn") - .TypeConstraint("T2", {"tensor(float16)"}, "Scale and output must be float16") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - // Output has same shape as input but FP16 type - auto input_type = ctx.getInputType(0); - if (input_type != nullptr && input_type->has_tensor_type()) { - auto output_type = ctx.getOutputType(0); - output_type->mutable_tensor_type()->set_elem_type( - ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - if (input_type->tensor_type().has_shape()) { - *output_type->mutable_tensor_type()->mutable_shape() = - input_type->tensor_type().shape(); - } - } - }); - - return schemas; -} - -void CreateFP8CustomOpModel(const PathString& model_name, const std::string& graph_name) { - // Create custom schema registry for TRT operators - auto custom_schema_registry = std::make_shared(); - - // Register TRT FP8 operator schemas - auto trt_schemas = CreateTRTFP8Schemas(); - auto status = custom_schema_registry->RegisterOpSet(trt_schemas, "trt", 1, 1); - ASSERT_TRUE(status.IsOK()) << "Failed to register TRT schemas: " << status.ErrorMessage(); - - IOnnxRuntimeOpSchemaRegistryList registries = {custom_schema_registry}; - - // Set the opset version: ONNX domain to 19, TRT domain to 1 - const std::unordered_map domain_to_version = {{"", 19}, {"trt", 1}}; - onnxruntime::Model model(graph_name, false, ModelMetaData(), PathString(), - registries, domain_to_version, {}, - DefaultLoggingManager().DefaultLogger()); - - auto& graph = model.MainGraph(); - - // Define input dimensions and data type - std::vector dims = {4, 64}; // 4x64 matrix - auto dtype = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; - - // Create input tensor type (FLOAT16) - ONNX_NAMESPACE::TypeProto input_tensor; - input_tensor.mutable_tensor_type()->set_elem_type(dtype); - for (auto dim : dims) { - input_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - - // Create input node arg - auto& input_arg = graph.GetOrCreateNodeArg("X", &input_tensor); - - // Create quantized tensor type (FLOAT8E4M3FN for per-tensor quantization) - ONNX_NAMESPACE::TypeProto quantized_tensor; - quantized_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN); - for (auto dim : dims) { - quantized_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - auto& quantized_arg = graph.GetOrCreateNodeArg("X_quantized", &quantized_tensor); - - // Create output tensor type (FLOAT16 - after dequantization) - ONNX_NAMESPACE::TypeProto output_tensor; - output_tensor.mutable_tensor_type()->set_elem_type(dtype); - for (auto dim : dims) { - output_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - auto& output_arg = graph.GetOrCreateNodeArg("Y", &output_tensor); - - // Create scale initializer (scalar FLOAT16 for per-tensor quantization) - ONNX_NAMESPACE::TensorProto scale_initializer; - scale_initializer.set_name("scale"); - scale_initializer.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - // For FLOAT16, we need to convert the float value to MLFloat16 and store as raw data - MLFloat16 scale_value(0.0078125f); - scale_initializer.set_raw_data(&scale_value, sizeof(MLFloat16)); - graph.AddInitializedTensor(scale_initializer); - - // Get the scale node arg - auto* scale_arg = graph.GetNodeArg("scale"); - - // Add TRT_FP8QuantizeLinear node from "trt" domain - // Inputs: X (float16), scale (float16) - // Output: X_quantized (float8e4m3fn) - std::vector quant_inputs = {&input_arg, scale_arg}; - std::vector quant_outputs = {&quantized_arg}; - graph.AddNode("trt_fp8_quantize_node", "TRT_FP8QuantizeLinear", "TRT FP8 Quantize FP16 input to FP8", - quant_inputs, quant_outputs, nullptr, "trt"); - - // Add TRT_FP8DequantizeLinear node from "trt" domain - // Inputs: X_quantized (float8e4m3fn), scale (float16) - // Output: Y (float16) - std::vector dequant_inputs = {&quantized_arg, scale_arg}; - std::vector dequant_outputs = {&output_arg}; - graph.AddNode("trt_fp8_dequantize_node", "TRT_FP8DequantizeLinear", "TRT FP8 Dequantize FP8 to FP16", - dequant_inputs, dequant_outputs, nullptr, "trt"); - - // Set graph inputs and outputs explicitly - graph.SetInputs({&input_arg}); - graph.SetOutputs({&output_arg}); - - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_STATUS_OK(Model::Save(model, model_name)); -} -#endif // !defined(DISABLE_FLOAT8_TYPES) - -#if !defined(DISABLE_FLOAT4_TYPES) && !defined(DISABLE_FLOAT8_TYPES) -// Helper function to create operator schemas for TRT FP4 Dynamic Quantize custom-op -static std::vector CreateTRTFP4Schemas() { - std::vector schemas; - - // TRT_FP4DynamicQuantize schema - schemas.emplace_back(); - ONNX_NAMESPACE::OpSchema& fp4_quant_schema = schemas.back(); - fp4_quant_schema - .SetName("TRT_FP4DynamicQuantize") - .SetDomain("trt") - .SinceVersion(1) - .SetDoc("TensorRT FP4 Dynamic Quantization - quantizes FP16 input to FP4 with block-wise quantization") - .Attr("axis", "Axis along which to quantize", ONNX_NAMESPACE::AttributeProto::INT, static_cast(-1)) - .Attr("block_size", "Block size for quantization", ONNX_NAMESPACE::AttributeProto::INT, static_cast(16)) - .Attr("scale_type", "Scale data type", ONNX_NAMESPACE::AttributeProto::INT, static_cast(17)) - .Input(0, "X", "Input tensor in FP16", "T1") - .Input(1, "scale", "Scale for quantization in FP16", "T1") - .Output(0, "Y_quantized", "Quantized output tensor in FP4", "T2") - .Output(1, "Y_scale", "Computed scales in FP8", "T3") - .TypeConstraint("T1", {"tensor(float16)"}, "Input and scale must be float16") - .TypeConstraint("T2", {"tensor(float4e2m1)"}, "Quantized output must be float4e2m1") - .TypeConstraint("T3", {"tensor(float8e4m3fn)"}, "Scale output must be float8e4m3fn") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - // Output 0 (Y_quantized) has same shape as input but FP4 type - auto input_type = ctx.getInputType(0); - if (input_type != nullptr && input_type->has_tensor_type()) { - auto output_type_0 = ctx.getOutputType(0); - output_type_0->mutable_tensor_type()->set_elem_type( - ONNX_NAMESPACE::TensorProto_DataType_FLOAT4E2M1); - if (input_type->tensor_type().has_shape()) { - *output_type_0->mutable_tensor_type()->mutable_shape() = - input_type->tensor_type().shape(); - } - - // Output 1 (Y_scale) shape depends on block_size and axis - // For simplicity, we'll just set the type and let runtime handle shape - auto output_type_1 = ctx.getOutputType(1); - output_type_1->mutable_tensor_type()->set_elem_type( - ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN); - } - }); - - return schemas; -} - -void CreateFP4CustomOpModel(const PathString& model_name, const std::string& graph_name) { - // Create custom schema registry for TRT operators - auto custom_schema_registry = std::make_shared(); - - // Register TRT FP4 operator schemas - auto trt_schemas = CreateTRTFP4Schemas(); - auto status = custom_schema_registry->RegisterOpSet(trt_schemas, "trt", 1, 1); - ASSERT_TRUE(status.IsOK()) << "Failed to register TRT schemas: " << status.ErrorMessage(); - - // Create registry list - must be std::list, not std::vector - IOnnxRuntimeOpSchemaRegistryList registries = {custom_schema_registry}; - - // Set the opset version: ONNX domain to 23, TRT domain to 1 - const std::unordered_map domain_to_version = {{"", 23}, {"trt", 1}}; - onnxruntime::Model model(graph_name, false, ModelMetaData(), PathString(), - registries, domain_to_version, {}, - DefaultLoggingManager().DefaultLogger()); - auto& graph = model.MainGraph(); - - // Define input dimensions and data type - std::vector input_dims = {64, 64}; // 64x64 matrix - auto dtype = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; - - // Create input tensor type (FLOAT16) - ONNX_NAMESPACE::TypeProto input_tensor; - input_tensor.mutable_tensor_type()->set_elem_type(dtype); - for (auto dim : input_dims) { - input_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - - // Create input node arg for data - auto& input_arg = graph.GetOrCreateNodeArg("X", &input_tensor); - - // Create scale initializer (scalar FLOAT16) - ONNX_NAMESPACE::TensorProto scale_initializer; - scale_initializer.set_name("scale"); - scale_initializer.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - MLFloat16 scale_value(0.1234f); - scale_initializer.set_raw_data(&scale_value, sizeof(MLFloat16)); - graph.AddInitializedTensor(scale_initializer); - auto* scale_arg = graph.GetNodeArg("scale"); - - // Create output 1: quantized tensor (FLOAT4E2M1) [64, 64] - ONNX_NAMESPACE::TypeProto quantized_output_tensor; - quantized_output_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT4E2M1); - for (auto dim : input_dims) { - quantized_output_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - auto& output_quantized = graph.GetOrCreateNodeArg("X_quantized", &quantized_output_tensor); - - // Create output 2: scale output (FLOAT8) [64, 4] - // With block_size=16 and last dim=64, we get 64/16=4 blocks per row - std::vector scale_output_dims = {64, 4}; - ONNX_NAMESPACE::TypeProto scale_output_tensor; - scale_output_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN); - for (auto dim : scale_output_dims) { - scale_output_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - auto& output_scale = graph.GetOrCreateNodeArg("X_scale", &scale_output_tensor); - - // Create attributes - NodeAttributes attrs; - - // axis attribute - ONNX_NAMESPACE::AttributeProto attr_axis; - attr_axis.set_name("axis"); - attr_axis.set_type(onnx::AttributeProto_AttributeType_INT); - attr_axis.set_i(-1); - attrs["axis"] = attr_axis; - - // block_size attribute - ONNX_NAMESPACE::AttributeProto attr_block_size; - attr_block_size.set_name("block_size"); - attr_block_size.set_type(onnx::AttributeProto_AttributeType_INT); - attr_block_size.set_i(16); - attrs["block_size"] = attr_block_size; - - // scale_type attribute - ONNX_NAMESPACE::AttributeProto attr_scale_type; - attr_scale_type.set_name("scale_type"); - attr_scale_type.set_type(onnx::AttributeProto_AttributeType_INT); - attr_scale_type.set_i(17); - attrs["scale_type"] = attr_scale_type; - - // Add TRT_FP4DynamicQuantize node from "trt" domain - // Inputs: X (float16), scale (float16 initializer) - // Outputs: X_quantized (float4e2m1), X_scale (float8) - std::vector inputs = {&input_arg, scale_arg}; - std::vector outputs = {&output_quantized, &output_scale}; - graph.AddNode("trt_fp4_dyn_quant", "TRT_FP4DynamicQuantize", "TRT FP4 Dynamic Quantize node", - inputs, outputs, &attrs, "trt"); - - // Create scale initializer for DequantizeLinear (FLOAT16 scalar) - ONNX_NAMESPACE::TensorProto dequant_scale_initializer; - dequant_scale_initializer.set_name("dequant_scale"); - dequant_scale_initializer.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); - MLFloat16 dequant_scale_value(0.0625f); - dequant_scale_initializer.set_raw_data(&dequant_scale_value, sizeof(MLFloat16)); - graph.AddInitializedTensor(dequant_scale_initializer); - - // Get the dequant_scale node arg - auto* dequant_scale_arg = graph.GetNodeArg("dequant_scale"); - - // Create dequantized scale output (FLOAT16) [64, 4] - ONNX_NAMESPACE::TypeProto dequant_scale_output_tensor; - dequant_scale_output_tensor.mutable_tensor_type()->set_elem_type(dtype); - for (auto dim : scale_output_dims) { - dequant_scale_output_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - auto& output_dequant_scale = graph.GetOrCreateNodeArg("X_scale_dequantized", &dequant_scale_output_tensor); - - // Add first DequantizeLinear node (#1) for scale - // Inputs: X_scale (float8), dequant_scale (float16) - // Output: X_scale_dequantized (float16) - std::vector dequant_inputs = {&output_scale, dequant_scale_arg}; - std::vector dequant_outputs = {&output_dequant_scale}; - graph.AddNode("dequantize_scale_node", "DequantizeLinear", "Dequantize FP8 scale to FP16", - dequant_inputs, dequant_outputs); - - // Create final output: fully dequantized tensor (FLOAT16) [64, 64] - ONNX_NAMESPACE::TypeProto final_output_tensor; - final_output_tensor.mutable_tensor_type()->set_elem_type(dtype); - for (auto dim : input_dims) { - final_output_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(dim); - } - auto& output_final = graph.GetOrCreateNodeArg("X_dequantized", &final_output_tensor); - - // Create attributes for second DequantizeLinear node - NodeAttributes dequant_attrs; - - ONNX_NAMESPACE::AttributeProto dequant_attr_axis; - dequant_attr_axis.set_name("axis"); - dequant_attr_axis.set_type(onnx::AttributeProto_AttributeType_INT); - dequant_attr_axis.set_i(-1); - dequant_attrs["axis"] = dequant_attr_axis; - - ONNX_NAMESPACE::AttributeProto dequant_attr_block_size; - dequant_attr_block_size.set_name("block_size"); - dequant_attr_block_size.set_type(onnx::AttributeProto_AttributeType_INT); - dequant_attr_block_size.set_i(16); - dequant_attrs["block_size"] = dequant_attr_block_size; - - // Add second DequantizeLinear node (#2) for data - // Inputs: X_quantized (float4e2m1), X_scale_dequantized (float16) - // Output: X_dequantized (float16) - std::vector dequant_data_inputs = {&output_quantized, &output_dequant_scale}; - std::vector dequant_data_outputs = {&output_final}; - graph.AddNode("dequantize_data_node", "DequantizeLinear", "Dequantize FP4 data to FP16 with block quantization", - dequant_data_inputs, dequant_data_outputs, &dequant_attrs); - - // Set graph inputs and outputs explicitly - // Input: X only (scale is an initializer) - // Output: X_dequantized only (final dequantized result) - graph.SetInputs({&input_arg}); - graph.SetOutputs({&output_final}); - - ASSERT_STATUS_OK(graph.Resolve()); - ASSERT_STATUS_OK(Model::Save(model, model_name)); -} -#endif // !defined(DISABLE_FLOAT4_TYPES) && !defined(DISABLE_FLOAT8_TYPES) - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h b/onnxruntime/test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h deleted file mode 100644 index 94cd2729c2dbe..0000000000000 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h +++ /dev/null @@ -1,214 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// Licensed under the MIT License. -#pragma once - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include "core/graph/constants.h" -#include "core/common/path_string.h" -#include "core/framework/tensor.h" -#include "core/framework/ort_value.h" -#include "test/util/include/api_asserts.h" - -namespace onnxruntime { -namespace test { - -using RegisteredEpDeviceUniquePtr = std::unique_ptr>; - -struct Utils { - struct NvTensorRtRtxEpInfo { - const std::filesystem::path library_path = -#ifdef _WIN32 - "onnxruntime_providers_nv_tensorrt_rtx.dll"; -#else - "libonnxruntime_providers_nv_tensorrt_rtx.so"; -#endif - const std::string registration_name = kNvTensorRTRTXExecutionProvider; - }; - - static NvTensorRtRtxEpInfo nv_tensorrt_rtx_ep_info; - - // get the OrtEpDevice for the NV TensorRT RTX EP from the environment - static void GetEp(Ort::Env& env, const std::string& ep_name, const OrtEpDevice*& ep_device); - - // Register the NV TensorRT RTX EP library, get the OrtEpDevice for it, and return a unique pointer that will - // automatically unregister the EP library. - static void RegisterAndGetNvTensorRtRtxEp(Ort::Env& env, RegisteredEpDeviceUniquePtr& nv_tensorrt_rtx_ep); -}; - -[[maybe_unused]] static std::string PathToUTF8(const PathString& path) { -#ifdef _WIN32 - std::wstring_convert> converter; - return converter.to_bytes(path); -#else - return path.c_str(); -#endif -} - -[[maybe_unused]] static void clearFileIfExists(PathString path) { - if (std::filesystem::exists(path)) { - std::filesystem::remove(path); - } -} - -template -static void VerifyOutputs(const std::vector& fetches, const std::vector& expected_dims, - const std::vector& expected_values) { - ASSERT_EQ(1, fetches.size()); - auto& rtensor = fetches.front().Get(); - TensorShape expected_shape(expected_dims); - ASSERT_EQ(expected_shape, rtensor.Shape()); - const std::vector found(rtensor.Data(), rtensor.Data() + expected_values.size()); - ASSERT_EQ(expected_values, found); -} - -/** - * Create a simple model with dynamic or non-dynamic input shape. - * \param model_name - model name - * \param graph_name - graph name - * \param dims - input dimensions - * \param add_fast_gelu - add FastGelu node which makes the whole model partition into TRT EP and CUDA EP subgraphs. - * \param external_initializer_file - file name to save external initializers to - * - * input: "X", "Y" and "Z" - * you can specify input dimensions, for example (1, 3, 2), (1, 2) or (1, -1, -1)). Note: -1 means the dimension is dynamic. - * All three inputs have the same dimensions. - * output: "M" - * - * "X" "Y" - * \ / - * "Z" Add - * \ / - * Add - * / - * Add (+ float scalar "S") - * / - * "O" - * - * or - * - * "X" "Y" - * \ / - * "Z" Add - * \ / - * Add - * / - * FastGelu (This node will be placed on CUDA EP) - * / - * * Add (+ float scalar "S") - * / - * "O" - */ -void CreateBaseModel(const PathString& model_name, - std::string graph_name, - std::vector dims, - bool add_fast_gelu = false, - ONNX_NAMESPACE::TensorProto_DataType dtype = ONNX_NAMESPACE::TensorProto_DataType_FLOAT, - const PathString& external_initializer_file = {}); - -void CreateLargeLLMModel(const PathString& model_path, const PathString& external_data_path); - -Ort::IoBinding generate_io_binding( - Ort::Session& session, - std::map> shape_overwrites = {}, - OrtAllocator* allocator = nullptr); - -#if !defined(DISABLE_FLOAT8_TYPES) -/** - * Create a model with TRT_FP8QuantizeLinear -> TRT_FP8DequantizeLinear nodes for per-tensor quantization. - * - * The model uses opset version 19 for ONNX domain and opset version 1 for TRT domain. - * - * \param model_name - output model file name - * \param graph_name - name of the graph - * - * input: "X" [4, 64] (FLOAT16) - * output: "Y" [4, 64] (FLOAT16) - * - * "X" [4x64] - * (FLOAT16) - * | - * | - * TRT_FP8QuantizeLinear (trt domain) - * (+ scale, FLOAT16 initializer) - * | - * | - * "X_quantized" - * [4x64] - * (FLOAT8E4M3FN) - * | - * | - * TRT_FP8DequantizeLinear (trt domain) - * (+ scale, FLOAT16 initializer) - * | - * | - * "Y" - * [4x64] - * (FLOAT16) - */ -void CreateFP8CustomOpModel(const PathString& model_name, - const std::string& graph_name); -#endif // !defined(DISABLE_FLOAT8_TYPES) - -#if !defined(DISABLE_FLOAT4_TYPES) && !defined(DISABLE_FLOAT8_TYPES) -/** - * Create a model with TRT_FP4DynamicQuantize node followed by DequantizeLinear nodes. - * - * The model uses opset version 23 for ONNX domain and opset version 1 for TRT domain. - * - * \param model_name - output model file name - * \param graph_name - name of the graph - * - * input: "X" [64, 64] (FLOAT16) - * output: "X_dequantized" [64, 64] (FLOAT16) - * - * "X" [64x64] - * (FLOAT16) - * | - * | - * TRT_FP4DynamicQuantize (trt domain) - * (axis=-1, block_size=16, scale_type=17) - * (+ scale, FLOAT16 initializer) - * / \ - * / \ - * "X_quantized" "X_scale" - * [64x64] [64x4] - * (FLOAT4E2M1) (FLOAT8E4M3FN) - * | | - * | | - * | DequantizeLinear #1 - * | (+ dequant_scale, FLOAT16 initializer) - * | | - * | | - * | "X_scale_dequantized" - * | [64x4] - * | (FLOAT16) - * | / - * \ / - * \ / - * \ / - * DequantizeLinear #2 - * (axis=-1, block_size=16) - * | - * | - * "X_dequantized" [OUTPUT] - * [64x64] - * (FLOAT16) - */ -void CreateFP4CustomOpModel(const PathString& model_name, - const std::string& graph_name); -#endif // !defined(DISABLE_FLOAT4_TYPES) && !defined(DISABLE_FLOAT8_TYPES) - -} // namespace test -} // namespace onnxruntime diff --git a/onnxruntime/test/python/onnxruntime_test_python_nv_tensorrt_rtx_ep_tests.py b/onnxruntime/test/python/onnxruntime_test_python_nv_tensorrt_rtx_ep_tests.py deleted file mode 100644 index 23cb006b01ad9..0000000000000 --- a/onnxruntime/test/python/onnxruntime_test_python_nv_tensorrt_rtx_ep_tests.py +++ /dev/null @@ -1,472 +0,0 @@ -# Copyright (c) NVIDIA Corporation. All rights reserved. -# Licensed under the MIT License. -from __future__ import annotations - -import sys -import unittest -from collections.abc import Sequence - -import numpy as np -import torch -from autoep_helper import AutoEpTestCase -from helper import get_name -from numpy.testing import assert_almost_equal -from onnx import TensorProto, helper - -import onnxruntime as onnxrt -from onnxruntime.capi._pybind_state import OrtDevice as C_OrtDevice -from onnxruntime.capi._pybind_state import OrtValue as C_OrtValue -from onnxruntime.capi._pybind_state import OrtValueVector, SessionIOBinding - -# Highest ai.onnx opset that has been *released* by the installed onnx package. -# Avoid onnx.defs.onnx_opset_version(), which returns the in-development next opset -# (ORT's loader rejects un-released opsets unless ALLOW_RELEASED_ONNX_OPSET_ONLY=0). -_LAST_RELEASED_AI_ONNX_OPSET = max(v for (d, v) in helper.OP_SET_ID_VERSION_MAP if d == "ai.onnx") - - -class TestNvTensorRTRTXAutoEP(AutoEpTestCase): - """ - Test suite for the NvTensorRTRTX Execution Provider. - - This class contains tests for registering the NvTensorRTRTX EP, - selecting it using different policies, and running inference with various - I/O binding configurations. - """ - - ep_lib_path = "onnxruntime_providers_nv_tensorrt_rtx.dll" - ep_name = "NvTensorRTRTXExecutionProvider" - - def setUp(self): - if sys.platform != "win32": - self.skipTest("Skipping test because device discovery is only supported on Windows") - self.register_execution_provider_library(self.ep_name, self.ep_lib_path) - - def tearDown(self): - self.unregister_execution_provider_library(self.ep_name) - - def _create_ortvalue_input_on_gpu(self, device): - return onnxrt.OrtValue.ortvalue_from_numpy( - np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32), device, 0 - ) - - def _create_ortvalue_alternate_input_on_gpu(self, device): - return onnxrt.OrtValue.ortvalue_from_numpy( - np.array([[2.0, 4.0], [6.0, 8.0], [10.0, 12.0]], dtype=np.float32), - device, - 0, - ) - - def _create_uninitialized_ortvalue_input_on_gpu(self, device): - return onnxrt.OrtValue.ortvalue_from_shape_and_type([3, 2], np.float32, device, 0) - - def _create_numpy_input(self): - return np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) - - def _create_expected_output(self): - return np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32) - - def _create_expected_output_alternate(self): - return np.array([[2.0, 8.0], [18.0, 32.0], [50.0, 72.0]], dtype=np.float32) - - def torch_to_onnx_type(self, torch_dtype): - if torch_dtype == torch.float32: - return TensorProto.FLOAT - elif torch_dtype == torch.float16: - return TensorProto.FLOAT16 - elif torch_dtype == torch.bfloat16: - return TensorProto.BFLOAT16 - elif torch_dtype == torch.int8: - return TensorProto.int8 - elif torch_dtype == torch.int32: - return TensorProto.INT32 - elif torch_dtype == torch.int64: - return TensorProto.INT64 - else: - raise TypeError(f"Unsupported dtype: {torch_dtype}") - - def test_nv_tensorrt_rtx_ep_register_and_inference(self): - """ - Test registration of NvTensorRTRTX EP, adding its OrtDevice to the SessionOptions, and running inference. - """ - ep_devices = onnxrt.get_ep_devices() - nv_tensorrt_rtx_ep_device = next((d for d in ep_devices if d.ep_name == self.ep_name), None) - self.assertIsNotNone(nv_tensorrt_rtx_ep_device) - self.assertEqual(nv_tensorrt_rtx_ep_device.ep_vendor, "NVIDIA") - - hw_device = nv_tensorrt_rtx_ep_device.device - self.assertEqual(hw_device.type, onnxrt.OrtHardwareDeviceType.GPU) - - # Run sample model and check output - sess = onnxrt.InferenceSession(get_name("mul_1.onnx")) - - x = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) - input_name = sess.get_inputs()[0].name - res = sess.run([], {input_name: x}) - output_expected = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32) - np.testing.assert_allclose(res[0], output_expected, rtol=1e-05, atol=1e-08) - - def test_nv_tensorrt_rtx_ep_prefer_gpu_and_inference(self): - """ - Test selecting NvTensorRTRTX EP via the PREFER_GPU policy and running inference. - """ - # Set a policy to prefer GPU. NvTensorRTRTX should be selected. - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - - # Run sample model and check output - sess = onnxrt.InferenceSession(get_name("mul_1.onnx"), sess_options=sess_options) - - x = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) - input_name = sess.get_inputs()[0].name - res = sess.run([], {input_name: x}) - output_expected = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32) - np.testing.assert_allclose(res[0], output_expected, rtol=1e-05, atol=1e-08) - - def test_nv_tensorrt_rtx_ep_selection_delegate_and_inference(self): - """ - Test selecting NvTensorRTRTX EP via the custom EP selection delegate function and then run inference. - """ - - # User's custom EP selection function. - def my_delegate( - ep_devices: Sequence[onnxrt.OrtEpDevice], - model_metadata: dict[str, str], - runtime_metadata: dict[str, str], - max_selections: int, - ) -> Sequence[onnxrt.OrtEpDevice]: - self.assertGreater(len(model_metadata), 0) - self.assertGreaterEqual(len(ep_devices), 1) - self.assertGreaterEqual(max_selections, 2) - - nv_tensorrt_rtx_ep_device = next((d for d in ep_devices if d.ep_name == self.ep_name), None) - self.assertIsNotNone(nv_tensorrt_rtx_ep_device) - - # Select the NvTensorRTRTX device - return [nv_tensorrt_rtx_ep_device] - - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy_delegate(my_delegate) - self.assertTrue(sess_options.has_providers()) - - # Run sample model and check output - sess = onnxrt.InferenceSession(get_name("mul_1.onnx"), sess_options=sess_options) - - x = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) - input_name = sess.get_inputs()[0].name - res = sess.run([], {input_name: x}) - output_expected = np.array([[1.0, 4.0], [9.0, 16.0], [25.0, 36.0]], dtype=np.float32) - np.testing.assert_allclose(res[0], output_expected, rtol=1e-05, atol=1e-08) - - def test_bind_input_only(self): - """ - Test I/O binding with input data only. - """ - # Set a policy to prefer GPU. NvTensorRTRTX should be selected. - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - - input = self._create_ortvalue_input_on_gpu("cuda") - - session = onnxrt.InferenceSession(get_name("mul_1.onnx"), sess_options=sess_options) - io_binding = session.io_binding() - - # Bind input to the GPU - io_binding.bind_input("X", "cuda", 0, np.float32, [3, 2], input.data_ptr()) - - # Sync if different streams - io_binding.synchronize_inputs() - - # Bind output to CPU - io_binding.bind_output("Y") - - # Invoke Run - session.run_with_iobinding(io_binding) - - # Sync if different streams - io_binding.synchronize_outputs() - - # Get outputs over to CPU (the outputs which were bound to the GPU will get copied over to the host - # here) - ort_output = io_binding.copy_outputs_to_cpu()[0] - - # Validate results - self.assertTrue(np.array_equal(self._create_expected_output(), ort_output)) - - def test_bind_input_and_bind_output_with_ortvalues(self): - """ - Test I/O binding with OrtValues for both input and output. - """ - # Set a policy to prefer GPU. NvTensorRTRTX EP should be selected. - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - - session = onnxrt.InferenceSession(get_name("mul_1.onnx"), sess_options=sess_options) - io_binding = session.io_binding() - - # Bind ortvalue as input - input_ortvalue = self._create_ortvalue_input_on_gpu("cuda") - io_binding.bind_ortvalue_input("X", input_ortvalue) - - # Bind ortvalue as output - output_ortvalue = self._create_uninitialized_ortvalue_input_on_gpu("cuda") - io_binding.bind_ortvalue_output("Y", output_ortvalue) - - # Sync if different streams - io_binding.synchronize_inputs() - - # Invoke Run - session.run_with_iobinding(io_binding) - - # Sync if different streams - io_binding.synchronize_outputs() - - # Inspect contents of output_ortvalue and make sure that it has the right contents - self.assertTrue(np.array_equal(self._create_expected_output(), output_ortvalue.numpy())) - - # Bind another ortvalue as input - input_ortvalue_2 = self._create_ortvalue_alternate_input_on_gpu("cuda") - io_binding.bind_ortvalue_input("X", input_ortvalue_2) - - # Sync if different streams - io_binding.synchronize_inputs() - - # Invoke Run - session.run_with_iobinding(io_binding) - - # Sync if different streams - io_binding.synchronize_outputs() - - # Inspect contents of output_ortvalue and make sure that it has the right contents - self.assertTrue(np.array_equal(self._create_expected_output_alternate(), output_ortvalue.numpy())) - - def test_bind_input_and_non_preallocated_output(self): - """ - Test I/O binding with non-preallocated output. - """ - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - - session = onnxrt.InferenceSession(get_name("mul_1.onnx"), sess_options=sess_options) - io_binding = session.io_binding() - - input = self._create_ortvalue_input_on_gpu("cuda") - - # Bind input to the GPU - io_binding.bind_input("X", "cuda", 0, np.float32, [3, 2], input.data_ptr()) - - # Bind output to the GPU - io_binding.bind_output("Y", "cuda") - - # Sync if different streams - io_binding.synchronize_inputs() - - # Invoke Run - session.run_with_iobinding(io_binding) - - # Sync if different streams - io_binding.synchronize_outputs() - - # This call returns an OrtValue which has data allocated by ORT on the GPU - ort_outputs = io_binding.get_outputs() - self.assertEqual(len(ort_outputs), 1) - self.assertEqual(ort_outputs[0].device_name(), "cuda") - # Validate results (by copying results to CPU by creating a Numpy object) - self.assertTrue(np.array_equal(self._create_expected_output(), ort_outputs[0].numpy())) - - # We should be able to repeat the above process as many times as we want - try once more - ort_outputs = io_binding.get_outputs() - self.assertEqual(len(ort_outputs), 1) - self.assertEqual(ort_outputs[0].device_name(), "cuda") - # Validate results (by copying results to CPU by creating a Numpy object) - self.assertTrue(np.array_equal(self._create_expected_output(), ort_outputs[0].numpy())) - - input = self._create_ortvalue_alternate_input_on_gpu("cuda") - - # Change the bound input and validate the results in the same bound OrtValue - # Bind alternate input to the GPU - io_binding.bind_input( - "X", - "cuda", - 0, - np.float32, - [3, 2], - input.data_ptr(), - ) - - # Sync if different streams - io_binding.synchronize_inputs() - - # Invoke Run - session.run_with_iobinding(io_binding) - - # Sync if different streams - io_binding.synchronize_outputs() - - # This call returns an OrtValue which has data allocated by ORT on the GPU - ort_outputs = io_binding.get_outputs() - self.assertEqual(len(ort_outputs), 1) - self.assertEqual(ort_outputs[0].device_name(), "cuda") - # Validate results (by copying results to CPU by creating a Numpy object) - self.assertTrue(np.array_equal(self._create_expected_output_alternate(), ort_outputs[0].numpy())) - - def test_bind_input_and_preallocated_output(self): - """ - Test I/O binding with preallocated output. - """ - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - - input = self._create_ortvalue_input_on_gpu("cuda") - - session = onnxrt.InferenceSession(get_name("mul_1.onnx"), sess_options=sess_options) - io_binding = session.io_binding() - - # Bind input to the GPU - io_binding.bind_input("X", "cuda", 0, np.float32, [3, 2], input.data_ptr()) - - # Bind output to the GPU - output = self._create_uninitialized_ortvalue_input_on_gpu("cuda") - io_binding.bind_output("Y", "cuda", 0, np.float32, [3, 2], output.data_ptr()) - - # Sync if different streams - io_binding.synchronize_inputs() - - # Invoke Run - session.run_with_iobinding(io_binding) - - # Sync if different streams - io_binding.synchronize_outputs() - - # Get outputs over to CPU (the outputs which were bound to the GPU will get copied over to the host - # here) - ort_output_vals = io_binding.copy_outputs_to_cpu()[0] - # Validate results - self.assertTrue(np.array_equal(self._create_expected_output(), ort_output_vals)) - - # Validate if ORT actually wrote to pre-allocated buffer by copying the allocated buffer - # to the host and validating its contents - ort_output_vals_in_cpu = output.numpy() - # Validate results - self.assertTrue(np.array_equal(self._create_expected_output(), ort_output_vals_in_cpu)) - - def test_bind_input_types(self): - """ - Test I/O binding with various input data types. - """ - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - opset = _LAST_RELEASED_AI_ONNX_OPSET - device = C_OrtDevice(C_OrtDevice.cuda(), C_OrtDevice.default_memory(), 0) - - for dtype in [ - np.float32, - # np.float64, - np.int32, - # np.uint32, - np.int64, - # np.uint64, - # np.int16, - # np.uint16, - # np.int8, - np.uint8, - np.float16, - np.bool_, - ]: - with self.subTest(dtype=dtype, inner_device=str(device)): - x = np.arange(8).reshape((-1, 2)).astype(dtype) - proto_dtype = helper.np_dtype_to_tensor_dtype(x.dtype) - - X = helper.make_tensor_value_info("X", proto_dtype, [None, x.shape[1]]) # noqa: N806 - Y = helper.make_tensor_value_info("Y", proto_dtype, [None, x.shape[1]]) # noqa: N806 - - # inference - node_add = helper.make_node("Identity", ["X"], ["Y"]) - - # graph - graph_def = helper.make_graph([node_add], "lr", [X], [Y], []) - model_def = helper.make_model( - graph_def, - producer_name="dummy", - ir_version=7, - producer_version="0", - opset_imports=[helper.make_operatorsetid("", opset)], - ) - - sess = onnxrt.InferenceSession(model_def.SerializeToString(), sess_options=sess_options) - - bind = SessionIOBinding(sess._sess) - ort_value = C_OrtValue.ortvalue_from_numpy(x, device) - bind.bind_ortvalue_input("X", ort_value) - bind.bind_output("Y", device) - sess._sess.run_with_iobinding(bind, None) - ortvaluevector = bind.get_outputs() - self.assertIsInstance(ortvaluevector, OrtValueVector) - ortvalue = bind.get_outputs()[0] - y = ortvalue.numpy() - assert_almost_equal(x, y) - - bind = SessionIOBinding(sess._sess) - bind.bind_input("X", device, dtype, x.shape, ort_value.data_ptr()) - bind.bind_output("Y", device) - sess._sess.run_with_iobinding(bind, None) - ortvalue = bind.get_outputs()[0] - y = ortvalue.numpy() - assert_almost_equal(x, y) - - def test_bind_onnx_types_from_torch(self): - """ - Test I/O binding with various input data types. - """ - sess_options = onnxrt.SessionOptions() - sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_GPU) - self.assertTrue(sess_options.has_providers()) - opset = _LAST_RELEASED_AI_ONNX_OPSET - - for dtype in [ - torch.float32, - torch.float16, - torch.bfloat16, - torch.int32, - torch.int64, - ]: - with self.subTest(dtype=dtype): - proto_dtype = self.torch_to_onnx_type(dtype) - - x_ = helper.make_tensor_value_info("X", proto_dtype, [None]) - y_ = helper.make_tensor_value_info("Y", proto_dtype, [None]) - node_add = helper.make_node("Identity", ["X"], ["Y"]) - graph_def = helper.make_graph([node_add], "lr", [x_], [y_], []) - model_def = helper.make_model( - graph_def, - producer_name="dummy", - ir_version=10, - producer_version="0", - opset_imports=[helper.make_operatorsetid("", opset)], - ) - sess = onnxrt.InferenceSession(model_def.SerializeToString(), sess_options=sess_options) - - dev = "cuda" if torch.cuda.is_available() else "cpu" - device = ( - C_OrtDevice(C_OrtDevice.cuda(), C_OrtDevice.default_memory(), 0) - if dev == "cuda" - else C_OrtDevice(C_OrtDevice.cpu(), C_OrtDevice.default_memory(), 0) - ) - - x = torch.arange(8, dtype=dtype, device=dev) - y = torch.empty(8, dtype=dtype, device=dev) - - bind = SessionIOBinding(sess._sess) - bind.bind_input("X", device, proto_dtype, x.shape, x.data_ptr()) - bind.bind_output("Y", device, proto_dtype, y.shape, y.data_ptr()) - sess._sess.run_with_iobinding(bind, None) - self.assertTrue(torch.equal(x, y)) - - -if __name__ == "__main__": - unittest.main(verbosity=1) diff --git a/onnxruntime/test/shared_lib/test_session_options.cc b/onnxruntime/test/shared_lib/test_session_options.cc index 7e29c6f9a71d1..51eb0f718749c 100644 --- a/onnxruntime/test/shared_lib/test_session_options.cc +++ b/onnxruntime/test/shared_lib/test_session_options.cc @@ -97,20 +97,21 @@ TEST(CApiTest, session_options_provider_interface_fail_add_cuda) { } #endif // defined(USE_CUDA_PROVIDER_INTERFACE) -#if defined(USE_NV_PROVIDER_INTERFACE) -// Test that loading NV EP when only the interface is built (but not the full EP) fails. -TEST(CApiTest, session_options_provider_interface_fail_add_nv) { +#if !defined(ORT_MINIMAL_BUILD) +// Test that appending the deprecated provider-bridge NV TensorRT RTX EP by name reports the plugin EP flow. +TEST(CApiTest, session_options_append_removed_nv_tensorrt_rtx_ep) { const OrtApi& api = Ort::GetApi(); - Ort::SessionOptions session_options; - Ort::Status status = Ort::Status{api.SessionOptionsAppendExecutionProvider(session_options, - kNvTensorRTRTXExecutionProvider, - nullptr, nullptr, 0)}; - ASSERT_TRUE(!status.IsOK()); - EXPECT_EQ(status.GetErrorCode(), ORT_FAIL); - EXPECT_THAT(status.GetErrorMessage(), testing::HasSubstr("Failed to load")); + for (const char* provider_name : {"NvTensorRtRtx", kNvTensorRTRTXExecutionProvider}) { + Ort::SessionOptions session_options; + Ort::Status status = Ort::Status{api.SessionOptionsAppendExecutionProvider(session_options, provider_name, + nullptr, nullptr, 0)}; + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.GetErrorCode(), ORT_INVALID_ARGUMENT); + EXPECT_THAT(status.GetErrorMessage(), testing::HasSubstr("Use the NVIDIA TensorRT RTX plugin EP")); + } } -#endif // defined(USE_OPENVINO_PROVIDER_INTERFACE) +#endif // !defined(ORT_MINIMAL_BUILD) #if defined(USE_TENSORRT_PROVIDER_INTERFACE) // Test that loading TensorRT EP when only the interface is built (but not the full EP) fails. diff --git a/onnxruntime/test/unittest_util/base_tester.cc b/onnxruntime/test/unittest_util/base_tester.cc index 02086dce49aec..92ad4948525ee 100644 --- a/onnxruntime/test/unittest_util/base_tester.cc +++ b/onnxruntime/test/unittest_util/base_tester.cc @@ -668,12 +668,6 @@ void BaseTester::RunWithConfig(size_t* number_of_pre_packed_weights_counter, #endif kDnnlExecutionProvider, kTensorrtExecutionProvider, -#ifdef USE_NV - // Only include NV TRT RTX EP when is ORT is built with the provider-bridge - // version of the EP (i.e., USE_NV is defined). This allows use of the plugin EP version of the EP - // when ORT is not built any provider-bridge EPs. - kNvTensorRTRTXExecutionProvider, -#endif kOpenVINOExecutionProvider, kDmlExecutionProvider, kAclExecutionProvider, @@ -733,8 +727,6 @@ void BaseTester::RunWithConfig(size_t* number_of_pre_packed_weights_counter, execution_provider = DefaultDnnlExecutionProvider(); else if (provider_type == onnxruntime::kOpenVINOExecutionProvider) execution_provider = DefaultOpenVINOExecutionProvider(); - else if (provider_type == onnxruntime::kNvTensorRTRTXExecutionProvider) - execution_provider = DefaultNvTensorRTRTXExecutionProvider(); else if (provider_type == onnxruntime::kTensorrtExecutionProvider) execution_provider = DefaultTensorrtExecutionProvider(); else if (provider_type == onnxruntime::kNnapiExecutionProvider) @@ -855,9 +847,6 @@ void BaseTester::ExecuteModelForEps( constexpr std::array kEpsThatCompileNodes{ kOpenVINOExecutionProvider, kTensorrtExecutionProvider, // uses kernel registry for Memcpy* nodes only -#ifdef USE_NV - kNvTensorRTRTXExecutionProvider, // uses kernel registry for Memcpy* nodes only -#endif kNnapiExecutionProvider, kVSINPUExecutionProvider, kCoreMLExecutionProvider, diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc index 86cc5719b9c2b..ecd0c3b43c6d1 100644 --- a/onnxruntime/test/util/default_providers.cc +++ b/onnxruntime/test/util/default_providers.cc @@ -94,14 +94,6 @@ std::unique_ptr DefaultTensorrtExecutionProvider() { return nullptr; } -std::unique_ptr DefaultNvTensorRTRTXExecutionProvider() { -#ifdef USE_NV - if (auto factory = NvProviderFactoryCreator::Create(0)) - return factory->CreateProvider(); -#endif - return nullptr; -} - std::unique_ptr TensorrtExecutionProviderWithOptions(const OrtTensorRTProviderOptions* params) { #ifdef USE_TENSORRT if (auto factory = TensorrtProviderFactoryCreator::Create(params)) diff --git a/onnxruntime/test/util/include/default_providers.h b/onnxruntime/test/util/include/default_providers.h index 0bf36822d4e85..574a66de30df5 100644 --- a/onnxruntime/test/util/include/default_providers.h +++ b/onnxruntime/test/util/include/default_providers.h @@ -44,7 +44,6 @@ std::unique_ptr DefaultDnnlExecutionProvider(); std::unique_ptr DnnlExecutionProviderWithOptions(const OrtDnnlProviderOptions* provider_options); // std::unique_ptr DefaultTvmExecutionProvider(); std::unique_ptr DefaultTensorrtExecutionProvider(); -std::unique_ptr DefaultNvTensorRTRTXExecutionProvider(); std::unique_ptr TensorrtExecutionProviderWithOptions(const OrtTensorRTProviderOptions* params); std::unique_ptr TensorrtExecutionProviderWithOptions(const OrtTensorRTProviderOptionsV2* params); std::unique_ptr DefaultMIGraphXExecutionProvider(); diff --git a/setup.py b/setup.py index fcde855cf20b0..d7b988d8423bc 100644 --- a/setup.py +++ b/setup.py @@ -198,19 +198,6 @@ def _rewrite_ld_preload_tensorrt(self, to_preload): f.write(" import os\n") f.write(' os.environ["ORT_TENSORRT_UNAVAILABLE"] = "1"\n') - def _rewrite_ld_preload_nv_tensorrt_rtx(self, to_preload): - with open("onnxruntime/capi/_ld_preload.py", "a", encoding="ascii") as f: - if len(to_preload) > 0: - f.write("from ctypes import CDLL, RTLD_GLOBAL\n") - f.write("try:\n") - f.writelines( - ' _{} = CDLL("{}", mode=RTLD_GLOBAL)\n'.format(library.split(".")[0], library) - for library in to_preload - ) - f.write("except OSError:\n") - f.write(" import os\n") - f.write(' os.environ["ORT_NV_TENSORRT_RTX_UNAVAILABLE"] = "1"\n') - def run(self): if is_manylinux: source = "onnxruntime/capi/onnxruntime_pybind11_state.so" @@ -221,7 +208,6 @@ def run(self): to_preload = [] to_preload_cuda = [] to_preload_tensorrt = [] - to_preload_nv_tensorrt_rtx = [] to_preload_cann = [] cuda_dependencies = [ @@ -298,7 +284,6 @@ def run(self): self._rewrite_ld_preload(to_preload) self._rewrite_ld_preload_cuda(to_preload_cuda) self._rewrite_ld_preload_tensorrt(to_preload_tensorrt) - self._rewrite_ld_preload_tensorrt(to_preload_nv_tensorrt_rtx) self._rewrite_ld_preload(to_preload_cann) elif platform.system() == "Linux": # Non-manylinux Linux builds: preload libcudart so that undefined CUDA symbols @@ -348,7 +333,6 @@ def finalize_options(self): providers_cuda = "onnxruntime_providers_cuda" providers_tensorrt_or_migraphx = "onnxruntime_providers_" + ("migraphx" if is_migraphx else "tensorrt") -providers_nv_tensorrt_rtx = "onnxruntime_providers_nv_tensorrt_rtx" providers_openvino = "onnxruntime_providers_openvino" providers_cann = "onnxruntime_providers_cann" providers_qnn = "onnxruntime_providers_qnn" @@ -356,14 +340,12 @@ def finalize_options(self): if platform.system() == "Linux": providers_cuda = "lib" + providers_cuda + ".so" providers_tensorrt_or_migraphx = "lib" + providers_tensorrt_or_migraphx + ".so" - providers_nv_tensorrt_rtx = "lib" + providers_nv_tensorrt_rtx + ".so" providers_openvino = "lib" + providers_openvino + ".so" providers_cann = "lib" + providers_cann + ".so" providers_qnn = "lib" + providers_qnn + ".so" elif platform.system() == "Windows": providers_cuda = providers_cuda + ".dll" providers_tensorrt_or_migraphx = providers_tensorrt_or_migraphx + ".dll" - providers_nv_tensorrt_rtx = providers_nv_tensorrt_rtx + ".dll" providers_openvino = providers_openvino + ".dll" providers_cann = providers_cann + ".dll" providers_qnn = providers_qnn + ".dll" @@ -402,7 +384,6 @@ def finalize_options(self): libs.extend(["libonnxruntime_providers_openvino.so"]) libs.extend(["libonnxruntime_providers_vitisai.so"]) libs.append(providers_cuda) - libs.append(providers_nv_tensorrt_rtx) libs.append(providers_tensorrt_or_migraphx) libs.append(providers_cann) libs.append(providers_qnn) @@ -421,10 +402,6 @@ def finalize_options(self): ] dl_libs.extend(qnn_deps) libs.extend(qnn_deps) - # NV TensorRT RTX - nv_tensorrt_rtx_deps = ["libtensorrt_rtx.so", "libtensorrt_onnxparser_rtx.so"] - dl_libs.extend(nv_tensorrt_rtx_deps) - libs.extend(nv_tensorrt_rtx_deps) if nightly_build: libs.extend(["libonnxruntime_pywrapper.so"]) elif platform.system() == "Darwin": @@ -451,7 +428,6 @@ def finalize_options(self): "libiomp5md.dll", providers_cuda, providers_tensorrt_or_migraphx, - providers_nv_tensorrt_rtx, providers_cann, "onnxruntime.dll", ] @@ -459,7 +435,6 @@ def finalize_options(self): libs.extend(["onnxruntime_providers_shared.dll"]) libs.extend(["onnxruntime_providers_dnnl.dll"]) libs.extend(["onnxruntime_providers_tensorrt.dll"]) - libs.extend(["onnxruntime_providers_nv_tensorrt_rtx.dll"]) libs.extend(["onnxruntime_providers_openvino.dll"]) libs.extend(["onnxruntime_providers_cuda.dll"]) libs.extend(["onnxruntime_providers_migraphx.dll"]) @@ -510,10 +485,6 @@ def finalize_options(self): "migraphx_tf.dll", ] libs.extend(migraphx_deps) - # NV TensorRT RTX Libs - nv_tensorrt_rtx_deps = ["tensorrt_onnxparser_rtx_*.dll", "tensorrt_rtx_*.dll"] - libs.extend(nv_tensorrt_rtx_deps) - if is_manylinux: if is_openvino: ov_libs = [ @@ -555,14 +526,12 @@ def finalize_options(self): examples_names = ["mul_1.onnx", "logreg_iris.onnx", "sigmoid.onnx"] examples = [path.join("datasets", x) for x in examples_names] -# Extra files such as EULA and ThirdPartyNotices (and Qualcomm License, only for QNN release packages)(and TensorRT RTX License and Acknowledgements for TRT RTX EP) +# Extra files such as EULA and ThirdPartyNotices (and Qualcomm License, only for QNN release packages) extra = [ "LICENSE", "ThirdPartyNotices.txt", "Privacy.md", "Qualcomm_LICENSE.pdf", - "TRT_RTX_LICENSE.txt", - "TRT_RTX_Acknowledgements.txt", ] # Description @@ -823,13 +792,6 @@ def reformat_run_count(count_str): with open(requirements_path) as f: install_requires = f.read().splitlines() -# Adding CUDA Runtime as dependency for NV TensorRT RTX python wheel -if package_name == "onnxruntime-trt-rtx": - major = cuda_major_version or "12" # Default to CUDA 12 - # CUDA 13 dropped the "-cuNN" suffix from the CUDA Runtime package name. - runtime_pkg = "nvidia-cuda-runtime" if int(major) >= 13 else f"nvidia-cuda-runtime-cu{major}" - install_requires.append(f"{runtime_pkg}~={major}.0") - def save_build_and_package_info(package_name, version_number, cuda_version, qnn_version): sys.path.append(path.join(path.dirname(__file__), "onnxruntime", "python")) diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 7205f1903cdd5..645ac22326fd1 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -442,7 +442,6 @@ def generate_build_tree( cudnn_home, nccl_home, tensorrt_home, - tensorrt_rtx_home, migraphx_home, acl_home, acl_libs, @@ -546,13 +545,11 @@ def generate_build_tree( "-Donnxruntime_ENABLE_MICROSOFT_INTERNAL=" + ("ON" if args.enable_msinternal else "OFF"), "-Donnxruntime_USE_VITISAI=" + ("ON" if args.use_vitisai else "OFF"), "-Donnxruntime_USE_TENSORRT=" + ("ON" if args.use_tensorrt else "OFF"), - "-Donnxruntime_USE_NV=" + ("ON" if args.use_nv_tensorrt_rtx else "OFF"), "-Donnxruntime_USE_TENSORRT_BUILTIN_PARSER=" + ("ON" if args.use_tensorrt_builtin_parser and not args.use_tensorrt_oss_parser else "OFF"), # interface variables are used only for building onnxruntime/onnxruntime_shared.dll but not EPs "-Donnxruntime_USE_TENSORRT_INTERFACE=" + ("ON" if args.enable_generic_interface else "OFF"), "-Donnxruntime_USE_CUDA_INTERFACE=" + ("ON" if args.enable_generic_interface else "OFF"), - "-Donnxruntime_USE_NV_INTERFACE=" + ("ON" if args.enable_generic_interface else "OFF"), "-Donnxruntime_USE_OPENVINO_INTERFACE=" + ("ON" if args.enable_generic_interface else "OFF"), "-Donnxruntime_USE_VITISAI_INTERFACE=" + ("ON" if args.enable_generic_interface else "OFF"), "-Donnxruntime_USE_QNN_INTERFACE=" + ("ON" if args.enable_generic_interface else "OFF"), @@ -845,8 +842,6 @@ def generate_build_tree( if args.use_tensorrt: cmake_args.append("-Donnxruntime_TENSORRT_HOME=" + tensorrt_home) - if args.use_nv_tensorrt_rtx: - cmake_args.append("-Donnxruntime_TENSORRT_RTX_HOME=" + tensorrt_rtx_home) if args.use_cuda: nvcc_threads = number_of_nvcc_threads(args) @@ -1883,8 +1878,6 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): dll_path_list = [] if args.use_tensorrt: dll_path_list.append(os.path.join(args.tensorrt_home, "lib")) - if args.use_nv_tensorrt_rtx: - dll_path_list.append(os.path.join(args.tensorrt_rtx_home, "lib")) dll_path = None if len(dll_path_list) > 0: @@ -2225,10 +2218,6 @@ def build_python_wheel( args.append(f"--qnn_version={qnn_version}") elif use_azure: args.append("--use_azure") - elif wheel_name_suffix == "trt-rtx": - cuda_version = cuda_version or parse_cuda_version_from_json(cuda_home) - if cuda_version: - args.append(f"--cuda_version={cuda_version}") run_subprocess(args, cwd=cwd) @@ -2629,9 +2618,6 @@ def main(): # if using tensorrt, setup tensorrt paths tensorrt_home = "" - tensorrt_rtx_home = "" - if args.use_nv_tensorrt_rtx: - tensorrt_rtx_home = args.tensorrt_rtx_home if args.use_tensorrt: tensorrt_home = setup_tensorrt_vars(args) @@ -2768,7 +2754,6 @@ def main(): cudnn_home, nccl_home, tensorrt_home, - tensorrt_rtx_home, migraphx_home, acl_home, acl_libs, diff --git a/tools/ci_build/build_args.py b/tools/ci_build/build_args.py index 8b97e1e568bf4..073a4f8012b45 100644 --- a/tools/ci_build/build_args.py +++ b/tools/ci_build/build_args.py @@ -736,12 +736,6 @@ def add_execution_provider_args(parser: argparse.ArgumentParser) -> None: ) trt_group.add_argument("--use_tensorrt_oss_parser", action="store_true", help="Use TensorRT OSS ONNX parser.") trt_group.add_argument("--tensorrt_home", help="Path to TensorRT installation directory.") - trt_group.add_argument("--tensorrt_rtx_home", help="Path to TensorRT RTX installation directory.") - - # --- Nv --- - nv_group = parser.add_argument_group("Nv Execution Provider") - nv_group.add_argument("--use_nv_tensorrt_rtx", action="store_true", help="Enable Nv EP.") - # --- DirectML --- dml_group = parser.add_argument_group("DirectML Execution Provider (Windows)") dml_group.add_argument("--use_dml", action="store_true", help="Enable DirectML EP (Windows).") diff --git a/tools/ci_build/gen_def.py b/tools/ci_build/gen_def.py index 15965a1fce193..159538722ed99 100755 --- a/tools/ci_build/gen_def.py +++ b/tools/ci_build/gen_def.py @@ -102,7 +102,6 @@ def parse_arguments(): "tensorrt", "azure", "webgpu", - "nv_tensorrt_rtx", ): file.write(f"#include \n") file.write("void* GetFunctionEntryByName(const char* name){\n")