From d9878a3bc4b3ff1e6cba84ac8866d8c91ddae311 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 28 May 2025 21:17:30 +0530 Subject: [PATCH 01/25] implement GetEPContextNodes() --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 21 +++++++++++++++++-- .../nv_tensorrt_rtx/nv_execution_provider.h | 3 +++ .../nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 7 +++++-- .../nv_tensorrt_rtx/onnx_ctx_model_helper.h | 3 ++- 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index cc9d9f3da1d81..577da9bc43415 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -1785,6 +1785,10 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, if (exclude_ops_set.find(node->OpType()) != exclude_ops_set.end()) { supported_node = false; } + // Exclude contrib ops + if (node->Domain() == kMSDomain) { + supported_node = false; + } if (supported_node) { if (new_subgraph) { @@ -2135,6 +2139,18 @@ static bool IsIOBindingRequired(TRTState* const trt_state, const Ort::KernelCont return require_io_binding; } +const InlinedVector NvExecutionProvider::GetEpContextNodes() const { + InlinedVector ep_context_nodes; + for (auto& model : ep_context_nodes_) { + auto& graph = model->MainGraph(); + for (int i = 0; i < graph.MaxNodeIndex(); i++) { + auto node = graph.GetNode(i); + ep_context_nodes.push_back(node); + } + } + return ep_context_nodes; +} + Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& graph_body_viewer, const Node& fused_node, std::unordered_map& input_map, @@ -2424,8 +2440,9 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr ep_context_embed_mode_, compute_capability_hw_compat, model_path_, - GetLogger())}; - DumpCtxModel(model_proto.get(), ctx_model_path_); + GetLogger(), + ep_context_nodes_, fused_node.Name())}; + // DumpCtxsModel(model_proto.get(), ctx_model_path_); } } diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 83b89a2e9d1fb..49c587983a524 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -283,6 +283,8 @@ class NvExecutionProvider : public IExecutionProvider { bool serialize_refitted_engine, bool detailed_build_log); + const InlinedVector GetEpContextNodes() const override; + private: mutable NvExecutionProviderInfo info_; bool external_stream_ = false; @@ -317,6 +319,7 @@ class NvExecutionProvider : public IExecutionProvider { std::string cache_prefix_; std::string op_types_to_exclude_; int nv_profile_index_ = 0; + std::vector> ep_context_nodes_; // The format is as for TENSORRT_VERSION: (MAJOR * 100 + MINOR) * 100 + PATCH int32_t trt_version_; 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 index 21d964b0c341f..698a3a510376f 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -72,7 +72,8 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const int64_t embed_mode, const std::string compute_capability, const std::string onnx_model_path, - const logging::Logger* logger) { + const logging::Logger* logger, + std::vector>& ep_context_nodes, const std::string& node_name) { auto model_build = graph_viewer.CreateModel(*logger); auto& graph_build = model_build->MainGraph(); @@ -126,7 +127,7 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, node_attributes->emplace(ONNX_MODEL_FILENAME, *attr_3); // Create EP context node - graph_build.AddNode(EPCONTEXT_OP, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); + graph_build.AddNode(node_name, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); ORT_ENFORCE(graph_build.Resolve().IsOK()); // Serialize modelproto to string @@ -137,6 +138,8 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, new_graph_viewer->ToProto(*model_proto->mutable_graph(), true, true); model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + ep_context_nodes.push_back(std::move(model_build)); + return model_proto.release(); } 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 index f0a05c42414e5..c5a56b28775ab 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -34,7 +34,8 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const int64_t embed_mode, const std::string compute_capability, const std::string onnx_model_path, - const logging::Logger* logger); + const logging::Logger* logger, + std::vector>& ep_context_nodes, const std::string& node_name); std::string GetCtxModelPath(const std::string& ep_context_file_path, const std::string& original_model_path); bool IsAbsolutePath(const std::string& path_string); From bde3ce566f8d8ad9ff49c9e055e82282956d5adc Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 28 May 2025 21:23:50 +0530 Subject: [PATCH 02/25] clean up --- .../providers/nv_tensorrt_rtx/nv_execution_provider.cc | 4 ++-- .../providers/nv_tensorrt_rtx/nv_execution_provider.h | 2 +- .../providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 9 +++++---- .../providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h | 2 +- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 577da9bc43415..f14a53348869a 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -2141,7 +2141,7 @@ static bool IsIOBindingRequired(TRTState* const trt_state, const Ort::KernelCont const InlinedVector NvExecutionProvider::GetEpContextNodes() const { InlinedVector ep_context_nodes; - for (auto& model : ep_context_nodes_) { + for (auto& model : ep_context_models_) { auto& graph = model->MainGraph(); for (int i = 0; i < graph.MaxNodeIndex(); i++) { auto node = graph.GetNode(i); @@ -2441,7 +2441,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr compute_capability_hw_compat, model_path_, GetLogger(), - ep_context_nodes_, fused_node.Name())}; + ep_context_models_, fused_node.Name())}; // DumpCtxsModel(model_proto.get(), ctx_model_path_); } } diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 49c587983a524..ae16b6626a138 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -319,7 +319,7 @@ class NvExecutionProvider : public IExecutionProvider { std::string cache_prefix_; std::string op_types_to_exclude_; int nv_profile_index_ = 0; - std::vector> ep_context_nodes_; + std::vector> ep_context_models_; // The format is as for TENSORRT_VERSION: (MAJOR * 100 + MINOR) * 100 + PATCH int32_t trt_version_; 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 index 698a3a510376f..aa108e680c66f 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -73,7 +73,7 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string compute_capability, const std::string onnx_model_path, const logging::Logger* logger, - std::vector>& ep_context_nodes, const std::string& node_name) { + std::vector>& ep_context_models, const std::string& ep_context_node_name) { auto model_build = graph_viewer.CreateModel(*logger); auto& graph_build = model_build->MainGraph(); @@ -127,9 +127,12 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, node_attributes->emplace(ONNX_MODEL_FILENAME, *attr_3); // Create EP context node - graph_build.AddNode(node_name, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); + graph_build.AddNode(ep_context_node_name, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); ORT_ENFORCE(graph_build.Resolve().IsOK()); + // A model with one EP context node is created for a supported subgraph + ep_context_models.push_back(std::move(model_build)); + // Serialize modelproto to string auto new_graph_viewer = graph_build.CreateGraphViewer(); auto& metadata = graph_viewer.GetGraph().GetModel().MetaData(); @@ -138,8 +141,6 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, new_graph_viewer->ToProto(*model_proto->mutable_graph(), true, true); model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); - ep_context_nodes.push_back(std::move(model_build)); - return model_proto.release(); } 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 index c5a56b28775ab..d650e0bf10827 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -35,7 +35,7 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string compute_capability, const std::string onnx_model_path, const logging::Logger* logger, - std::vector>& ep_context_nodes, const std::string& node_name); + std::vector>& ep_context_models, const std::string& ep_context_node_name); std::string GetCtxModelPath(const std::string& ep_context_file_path, const std::string& original_model_path); bool IsAbsolutePath(const std::string& path_string); From 7165dfe3c443e15eaa513bc95add5761277873e8 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Thu, 29 May 2025 15:20:51 +0530 Subject: [PATCH 03/25] rebase to latest --- .../core/providers/nv_tensorrt_rtx/nv_execution_provider.cc | 4 ---- 1 file changed, 4 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index f14a53348869a..85c03253b9074 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -1785,10 +1785,6 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, if (exclude_ops_set.find(node->OpType()) != exclude_ops_set.end()) { supported_node = false; } - // Exclude contrib ops - if (node->Domain() == kMSDomain) { - supported_node = false; - } if (supported_node) { if (new_subgraph) { From 7b1f5bc9a64b291a259086e913afe3f7d4c71337 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 4 Jun 2025 22:31:06 +0530 Subject: [PATCH 04/25] remove ctx model to just add node --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 20 +++++++++--------- .../nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 21 +++++-------------- .../nv_tensorrt_rtx/onnx_ctx_model_helper.h | 4 ++-- 3 files changed, 17 insertions(+), 28 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 85c03253b9074..4038f7a46a70f 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -2429,16 +2429,16 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr ep_cache_context_attr_ = std::filesystem::path(engine_cache_relative_path_to_context_model_dir).append(cache_file_name.string()).string(); } std::string compute_capability_hw_compat = compute_capability_ + "+"; - std::unique_ptr model_proto{CreateCtxModel(graph_body_viewer, - ep_cache_context_attr_, - reinterpret_cast(serialized_engine->data()), - serialized_engine->size(), - ep_context_embed_mode_, - compute_capability_hw_compat, - model_path_, - GetLogger(), - ep_context_models_, fused_node.Name())}; - // DumpCtxsModel(model_proto.get(), ctx_model_path_); + + ep_context_models_.push_back(CreateCtxNode(graph_body_viewer, + ep_cache_context_attr_, + reinterpret_cast(serialized_engine->data()), + serialized_engine->size(), + ep_context_embed_mode_, + compute_capability_hw_compat, + model_path_, + GetLogger(), + fused_node.Name())); } } 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 index aa108e680c66f..e7d29e58514fa 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -63,9 +63,9 @@ void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, } /* - * Create "EP context node" model where engine information is embedded + * Create EP context node where engine information is embedded */ -ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, +std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewer, const std::string engine_cache_path, char* engine_data, size_t size, @@ -73,8 +73,8 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string compute_capability, const std::string onnx_model_path, const logging::Logger* logger, - std::vector>& ep_context_models, const std::string& ep_context_node_name) { - auto model_build = graph_viewer.CreateModel(*logger); + const std::string& ep_context_node_name) { + auto model_build = Model::Create("nv_trt_rtx_ep_context_model", false, *logger); auto& graph_build = model_build->MainGraph(); // Get graph inputs and outputs @@ -130,18 +130,7 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, graph_build.AddNode(ep_context_node_name, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); ORT_ENFORCE(graph_build.Resolve().IsOK()); - // A model with one EP context node is created for a supported subgraph - ep_context_models.push_back(std::move(model_build)); - - // Serialize modelproto to string - auto new_graph_viewer = graph_build.CreateGraphViewer(); - auto& metadata = graph_viewer.GetGraph().GetModel().MetaData(); - auto model = new_graph_viewer->CreateModel(*logger, metadata); - auto model_proto = model->ToProto(); - new_graph_viewer->ToProto(*model_proto->mutable_graph(), true, true); - model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); - - return model_proto.release(); + return model_build; } /* 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 index d650e0bf10827..fc6c00cb75731 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -27,7 +27,7 @@ static const std::string EPCONTEXT_WARNING = bool GraphHasCtxNode(const GraphViewer& graph_viewer); const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer); std::filesystem::path GetPathOrParentPathOfCtxModel(const std::string& ep_context_file_path); -ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, +std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewer, const std::string engine_cache_path, char* engine_data, size_t size, @@ -35,7 +35,7 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string compute_capability, const std::string onnx_model_path, const logging::Logger* logger, - std::vector>& ep_context_models, const std::string& ep_context_node_name); + const std::string& ep_context_node_name); std::string GetCtxModelPath(const std::string& ep_context_file_path, const std::string& original_model_path); bool IsAbsolutePath(const std::string& path_string); From 03c42fd32bdfeb63a4f7c0aef712ee9e4d4bd5d2 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Mon, 9 Jun 2025 18:12:40 +0530 Subject: [PATCH 05/25] update GetCapabilities for multiple EP Context Nodes --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 34 ++++++++++++------- .../nv_tensorrt_rtx/onnx_ctx_model_helper.h | 4 +-- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 4038f7a46a70f..f835f0956713c 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -1708,21 +1708,32 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, #endif model_path_[sizeof(model_path_) - 1] = '\0'; - // If the model consists of only a single "EPContext" contrib op, 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 the ComputeCapability here. - if (graph.NumberOfNodes() == 1 && GraphHasCtxNode(graph)) { - SubGraph_t supported_node_vector = {{0}, true}; - std::unique_ptr sub_graph = GetSubGraph(supported_node_vector, graph, TRTGenerateId(graph, std::to_string(trt_version_), std::to_string(cuda_version_)), 0); - result.push_back(ComputeCapability::Create(std::move(sub_graph))); - 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_)); - // Get supported node list from TensorRT parser - const int number_of_ort_nodes = graph.NumberOfNodes(); + // 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. + if (GraphHasCtxNode(graph)) { + int subgraph_idx = 0; + for (size_t i = 0; i < static_cast(number_of_ort_nodes); i++) { + const auto& node = graph.GetNode(node_index[i]); + const bool is_context_node = node && !node->OpType().empty() && node->OpType() == "EPContext"; + if (is_context_node) { + SubGraph_t supported_node_vector(std::make_pair(std::vector{i}, 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 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); @@ -1741,7 +1752,6 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, auto exclude_ops_set = get_exclude_ops_set(op_types_to_exclude_); SubGraphCollection_t parser_nodes_vector, supported_nodes_vector; - const std::vector& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); bool new_subgraph = true; /* Iterate all the nodes and exclude the node if: 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 index fc6c00cb75731..21ac474f5dede 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -20,9 +20,7 @@ static const std::string COMPUTE_CAPABILITY = "hardware_architecture"; static const std::string ONNX_MODEL_FILENAME = "onnx_model_filename"; static const std::string EPCONTEXT_OP_DOMAIN = "com.microsoft"; static const std::string EPCONTEXT_WARNING = - "It's suggested to set the ORT graph optimization level to 0 and \ - make \"embed_mode\" to 0 (\"ep_cache_context\" is the cache path)\ - for the best model loading time"; + "It's suggested to set the ORT graph optimization level to 0 for the best performance"; bool GraphHasCtxNode(const GraphViewer& graph_viewer); const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer); From 43ac5d51dcaf5a9e15f77ed440dcddc32a845898 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 11 Jun 2025 10:39:37 +0530 Subject: [PATCH 06/25] fix lint --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 16 ++++++++-------- .../nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 16 ++++++++-------- .../nv_tensorrt_rtx/onnx_ctx_model_helper.h | 16 ++++++++-------- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index f835f0956713c..185a30859f16c 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -2441,14 +2441,14 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr std::string compute_capability_hw_compat = compute_capability_ + "+"; ep_context_models_.push_back(CreateCtxNode(graph_body_viewer, - ep_cache_context_attr_, - reinterpret_cast(serialized_engine->data()), - serialized_engine->size(), - ep_context_embed_mode_, - compute_capability_hw_compat, - model_path_, - GetLogger(), - fused_node.Name())); + ep_cache_context_attr_, + reinterpret_cast(serialized_engine->data()), + serialized_engine->size(), + ep_context_embed_mode_, + compute_capability_hw_compat, + model_path_, + GetLogger(), + fused_node.Name())); } } 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 index e7d29e58514fa..001742bcc2266 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -66,14 +66,14 @@ void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, * Create EP context node where engine information is embedded */ std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewer, - 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 logging::Logger* logger, - const std::string& ep_context_node_name) { + 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 logging::Logger* logger, + const std::string& ep_context_node_name) { auto model_build = Model::Create("nv_trt_rtx_ep_context_model", false, *logger); auto& graph_build = model_build->MainGraph(); 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 index 21ac474f5dede..3c18fd2b59789 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -26,14 +26,14 @@ bool GraphHasCtxNode(const GraphViewer& graph_viewer); const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer); std::filesystem::path GetPathOrParentPathOfCtxModel(const std::string& ep_context_file_path); std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewer, - 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 logging::Logger* logger, - const std::string& ep_context_node_name); + 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 logging::Logger* logger, + const std::string& ep_context_node_name); std::string GetCtxModelPath(const std::string& ep_context_file_path, const std::string& original_model_path); bool IsAbsolutePath(const std::string& path_string); From 313f4ceaef6f28797d1875a3b0a6096c8bf993bf Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 18 Jun 2025 08:09:03 +0530 Subject: [PATCH 07/25] add support for TRT external weights API --- cmake/deps.txt | 4 +- .../nv_tensorrt_rtx/nv_provider_options.h | 2 + .../nv_tensorrt_rtx/nv_execution_provider.cc | 257 +++++++++++++++++- .../nv_tensorrt_rtx/nv_execution_provider.h | 7 +- .../nv_execution_provider_info.cc | 13 + .../nv_execution_provider_info.h | 2 + .../nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 4 + .../nv_tensorrt_rtx/onnx_ctx_model_helper.h | 6 + 8 files changed, 278 insertions(+), 17 deletions(-) diff --git a/cmake/deps.txt b/cmake/deps.txt index 7bac2a0fbced9..1bd50fb5749c7 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -35,8 +35,8 @@ microsoft_wil;https://github.com/microsoft/wil/archive/refs/tags/v1.0.230629.1.z mimalloc;https://github.com/microsoft/mimalloc/archive/refs/tags/v2.1.1.zip;d5ee7d34223d0567892db5179849939c8769dc41 mp11;https://github.com/boostorg/mp11/archive/refs/tags/boost-1.82.0.zip;9bc9e01dffb64d9e0773b2e44d2f22c51aace063 onnx;https://github.com/onnx/onnx/archive/refs/tags/v1.18.0.zip;f156d032a3af91b66d554e11158b33ca77bbb1f2 -# Use the latest commit of 10.9-GA -onnx_tensorrt;https://github.com/onnx/onnx-tensorrt/archive/d5dce67db7c2e64b07e055571f5ec06f7f254de2.zip;01114d3b67650857281fa50faa2e412130a63b69 +# Side branch +onnx_tensorrt;https://github.com/gedoensmax/onnx-tensorrt/archive/1ff5201c5eb209a4a58330b9e8cae97f153c93a4.zip;60af1687b0e03c2c920197f4e059cb666aefdd9c protobuf;https://github.com/protocolbuffers/protobuf/archive/refs/tags/v21.12.zip;7cf2733949036c7d52fda017badcab093fe73bfa protoc_win64;https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-win64.zip;b4521f7ada5b260380f94c4bd7f1b7684c76969a protoc_win32;https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-win32.zip;3688010318192c46ce73213cdfb6b3e5656da874 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 index 11cc6f131dab3..e4de1d3a6d647 100644 --- a/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h +++ b/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h @@ -35,6 +35,8 @@ constexpr const char* kCudaGraphEnable = "nv_cuda_graph_enable"; constexpr const char* kONNXBytestream = "nv_onnx_bytestream"; constexpr const char* kONNXBytestreamSize = "nv_onnx_bytestream_size"; constexpr const char* kMultiProfileEnable = "nv_multi_profile_enable"; +constexpr const char* kExternalDataBytestream = "nv_external_data_bytestream"; +constexpr const char* kExternalDataBytestreamSize = "nv_external_data_bytestream_size"; } // namespace provider_option_names namespace run_option_names { diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 185a30859f16c..c100d4754f0ce 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -866,6 +866,14 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) "When providing either 'trt_onnx_bytestream_size' or " "'trt_onnx_bytestream' both have to be provided")); } + 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 'trt_external_data_bytestream_size' or " + "'trt_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; @@ -1029,6 +1037,7 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) << ", 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_op_types_to_exclude: " << op_types_to_exclude_; } @@ -1502,7 +1511,34 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // 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*/); + // graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/); + + // Set export initializers to false so that we can succesfully serialize. + std::vector names; + std::vector bytes; + std::vector sizes; + + auto allInitializers = graph_viewer->GetAllInitializedTensors(); + + for (auto entry : allInitializers) { + auto name = entry.first; + auto* tp = entry.second; + + // std::cout << "ORT: Saving initializers in mem: " << tp->name() << ", has raw data? " << tp->has_raw_data() << std::endl; + + // TODO: Handle non-raw-data? + if (tp->has_raw_data()) { + names.push_back(tp->name().c_str()); + bytes.push_back(tp->raw_data().c_str()); + sizes.push_back(tp->raw_data().size()); + } + // else { + // std::cout << "ORT: Tensor has no raw data: " << tp->name() << std::endl; + // } + } + graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, false); + + model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); std::string string_buf; @@ -1525,7 +1561,19 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t { auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); - auto is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); + // auto is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); + + bool loadSuccess = trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); + + bool loadInit = true; + for (int i=0; iloadInitializer(names[i], bytes[i], sizes[i]); + } + if (!(loadSuccess && loadInit)) { + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TRT Parser load failed")); + } + // auto is_model_supported = trt_parser->supportsModelV3(); + bool is_model_supported = trt_parser->parseModelProto(); // Note: Calling getNbSubgraphs or getSubgraphNodes before calling supportsModelV2 results in undefined behavior. auto num_subgraphs = trt_parser->getNbSubgraphs(); @@ -1946,9 +1994,12 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, bool path_check, 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 serialize_refitted_engine, - bool detailed_build_log) { + bool detailed_build_log, + const GraphViewer* graph_body_viewer) { bool refit_from_file = onnx_model_bytestream == nullptr && onnx_model_bytestream_size == 0; std::filesystem::path onnx_model_path{onnx_model_folder_path}; if (refit_from_file) { @@ -1992,18 +2043,157 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Nv EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } + if (refitter->refitCudaEngine()) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + } } else { + int required_weights = refitter->getAllWeights(0, nullptr); + std::vector refit_names(required_weights); + refitter->getAllWeights(required_weights, refit_names.data()); + LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refitting from byte array"; - if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { + // if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Weights required for a full refit " << required_weights; + + // predeclare variables here to keep them alive for the refitter + std::vector names; + names.reserve(required_weights); + std::vector bytes; + bytes.reserve(required_weights); + std::vector bytes_copy; + bytes_copy.reserve(required_weights); + std::vector sizes; + sizes.reserve(required_weights); + auto onnx_model = ModelProto::Create(); + TensorProtos* allInitializers_byte_stream; + const InitializedTensorSet* allInitializers_graph_body; + + // load graph structure + bool refitloadSuccess = parser_refitter->loadModelProto(onnx_model_bytestream, onnx_model_bytestream_size, nullptr); + if (!refitloadSuccess) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "Nv EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); } - } - if (refitter->refitCudaEngine()) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully refitted the weight-stripped engine."; - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + // } + // if (refitter->refitCudaEngine()) { + // LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully refitted the weight-stripped engine."; + // } else { + // return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + // "Nv EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + + // conditional branch loading additional initializers if needed + if (onnx_external_data_bytestream_size) { + // This code path is leveraged when loading a context file as byte array + // and providing the corresponding ONNX as byte array through provider options + + 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."); + } + + + + auto const& graph = onnx_model->mutable_graph(); + allInitializers_byte_stream = graph->mutable_initializer(); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializers that were found " << allInitializers_byte_stream->size(); + struct DataInfo { + int64_t offset; + size_t length; + }; + 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(); + bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); + if (weight_is_refittable) { + if (proto.has_data_location()) { + if (proto.data_location() == TensorProto_DataLocation_EXTERNAL) { + DataInfo external_data_info = {}; + 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()) { + external_data_info.offset = std::stoll(*current_value); + } else if (*current_key == kLength && !current_value->empty()) { + external_data_info.length = std::stoul(*current_value); + } + } + names.push_back(proto.name().c_str()); + bytes.push_back(static_cast(onnx_external_data_bytestream) + external_data_info.offset); + sizes.push_back(external_data_info.length); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[TensorRT EP] Proto: " + proto_name + " has default as data locationn which is not supported"); + } + } else { + if (!proto.has_raw_data()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[TensorRT EP] Proto: " + proto_name + " has no raw data"); + } + auto& raw_data = proto.raw_data(); + + names.push_back(proto.name().c_str()); + bytes.push_back(raw_data.c_str()); + sizes.push_back(raw_data.size()); + } + } else { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " cannot be refitted"; + } + } + } else if (graph_body_viewer) { + // This path will only be used if + // 1. An ONNX was provided as byte array including initializers + // 2. An ONNX was provided as byte array and it's initializers were provided using AddExternalInitializers* API + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting using initializers of the current graph in memory"; + + allInitializers_graph_body = &graph_body_viewer->GetAllInitializedTensors(); + + for (auto& entry : *allInitializers_graph_body) { + auto* tp = entry.second; + auto& proto_name = tp->name(); + // TODO: Handle non-raw-data? + bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); + if (tp->has_raw_data() && weight_is_refittable) { + names.push_back(proto_name.c_str()); + bytes.push_back(tp->raw_data().c_str()); + sizes.push_back(tp->raw_data().size()); + } + } + } else { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] No external initializers are present"; + } + // provide dedicated initializers if ONNX serialization is not complete + if (!names.empty()) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Number of initializers submitted to refitter " << names.size(); + bool refloadInit = true; + for(int i=0; iloadInitializer(names[i], bytes[i], sizes[i]); + } + + // bool refloadInit = parser_refitter->loadInitializer(names.data(), bytes.data(), sizes.data(), names.size()); + if (!refloadInit) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + } + } + + bool refparseModelProto = parser_refitter->refitModelProto(); + if (!refparseModelProto) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + } + if (refitter->refitCudaEngine()) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + } } // serialize the refitted engine to disk @@ -2165,12 +2355,31 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // Reconstruct graph proto from fused node's function body auto model = graph_body_viewer.CreateModel(*GetLogger()); auto model_proto = model->ToProto(); + // Set export initializers to false so that we can succesfully serialize. + + std::vector names; + std::vector bytes; + std::vector sizes; + + auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); + + for (auto entry : allInitializers) { + auto name = entry.first; + auto* tp = entry.second; + // TODO: Handle non-raw-data? + if (tp->has_raw_data()) { + names.push_back(tp->name().c_str()); + bytes.push_back(tp->raw_data().c_str()); + sizes.push_back(tp->raw_data().size()); + } + } + // 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*/); + graph_body_viewer.ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, false); model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); std::string string_buf; model_proto->SerializeToString(string_buf); @@ -2187,7 +2396,21 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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)); - trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); + // trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); + bool loadSuccess = trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); + // bool loadInit = trt_parser->loadInitializer(names.data(), bytes.data(), sizes.data(), names.size()); + bool loadInit = true; + for (int i=0; iloadInitializer(names[i], bytes[i], sizes[i]); + } + if (!(loadSuccess && loadInit)) { + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TRT Parser load failed")); + } + bool parseModelProto = trt_parser->parseModelProto(); + if (!parseModelProto) { + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TRT Parser failed")); + } + if (max_workspace_size_ > 0) { trt_config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_size_); } @@ -2462,9 +2685,12 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr false /* path check for security */, onnx, onnx_size, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, trt_engine.get(), false /* serialize refitted engine to disk */, - detailed_build_log_); + detailed_build_log_, + &graph_body_viewer); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); } @@ -2517,6 +2743,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // Create function state // TODO: remove default capture NodeComputeInfo compute_info; + auto* graph_body_viewer_ptr = &graph_body_viewer; 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(), @@ -2527,7 +2754,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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_, cache_suffix}; + auxiliary_streams_, cuda_graph_enable_, is_dynamic_shape_context, cache_prefix_, cache_suffix, graph_body_viewer_ptr}; *state = p.release(); return 0; }; @@ -2785,6 +3012,8 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra 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); if (status != Status::OK()) { diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index ae16b6626a138..079ea8e17b1c5 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -279,9 +279,12 @@ class NvExecutionProvider : public IExecutionProvider { bool path_check, 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 serialize_refitted_engine, - bool detailed_build_log); + bool detailed_build_log, + const GraphViewer* graph_body_viewer = nullptr); const InlinedVector GetEpContextNodes() const override; @@ -301,6 +304,8 @@ class NvExecutionProvider : public IExecutionProvider { std::string onnx_model_folder_path_; const void* onnx_model_bytestream_; size_t onnx_model_bytestream_size_; + 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_; 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 index f90bf24ef4975..411a5d6447da2 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc @@ -17,6 +17,7 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi NvExecutionProviderInfo info{}; void* user_compute_stream = nullptr; void* onnx_bytestream = nullptr; + void* external_data_bytestream = nullptr; ORT_THROW_IF_ERROR( ProviderOptionsParser{} .AddValueParser( @@ -58,11 +59,21 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi return Status::OK(); }) .AddAssignmentToReference(nv::provider_option_names::kONNXBytestreamSize, info.onnx_bytestream_size) + .AddValueParser( + nv::provider_option_names::kExternalDataBytestream, + [&external_data_bytestream](const std::string& value_str) -> Status { + size_t address; + ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, address)); + external_data_bytestream = reinterpret_cast(address); + return Status::OK(); + }) + .AddAssignmentToReference(nv::provider_option_names::kExternalDataBytestreamSize, info.external_data_bytestream_size) .Parse(options)); // add new provider option here. info.user_compute_stream = user_compute_stream; info.has_user_compute_stream = (user_compute_stream != nullptr); 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 @@ -112,6 +123,8 @@ ProviderOptions NvExecutionProviderInfo::ToProviderOptions(const NvExecutionProv {nv::provider_option_names::kCudaGraphEnable, MakeStringWithClassicLocale(info.cuda_graph_enable)}, {nv::provider_option_names::kONNXBytestream, MakeStringWithClassicLocale(info.onnx_bytestream)}, {nv::provider_option_names::kONNXBytestreamSize, MakeStringWithClassicLocale(info.onnx_bytestream_size)}, + {nv::provider_option_names::kExternalDataBytestream, MakeStringWithClassicLocale(info.external_data_bytestream)}, + {nv::provider_option_names::kExternalDataBytestreamSize, MakeStringWithClassicLocale(info.external_data_bytestream_size)}, }; return options; } 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 index 4d6c6fe116076..89f92b7723db1 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h @@ -31,6 +31,8 @@ struct NvExecutionProviderInfo { std::string onnx_model_folder_path{""}; const void* onnx_bytestream{nullptr}; size_t onnx_bytestream_size{0}; + 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}; 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 index 001742bcc2266..e92576590764a 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -294,6 +294,8 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph make_secure_path_checks, onnx_model_bytestream_, onnx_model_bytestream_size_, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, (*trt_engine_).get(), false /* serialize refitted engine to disk */, detailed_build_log_); @@ -363,6 +365,8 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph make_secure_path_checks, onnx_model_bytestream_, onnx_model_bytestream_size_, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, (*trt_engine_).get(), true /* serialize refitted engine to disk */, detailed_build_log_); 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 index 3c18fd2b59789..dc9d9c7b5ae39 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -54,6 +54,8 @@ class TensorRTCacheModelHandler { 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), @@ -63,6 +65,8 @@ class TensorRTCacheModelHandler { 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); @@ -80,6 +84,8 @@ class TensorRTCacheModelHandler { 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 From a7eadab6230b55915ccb54bf1f164a0896cdea4f Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Mon, 14 Jul 2025 21:52:19 +0000 Subject: [PATCH 08/25] add new changes --- cmake/deps.txt | 4 +- .../nv_tensorrt_rtx/nv_execution_provider.cc | 259 ++++++++---------- .../nv_tensorrt_rtx/nv_execution_provider.h | 9 + 3 files changed, 125 insertions(+), 147 deletions(-) diff --git a/cmake/deps.txt b/cmake/deps.txt index 1bd50fb5749c7..7bac2a0fbced9 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -35,8 +35,8 @@ microsoft_wil;https://github.com/microsoft/wil/archive/refs/tags/v1.0.230629.1.z mimalloc;https://github.com/microsoft/mimalloc/archive/refs/tags/v2.1.1.zip;d5ee7d34223d0567892db5179849939c8769dc41 mp11;https://github.com/boostorg/mp11/archive/refs/tags/boost-1.82.0.zip;9bc9e01dffb64d9e0773b2e44d2f22c51aace063 onnx;https://github.com/onnx/onnx/archive/refs/tags/v1.18.0.zip;f156d032a3af91b66d554e11158b33ca77bbb1f2 -# Side branch -onnx_tensorrt;https://github.com/gedoensmax/onnx-tensorrt/archive/1ff5201c5eb209a4a58330b9e8cae97f153c93a4.zip;60af1687b0e03c2c920197f4e059cb666aefdd9c +# Use the latest commit of 10.9-GA +onnx_tensorrt;https://github.com/onnx/onnx-tensorrt/archive/d5dce67db7c2e64b07e055571f5ec06f7f254de2.zip;01114d3b67650857281fa50faa2e412130a63b69 protobuf;https://github.com/protocolbuffers/protobuf/archive/refs/tags/v21.12.zip;7cf2733949036c7d52fda017badcab093fe73bfa protoc_win64;https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-win64.zip;b4521f7ada5b260380f94c4bd7f1b7684c76969a protoc_win32;https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-win32.zip;3688010318192c46ce73213cdfb6b3e5656da874 diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index c100d4754f0ce..e04f342c22bd1 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -871,8 +871,8 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) 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 'trt_external_data_bytestream_size' or " - "'trt_external_data_bytestream' both have to be provided")); + "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; @@ -1511,32 +1511,19 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // 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*/); - // Set export initializers to false so that we can succesfully serialize. - std::vector names; - std::vector bytes; - std::vector sizes; + // get initializer data + std::vector userWeights; auto allInitializers = graph_viewer->GetAllInitializedTensors(); - - for (auto entry : allInitializers) { - auto name = entry.first; - auto* tp = entry.second; - - // std::cout << "ORT: Saving initializers in mem: " << tp->name() << ", has raw data? " << tp->has_raw_data() << std::endl; - - // TODO: Handle non-raw-data? - if (tp->has_raw_data()) { - names.push_back(tp->name().c_str()); - bytes.push_back(tp->raw_data().c_str()); - sizes.push_back(tp->raw_data().size()); - } - // else { - // std::cout << "ORT: Tensor has no raw data: " << tp->name() << std::endl; - // } + for (auto entry : allInitializers){ + auto* tp = entry.second; + if (tp->has_raw_data()){ + userWeights.push_back( + TensorrtUserWeights{tp->name(), tp->raw_data(), (int64_t)tp->raw_data().size()}); + } } - graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, false); + graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, false /*include raw initializers*/); model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); @@ -1557,23 +1544,18 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t auto network_flags = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kSTRONGLY_TYPED); 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)); - // auto is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); - - bool loadSuccess = trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); + trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - bool loadInit = true; - for (int i=0; iloadInitializer(names[i], bytes[i], sizes[i]); - } - if (!(loadSuccess && loadInit)) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TRT Parser load failed")); - } - // auto is_model_supported = trt_parser->supportsModelV3(); - bool is_model_supported = trt_parser->parseModelProto(); + for (auto const& userWeight : userWeights){ + trt_parser->loadInitializer(userWeight.name.c_str(), static_cast(userWeight.data.c_str()), userWeight.size); + } + is_model_supported = trt_parser->parseModelProto(); // Note: Calling getNbSubgraphs or getSubgraphNodes before calling supportsModelV2 results in undefined behavior. auto num_subgraphs = trt_parser->getNbSubgraphs(); @@ -2001,6 +1983,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, bool detailed_build_log, const GraphViewer* graph_body_viewer) { bool refit_from_file = onnx_model_bytestream == nullptr && onnx_model_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()) { @@ -2023,7 +2006,6 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, "The ONNX model path has '..'. For security purpose, it's not " "allowed to point outside the directory."); } - 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() + @@ -2031,80 +2013,67 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, } } } - // 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)); - 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())) { + + bool refit_with_external_data = onnx_external_data_bytestream != nullptr && onnx_external_data_bytestream_size != 0; + + // New refit APIs + if (refit_with_external_data || graph_body_viewer) { + if (refit_from_file) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "TensorRT EP's refit with external data must be called with a valid ONNX model bytestream"); } - if (refitter->refitCudaEngine()) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; - } else { + + if (!parser_refitter->loadModelProto(onnx_model_bytestream, onnx_model_bytestream_size, nullptr)){ return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "TensorRT EP's IParserRefitter could not load model from provided onnx_model_bytestream"); } - } else { + + // Extract weight information from the Refitter int required_weights = refitter->getAllWeights(0, nullptr); std::vector refit_names(required_weights); refitter->getAllWeights(required_weights, refit_names.data()); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refitting from byte array"; - // if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Weights required for a full refit " << required_weights; - - // predeclare variables here to keep them alive for the refitter + // Vectors to keep track of data pointers std::vector names; names.reserve(required_weights); std::vector bytes; bytes.reserve(required_weights); - std::vector bytes_copy; - bytes_copy.reserve(required_weights); std::vector sizes; sizes.reserve(required_weights); - auto onnx_model = ModelProto::Create(); - TensorProtos* allInitializers_byte_stream; - const InitializedTensorSet* allInitializers_graph_body; - // load graph structure - bool refitloadSuccess = parser_refitter->loadModelProto(onnx_model_bytestream, onnx_model_bytestream_size, nullptr); - if (!refitloadSuccess) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + /* + TODO: Instead of passing in a single pointer and reconstructing the entire ONNX model, is it better to pre-partition the weights? + This function will now have parameters void ** data, const char ** names, int64_t * sizes, int64_t num_weights + for (int i = 0; i < num_weights; i++) + { + parser_refitter->loadInitializer(data[i], names[i], sizes[i]); } - // } - // if (refitter->refitCudaEngine()) { - // LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully refitted the weight-stripped engine."; - // } else { - // return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - // "Nv EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + */ - // conditional branch loading additional initializers if needed - if (onnx_external_data_bytestream_size) { - // This code path is leveraged when loading a context file as byte array - // and providing the corresponding ONNX as byte array through provider options + if (refit_with_external_data) + { + 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); + 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."); + "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) << "[TensorRT EP] Initializers that were found " << allInitializers_byte_stream->size(); - struct DataInfo { - int64_t offset; - size_t length; - }; + + // Loop through all initializers 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(); @@ -2112,52 +2081,53 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, if (weight_is_refittable) { if (proto.has_data_location()) { if (proto.data_location() == TensorProto_DataLocation_EXTERNAL) { - DataInfo external_data_info = {}; + // 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()) { - external_data_info.offset = std::stoll(*current_value); + offset = std::stoll(*current_value); } else if (*current_key == kLength && !current_value->empty()) { - external_data_info.length = std::stoul(*current_value); + length = std::stoul(*current_value); } } names.push_back(proto.name().c_str()); - bytes.push_back(static_cast(onnx_external_data_bytestream) + external_data_info.offset); - sizes.push_back(external_data_info.length); + bytes.push_back(static_cast(onnx_external_data_bytestream) + offset); + sizes.push_back(length); } else { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[TensorRT EP] Proto: " + proto_name + " has default as data locationn which is not supported"); + "[TensorRT EP] Proto: " + proto_name + " has default as data location which is not supported"); } } else { if (!proto.has_raw_data()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[TensorRT EP] Proto: " + proto_name + " has no raw data"); + "[TensorRT EP] Proto: " + proto_name + " has no raw data"); } auto& raw_data = proto.raw_data(); - names.push_back(proto.name().c_str()); bytes.push_back(raw_data.c_str()); sizes.push_back(raw_data.size()); } } else { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " cannot be refitted"; + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " was not marked as refittable"; } } - } else if (graph_body_viewer) { + } + else { // graph_body_viewer path. // This path will only be used if // 1. An ONNX was provided as byte array including initializers // 2. An ONNX was provided as byte array and it's initializers were provided using AddExternalInitializers* API LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting using initializers of the current graph in memory"; - allInitializers_graph_body = &graph_body_viewer->GetAllInitializedTensors(); + auto allInitializers_graph_body = &graph_body_viewer->GetAllInitializedTensors(); for (auto& entry : *allInitializers_graph_body) { auto* tp = entry.second; auto& proto_name = tp->name(); - // TODO: Handle non-raw-data? bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); if (tp->has_raw_data() && weight_is_refittable) { names.push_back(proto_name.c_str()); @@ -2165,35 +2135,50 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, sizes.push_back(tp->raw_data().size()); } } - } else { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] No external initializers are present"; } - // provide dedicated initializers if ONNX serialization is not complete + + // Load extracted initializers into the parser if (!names.empty()) { LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Number of initializers submitted to refitter " << names.size(); - bool refloadInit = true; - for(int i=0; iloadInitializer(names[i], bytes[i], sizes[i]); + for (size_t i = 0; i < names.size(); i++) + { + bool refloadInit = parser_refitter->loadInitializer(names[i], bytes[i], sizes[i]); + if (!refloadInit) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + } } + } + // Perform refit. + if(!parser_refitter->refitModelProto()) + { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); + } + refit_complete = true; - // bool refloadInit = parser_refitter->loadInitializer(names.data(), bytes.data(), sizes.data(), names.size()); - if (!refloadInit) { + } + // If new refit flow was not completed, then fallback to refit_from_file. + if (!refit_complete){ + if (refit_from_file) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT 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, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + } + }else { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from byte array"; + if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); } } - - bool refparseModelProto = parser_refitter->refitModelProto(); - if (!refparseModelProto) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); - } - if (refitter->refitCudaEngine()) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); - } + } + if (refitter->refitCudaEngine()) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } // serialize the refitted engine to disk @@ -2355,26 +2340,20 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // Reconstruct graph proto from fused node's function body auto model = graph_body_viewer.CreateModel(*GetLogger()); auto model_proto = model->ToProto(); - // Set export initializers to false so that we can succesfully serialize. - std::vector names; - std::vector bytes; - std::vector sizes; + // exclude weights if external + auto userWeights = std::make_unique>(); auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); - - for (auto entry : allInitializers) { - auto name = entry.first; - auto* tp = entry.second; - // TODO: Handle non-raw-data? - if (tp->has_raw_data()) { - names.push_back(tp->name().c_str()); - bytes.push_back(tp->raw_data().c_str()); - sizes.push_back(tp->raw_data().size()); - } + for (auto entry : allInitializers){ + auto name = entry.first; + auto* tp = entry.second; + if (tp->has_raw_data()){ + userWeights->push_back( + TensorrtUserWeights{tp->name(), tp->raw_data(), (int64_t)tp->raw_data().size()}); + } } - // 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 @@ -2396,20 +2375,11 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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)); - // trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); - bool loadSuccess = trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - // bool loadInit = trt_parser->loadInitializer(names.data(), bytes.data(), sizes.data(), names.size()); - bool loadInit = true; - for (int i=0; iloadInitializer(names[i], bytes[i], sizes[i]); - } - if (!(loadSuccess && loadInit)) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TRT Parser load failed")); - } - bool parseModelProto = trt_parser->parseModelProto(); - if (!parseModelProto) { - ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TRT Parser failed")); + trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); + for (auto const& userWeight : *userWeights){ + trt_parser->loadInitializer(userWeight.name.c_str(), static_cast(userWeight.data.c_str()), userWeight.size); } + trt_parser->parseModelProto(); if (max_workspace_size_ > 0) { trt_config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_size_); @@ -2689,8 +2659,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr onnx_external_data_bytestream_size_, trt_engine.get(), false /* serialize refitted engine to disk */, - detailed_build_log_, - &graph_body_viewer); + detailed_build_log_); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); } @@ -2734,6 +2703,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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)); + weights_.emplace(fused_node.Name(), std::move(userWeights)); 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); @@ -2743,7 +2713,6 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // Create function state // TODO: remove default capture NodeComputeInfo compute_info; - auto* graph_body_viewer_ptr = &graph_body_viewer; 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(), @@ -2754,7 +2723,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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_, cache_suffix, graph_body_viewer_ptr}; + auxiliary_streams_, cuda_graph_enable_, is_dynamic_shape_context, cache_prefix_, cache_suffix, &weights_[context->node_name]}; *state = p.release(); return 0; }; diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 079ea8e17b1c5..b58dc33bef918 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -153,6 +153,13 @@ struct TensorParams { } }; +// Struct to hold user weights when ModelProtos are serialized with external data +struct TensorrtUserWeights { + std::string name{}; + std::string data{}; + int64_t size{}; +}; + // Information to construct kernel function state. struct TensorrtFuncState { AllocateFunc test_allocate_func = nullptr; @@ -190,6 +197,7 @@ struct TensorrtFuncState { bool skip_io_binding_allowed = false; // Indicates if input/output binding can be skipped IAllocatorUniquePtr context_memory = nullptr; size_t context_memory_size = 0; + std::unique_ptr> *userWeights = nullptr; }; // Minimum information to construct kernel function state for direct engine load code path @@ -364,6 +372,7 @@ class NvExecutionProvider : public IExecutionProvider { 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_; + std::unordered_map>> weights_; // User provided weights // for external stream, we need to create its cudnn/cublass handle before cuda EP enable cuda graph capture cudnnHandle_t external_cudnn_handle_ = nullptr; From 23c639384d2a4679df02adef7c4cc33f5dd3467f Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 23 Jul 2025 11:06:11 +0000 Subject: [PATCH 09/25] update external initializer fix --- .../nv_tensorrt_rtx/nv_provider_options.h | 1 + .../nv_tensorrt_rtx/nv_execution_provider.cc | 258 +++++++++--------- .../nv_tensorrt_rtx/nv_execution_provider.h | 29 +- .../nv_execution_provider_info.cc | 4 + .../nv_execution_provider_info.h | 1 + tools/ci_build/build.py | 4 +- 6 files changed, 158 insertions(+), 139 deletions(-) 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 index e4de1d3a6d647..620cec3bac594 100644 --- a/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h +++ b/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h @@ -35,6 +35,7 @@ constexpr const char* kCudaGraphEnable = "nv_cuda_graph_enable"; constexpr const char* kONNXBytestream = "nv_onnx_bytestream"; constexpr const char* kONNXBytestreamSize = "nv_onnx_bytestream_size"; constexpr const char* kMultiProfileEnable = "nv_multi_profile_enable"; +constexpr const char* kUseExternalDataInitializer = "nv_use_external_data_initializer"; constexpr const char* kExternalDataBytestream = "nv_external_data_bytestream"; constexpr const char* kExternalDataBytestreamSize = "nv_external_data_bytestream_size"; diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index e04f342c22bd1..72b26a41785cf 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -866,6 +866,7 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) "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) || @@ -1038,6 +1039,7 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) << ", 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_; } @@ -1512,18 +1514,25 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // 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. - // get initializer data + // save user provided external data in memory instead of writing to ModelProto + // needed for models > 2GB std::vector userWeights; - auto allInitializers = graph_viewer->GetAllInitializedTensors(); - for (auto entry : allInitializers){ + if(use_external_data_initializer_) { + auto allInitializers = graph_viewer->GetAllInitializedTensors(); + for (auto &entry : allInitializers) { auto* tp = entry.second; - if (tp->has_raw_data()){ - userWeights.push_back( - TensorrtUserWeights{tp->name(), tp->raw_data(), (int64_t)tp->raw_data().size()}); + if (tp->has_raw_data()) { + userWeights.emplace_back(tp->name(), tp->raw_data()); + } else if (utils::HasExternalDataInMemory(*tp)) { + std::unique_ptr full_init; + ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tp, full_init)); + userWeights.emplace_back(full_init->name(), full_init->raw_data()); } + } } - graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, false /*include raw initializers*/); + + 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); @@ -1550,12 +1559,16 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t { auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); - trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - - for (auto const& userWeight : userWeights){ - trt_parser->loadInitializer(userWeight.name.c_str(), static_cast(userWeight.data.c_str()), userWeight.size); + if (use_external_data_initializer_) { + 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 { + is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); } - is_model_supported = trt_parser->parseModelProto(); // Note: Calling getNbSubgraphs or getSubgraphNodes before calling supportsModelV2 results in undefined behavior. auto num_subgraphs = trt_parser->getNbSubgraphs(); @@ -1980,9 +1993,10 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, size_t onnx_external_data_bytestream_size, nvinfer1::ICudaEngine* trt_engine, bool serialize_refitted_engine, - bool detailed_build_log, - const GraphViewer* graph_body_viewer) { + 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) { @@ -2006,6 +2020,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, "The ONNX model path has '..'. For security purpose, it's not " "allowed to point outside the directory."); } + 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() + @@ -2013,164 +2028,136 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, } } } + // 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)); - bool refit_with_external_data = onnx_external_data_bytestream != nullptr && onnx_external_data_bytestream_size != 0; - // New refit APIs - if (refit_with_external_data || graph_body_viewer) { + if (refit_with_external_data) { + // A valid model bytestream must be passed. if (refit_from_file) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's refit with external data must be called with a valid ONNX model bytestream"); + "TensorRT 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)){ + if (!parser_refitter->loadModelProto(onnx_model_bytestream, onnx_model_bytestream_size, nullptr)) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP's IParserRefitter could not load model from provided onnx_model_bytestream"); } - // Extract weight information from the Refitter + // Extract weight information from the Refitter. int required_weights = refitter->getAllWeights(0, nullptr); std::vector refit_names(required_weights); refitter->getAllWeights(required_weights, refit_names.data()); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitter requires " << required_weights << " weights"; - // Vectors to keep track of data pointers - std::vector names; + // 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); - /* - TODO: Instead of passing in a single pointer and reconstructing the entire ONNX model, is it better to pre-partition the weights? - This function will now have parameters void ** data, const char ** names, int64_t * sizes, int64_t num_weights - for (int i = 0; i < num_weights; i++) - { - parser_refitter->loadInitializer(data[i], names[i], sizes[i]); - } - */ - - if (refit_with_external_data) - { - 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."); - } + auto onnx_model = ModelProto::Create(); + TensorProtos* allInitializers_byte_stream; - // Extract graph and initializer information. - auto const& graph = onnx_model->mutable_graph(); - allInitializers_byte_stream = graph->mutable_initializer(); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializers that were found " << allInitializers_byte_stream->size(); - - // Loop through all initializers - 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(); - bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); - if (weight_is_refittable) { - 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); - } + // 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) << "[TensorRT 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(); + bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); + if (weight_is_refittable) { + 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().c_str()); - bytes.push_back(static_cast(onnx_external_data_bytestream) + offset); - sizes.push_back(length); - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[TensorRT EP] Proto: " + proto_name + " has default as data location which is not supported"); } + names.push_back(proto.name()); + bytes.push_back(static_cast(onnx_external_data_bytestream) + offset); + sizes.push_back(length); } else { - if (!proto.has_raw_data()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[TensorRT EP] Proto: " + proto_name + " has no raw data"); - } - auto& raw_data = proto.raw_data(); - names.push_back(proto.name().c_str()); - bytes.push_back(raw_data.c_str()); - sizes.push_back(raw_data.size()); + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[TensorRT 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(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " was not marked as refittable"; + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Proto: " + proto_name + " has no raw nor external data."; + ++missing_initializer_data; } + } else { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " was not marked as refittable"; } } - else { // graph_body_viewer path. - // This path will only be used if - // 1. An ONNX was provided as byte array including initializers - // 2. An ONNX was provided as byte array and it's initializers were provided using AddExternalInitializers* API - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting using initializers of the current graph in memory"; - - auto allInitializers_graph_body = &graph_body_viewer->GetAllInitializedTensors(); - - for (auto& entry : *allInitializers_graph_body) { - auto* tp = entry.second; - auto& proto_name = tp->name(); - bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); - if (tp->has_raw_data() && weight_is_refittable) { - names.push_back(proto_name.c_str()); - bytes.push_back(tp->raw_data().c_str()); - sizes.push_back(tp->raw_data().size()); - } - } + if (missing_initializer_data) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[TensorRT EP] RefitEngine is missing " + std::to_string(missing_initializer_data) + " initializers."); } // Load extracted initializers into the parser if (!names.empty()) { LOGS_DEFAULT(VERBOSE) << "[TensorRT 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], bytes[i], sizes[i]); + 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, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + "TensorRT 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, - "TensorRT EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); + if (!parser_refitter->refitModelProto()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); } refit_complete = true; - } + // If new refit flow was not completed, then fallback to refit_from_file. - if (!refit_complete){ + if (!refit_complete) { if (refit_from_file) { LOGS_DEFAULT(VERBOSE) << "[TensorRT 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, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } - }else { + } else { LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from byte array"; if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); } } } @@ -2178,7 +2165,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; } else { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } // serialize the refitted engine to disk @@ -2187,11 +2174,12 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, nvinfer1::IHostMemory* serialized_engine = trt_engine->serialize(); std::ofstream engine_file(refitted_engine_cache, std::ios::binary | std::ios::out); engine_file.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Serialize the refitted engine to " << refitted_engine_cache; + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialize the refitted engine to " << refitted_engine_cache; } 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) { @@ -2344,21 +2332,25 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // exclude weights if external auto userWeights = std::make_unique>(); - auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); - for (auto entry : allInitializers){ - auto name = entry.first; + if (use_external_data_initializer_) { + auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); + for (auto& entry : allInitializers){ auto* tp = entry.second; if (tp->has_raw_data()){ - userWeights->push_back( - TensorrtUserWeights{tp->name(), tp->raw_data(), (int64_t)tp->raw_data().size()}); + userWeights->emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data())); + } else if (utils::HasExternalDataInMemory(*tp)) { + std::unique_ptr full_init; + ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tp, full_init)); + userWeights->emplace_back(TensorrtUserWeights(full_init->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*/, false); + 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); @@ -2375,11 +2367,17 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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)); - trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - for (auto const& userWeight : *userWeights){ - trt_parser->loadInitializer(userWeight.name.c_str(), static_cast(userWeight.data.c_str()), userWeight.size); + + if (use_external_data_initializer_) { + 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 { + trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); } - trt_parser->parseModelProto(); if (max_workspace_size_ > 0) { trt_config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_size_); @@ -2647,14 +2645,12 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr if (weight_stripped_engine_refit_) { LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refit engine from main ONNX file after engine build"; - char* onnx = string_buf.data(); - size_t onnx_size = string_buf.size(); auto status = RefitEngine(model_path_, onnx_model_folder_path_, engine_cache_path, false /* path check for security */, - onnx, - onnx_size, + onnx_model_bytestream_, + onnx_model_bytestream_size_, onnx_external_data_bytestream_, onnx_external_data_bytestream_size_, trt_engine.get(), diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index b58dc33bef918..c9eda82361658 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -153,11 +153,26 @@ struct TensorParams { } }; -// Struct to hold user weights when ModelProtos are serialized with external data -struct TensorrtUserWeights { - std::string name{}; - std::string data{}; - int64_t size{}; +// 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_(data) {}; + + const char* Name() const { + return name_.c_str(); + }; + + const void* Data() const { + return static_cast(data_.data()); + } + + const int64_t Size() const { + return static_cast(data_.size()); + } + + private: + std::string name_{}; + std::string data_{}; }; // Information to construct kernel function state. @@ -291,8 +306,7 @@ class NvExecutionProvider : public IExecutionProvider { size_t onnx_external_data_bytestream_size, nvinfer1::ICudaEngine* trt_engine, bool serialize_refitted_engine, - bool detailed_build_log, - const GraphViewer* graph_body_viewer = nullptr); + bool detailed_build_log); const InlinedVector GetEpContextNodes() const override; @@ -312,6 +326,7 @@ class NvExecutionProvider : public IExecutionProvider { 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; 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 index 411a5d6447da2..2d97f8f037b23 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc @@ -49,6 +49,7 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi .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) .AddValueParser( nv::provider_option_names::kONNXBytestream, @@ -105,6 +106,8 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi ORT_THROW("Invalid ", kOrtSessionOptionEpContextEmbedMode, " must 0 or 1"); } + // info.use_external_data_initializer = true; + return info; } @@ -123,6 +126,7 @@ ProviderOptions NvExecutionProviderInfo::ToProviderOptions(const NvExecutionProv {nv::provider_option_names::kCudaGraphEnable, MakeStringWithClassicLocale(info.cuda_graph_enable)}, {nv::provider_option_names::kONNXBytestream, MakeStringWithClassicLocale(info.onnx_bytestream)}, {nv::provider_option_names::kONNXBytestreamSize, MakeStringWithClassicLocale(info.onnx_bytestream_size)}, + {nv::provider_option_names::kUseExternalDataInitializer, MakeStringWithClassicLocale(info.use_external_data_initializer)}, {nv::provider_option_names::kExternalDataBytestream, MakeStringWithClassicLocale(info.external_data_bytestream)}, {nv::provider_option_names::kExternalDataBytestreamSize, MakeStringWithClassicLocale(info.external_data_bytestream_size)}, }; 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 index 89f92b7723db1..b826925361b05 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.h @@ -31,6 +31,7 @@ struct NvExecutionProviderInfo { 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}; diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index dd3e096c0334b..bf89ff2010ec5 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1708,8 +1708,10 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): run_ios_tests(args, source_dir, config, cwd) continue dll_path_list = [] - if args.use_tensorrt or args.use_nv_tensorrt_rtx: + 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: From 7b1320ee365235b1e9931b6d50b8eae2c672f707 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Thu, 24 Jul 2025 02:35:08 +0000 Subject: [PATCH 10/25] fix EP name --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 66 +++++++++---------- .../nv_execution_provider_info.cc | 2 - 2 files changed, 33 insertions(+), 35 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 72b26a41785cf..4b1cf65fd5e4b 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -487,7 +487,7 @@ Status BindContextInput(Ort::KernelContext& ctx, if (!trt_context->setTensorAddress(input_name, &shape_tensor_values[input_name][0])) { std::string error_input_name = input_name; std::string error_msg = - "Nv EP failed to call nvinfer1::IExecutionContext::setTensorAddress() for shape input '" + + "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)); } @@ -510,7 +510,7 @@ Status BindContextInput(Ort::KernelContext& ctx, if (!trt_context->setTensorAddress(input_name, &shape_tensor_values_int64[input_name][0])) { std::string error_input_name = input_name; std::string error_msg = - "Nv EP failed to call nvinfer1::IExecutionContext::setTensorAddress() for shape input '" + + "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)); } @@ -532,7 +532,7 @@ Status BindContextInput(Ort::KernelContext& ctx, if (!trt_context->setInputShape(input_name, dims)) { std::string error_input_name = input_name; ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP failed to call nvinfer1::IExecutionContext::setInputShape() for input '" + error_input_name + "'")); + "NvTensorRTRTX EP failed to call nvinfer1::IExecutionContext::setInputShape() for input '" + error_input_name + "'")); } // Bind "execution tensor" input buffer @@ -553,7 +553,7 @@ Status BindContextInput(Ort::KernelContext& ctx, CASE_GET_CAST_INPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE, double, float) default: { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP input onnx tensor data type: " + std::to_string(tensor_type) + " not supported."); + "NvTensorRTRTX EP input onnx tensor data type: " + std::to_string(tensor_type) + " not supported."); } } trt_context->setTensorAddress(input_name, data); @@ -644,7 +644,7 @@ Status BindContextOutput(Ort::KernelContext& ctx, CASE_GET_CAST_OUTPUT_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE, double, float) default: { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP output tensor data type: " + std::to_string(output_type) + " not supported."); + "NvTensorRTRTX EP output tensor data type: " + std::to_string(output_type) + " not supported."); } } trt_context->setTensorAddress(output_name, buffers[output_name]); @@ -707,7 +707,7 @@ Status BindKernelOutput(Ort::KernelContext& ctx, CASE_CAST_TENSOR(ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE, float, double) default: { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP output tensor data type: " + std::to_string(output_type) + " not supported."); + "NvTensorRTRTX EP output tensor data type: " + std::to_string(output_type) + " not supported."); } } return Status::OK(); @@ -988,13 +988,13 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) LIBTYPE handle = OPENLIB(engine_decryption_lib_path_.c_str()); if (handle == nullptr) { ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP could not open shared library from " + engine_decryption_lib_path_)); + "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, - "Nv EP could not find decryption function in shared library from " + engine_decryption_lib_path_)); + "NvTensorRTRTX EP could not find decryption function in shared library from " + engine_decryption_lib_path_)); } } @@ -2040,19 +2040,19 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, // A valid model bytestream must be passed. if (refit_from_file) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's refit with external data must be called with a valid ONNX model bytestream"); + "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, - "TensorRT EP's IParserRefitter could not load model from provided onnx_model_bytestream"); + "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(required_weights); refitter->getAllWeights(required_weights, refit_names.data()); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitter requires " << required_weights << " weights"; + LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refitter requires " << required_weights << " weights"; // Vectors to keep track of data pointers. std::vector names; @@ -2076,7 +2076,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, // Extract graph and initializer information. auto const& graph = onnx_model->mutable_graph(); allInitializers_byte_stream = graph->mutable_initializer(); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializers that were found " << allInitializers_byte_stream->size(); + LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Initializers that were found " << allInitializers_byte_stream->size(); // Loop through all initializers int missing_initializer_data = 0; @@ -2106,7 +2106,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, sizes.push_back(length); } else { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[TensorRT EP] Proto: " + proto_name + " expected to have external datalocation, but default datalocation was provided instead."); + "[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(); @@ -2114,33 +2114,33 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, bytes.push_back(raw_data.c_str()); sizes.push_back(raw_data.size()); } else { - LOGS_DEFAULT(WARNING) << "[TensorRT EP] Proto: " + proto_name + " has no raw nor external data."; + LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Proto: " + proto_name + " has no raw nor external data."; ++missing_initializer_data; } } else { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " was not marked as refittable"; + 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, - "[TensorRT EP] RefitEngine is missing " + std::to_string(missing_initializer_data) + " initializers."); + "[NvTensorRTRTX EP] RefitEngine is missing " + std::to_string(missing_initializer_data) + " initializers."); } // Load extracted initializers into the parser if (!names.empty()) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Number of initializers submitted to refitter " << names.size(); + 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, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); + "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, - "TensorRT EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); + "NvTensorRTRTX EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); } refit_complete = true; } @@ -2148,24 +2148,24 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, // If new refit flow was not completed, then fallback to refit_from_file. if (!refit_complete) { if (refit_from_file) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from file on disk: " << onnx_model_path.string(); + 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, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "NvTensorRTRTX EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } } else { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from byte array"; + 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, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); + "NvTensorRTRTX EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); } } } if (refitter->refitCudaEngine()) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Successfully refitted the weight-stripped engine."; + LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully refitted the weight-stripped engine."; } else { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "NvTensorRTRTX EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } // serialize the refitted engine to disk @@ -2174,7 +2174,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, nvinfer1::IHostMemory* serialized_engine = trt_engine->serialize(); std::ofstream engine_file(refitted_engine_cache, std::ios::binary | std::ios::out); engine_file.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialize the refitted engine to " << refitted_engine_cache; + LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Serialize the refitted engine to " << refitted_engine_cache; } return Status::OK(); } @@ -2611,12 +2611,12 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr std::unique_ptr serialized_engine{trt_builder->buildSerializedNetwork(*trt_network, *trt_config)}; if (serialized_engine == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP failed to create engine from network for fused node: " + fused_node.Name()); + "NvTensorRTRTX EP failed to create engine from network for fused node: " + fused_node.Name()); } trt_engine = std::unique_ptr(runtime_->deserializeCudaEngine(serialized_engine->data(), serialized_engine->size())); if (trt_engine == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP failed to deserialize engine for fused node: " + fused_node.Name()); + "NvTensorRTRTX EP failed to deserialize engine for fused node: " + fused_node.Name()); } if (detailed_build_log_) { auto engine_build_stop = std::chrono::steady_clock::now(); @@ -2667,7 +2667,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); if (!trt_context) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP could not build execution context for fused node: " + fused_node.Name()); + "NvTensorRTRTX EP could not build execution context for fused node: " + fused_node.Name()); } bool is_dynamic_shape_context = false; @@ -2767,7 +2767,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr if (multi_profile_enable_ == true) { if (!trt_context->setOptimizationProfileAsync(nv_profile_index_, stream)) - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Nv EP select an optimization profile for the current context failed"); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "NvTensorRTRTX EP select an optimization profile for the current context failed"); } // Check before using trt_engine @@ -2881,7 +2881,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // Run TRT inference if (!trt_context->enqueueV3(stream)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Nv EP execution context enqueue failed."); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "NvTensorRTRTX EP execution context enqueue failed."); } /* @@ -2992,7 +2992,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra trt_context = std::unique_ptr(trt_engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); if (!trt_context) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "Nv EP could not build execution context for fused node: " + fused_node.Name()); + "NvTensorRTRTX EP could not build execution context for fused node: " + fused_node.Name()); } bool is_dynamic_shape_context = false; @@ -3197,7 +3197,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra // Run TRT inference if (!trt_context->enqueueV3(stream)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Nv EP execution context enqueue failed."); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "NvTensorRTRTX EP execution context enqueue failed."); } /* 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 index 2d97f8f037b23..62c63fdcdd979 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc @@ -106,8 +106,6 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi ORT_THROW("Invalid ", kOrtSessionOptionEpContextEmbedMode, " must 0 or 1"); } - // info.use_external_data_initializer = true; - return info; } From 3b039fa74c9c20eb697951d01afa71c25cd9f572 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Thu, 31 Jul 2025 12:45:55 +0200 Subject: [PATCH 11/25] reorganize unittest helpers --- .../nv_tensorrt_rtx/nv_basic_test.cc | 174 +----------------- .../test_nv_trt_rtx_ep_util.cc | 118 ++++++++++++ .../nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h | 86 ++++++++- 3 files changed, 204 insertions(+), 174 deletions(-) diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc index 19505da1bbe56..7eb1cdab9bb81 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc @@ -5,21 +5,13 @@ #include "core/session/inference_session.h" #include "test/providers/provider_test_utils.h" #include "test/framework/test_utils.h" -#include "gtest/gtest.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 "test/util/include/api_asserts.h" -#include "test/util/include/asserts.h" -#include -#include -#include -#include -#include #include -#include #include using namespace std; @@ -59,170 +51,6 @@ class NvExecutionProviderTest : public ::testing::Test { using NvExecutionProviderTestTypes = ::testing::Types; // double, TYPED_TEST_SUITE(NvExecutionProviderTest, NvExecutionProviderTestTypes); -std::string PathToUTF8(const PathString& path) { -#ifdef WIN32 - std::wstring_convert> converter; - return converter.to_bytes(path); -#else - return path.c_str(); -#endif -} - -void clearFileIfExists(PathString path) { - if (std::filesystem::exists(path)) { - std::filesystem::remove(path); - } -} - -template -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. - * - * 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" - */ -static 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) { - 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()); - status = onnxruntime::Model::Save(model, model_name); - ASSERT_TRUE(status.IsOK()); -} - -static Ort::IoBinding generate_io_binding(Ort::Session& session, std::map> shape_overwrites = {}) { - Ort::IoBinding binding(session); - auto allocator = Ort::AllocatorWithDefaultOptions(); - 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(), allocator.GetInfo()); - } - return binding; -} TEST(NvExecutionProviderTest, ContextEmbedAndReload) { PathString model_name = ORT_TSTR("nv_execution_provider_test.onnx"); 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 index f0ce5c0b296ca..53b2a3627bf0f 100644 --- 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 @@ -12,6 +12,11 @@ #include "core/session/onnxruntime_cxx_api.h" #include "test/util/include/api_asserts.h" +#include "core/graph/onnx_protobuf.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/framework/test_utils.h" namespace onnxruntime { namespace test { @@ -52,6 +57,119 @@ void Utils::RegisterAndGetNvTensorRtRtxEp(Ort::Env& env, RegisteredEpDeviceUniqu }); } +void CreateBaseModel(const PathString& model_name, + std::string graph_name, + std::vector dims, + bool add_fast_gelu, + ONNX_NAMESPACE::TensorProto_DataType dtype) { + 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()); + status = onnxruntime::Model::Save(model, model_name); + ASSERT_TRUE(status.IsOK()); +} + +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, info->alloc_type, 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; +} + } // 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 index ef14d3cb382c0..4369fe5e692c2 100644 --- 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 @@ -5,9 +5,19 @@ #include #include +#include +#include + +#include +#include +#include +#include +#include -#include "core/session/onnxruntime_cxx_api.h" #include "core/graph/constants.h" +#include "core/common/path_string.h" +#include "core/framework/tensor.h" +#include "test/util/include/api_asserts.h" namespace onnxruntime { namespace test { @@ -34,5 +44,79 @@ struct Utils { // automatically unregister the EP library. static void RegisterAndGetNvTensorRtRtxEp(Ort::Env& env, RegisteredEpDeviceUniquePtr& nv_tensorrt_rtx_ep); }; + +static std::string PathToUTF8(const PathString& path) { +#ifdef WIN32 + std::wstring_convert> converter; + return converter.to_bytes(path); +#else + return path.c_str(); +#endif +} + +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. + * + * 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); + +Ort::IoBinding generate_io_binding( + Ort::Session& session, + std::map> shape_overwrites = {}, + OrtAllocator* allocator = nullptr); + } // namespace test } // namespace onnxruntime From 28d211e6558abe34e2867f6a3567820a596c43e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Thu, 31 Jul 2025 18:38:03 +0200 Subject: [PATCH 12/25] fix type tests --- .../nv_tensorrt_rtx/nv_basic_test.cc | 78 +++++++++++-------- .../test_nv_trt_rtx_ep_util.cc | 2 + .../nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h | 5 +- 3 files changed, 50 insertions(+), 35 deletions(-) diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc index 7eb1cdab9bb81..de13ef547eac6 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc @@ -22,36 +22,6 @@ namespace onnxruntime { namespace test { -template -class NvExecutionProviderTest : public ::testing::Test { - protected: - std::string getTypeAsName() { - std::string dtype_name = ""; - if constexpr (std::is_same::value) { - dtype_name = "fp64"; - } else if constexpr (std::is_same::value) { - dtype_name = "fp32"; - } else if constexpr (std::is_same::value) { - dtype_name = "bf16"; - } else if constexpr (std::is_same::value) { - dtype_name = "fp16"; - } else if constexpr (std::is_same::value) { - dtype_name = "int8"; - } else if constexpr (std::is_same::value) { - dtype_name = "uint8"; - } else if constexpr (std::is_same::value) { - dtype_name = "int32"; - } else if constexpr (std::is_same::value) { - dtype_name = "int64"; - } - return dtype_name; - } -}; - -using NvExecutionProviderTestTypes = ::testing::Types; // double, -TYPED_TEST_SUITE(NvExecutionProviderTest, NvExecutionProviderTestTypes); - - 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"); @@ -196,15 +166,44 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDataDynamic) { } } -TYPED_TEST(NvExecutionProviderTest, IOTypeTests) { - std::string dtype_name = this->getTypeAsName(); +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) { + 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); std::string graph_name = "test" + dtype_name; std::vector dims = {1, -1, -1}; - CreateBaseModel(model_name, graph_name, dims); + CreateBaseModel(model_name, graph_name, dims, false, GetParam()); auto env = Ort::Env(); auto logging_level = OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING; @@ -222,6 +221,19 @@ TYPED_TEST(NvExecutionProviderTest, IOTypeTests) { } } +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); }); + #if defined(WIN32) static bool SessionHasEp(Ort::Session& session, const char* ep_name) { // Access the underlying InferenceSession. 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 index 53b2a3627bf0f..be1cb4efc3942 100644 --- 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 @@ -20,6 +20,7 @@ namespace onnxruntime { namespace test { +#ifdef _WIN32 Utils::NvTensorRtRtxEpInfo Utils::nv_tensorrt_rtx_ep_info; @@ -56,6 +57,7 @@ void Utils::RegisterAndGetNvTensorRtRtxEp(Ort::Env& env, RegisteredEpDeviceUniqu c_api.UnregisterExecutionProviderLibrary(env, nv_tensorrt_rtx_ep_info.registration_name.c_str()); }); } +#endif // _WIN32 void CreateBaseModel(const PathString& model_name, std::string graph_name, 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 index 4369fe5e692c2..f07dddd008aa7 100644 --- 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 @@ -17,6 +17,7 @@ #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 { @@ -45,7 +46,7 @@ struct Utils { static void RegisterAndGetNvTensorRtRtxEp(Ort::Env& env, RegisteredEpDeviceUniquePtr& nv_tensorrt_rtx_ep); }; -static std::string PathToUTF8(const PathString& path) { +[[maybe_unused]] static std::string PathToUTF8(const PathString& path) { #ifdef WIN32 std::wstring_convert> converter; return converter.to_bytes(path); @@ -54,7 +55,7 @@ static std::string PathToUTF8(const PathString& path) { #endif } -static void clearFileIfExists(PathString path) { +[[maybe_unused]] static void clearFileIfExists(PathString path) { if (std::filesystem::exists(path)) { std::filesystem::remove(path); } From bd3d4edf37f777cf9be87bc88c1a46595b0a65f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Thu, 31 Jul 2025 21:52:30 +0200 Subject: [PATCH 13/25] basic EP context support expand EP context tests --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 143 +++++++-------- .../nv_tensorrt_rtx/nv_execution_provider.h | 11 +- .../nv_execution_provider_info.cc | 3 +- .../nv_execution_provider_utils.h | 13 +- .../nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 170 +++++++----------- .../nv_tensorrt_rtx/onnx_ctx_model_helper.h | 33 ++-- .../nv_tensorrt_rtx/nv_basic_test.cc | 47 ++--- .../nv_tensorrt_rtx/nv_ep_context_test.cc | 131 ++++++++++++++ .../test_nv_trt_rtx_ep_util.cc | 14 +- .../nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h | 7 +- 10 files changed, 310 insertions(+), 262 deletions(-) create mode 100644 onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 4b1cf65fd5e4b..786da87a379c4 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -836,7 +836,11 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) cudaDeviceProp prop; CUDA_CALL_THROW(cudaGetDeviceProperties(&prop, device_id_)); - compute_capability_ = GetComputeCapacity(prop); + if (prop.major < 8 || prop.major == 9 || prop.major == 10) { + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[NvTensorRTRTX EP] The execution provider only supports RTX devices with compute capabilities =< 80.")); + } + compute_capability_ = GetComputeCapability(prop); if (info.has_user_compute_stream) { external_stream_ = true; stream_ = static_cast(info.user_compute_stream); @@ -872,8 +876,8 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) 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")); + "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; @@ -1518,9 +1522,9 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // needed for models > 2GB std::vector userWeights; - if(use_external_data_initializer_) { + if (use_external_data_initializer_) { auto allInitializers = graph_viewer->GetAllInitializedTensors(); - for (auto &entry : allInitializers) { + for (auto& entry : allInitializers) { auto* tp = entry.second; if (tp->has_raw_data()) { userWeights.emplace_back(tp->name(), tp->raw_data()); @@ -1534,7 +1538,6 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t 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; @@ -1561,12 +1564,11 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t if (use_external_data_initializer_) { trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - for (auto const& userWeight : userWeights){ + for (auto const& userWeight : userWeights) { trt_parser->loadInitializer(userWeight.Name(), userWeight.Data(), userWeight.Size()); } is_model_supported = trt_parser->parseModelProto(); - } - else { + } else { is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); } @@ -1760,13 +1762,14 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, // 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. - if (GraphHasCtxNode(graph)) { + size_t node_idx = 0; + if (GraphHasCtxNode(graph, node_idx)) { int subgraph_idx = 0; - for (size_t i = 0; i < static_cast(number_of_ort_nodes); i++) { - const auto& node = graph.GetNode(node_index[i]); - const bool is_context_node = node && !node->OpType().empty() && node->OpType() == "EPContext"; + 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) { - SubGraph_t supported_node_vector(std::make_pair(std::vector{i}, true)); + 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))); @@ -1985,16 +1988,13 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, */ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, std::string& onnx_model_folder_path, - std::string& weight_stripped_engine_cath_path, bool path_check, 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 serialize_refitted_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; @@ -2168,18 +2168,9 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, "NvTensorRTRTX EP's IRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); } - // serialize the refitted engine to disk - if (serialize_refitted_engine) { - std::string refitted_engine_cache = GetWeightRefittedEnginePath(weight_stripped_engine_cath_path); - nvinfer1::IHostMemory* serialized_engine = trt_engine->serialize(); - std::ofstream engine_file(refitted_engine_cache, std::ios::binary | std::ios::out); - engine_file.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Serialize the refitted engine to " << refitted_engine_cache; - } 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) { @@ -2202,8 +2193,10 @@ common::Status NvExecutionProvider::Compile(const std::vector } Status status; - if (GraphHasCtxNode(graph_body_viewer)) { + 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, @@ -2310,10 +2303,8 @@ static bool IsIOBindingRequired(TRTState* const trt_state, const Ort::KernelCont const InlinedVector NvExecutionProvider::GetEpContextNodes() const { InlinedVector ep_context_nodes; - for (auto& model : ep_context_models_) { - auto& graph = model->MainGraph(); - for (int i = 0; i < graph.MaxNodeIndex(); i++) { - auto node = graph.GetNode(i); + if (ep_context_model_) { + for (auto* node : ep_context_model_->MainGraph().Nodes()) { ep_context_nodes.push_back(node); } } @@ -2334,9 +2325,9 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr if (use_external_data_initializer_) { auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); - for (auto& entry : allInitializers){ + for (auto& entry : allInitializers) { auto* tp = entry.second; - if (tp->has_raw_data()){ + if (tp->has_raw_data()) { userWeights->emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data())); } else if (utils::HasExternalDataInMemory(*tp)) { std::unique_ptr full_init; @@ -2370,12 +2361,11 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr if (use_external_data_initializer_) { trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); - for (auto const& userWeight : *userWeights){ + for (auto const& userWeight : *userWeights) { trt_parser->loadInitializer(userWeight.Name(), userWeight.Data(), userWeight.Size()); } trt_parser->parseModelProto(); - } - else { + } else { trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); } @@ -2542,7 +2532,6 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr ; } } - std::string trt_node_name_with_precision = fused_node.Name() + "_strong_typed"; // enable sparse weights if (sparsity_enable_) { @@ -2571,32 +2560,6 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr std::unique_ptr trt_engine; std::unique_ptr trt_context; - std::string cache_path = ""; - std::string cache_suffix = ""; - // Customize cache prefix if assigned - if (!cache_prefix_.empty()) { - // Generate cache suffix in case user would like to customize cache prefix - cache_suffix = "_" + GetCacheSuffix(fused_node.Name(), trt_node_name_with_precision); - cache_path = GetCachePath(cache_path_, cache_prefix_) + cache_suffix; - } else { - cache_path = GetCachePath(cache_path_, trt_node_name_with_precision); - } - - // Name the engine cache based on GPU compute capacity and reduce the chance of loading an incompatible cache - // Note: Engine cache generated on a GPU with large memory might not be loadable on a GPU with smaller memory, even if they share the same compute capacity - const std::string cache_path_prefix = cache_path; - std::string engine_cache_path = cache_path_prefix + ".engine"; - const std::string encrypted_engine_cache_path = engine_cache_path + ".encrypted"; - const std::string profile_cache_path = cache_path_prefix + ".profile"; - - // If weight-stripped engine is enabled and refitted engine cache is not present, - // TRT EP will use the engine cache with ".stripped.engine" appended to the end. - const std::filesystem::path engine_cache_fs_path = engine_cache_path; - if (weight_stripped_engine_enable_ && !std::filesystem::exists(engine_cache_fs_path)) { - engine_cache_path = cache_path_prefix + ".stripped.engine"; - weight_stripped_engine_refit_ = true; - } - // Generate file name for dumping ep context model if (dump_ep_context_model_ && ctx_model_path_.empty()) { ctx_model_path_ = GetCtxModelPath(ep_context_file_path_, model_path_); @@ -2620,26 +2583,43 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr } if (detailed_build_log_) { auto engine_build_stop = std::chrono::steady_clock::now(); - LOGS_DEFAULT(INFO) << "TensorRT engine build for " << trt_node_name_with_precision << " took: " << std::chrono::duration_cast(engine_build_stop - engine_build_start).count() << "ms" << std::endl; + 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 - if (ep_cache_context_attr_.empty()) { - auto cache_file_name = std::filesystem::path(engine_cache_path).filename(); - ep_cache_context_attr_ = std::filesystem::path(engine_cache_relative_path_to_context_model_dir).append(cache_file_name.string()).string(); + + 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"; + ; + } + auto cache_file_name = std::filesystem::path(cache_path).filename(); + cache_path = std::filesystem::path(engine_cache_relative_path_to_context_model_dir).append(cache_file_name.string()).string(); + // 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()); } - std::string compute_capability_hw_compat = compute_capability_ + "+"; - ep_context_models_.push_back(CreateCtxNode(graph_body_viewer, - ep_cache_context_attr_, - reinterpret_cast(serialized_engine->data()), - serialized_engine->size(), - ep_context_embed_mode_, - compute_capability_hw_compat, - model_path_, - GetLogger(), - fused_node.Name())); + 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()); + } } } @@ -2647,14 +2627,12 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Refit engine from main ONNX file after engine build"; auto status = RefitEngine(model_path_, onnx_model_folder_path_, - engine_cache_path, false /* path check for security */, onnx_model_bytestream_, onnx_model_bytestream_size_, onnx_external_data_bytestream_, onnx_external_data_bytestream_size_, trt_engine.get(), - false /* serialize refitted engine to disk */, detailed_build_log_); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); @@ -2714,12 +2692,12 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr *p = {context->allocate_func, context->release_func, context->allocator_handle, context->node_name, builder_.get(), &parsers_[context->node_name], &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_, trt_node_name_with_precision, + 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_, cache_suffix, &weights_[context->node_name]}; + auxiliary_streams_, cuda_graph_enable_, is_dynamic_shape_context, cache_prefix_, &weights_[context->node_name]}; *state = p.release(); return 0; }; @@ -2958,6 +2936,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr } 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, @@ -2980,7 +2959,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const Gra onnx_external_data_bytestream_, onnx_external_data_bytestream_size_, detailed_build_log_); - auto status = trt_cache_model_handler.GetEpContextFromGraph(graph_body_viewer); + 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()); } diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index c9eda82361658..33984ecf2b553 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -166,7 +166,7 @@ class TensorrtUserWeights { return static_cast(data_.data()); } - const int64_t Size() const { + int64_t Size() const { return static_cast(data_.size()); } @@ -190,7 +190,6 @@ struct TensorrtFuncState { std::vector> output_info; std::unordered_map>>> input_shape_ranges; std::mutex* tensorrt_mu_ptr = nullptr; - std::string trt_node_name_with_precision; bool engine_cache_enable = false; std::string engine_cache_path; nvinfer1::IRuntime* runtime = nullptr; @@ -298,14 +297,12 @@ class NvExecutionProvider : public IExecutionProvider { static common::Status RefitEngine(std::string onnx_model_filename, std::string& onnx_model_folder_path, - std::string& weight_stripped_engine_cath_path, bool path_check, 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 serialize_refitted_engine, bool detailed_build_log); const InlinedVector GetEpContextNodes() const override; @@ -347,7 +344,7 @@ class NvExecutionProvider : public IExecutionProvider { std::string cache_prefix_; std::string op_types_to_exclude_; int nv_profile_index_ = 0; - std::vector> ep_context_models_; + std::unique_ptr ep_context_model_; // The format is as for TENSORRT_VERSION: (MAJOR * 100 + MINOR) * 100 + PATCH int32_t trt_version_; @@ -362,7 +359,6 @@ class NvExecutionProvider : public IExecutionProvider { std::string ep_context_file_path_; int ep_context_embed_mode_ = 0; std::string ctx_model_path_; - std::string ep_cache_context_attr_; std::string engine_cache_relative_path_to_context_model_dir; std::unordered_set control_flow_op_set_ = {"If", "Loop", "Scan"}; @@ -387,7 +383,7 @@ class NvExecutionProvider : public IExecutionProvider { 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_; - std::unordered_map>> weights_; // User provided weights + std::unordered_map>> weights_; // User provided weights // for external stream, we need to create its cudnn/cublass handle before cuda EP enable cuda graph capture cudnnHandle_t external_cudnn_handle_ = nullptr; @@ -582,6 +578,7 @@ class NvExecutionProvider : public IExecutionProvider { * 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, 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 index 62c63fdcdd979..090ae1b20e0ea 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc @@ -85,7 +85,8 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi info.dump_ep_context_model = false; } else if (ep_context_enable == "1") { info.dump_ep_context_model = true; - info.weight_stripped_engine_enable = 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"); } 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 index ea586ba445ba2..c564fe65c3d5c 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_utils.h @@ -386,22 +386,11 @@ std::string GetCachePath(const std::string& root, const std::string& name) { * Get compute capability * */ -std::string GetComputeCapacity(const cudaDeviceProp& prop) { +std::string GetComputeCapability(const cudaDeviceProp& prop) { const std::string compute_capability = std::to_string(prop.major * 10 + prop.minor); return compute_capability; } -/* - * Get Timing by compute capability - * - */ -std::string GetTimingCachePath(const std::string& root, std::string& compute_cap) { - // append compute capability of the GPU as this invalidates the cache and TRT will throw when loading the cache - const std::string timing_cache_name = "NvExecutionProvider_cache_sm" + - compute_cap + ".timing"; - return GetCachePath(root, timing_cache_name); -} - /* * Get cache by type * 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 index e92576590764a..9e0db4ac377f3 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -20,10 +20,11 @@ extern TensorrtLogger& GetTensorrtLogger(bool verbose_log); * * Note: Please see more details about "EPContext" contrib op in contrib_defs.cc */ -bool GraphHasCtxNode(const GraphViewer& graph_viewer) { +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) { + node_idx = i; return true; } } @@ -65,18 +66,16 @@ void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, /* * Create EP context node where engine information is embedded */ -std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewer, - 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 logging::Logger* logger, - const std::string& ep_context_node_name) { - auto model_build = Model::Create("nv_trt_rtx_ep_context_model", false, *logger); - auto& graph_build = model_build->MainGraph(); - +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()) { @@ -90,47 +89,72 @@ std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewe } // Create EP context node attributes - auto attr_0 = ONNX_NAMESPACE::AttributeProto::Create(); // embed_mode - auto attr_1 = ONNX_NAMESPACE::AttributeProto::Create(); // ep_cache_context - auto attr_2 = ONNX_NAMESPACE::AttributeProto::Create(); // hardware_architecture - auto attr_3 = ONNX_NAMESPACE::AttributeProto::Create(); // onnx_model_filename + 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(); std::string engine_data_str = ""; - attr_0->set_name(EMBED_MODE); - attr_0->set_type(onnx::AttributeProto_AttributeType_INT); - attr_0->set_i(embed_mode); - attr_1->set_name(EP_CACHE_CONTEXT); - attr_1->set_type(onnx::AttributeProto_AttributeType_STRING); + 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_1->set_s(engine_data_str); + attr_ep_cache_context->set_s(engine_data_str); // TODO(maximilianm) we might want to disable this warning as we only support weightless engines that are really small // the reason we had this was that the field will be hashed and storing a large bytestream has significant overhead - LOGS_DEFAULT(WARNING) << EPCONTEXT_WARNING; } else { - attr_1->set_s(engine_cache_path); + attr_ep_cache_context->set_s(engine_cache_path); + 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_2->set_name(COMPUTE_CAPABILITY); - attr_2->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_2->set_s(compute_capability); - attr_3->set_name(ONNX_MODEL_FILENAME); - attr_3->set_type(onnx::AttributeProto_AttributeType_STRING); - attr_3->set_s(std::filesystem::path(onnx_model_path).filename().string()); + + 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(std::filesystem::path(onnx_model_path).filename().string()); + + 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)); auto node_attributes = ONNX_NAMESPACE::NodeAttributes::Create(); constexpr int num_attributes = 4; node_attributes->reserve(num_attributes); - node_attributes->emplace(EMBED_MODE, *attr_0); - node_attributes->emplace(EP_CACHE_CONTEXT, *attr_1); - node_attributes->emplace(COMPUTE_CAPABILITY, *attr_2); - node_attributes->emplace(ONNX_MODEL_FILENAME, *attr_3); + 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); // 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 model_build; + return Status::OK(); } /* @@ -199,17 +223,6 @@ std::string GetCtxModelPath(const std::string& ep_context_file_path, return ctx_model_path; } -/* - * Dump "EP context" model - * - */ -void DumpCtxModel(ONNX_NAMESPACE::ModelProto* model_proto, - const std::string& ctx_model_path) { - std::fstream dump(ctx_model_path, std::ios::out | std::ios::trunc | std::ios::binary); - model_proto->SerializeToOstream(dump); - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Dumped " + ctx_model_path; -} - bool IsAbsolutePath(const std::string& path_string) { #ifdef _WIN32 onnxruntime::PathString ort_path_string = onnxruntime::ToPathString(path_string); @@ -241,38 +254,12 @@ bool IsRelativePathToParentPath(const std::string& path_string) { #endif } -/* - * Get the weight-refitted engine cache path from a weight-stripped engine cache path - * - * Weight-stipped engine: - * An engine with weights stripped and its size is smaller than a regualr engine. - * The cache name of weight-stripped engine is NvExecutionProvider_TRTKernel_XXXXX.stripped.engine - * - * Weight-refitted engine: - * An engine that its weights have been refitted and it's simply a regular engine. - * The cache name of weight-refitted engine is NvExecutionProvider_TRTKernel_XXXXX.engine - */ -std::string GetWeightRefittedEnginePath(std::string stripped_engine_cache) { - std::filesystem::path stripped_engine_cache_path(stripped_engine_cache); - std::string refitted_engine_cache_path = stripped_engine_cache_path.stem().stem().string() + ".engine"; - return refitted_engine_cache_path; -} - -bool IsWeightStrippedEngineCache(std::filesystem::path& engine_cache_path) { - // The weight-stripped engine cache has the naming of xxx.stripped.engine - return engine_cache_path.stem().extension().string() == ".stripped"; -} - -Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph_viewer) { - if (!ValidateEPCtxNode(graph_viewer)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "It's not a valid EP Context node"); - } - auto node = graph_viewer.GetNode(0); - auto& attrs = node->GetAttributes(); +Status TensorRTCacheModelHandler::GetEpContextFromGraph(const Node& node) { + auto& attrs = node.GetAttributes(); const int64_t embed_mode = attrs.at(EMBED_MODE).i(); // Only make path checks if model not provided as byte buffer - bool make_secure_path_checks = !GetModelPath(graph_viewer).empty(); + bool make_secure_path_checks = ep_context_model_path_.empty(); if (embed_mode) { // Get engine from byte stream. @@ -287,17 +274,14 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph if (weight_stripped_engine_refit_) { const std::string onnx_model_filename = attrs.at(ONNX_MODEL_FILENAME).s(); - std::string placeholder; auto status = NvExecutionProvider::RefitEngine(onnx_model_filename, onnx_model_folder_path_, - placeholder, make_secure_path_checks, onnx_model_bytestream_, onnx_model_bytestream_size_, onnx_external_data_bytestream_, onnx_external_data_bytestream_size_, (*trt_engine_).get(), - false /* serialize refitted engine to disk */, detailed_build_log_); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); @@ -322,21 +306,6 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph auto engine_cache_path = ctx_model_dir.append(cache_path); LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] GetEpContextFromGraph engine_cache_path: " + engine_cache_path.string(); - // If it's a weight-stripped engine cache, it needs to be refitted even though the refit flag is not enabled - if (!weight_stripped_engine_refit_) { - weight_stripped_engine_refit_ = IsWeightStrippedEngineCache(engine_cache_path); - } - - // If the serialized refitted engine is present, use it directly without refitting the engine again - if (weight_stripped_engine_refit_) { - const std::filesystem::path refitted_engine_cache_path = GetWeightRefittedEnginePath(engine_cache_path.string()); - if (std::filesystem::exists(refitted_engine_cache_path)) { - LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] " + refitted_engine_cache_path.string() + " exists."; - engine_cache_path = refitted_engine_cache_path.string(); - weight_stripped_engine_refit_ = false; - } - } - 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() + @@ -361,14 +330,12 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph std::string weight_stripped_engine_cache = engine_cache_path.string(); auto status = NvExecutionProvider::RefitEngine(onnx_model_filename, onnx_model_folder_path_, - weight_stripped_engine_cache, make_secure_path_checks, onnx_model_bytestream_, onnx_model_bytestream_size_, onnx_external_data_bytestream_, onnx_external_data_bytestream_size_, (*trt_engine_).get(), - true /* serialize refitted engine to disk */, detailed_build_log_); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); @@ -381,11 +348,8 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph /* * The sanity check for EP context contrib op. */ -bool TensorRTCacheModelHandler::ValidateEPCtxNode(const GraphViewer& graph_viewer) { - assert(graph_viewer.NumberOfNodes() == 1); - assert(graph_viewer.GetNode(0)->OpType() == EPCONTEXT_OP); - auto node = graph_viewer.GetNode(0); - auto& attrs = node->GetAttributes(); +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) { @@ -410,7 +374,7 @@ bool TensorRTCacheModelHandler::ValidateEPCtxNode(const GraphViewer& graph_viewe const int64_t embed_mode = attrs.at(EMBED_MODE).i(); if (embed_mode == 1) { // engine binary data - LOGS_DEFAULT(WARNING) << EPCONTEXT_WARNING; + // LOGS_DEFAULT(WARNING) << EPCONTEXT_WARNING; } return true; 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 index dc9d9c7b5ae39..7c52f26cc9177 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "core/providers/nv_tensorrt_rtx/nv_includes.h" #include "core/providers/shared_library/provider_api.h" @@ -14,32 +15,32 @@ 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 EPCONTEXT_OP_DOMAIN = "com.microsoft"; -static const std::string EPCONTEXT_WARNING = - "It's suggested to set the ORT graph optimization level to 0 for the best performance"; -bool GraphHasCtxNode(const GraphViewer& graph_viewer); +bool GraphHasCtxNode(const GraphViewer& graph_viewer, size_t& node_idx); const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer); std::filesystem::path GetPathOrParentPathOfCtxModel(const std::string& ep_context_file_path); -std::unique_ptr CreateCtxNode(const GraphViewer& graph_viewer, - 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 logging::Logger* logger, - const std::string& ep_context_node_name); +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); std::string GetCtxModelPath(const std::string& ep_context_file_path, const std::string& original_model_path); bool IsAbsolutePath(const std::string& path_string); bool IsRelativePathToParentPath(const std::string& path_string); -void DumpCtxModel(ONNX_NAMESPACE::ModelProto* model_proto, - const std::string& ctx_model_path); void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, char* engine_data, size_t size); @@ -71,9 +72,9 @@ class TensorRTCacheModelHandler { } ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TensorRTCacheModelHandler); - bool ValidateEPCtxNode(const GraphViewer& graph_viewer); + bool ValidateEPCtxNode(const Node& node); - Status GetEpContextFromGraph(const GraphViewer& graph_viewer); + Status GetEpContextFromGraph(const Node& node); private: std::unique_ptr* trt_engine_; diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc index de13ef547eac6..35a63d0ac6333 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc @@ -31,11 +31,6 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReload) { std::vector dims = {1, 3, 2}; CreateBaseModel(model_name, graph_name, dims); - - auto env = Ort::Env(); - auto logging_level = OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING; - env.UpdateEnvWithCustomLogLevel(logging_level); - // AOT time { auto start = std::chrono::high_resolution_clock::now(); @@ -44,7 +39,7 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReload) { so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name.c_str(), so); + 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; @@ -59,7 +54,7 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReload) { Ort::RunOptions run_options; so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name_ctx.c_str(), so); + 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; @@ -78,10 +73,6 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDynamic) { CreateBaseModel(model_name, graph_name, dims); - auto env = Ort::Env(); - auto logging_level = OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING; - env.UpdateEnvWithCustomLogLevel(logging_level); - // AOT time { auto start = std::chrono::high_resolution_clock::now(); @@ -90,7 +81,7 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDynamic) { so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name.c_str(), so); + 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; @@ -105,7 +96,7 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDynamic) { Ort::RunOptions run_options; so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name_ctx.c_str(), so); + 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; @@ -127,10 +118,6 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDataDynamic) { CreateBaseModel(model_name, graph_name, dims); - auto env = Ort::Env(); - auto logging_level = OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING; - env.UpdateEnvWithCustomLogLevel(logging_level); - // AOT time { auto start = std::chrono::high_resolution_clock::now(); @@ -139,7 +126,7 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDataDynamic) { so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); so.AddConfigEntry(kOrtSessionOptionEpContextFilePath, model_name_ctx_str.c_str()); so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name.c_str(), so); + 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; @@ -154,7 +141,7 @@ TEST(NvExecutionProviderTest, ContextEmbedAndReloadDataDynamic) { Ort::RunOptions run_options; so.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1"); so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name_ctx.c_str(), so); + 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; @@ -196,25 +183,21 @@ class TypeTests : public ::testing::TestWithParam dims = {1, -1, -1}; + const std::string graph_name = "test" + dtype_name; + const std::vector dims = {1, 5, 10}; CreateBaseModel(model_name, graph_name, dims, false, GetParam()); - auto env = Ort::Env(); - auto logging_level = OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING; - env.UpdateEnvWithCustomLogLevel(logging_level); - // AOT time { Ort::SessionOptions so; Ort::RunOptions run_options; so.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); - Ort::Session session_object(env, model_name.c_str(), so); + 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); @@ -260,20 +243,16 @@ TEST(NvExecutionProviderTest, AutoEp_PreferGpu) { CreateBaseModel(model_name, graph_name, dims); - auto env = Ort::Env(); - auto logging_level = OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING; - env.UpdateEnvWithCustomLogLevel(logging_level); - { - env.RegisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider, ORT_TSTR("onnxruntime_providers_nv_tensorrt_rtx.dll")); + ort_env->RegisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider, ORT_TSTR("onnxruntime_providers_nv_tensorrt_rtx.dll")); Ort::SessionOptions so; so.SetEpSelectionPolicy(OrtExecutionProviderDevicePolicy_PREFER_GPU); - Ort::Session session_object(env, model_name.c_str(), so); + Ort::Session session_object(*ort_env, model_name.c_str(), so); EXPECT_TRUE(SessionHasEp(session_object, kNvTensorRTRTXExecutionProvider)); } - env.UnregisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider); + ort_env->UnregisterExecutionProviderLibrary(kNvTensorRTRTXExecutionProvider); } TEST(NvExecutionProviderTest, GetSharedAllocator) { 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 new file mode 100644 index 0000000000000..f9f832096daaf --- /dev/null +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc @@ -0,0 +1,131 @@ +// 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/common/path_utils.h" +#include "core/graph/onnx_protobuf.h" +#include "core/session/inference_session.h" +#include "test/providers/provider_test_utils.h" +#include "test/framework/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 { + +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; + const std::string to_string() const { + return "embed_mode_" + std::to_string(embed_mode) + "_bytestream_io_" + std::to_string(bytestream_io); + } +}; +class CompileApiTest + : public testing::TestWithParam { + public: + const CompileParam& GetCompileParam() const { + return GetParam(); + } +}; + +TEST_P(CompileApiTest, EmbedMode) { + const auto& test_param = GetCompileParam(); + const std::string test_name = test_param.to_string(); + 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, true); + + Ort::SessionOptions session_options; +#ifdef _WIN32 + /// 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. + RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; + Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); + const OrtEpDevice* const* ep_devices = nullptr; + size_t num_devices = 0; + ASSERT_ORTSTATUS_OK(Ort::GetApi().GetEpDevices(*ort_env, &ep_devices, &num_devices)); + ASSERT_ORTSTATUS_OK(Ort::GetApi().SessionOptionsAppendExecutionProvider_V2(session_options, *ort_env, + ep_devices, 1, nullptr, nullptr, 0)); +#else + session_options.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); +#endif + + 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 + auto status = Ort::CompileModel(*ort_env, model_compile_options); + if (!status.IsOK()) { + std::cerr << status.GetErrorMessage() << std::endl; + } + ASSERT_TRUE(status.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); +} + +INSTANTIATE_TEST_SUITE_P( + NvExecutionProviderTest, CompileApiTest, + ::testing::Values( + CompileParam{true, false}, + CompileParam{false, false}, + CompileParam{true, true}, + CompileParam{false, true}), + [](const testing::TestParamInfo& info) { + return info.param.to_string(); + }); + + +} // namespace test +} // namespace onnxruntime 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 index be1cb4efc3942..379ab21f20ff1 100644 --- 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 @@ -3,7 +3,6 @@ // Licensed under the MIT License. // registration/selection is only supported on windows as there's no device discovery on other platforms -#ifdef _WIN32 #include "test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h" @@ -13,6 +12,7 @@ #include "core/session/onnxruntime_cxx_api.h" #include "test/util/include/api_asserts.h" #include "core/graph/onnx_protobuf.h" +#include "core/graph/model_saving_options.h" #include "test/util/include/scoped_env_vars.h" #include "test/common/trt_op_test_utils.h" #include "test/providers/provider_test_utils.h" @@ -63,7 +63,8 @@ void CreateBaseModel(const PathString& model_name, std::string graph_name, std::vector dims, bool add_fast_gelu, - ONNX_NAMESPACE::TensorProto_DataType dtype) { + 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; @@ -126,7 +127,12 @@ void CreateBaseModel(const PathString& model_name, auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); - status = onnxruntime::Model::Save(model, model_name); + 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()); } @@ -174,5 +180,3 @@ Ort::IoBinding generate_io_binding( } // namespace test } // namespace onnxruntime - -#endif // _WIN32 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 index f07dddd008aa7..eef608fc20eb4 100644 --- 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 @@ -13,6 +13,7 @@ #include #include #include +#include #include "core/graph/constants.h" #include "core/common/path_string.h" @@ -63,7 +64,7 @@ struct Utils { template static void VerifyOutputs(const std::vector& fetches, const std::vector& expected_dims, - const std::vector& expected_values) { + const std::vector& expected_values) { ASSERT_EQ(1, fetches.size()); auto& rtensor = fetches.front().Get(); TensorShape expected_shape(expected_dims); @@ -78,6 +79,7 @@ static void VerifyOutputs(const std::vector& fetches, const std::vecto * \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. @@ -112,7 +114,8 @@ 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); + ONNX_NAMESPACE::TensorProto_DataType dtype = ONNX_NAMESPACE::TensorProto_DataType_FLOAT, + const PathString& external_initializer_file = {}); Ort::IoBinding generate_io_binding( Ort::Session& session, From d5151d731132f0499fabbf60727f5c3a3adf1023 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Fri, 1 Aug 2025 13:30:02 +0200 Subject: [PATCH 14/25] large model test fix large model unit test --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 19 +- .../nv_tensorrt_rtx/nv_ep_context_test.cc | 113 ++++++- .../test_nv_trt_rtx_ep_util.cc | 286 ++++++++++++++++++ .../nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h | 2 + 4 files changed, 406 insertions(+), 14 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 786da87a379c4..7a088b2e6f45f 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -1464,9 +1464,20 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t SetGraphOuterScopeValuesAndInputs(graph_build, graph.GetGraph()); SetAllGraphInputs(graph_build); } - - ORT_ENFORCE(graph_build.Resolve().IsOK()); - + // for (auto& tensor:graph.GetAllInitializedTensors()) { + // if (utils::HasExternalDataInMemory(*tensor.second)) { + // std::unique_ptr full_init; + // ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tensor.second, full_init)); + // (*tensor.second).clear_external_data(); + // (*tensor.second).set_raw_data(full_init->raw_data()); + // } + // } + + auto status = graph_build.Resolve(); + if (!status.IsOK()) { + LOGS_DEFAULT(ERROR) << status.ErrorMessage(); + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "ONNX graph resolve failed: " + status.ErrorMessage())); + } // Add parent graph output to the subgraph int i = 0; std::vector subgraph_outputs; @@ -1521,9 +1532,9 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // save user provided external data in memory instead of writing to ModelProto // needed for models > 2GB std::vector userWeights; + auto allInitializers = graph_viewer->GetAllInitializedTensors(); if (use_external_data_initializer_) { - auto allInitializers = graph_viewer->GetAllInitializedTensors(); for (auto& entry : allInitializers) { auto* tp = entry.second; if (tp->has_raw_data()) { 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 index f9f832096daaf..4b5b9e4be4522 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc @@ -56,16 +56,17 @@ class CompileApiTest } }; -TEST_P(CompileApiTest, EmbedMode) { - const auto& test_param = GetCompileParam(); - const std::string test_name = test_param.to_string(); +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, true); + CreateBaseModel(model_name, graph_name, dims, !fully_supported_model); Ort::SessionOptions session_options; #ifdef _WIN32 @@ -97,11 +98,7 @@ TEST_P(CompileApiTest, EmbedMode) { model_compile_options.SetOutputModelPath(model_name_ctx.c_str()); } // AOT time - auto status = Ort::CompileModel(*ort_env, model_compile_options); - if (!status.IsOK()) { - std::cerr << status.GetErrorMessage() << std::endl; - } - ASSERT_TRUE(status.IsOK()); + ASSERT_TRUE(Ort::CompileModel(*ort_env, model_compile_options).IsOK()); // JIT time Ort::Session session_object{nullptr}; @@ -115,6 +112,103 @@ TEST_P(CompileApiTest, EmbedMode) { 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(); + PathString model_name = path_utils::MakePathString("nv_execution_provider_compile_large.onnx"); + PathString external_data_name = path_utils::MakePathString("nv_execution_provider_compile_large.onnx_data"); + PathString model_name_ctx = path_utils::MakePathString("nv_execution_provider_compile_large_ctx.onnx"); + PathString model_name_ctx_data = path_utils::MakePathString("nv_execution_provider_compile_large_ctx.onnx_data"); + clearFileIfExists(model_name_ctx); + std::string graph_name = "test"; + std::vector dims = {1, 3, 2}; + if (!std::filesystem::exists(model_name) || !std::filesystem::exists(external_data_name)) { + CreateLargeLLMModel(model_name, external_data_name); + } + + Ort::SessionOptions session_options; + std::vector option_keys; + std::vector option_values; + if (test_param.bytestream_io) { + option_keys = {onnxruntime::nv::provider_option_names::kUseExternalDataInitializer}; + option_values = {"0"}; + } + ASSERT_EQ(option_keys.size(), option_values.size()); +#ifdef _WIN32 + /// 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. + RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; + Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); + const OrtEpDevice* const* ep_devices = nullptr; + const OrtEpDevice* const* selected_device = nullptr; + size_t num_devices = 0; + ASSERT_ORTSTATUS_OK(Ort::GetApi().GetEpDevices(*ort_env, &ep_devices, &num_devices)); + for (int i = 0; i < num_devices; i++) { + if (ep_devices[i]->ep_name == kNvTensorRTRTXExecutionProvider) { + selected_device = &ep_devices[i]; + } + } + ASSERT_ORTSTATUS_OK(Ort::GetApi().SessionOptionsAppendExecutionProvider_V2(session_options, *ort_env, selected_device, 1, + option_keys.data(), option_values.data(), option_keys.size())); +#else + std::unordered_map option_map; + for (size_t i = 0; i < option_keys.size(); ++i) { + option_map[option_keys[i]] = option_values[i]; + } + session_options.AppendExecutionProvider(onnxruntime::kNvTensorRTRTXExecutionProvider, option_map); +#endif + + Ort::ModelCompilationOptions model_compile_options(*ort_env, session_options); + // with embed mode == 1 the resulting file will be over the 2GB proto limit + model_compile_options.SetEpContextEmbedMode(0); + + 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( @@ -126,6 +220,5 @@ INSTANTIATE_TEST_SUITE_P( return info.param.to_string(); }); - } // namespace test } // namespace onnxruntime 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 index 379ab21f20ff1..ca9866833d9ca 100644 --- 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 @@ -7,10 +7,12 @@ #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/onnx_protobuf.h" #include "core/graph/model_saving_options.h" #include "test/util/include/scoped_env_vars.h" @@ -136,6 +138,290 @@ void CreateBaseModel(const PathString& 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; + if (dtype == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + std::vector data(total_size); + std::default_random_engine rng; + 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::default_random_engine rng; + 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, 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 index eef608fc20eb4..88c552dac2394 100644 --- 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 @@ -117,6 +117,8 @@ void CreateBaseModel(const PathString& model_name, 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 = {}, From 8de13f9b148edfc724cbdffd1fe9433e143a4b7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Fri, 1 Aug 2025 17:32:40 +0200 Subject: [PATCH 15/25] remove support for weightless --- .../nv_tensorrt_rtx/nv_provider_options.h | 4 ---- .../nv_execution_provider_info.cc | 24 +------------------ 2 files changed, 1 insertion(+), 27 deletions(-) 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 index 620cec3bac594..dc27204017caa 100644 --- a/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h +++ b/include/onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_options.h @@ -32,12 +32,8 @@ 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 = "nv_cuda_graph_enable"; -constexpr const char* kONNXBytestream = "nv_onnx_bytestream"; -constexpr const char* kONNXBytestreamSize = "nv_onnx_bytestream_size"; constexpr const char* kMultiProfileEnable = "nv_multi_profile_enable"; constexpr const char* kUseExternalDataInitializer = "nv_use_external_data_initializer"; -constexpr const char* kExternalDataBytestream = "nv_external_data_bytestream"; -constexpr const char* kExternalDataBytestreamSize = "nv_external_data_bytestream_size"; } // namespace provider_option_names namespace run_option_names { 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 index 090ae1b20e0ea..0658f26e4fdbe 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc @@ -51,24 +51,6 @@ NvExecutionProviderInfo NvExecutionProviderInfo::FromProviderOptions(const Provi .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) - .AddValueParser( - nv::provider_option_names::kONNXBytestream, - [&onnx_bytestream](const std::string& value_str) -> Status { - size_t address; - ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, address)); - onnx_bytestream = reinterpret_cast(address); - return Status::OK(); - }) - .AddAssignmentToReference(nv::provider_option_names::kONNXBytestreamSize, info.onnx_bytestream_size) - .AddValueParser( - nv::provider_option_names::kExternalDataBytestream, - [&external_data_bytestream](const std::string& value_str) -> Status { - size_t address; - ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, address)); - external_data_bytestream = reinterpret_cast(address); - return Status::OK(); - }) - .AddAssignmentToReference(nv::provider_option_names::kExternalDataBytestreamSize, info.external_data_bytestream_size) .Parse(options)); // add new provider option here. info.user_compute_stream = user_compute_stream; @@ -123,11 +105,7 @@ ProviderOptions NvExecutionProviderInfo::ToProviderOptions(const NvExecutionProv {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::kONNXBytestream, MakeStringWithClassicLocale(info.onnx_bytestream)}, - {nv::provider_option_names::kONNXBytestreamSize, MakeStringWithClassicLocale(info.onnx_bytestream_size)}, - {nv::provider_option_names::kUseExternalDataInitializer, MakeStringWithClassicLocale(info.use_external_data_initializer)}, - {nv::provider_option_names::kExternalDataBytestream, MakeStringWithClassicLocale(info.external_data_bytestream)}, - {nv::provider_option_names::kExternalDataBytestreamSize, MakeStringWithClassicLocale(info.external_data_bytestream_size)}, + {nv::provider_option_names::kUseExternalDataInitializer, MakeStringWithClassicLocale(info.use_external_data_initializer)} }; return options; } From e2b67a4eee2ba07244285201e9150b49e2e530d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Tue, 5 Aug 2025 11:36:34 +0200 Subject: [PATCH 16/25] reduce header usages, cleanup and unify usage of windows ifdef remove unused lines revert changes from main in grahp.cc --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 3 +- .../nv_tensorrt_rtx/nv_execution_provider.h | 3 +- .../nv_execution_provider_info.cc | 3 +- .../nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 2 - .../nv_tensorrt_rtx/nv_basic_test.cc | 5 +- .../nv_tensorrt_rtx/nv_ep_context_test.cc | 102 +++++++----------- .../nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.h | 4 +- 7 files changed, 46 insertions(+), 76 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 7a088b2e6f45f..f56b6e5aa7428 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -2688,7 +2688,6 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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)); - weights_.emplace(fused_node.Name(), std::move(userWeights)); 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); @@ -2708,7 +2707,7 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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_, &weights_[context->node_name]}; + auxiliary_streams_, cuda_graph_enable_, is_dynamic_shape_context, cache_prefix_}; *state = p.release(); return 0; }; diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 33984ecf2b553..4ca3cec12b1f8 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -204,6 +204,7 @@ struct TensorrtFuncState { 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; @@ -211,7 +212,6 @@ struct TensorrtFuncState { bool skip_io_binding_allowed = false; // Indicates if input/output binding can be skipped IAllocatorUniquePtr context_memory = nullptr; size_t context_memory_size = 0; - std::unique_ptr> *userWeights = nullptr; }; // Minimum information to construct kernel function state for direct engine load code path @@ -226,6 +226,7 @@ struct TensorrtShortFuncState { 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; 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 index 0658f26e4fdbe..527a37f6c2b57 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider_info.cc @@ -105,8 +105,7 @@ ProviderOptions NvExecutionProviderInfo::ToProviderOptions(const NvExecutionProv {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::kUseExternalDataInitializer, MakeStringWithClassicLocale(info.use_external_data_initializer)}}; return options; } } // 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 index 9e0db4ac377f3..e9420bce880bd 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -110,8 +110,6 @@ Status CreateCtxNode(const GraphViewer& graph_viewer, engine_data_str.assign(engine_data, size); } attr_ep_cache_context->set_s(engine_data_str); - // TODO(maximilianm) we might want to disable this warning as we only support weightless engines that are really small - // the reason we had this was that the field will be hashed and storing a large bytestream has significant overhead } else { attr_ep_cache_context->set_s(engine_cache_path); std::fstream engine_cache_file(engine_cache_path, std::ios::binary | std::ios::out); diff --git a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc index 35a63d0ac6333..2327bc2094d1a 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_basic_test.cc @@ -1,4 +1,3 @@ -// 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/graph/onnx_protobuf.h" @@ -217,7 +216,7 @@ INSTANTIATE_TEST_SUITE_P(NvExecutionProviderTest, TypeTests, ), [](const testing::TestParamInfo& info) { return getTypeAsName(info.param); }); -#if defined(WIN32) +#ifdef _WIN32 static bool SessionHasEp(Ort::Session& session, const char* ep_name) { // Access the underlying InferenceSession. const OrtSession* ort_session = session; @@ -399,7 +398,7 @@ TEST(NvExecutionProviderTest, DataTransfer) { device_tensor = Ort::Value(); } -#endif // defined(WIN32) +#endif } // 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 index 4b5b9e4be4522..3d3ab06d8ac39 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc @@ -1,28 +1,37 @@ -// 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/common/path_utils.h" -#include "core/graph/onnx_protobuf.h" -#include "core/session/inference_session.h" -#include "test/providers/provider_test_utils.h" #include "test/framework/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 +#include -using namespace std; -using namespace ONNX_NAMESPACE; -using namespace ::onnxruntime::logging; 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; +#ifdef _WIN32 + /// 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); +#else + session_options.AppendExecutionProvider(onnxruntime::kNvTensorRTRTXExecutionProvider, option_map); +#endif + return nv_tensorrt_rtx_ep; +} + std::vector readBinaryFile(const PathString& filename) { std::ifstream file(filename, std::ios::binary); if (!file.is_open()) { @@ -69,19 +78,8 @@ void SmallModelTest(CompileParam test_param, bool fully_supported_model) { CreateBaseModel(model_name, graph_name, dims, !fully_supported_model); Ort::SessionOptions session_options; -#ifdef _WIN32 - /// 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. - RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); - const OrtEpDevice* const* ep_devices = nullptr; - size_t num_devices = 0; - ASSERT_ORTSTATUS_OK(Ort::GetApi().GetEpDevices(*ort_env, &ep_devices, &num_devices)); - ASSERT_ORTSTATUS_OK(Ort::GetApi().SessionOptionsAppendExecutionProvider_V2(session_options, *ort_env, - ep_devices, 1, nullptr, nullptr, 0)); -#else - session_options.AppendExecutionProvider(kNvTensorRTRTXExecutionProvider, {}); -#endif + std::unordered_map option_map{}; + auto ep = AppendTrtEtxEP(session_options, option_map); Ort::ModelCompilationOptions model_compile_options(*ort_env, session_options); model_compile_options.SetEpContextEmbedMode(test_param.embed_mode); @@ -124,52 +122,28 @@ TEST_P(CompileApiTest, SmallSplitModel) { TEST_P(CompileApiTest, LargeModel) { const auto& test_param = GetCompileParam(); - PathString model_name = path_utils::MakePathString("nv_execution_provider_compile_large.onnx"); - PathString external_data_name = path_utils::MakePathString("nv_execution_provider_compile_large.onnx_data"); - PathString model_name_ctx = path_utils::MakePathString("nv_execution_provider_compile_large_ctx.onnx"); - PathString model_name_ctx_data = path_utils::MakePathString("nv_execution_provider_compile_large_ctx.onnx_data"); + // 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); - std::string graph_name = "test"; - std::vector dims = {1, 3, 2}; + 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::vector option_keys; - std::vector option_values; - if (test_param.bytestream_io) { - option_keys = {onnxruntime::nv::provider_option_names::kUseExternalDataInitializer}; - option_values = {"0"}; - } - ASSERT_EQ(option_keys.size(), option_values.size()); -#ifdef _WIN32 - /// 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. - RegisteredEpDeviceUniquePtr nv_tensorrt_rtx_ep; - Utils::RegisterAndGetNvTensorRtRtxEp(*ort_env, nv_tensorrt_rtx_ep); - const OrtEpDevice* const* ep_devices = nullptr; - const OrtEpDevice* const* selected_device = nullptr; - size_t num_devices = 0; - ASSERT_ORTSTATUS_OK(Ort::GetApi().GetEpDevices(*ort_env, &ep_devices, &num_devices)); - for (int i = 0; i < num_devices; i++) { - if (ep_devices[i]->ep_name == kNvTensorRTRTXExecutionProvider) { - selected_device = &ep_devices[i]; - } - } - ASSERT_ORTSTATUS_OK(Ort::GetApi().SessionOptionsAppendExecutionProvider_V2(session_options, *ort_env, selected_device, 1, - option_keys.data(), option_values.data(), option_keys.size())); -#else - std::unordered_map option_map; - for (size_t i = 0; i < option_keys.size(); ++i) { - option_map[option_keys[i]] = option_values[i]; - } - session_options.AppendExecutionProvider(onnxruntime::kNvTensorRTRTXExecutionProvider, option_map); -#endif + std::unordered_map option_map{{onnxruntime::nv::provider_option_names::kUseExternalDataInitializer, std::to_string(test_param.bytestream_io)}}; + auto ep = AppendTrtEtxEP(session_options, option_map); Ort::ModelCompilationOptions model_compile_options(*ort_env, session_options); - // with embed mode == 1 the resulting file will be over the 2GB proto limit - model_compile_options.SetEpContextEmbedMode(0); + model_compile_options.SetEpContextEmbedMode(test_param.embed_mode); void* output_context = nullptr; size_t output_context_size = 0; 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 index 88c552dac2394..0f011af8211ca 100644 --- 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 @@ -29,7 +29,7 @@ using RegisteredEpDeviceUniquePtr = std::unique_ptr> converter; return converter.to_bytes(path); #else From 3728ce0c929c33200177fbf9f2e42a7148cb70be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Fri, 8 Aug 2025 12:54:59 +0200 Subject: [PATCH 17/25] address review comments --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index f56b6e5aa7428..c8ef85df928ee 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -1773,19 +1773,19 @@ NvExecutionProvider::GetCapability(const GraphViewer& graph, // 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. - size_t node_idx = 0; - if (GraphHasCtxNode(graph, node_idx)) { - 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) { - 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))); - } + 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) { + 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; } @@ -2061,9 +2061,11 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, // Extract weight information from the Refitter. int required_weights = refitter->getAllWeights(0, nullptr); - std::vector refit_names(required_weights); - refitter->getAllWeights(required_weights, refit_names.data()); + 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; @@ -2094,8 +2096,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, 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(); - bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); - if (weight_is_refittable) { + 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. From 80574d29b82c68c1e566658fa8d0e7bd24d236d9 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Tue, 12 Aug 2025 13:13:45 +0530 Subject: [PATCH 18/25] fix engine cache path with EP context --- .../core/providers/nv_tensorrt_rtx/nv_execution_provider.cc | 2 -- .../core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc | 3 ++- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index c8ef85df928ee..82aef341484e0 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -2611,8 +2611,6 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr cache_path = GetCachePath(cache_path_, fused_node.Name()) + ".engine"; ; } - auto cache_file_name = std::filesystem::path(cache_path).filename(); - cache_path = std::filesystem::path(engine_cache_relative_path_to_context_model_dir).append(cache_file_name.string()).string(); // 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_) { 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 index e9420bce880bd..1f34a0f25877d 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/onnx_ctx_model_helper.cc @@ -111,7 +111,8 @@ Status CreateCtxNode(const GraphViewer& graph_viewer, } attr_ep_cache_context->set_s(engine_data_str); } else { - attr_ep_cache_context->set_s(engine_cache_path); + std::string engine_cache_filename = std::filesystem::path(engine_cache_path).filename().string(); + 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); From 3996d9b5f34b2c5fc9f255ebd164ed15110b898c Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Tue, 12 Aug 2025 17:36:05 +0530 Subject: [PATCH 19/25] fix unit test to add seed for random tensors --- .../test/providers/nv_tensorrt_rtx/test_nv_trt_rtx_ep_util.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index ca9866833d9ca..17182ab032f7a 100644 --- 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 @@ -151,15 +151,15 @@ ONNX_NAMESPACE::TensorProto CreateLargeWeight( // 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::default_random_engine rng; 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::default_random_engine rng; 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)); From a4f8c4540f9f62f4885fd66022b15a854120d561 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 13 Aug 2025 10:34:22 +0530 Subject: [PATCH 20/25] support sm86 and onwards RTX devices --- .../core/providers/nv_tensorrt_rtx/nv_execution_provider.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 82aef341484e0..945a3cd2d797f 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -836,9 +836,9 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) cudaDeviceProp prop; CUDA_CALL_THROW(cudaGetDeviceProperties(&prop, device_id_)); - if (prop.major < 8 || prop.major == 9 || prop.major == 10) { + if (prop.major < 8 || (prop.major == 8 && prop.minor < 6) || prop.major == 9 || prop.major == 10) { ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[NvTensorRTRTX EP] The execution provider only supports RTX devices with compute capabilities =< 80.")); + "[NvTensorRTRTX EP] The execution provider only supports RTX devices with compute capabilities >= 86.")); } compute_capability_ = GetComputeCapability(prop); if (info.has_user_compute_stream) { From b0bab1c1fb6c616ecd720831634bedd8898b7e37 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Wed, 13 Aug 2025 14:49:40 +0530 Subject: [PATCH 21/25] update cc check --- .../core/providers/nv_tensorrt_rtx/nv_execution_provider.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 945a3cd2d797f..5832163f8ebc1 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -836,9 +836,10 @@ NvExecutionProvider::NvExecutionProvider(const NvExecutionProviderInfo& info) cudaDeviceProp prop; CUDA_CALL_THROW(cudaGetDeviceProperties(&prop, device_id_)); - if (prop.major < 8 || (prop.major == 8 && prop.minor < 6) || prop.major == 9 || prop.major == 10) { + 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.")); + "[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) { From f762d86094211001a002bf8e1e4294294478f550 Mon Sep 17 00:00:00 2001 From: Vishal Agarwal Date: Fri, 15 Aug 2025 23:56:33 +0530 Subject: [PATCH 22/25] fix lint --- .../core/providers/nv_tensorrt_rtx/nv_execution_provider.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 5832163f8ebc1..f01e7e89f7393 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -2066,7 +2066,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, 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())); + std::make_move_iterator(refit_names_prealocated.end())); // Vectors to keep track of data pointers. std::vector names; From b81a6d547b20fe045882122853d29928d35e666e Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Sun, 17 Aug 2025 14:31:08 +0200 Subject: [PATCH 23/25] do not copy memory to EP owned memory for raw initializers --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 49 +++++++++++-------- .../nv_tensorrt_rtx/nv_execution_provider.h | 21 ++++++-- 2 files changed, 46 insertions(+), 24 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index f01e7e89f7393..358440943dc16 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -1156,6 +1156,9 @@ nvinfer1::IBuilder* NvExecutionProvider::GetBuilder(TensorrtLogger& trt_logger) { 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(); @@ -1465,14 +1468,6 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t SetGraphOuterScopeValuesAndInputs(graph_build, graph.GetGraph()); SetAllGraphInputs(graph_build); } - // for (auto& tensor:graph.GetAllInitializedTensors()) { - // if (utils::HasExternalDataInMemory(*tensor.second)) { - // std::unique_ptr full_init; - // ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tensor.second, full_init)); - // (*tensor.second).clear_external_data(); - // (*tensor.second).set_raw_data(full_init->raw_data()); - // } - // } auto status = graph_build.Resolve(); if (!status.IsOK()) { @@ -1533,17 +1528,18 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // save user provided external data in memory instead of writing to ModelProto // needed for models > 2GB std::vector userWeights; - auto allInitializers = graph_viewer->GetAllInitializedTensors(); - if (use_external_data_initializer_) { + const InitializedTensorSet& allInitializers = graph_viewer->GetAllInitializedTensors(); + userWeights.reserve(allInitializers.size()); for (auto& entry : allInitializers) { auto* tp = entry.second; - if (tp->has_raw_data()) { - userWeights.emplace_back(tp->name(), tp->raw_data()); + if (utils::HasRawData(*tp)) { + userWeights.emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data().data(), tp->raw_data().size())); } else if (utils::HasExternalDataInMemory(*tp)) { + // TODO(maximilianm) remove this memory copy inside `GetTensorProtoWithDataIfInMemory` by using https://github.com/microsoft/onnxruntime/pull/25761 std::unique_ptr full_init; ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tp, full_init)); - userWeights.emplace_back(full_init->name(), full_init->raw_data()); + userWeights.emplace_back(std::move(full_init->name()), std::move(full_init->raw_data())); } } } @@ -1575,11 +1571,15 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t 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_); } @@ -2049,6 +2049,7 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, // 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, @@ -2156,6 +2157,9 @@ common::Status NvExecutionProvider::RefitEngine(std::string onnx_model_filename, "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. @@ -2334,18 +2338,19 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr auto model_proto = model->ToProto(); // exclude weights if external - auto userWeights = std::make_unique>(); - + std::vector userWeights; if (use_external_data_initializer_) { - auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); + const InitializedTensorSet& allInitializers = graph_body_viewer.GetAllInitializedTensors(); + userWeights.reserve(allInitializers.size()); for (auto& entry : allInitializers) { auto* tp = entry.second; - if (tp->has_raw_data()) { - userWeights->emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data())); + if (utils::HasRawData(*tp)) { + userWeights.emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data().data(), tp->raw_data().size())); } else if (utils::HasExternalDataInMemory(*tp)) { + // TODO(maximilianm) remove this memory copy inside `GetTensorProtoWithDataIfInMemory` by using https://github.com/microsoft/onnxruntime/pull/25761 std::unique_ptr full_init; ORT_THROW_IF_ERROR(utils::GetTensorProtoWithDataIfInMemory(*tp, full_init)); - userWeights->emplace_back(TensorrtUserWeights(full_init->name(), full_init->raw_data())); + userWeights.emplace_back(TensorrtUserWeights(std::move(full_init->name()), std::move(full_init->raw_data()))); } } } @@ -2373,11 +2378,15 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr 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) { + 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_); } diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index 4ca3cec12b1f8..abbb9149b67d7 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -156,23 +156,36 @@ struct TensorParams { // 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_(data) {}; + 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 { - return static_cast(data_.data()); + if (!data_cpy_.empty()) { + return data_cpy_.data(); + } + return data_; } int64_t Size() const { - return static_cast(data_.size()); + if (!data_cpy_.empty()) { + return static_cast(data_cpy_.size()); + } + return static_cast(size_); } private: std::string name_{}; - std::string data_{}; + std::string data_cpy_{}; + void const* data_; + size_t size_; }; // Information to construct kernel function state. From d0926f8a09e056074a2aff332ef91d4c15be862a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Thu, 21 Aug 2025 00:20:30 +0200 Subject: [PATCH 24/25] use ort values i data is already loaded in memory --- .../nv_tensorrt_rtx/nv_execution_provider.cc | 23 +++++++++++++++++-- .../nv_tensorrt_rtx/nv_ep_context_test.cc | 16 +++++++++---- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc index 358440943dc16..4619faddba150 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.cc @@ -11,6 +11,7 @@ #include "core/common/common.h" #include "core/common/narrow.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" @@ -1529,14 +1530,23 @@ SubGraphCollection_t NvExecutionProvider::GetSupportedList(SubGraphCollection_t // needed for models > 2GB std::vector userWeights; if (use_external_data_initializer_) { + auto c_api = Ort::GetApi(); const InitializedTensorSet& allInitializers = graph_viewer->GetAllInitializedTensors(); userWeights.reserve(allInitializers.size()); for (auto& entry : allInitializers) { + OrtValue initializer_value; auto* tp = entry.second; if (utils::HasRawData(*tp)) { userWeights.emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data().data(), tp->raw_data().size())); + } else if (graph_viewer->GetOrtValueInitializer(tp->name(), initializer_value)) { + // the initializer was marked as external data by the ORT graph at load time since it was provided in memory + size_t size = 0; + const void* ptr = nullptr; + c_api.GetTensorSizeInBytes(&initializer_value, &size); + c_api.GetTensorData(&initializer_value, &ptr); + userWeights.emplace_back(tp->name(), ptr, size); } else if (utils::HasExternalDataInMemory(*tp)) { - // TODO(maximilianm) remove this memory copy inside `GetTensorProtoWithDataIfInMemory` by using https://github.com/microsoft/onnxruntime/pull/25761 + // 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(std::move(full_init->name()), std::move(full_init->raw_data())); @@ -2340,14 +2350,23 @@ Status NvExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphViewer& gr // exclude weights if external std::vector userWeights; if (use_external_data_initializer_) { + auto c_api = Ort::GetApi(); const InitializedTensorSet& allInitializers = graph_body_viewer.GetAllInitializedTensors(); userWeights.reserve(allInitializers.size()); for (auto& entry : allInitializers) { + OrtValue initializer_value; auto* tp = entry.second; if (utils::HasRawData(*tp)) { userWeights.emplace_back(TensorrtUserWeights(tp->name(), tp->raw_data().data(), tp->raw_data().size())); + } else if (graph_body_viewer.GetOrtValueInitializer(tp->name(), initializer_value)) { + // the initializer was marked as external data by the ORT graph at load time since it was provided in memory + size_t size = 0; + const void* ptr = nullptr; + c_api.GetTensorSizeInBytes(&initializer_value, &size); + c_api.GetTensorData(&initializer_value, &ptr); + userWeights.emplace_back(tp->name(), ptr, size); } else if (utils::HasExternalDataInMemory(*tp)) { - // TODO(maximilianm) remove this memory copy inside `GetTensorProtoWithDataIfInMemory` by using https://github.com/microsoft/onnxruntime/pull/25761 + // 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(TensorrtUserWeights(std::move(full_init->name()), std::move(full_init->raw_data()))); 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 index 3d3ab06d8ac39..ce49ae81c81c0 100644 --- a/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc +++ b/onnxruntime/test/providers/nv_tensorrt_rtx/nv_ep_context_test.cc @@ -53,8 +53,10 @@ std::vector readBinaryFile(const PathString& filename) { 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); + 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 @@ -78,7 +80,8 @@ void SmallModelTest(CompileParam test_param, bool fully_supported_model) { CreateBaseModel(model_name, graph_name, dims, !fully_supported_model); Ort::SessionOptions session_options; - std::unordered_map option_map{}; + 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); @@ -139,7 +142,9 @@ TEST_P(CompileApiTest, LargeModel) { } Ort::SessionOptions session_options; - std::unordered_map option_map{{onnxruntime::nv::provider_option_names::kUseExternalDataInitializer, std::to_string(test_param.bytestream_io)}}; + 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); @@ -189,7 +194,10 @@ INSTANTIATE_TEST_SUITE_P( CompileParam{true, false}, CompileParam{false, false}, CompileParam{true, true}, - CompileParam{false, 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(); }); From f6404304f773e69446f634a052f137b5470521d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Thu, 21 Aug 2025 00:26:20 +0200 Subject: [PATCH 25/25] remove unused var --- .../core/providers/nv_tensorrt_rtx/nv_execution_provider.h | 1 - 1 file changed, 1 deletion(-) diff --git a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h index abbb9149b67d7..e3dd38eb837ff 100644 --- a/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h +++ b/onnxruntime/core/providers/nv_tensorrt_rtx/nv_execution_provider.h @@ -397,7 +397,6 @@ class NvExecutionProvider : public IExecutionProvider { 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_; - std::unordered_map>> weights_; // User provided weights // for external stream, we need to create its cudnn/cublass handle before cuda EP enable cuda graph capture cudnnHandle_t external_cudnn_handle_ = nullptr;