From 00cecb73d57efbfa545a63af39023ddf6a1351b3 Mon Sep 17 00:00:00 2001 From: Jingyan Wang Date: Sat, 21 Dec 2024 06:58:22 +0000 Subject: [PATCH] Support EP context partition --- .../tensorrt/onnx_ctx_model_helper.cc | 271 +++++++++- .../tensorrt/onnx_ctx_model_helper.h | 20 +- .../tensorrt/tensorrt_execution_provider.cc | 484 ++++++++++++++++-- .../tensorrt/tensorrt_execution_provider.h | 70 +++ 4 files changed, 807 insertions(+), 38 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc index ef45d6c85d..427a197e7c 100644 --- a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc @@ -20,8 +20,16 @@ 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) { + // LOGS_DEFAULT(VERBOSE) << "*#* graph_viewer.MaxNodeIndex()=" << graph_viewer.MaxNodeIndex(); for (int i = 0; i < graph_viewer.MaxNodeIndex(); ++i) { auto node = graph_viewer.GetNode(i); + // if (!node) { + // LOGS_DEFAULT(VERBOSE) << "#*# Node at index " << i << " is null!"; + // continue; + // } + // if (!node->Name().empty()) { + // LOGS_DEFAULT(VERBOSE) << "*#* node->Name()=" << node->Name() << " node->OpType()=" << node->OpType(); + // } if (node != nullptr && node->OpType() == EPCONTEXT_OP) { return true; } @@ -29,6 +37,26 @@ bool GraphHasCtxNode(const GraphViewer& graph_viewer) { return false; } +int FindCtxNodeInGraph(const GraphViewer& graph_viewer) { +// Assumes there's only 1 context node in this subgraph (graph_viewer) + // LOGS_DEFAULT(VERBOSE) << "*#* graph_viewer.MaxNodeIndex()=" << graph_viewer.MaxNodeIndex(); + for (int i = 0; i < graph_viewer.MaxNodeIndex(); ++i) { + auto node = graph_viewer.GetNode(i); + // if (!node) { + // LOGS_DEFAULT(VERBOSE) << "#*# Node at index " << i << " is null!"; + // continue; + // } + // if (!node->Name().empty()) { + // LOGS_DEFAULT(VERBOSE) << "*#* node->Name()=" << node->Name() << " node->OpType()=" << node->OpType(); + // } + if (node != nullptr && node->OpType() == EPCONTEXT_OP) { + LOGS_DEFAULT(VERBOSE) << "*#* context node found at index=" << i; + return i; + } + } + return -1; +} + const std::filesystem::path& GetModelPath(const GraphViewer& graph_viewer) { // find the top level graph const Graph* cur_graph = &graph_viewer.GetGraph(); @@ -64,6 +92,10 @@ void UpdateCtxNodeModelEngineContext(ONNX_NAMESPACE::ModelProto* model_proto, /* * Create "EP context node" model where engine information is embedded */ +// ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, +// std::unique_ptr CreateCtxModel(const GraphViewer& graph_viewer, +// std::unique_ptr CreateCtxModel(const GraphViewer& graph_viewer, +// Status CreateCtxModel(const GraphViewer& graph_viewer, ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string engine_cache_path, char* engine_data, @@ -136,6 +168,93 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, return model_proto.release(); } +std::unique_ptr CreateCtxModel2(const GraphViewer& graph_viewer, + const std::string fused_subgraph_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) { + LOGS_DEFAULT(VERBOSE) << "*#* In CreateCtxModel2"; + auto model_build = graph_viewer.CreateModel(*logger); + auto& graph_build = model_build->MainGraph(); + + // Get graph inputs and outputs + std::vector inputs, outputs; + for (auto input : graph_viewer.GetInputs()) { + auto& n_input = graph_build.GetOrCreateNodeArg(input->Name(), input->TypeAsProto()); + inputs.push_back(&n_input); + } + + for (auto output : graph_viewer.GetOutputs()) { + auto& n_output = graph_build.GetOrCreateNodeArg(output->Name(), output->TypeAsProto()); + outputs.push_back(&n_output); + } + + // Create EP context node attributes + auto attr_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 + 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); + if (embed_mode) { + if (size > 0) { + engine_data_str.assign(engine_data, size); + } + attr_1->set_s(engine_data_str); + LOGS_DEFAULT(WARNING) << EPCONTEXT_WARNING; + } else { + attr_1->set_s(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()); + + 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); + + // Create EP context node + //graph_build.AddNode(EPCONTEXT_OP, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); + LOGS_DEFAULT(VERBOSE) << "*#* fused_subgraph_name=" << fused_subgraph_name; + graph_build.AddNode(fused_subgraph_name, EPCONTEXT_OP, "", inputs, outputs, node_attributes.get(), EPCONTEXT_OP_DOMAIN); + LOGS_DEFAULT(VERBOSE) << "*#* graph_build.GetNode(0)->Name()" << graph_build.GetNode(0)->Name(); + ORT_ENFORCE(graph_build.Resolve().IsOK()); + + // Serialize modelproto to string + // auto new_graph_viewer = graph_build.CreateGraphViewer(); + // std::unique_ptr model = new_graph_viewer->CreateModel(*logger); + // 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 std::unique_ptr(graph_build); + // return std::make_unique(graph_build); + // return std::make_unique(std::move(graph_build)); + // return std::move(graph_build); + // trt_ep_context_models.emplace("node name", std::move(model_build)); + // trt_ep_context_models.emplace_back(std::move(model_build)); + // return Status::OK(); + // return std::unique_ptr(model_build); + // return std::move(model_build); // Transfer ownership + return model_build; +} + /* * Return the directory where the ep context model locates */ @@ -266,11 +385,11 @@ bool IsWeightStrippedEngineCache(std::filesystem::path& engine_cache_path) { return engine_cache_path.stem().extension().string() == ".stripped"; } -Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph_viewer) { - if (!ValidateEPCtxNode(graph_viewer)) { +Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph_viewer, const int ctx_node_idx) { + if (!ValidateEPCtxNode(graph_viewer, ctx_node_idx)) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "It's not a valid EP Context node"); } - auto node = graph_viewer.GetNode(0); + auto node = graph_viewer.GetNode(ctx_node_idx); auto& attrs = node->GetAttributes(); const int64_t embed_mode = attrs.at(EMBED_MODE).i(); @@ -380,14 +499,26 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph /* * The sanity check for EP context contrib op. */ -bool TensorRTCacheModelHandler::ValidateEPCtxNode(const GraphViewer& graph_viewer) { +bool TensorRTCacheModelHandler::ValidateEPCtxNode(const GraphViewer& graph_viewer, const int ctx_node_idx) { assert(graph_viewer.NumberOfNodes() == 1); - assert(graph_viewer.GetNode(0)->OpType() == EPCONTEXT_OP); - auto node = graph_viewer.GetNode(0); + assert(graph_viewer.GetNode(ctx_node_idx)->OpType() == EPCONTEXT_OP); + auto node = graph_viewer.GetNode(ctx_node_idx); + LOGS_DEFAULT(VERBOSE) << "*#* node->Name()=" << node->Name(); auto& attrs = node->GetAttributes(); + // print node info + if (!node) { + LOGS_DEFAULT(VERBOSE) << "*#* node is null"; + return false; + } + LOGS_DEFAULT(VERBOSE) << (node && !node->Name().empty() + ? "*#* Node Name: " + node->Name() + : "*#* Node has empty name."); + LOGS_DEFAULT(VERBOSE) << (node && !node->OpType().empty() + ? "*#* Node Name: " + node->OpType() + : "*#* Node has empty OpType."); // Show the warning if compute capability is not matched - if (attrs.count(COMPUTE_CAPABILITY) > 0) { + if (attrs.find(COMPUTE_CAPABILITY)!=attrs.end() && attrs.count(COMPUTE_CAPABILITY) > 0) { std::string model_compute_capability = attrs.at(COMPUTE_CAPABILITY).s(); // Verify if engine was compiled with ampere+ hardware compatibility enabled if (model_compute_capability == "80+") { @@ -414,4 +545,130 @@ bool TensorRTCacheModelHandler::ValidateEPCtxNode(const GraphViewer& graph_viewe return true; } + +// Status SetGraphInputOutputInfo(const GraphViewer& graph_viewer, +// const onnxruntime::Node& fused_node, +// const logging::Logger& logger) { +// auto graph_initializers = graph_viewer.GetAllInitializedTensors(); +// for (auto graph_ini : graph_initializers) { +// initializer_inputs_.emplace(graph_ini.first); +// } +// auto input_defs = fused_node.InputDefs(); + +// ORT_RETURN_IF_ERROR(ParseGraphInputOrOutput(input_defs, input_names_, inputs_info_, +// model_input_index_map_, logger, true)); + +// auto output_defs = fused_node.OutputDefs(); +// ORT_RETURN_IF_ERROR(ParseGraphInputOrOutput(output_defs, output_names_, outputs_info_, +// model_output_index_map_, logger)); + +// return Status::OK(); +// } + +// Status SetupTrtInputOutput(const logging::Logger& logger) { +// LOGS(logger, VERBOSE) << "Setting up QNN input/output for graph: " << graph_info_->Name(); + +// auto result = SetupTensors(qnn_input_infos_, graph_info_->InputTensors()); + +// if (Status::OK() != result) { +// LOGS(logger, ERROR) << "Failed to setup QNN input output tensors for graph: " << graph_info_->Name(); +// return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to setup QNN input tensors!"); +// } + +// result = SetupTensors(qnn_output_infos_, graph_info_->OutputTensors(), false); +// if (Status::OK() != result) { +// LOGS(logger, ERROR) << "Failed to setup QNN input output tensors for graph: " << graph_info_->Name(); +// return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to setup QNN output tensors!"); +// } + +// return Status::OK(); +// } + +// Status LoadCachedTrtContextFromBuffer(char* buffer, uint64_t buffer_length, +// std::string node_name, +// std::unordered_map> trt_models) { +// LOGS_DEFAULT(VERBOSE) << buffer; +// return Status::OK(); +// } + +// Status GetEpContextFromMainNode(const onnxruntime::Node& main_context_node, +// const onnxruntime::PathString& ctx_onnx_model_path, +// std::unordered_map>& trt_models) { +// ORT_RETURN_IF_NOT(EPCONTEXT_OP == main_context_node.OpType(), "Should only filter in the EPContext node."); +// NodeAttrHelper node_helper(main_context_node); +// bool is_embed_mode = node_helper.Get(EMBED_MODE, true); +// if (is_embed_mode) { +// const std::string& context_binary = node_helper.Get(EP_CACHE_CONTEXT, ""); +// return LoadCachedTrtContextFromBuffer(const_cast(context_binary.c_str()), +// static_cast(context_binary.length()), +// main_context_node.Name(), +// trt_models); +// } + +// std::filesystem::path folder_path = std::filesystem::path(ctx_onnx_model_path).parent_path(); +// std::string external_trt_ctx_binary_file_name = node_helper.Get(EP_CACHE_CONTEXT, ""); +// ORT_RETURN_IF(external_trt_ctx_binary_file_name.empty(), "The file path in ep_cache_context should not be empty."); +// #ifdef _WIN32 +// onnxruntime::PathString external_qnn_context_binary_path = onnxruntime::ToPathString(external_qnn_ctx_binary_file_name); +// auto ctx_file_path = std::filesystem::path(external_qnn_context_binary_path.c_str()); +// ORT_RETURN_IF(ctx_file_path.is_absolute(), "External mode should set ep_cache_context field with a relative path, but it is an absolute path: ", +// external_qnn_ctx_binary_file_name); +// auto relative_path = ctx_file_path.lexically_normal().make_preferred().wstring(); +// if (relative_path.find(L"..", 0) != std::string::npos) { +// return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, "The file path in ep_cache_context field has '..'. It's not allowed to point outside the directory."); +// } + +// std::filesystem::path context_binary_path = folder_path.append(relative_path); +// #else +// ORT_RETURN_IF(external_trt_ctx_binary_file_name[0] == '/', +// "External mode should set ep_cache_context field with a relative path, but it is an absolute path: ", +// external_trt_ctx_binary_file_name); +// if (external_trt_ctx_binary_file_name.find("..", 0) != std::string::npos) { +// return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, "The file path in ep_cache_context field has '..'. It's not allowed to point outside the directory."); +// } +// std::filesystem::path context_binary_path = folder_path.append(external_trt_ctx_binary_file_name); +// std::string file_full_path = context_binary_path.string(); +// #endif +// if (!std::filesystem::is_regular_file(context_binary_path)) { +// return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, "The file path in ep_cache_context does not exist or is not accessible."); +// } + +// size_t buffer_size{0}; +// std::ifstream cache_file(context_binary_path.string().c_str(), std::ifstream::binary); +// ORT_RETURN_IF(!cache_file || !cache_file.good(), "Failed to open cache file."); + +// cache_file.seekg(0, cache_file.end); +// buffer_size = static_cast(cache_file.tellg()); +// ORT_RETURN_IF(0 == buffer_size, "Empty cache file encountered."); + +// cache_file.seekg(0, cache_file.beg); +// std::unique_ptr buffer = std::make_unique(buffer_size); +// ORT_RETURN_IF(nullptr == buffer, "Failed to allocate memory for cache file."); +// // Load file into buffer +// const auto& read_result = cache_file.read(buffer.get(), buffer_size); +// ORT_RETURN_IF(!read_result, "Failed to read contents from cached context file."); +// cache_file.close(); +// return LoadCachedTrtContextFromBuffer(const_cast(context_binary.c_str()), +// static_cast(context_binary.length()), +// main_context_node.Name(), +// trt_models); +// } + +// Status LoadTrtCtxFromOnnxGraph(const onnxruntime::GraphViewer& graph_viewer, +// const onnxruntime::PathString& ctx_onnx_model_path, +// std::unordered_map>& trt_models, +// const logging::Logger& logger) { +// ORT_RETURN_IF(graph_viewer.NumberOfNodes() != 1, "One filtered graph should has only one EPContext node!"); +// Status status = GetEpContextFromMainNode(*graph_viewer.Nodes().begin(), ctx_onnx_model_path, +// trt_models); + +// // This is the protocol with customer that status with INVALID_GRAPH will be generated if failed to load context model +// if (!status.IsOK()) { +// LOGS(logger, ERROR) << "Failed to load from EpContext model. " << status.ErrorMessage(); +// return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH, "Failed to load from EpContext model. ", status.ErrorMessage()); +// } + +// return Status::OK(); +// } + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h index 3af0143cbf..e652ecdf0b 100644 --- a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h @@ -24,8 +24,14 @@ static const std::string EPCONTEXT_WARNING = for the best model loading time"; bool GraphHasCtxNode(const GraphViewer& graph_viewer); +int FindCtxNodeInGraph(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 CreateCtxModel(const GraphViewer& graph_viewer, +// std::unique_ptr CreateCtxModel(const GraphViewer& graph_viewer, +// Status CreateCtxModel(const GraphViewer& graph_viewer, ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string engine_cache_path, char* engine_data, @@ -34,6 +40,16 @@ ONNX_NAMESPACE::ModelProto* CreateCtxModel(const GraphViewer& graph_viewer, const std::string compute_capability, const std::string onnx_model_path, const logging::Logger* logger); + +std::unique_ptr CreateCtxModel2(const GraphViewer& graph_viewer, + const std::string fused_subgraph_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); std::string GetCtxModelPath(const std::string& ep_context_file_path, const std::string& original_model_path); bool IsAbsolutePath(const std::string& path_string); @@ -67,9 +83,9 @@ class TensorRTCacheModelHandler { } ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TensorRTCacheModelHandler); - bool ValidateEPCtxNode(const GraphViewer& graph_viewer); + bool ValidateEPCtxNode(const GraphViewer& graph_viewer, const int ctx_node_idx); - Status GetEpContextFromGraph(const GraphViewer& graph_viewer); + Status GetEpContextFromGraph(const GraphViewer& graph_viewer, const int ctx_node_idx); private: std::unique_ptr* trt_engine_; diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 1a5cf6abab..6df42d5959 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -1344,6 +1344,7 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv timing_cache_enable_ = info.timing_cache_enable; force_timing_cache_match_ = info.force_timing_cache; detailed_build_log_ = info.detailed_build_log; + // EPCTX TODO combine session option with trt options dump_ep_context_model_ = info.dump_ep_context_model; ep_context_file_path_ = info.ep_context_file_path; ep_context_embed_mode_ = info.ep_context_embed_mode; @@ -1952,6 +1953,13 @@ std::unique_ptr TensorrtExecutionProvider::GetSubGraph(SubGraph node_set.insert(node_index[index]); } + std::ostringstream oss; + oss << "*#* node_set="; + for (const auto& value : node_set) { + oss << value << " "; + } + LOGS_DEFAULT(VERBOSE) << oss.str(); + // Get parent graph output names std::unordered_set graph_output_names; for (const auto* output_arg : graph.GetOutputs()) { @@ -2467,15 +2475,52 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, strncpy(model_path_, path_string.c_str(), sizeof(model_path_) - 1); #endif model_path_[sizeof(model_path_) - 1] = '\0'; + LOGS_DEFAULT(VERBOSE) << "*#* graph.NumberOfNodes()=" << graph.NumberOfNodes(); + LOGS_DEFAULT(VERBOSE) << "*#* TRTGenerateId(graph)" << TRTGenerateId(graph); // 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), 0); - result.push_back(ComputeCapability::Create(std::move(sub_graph))); - return result; + if (GraphHasCtxNode(graph)) { + if (graph.NumberOfNodes() == 1) { + SubGraph_t supported_node_vector = {{0}, true}; + std::unique_ptr sub_graph = GetSubGraph(supported_node_vector, graph, TRTGenerateId(graph), 0); + result.push_back(ComputeCapability::Create(std::move(sub_graph))); + return result; + } else { + SubGraphCollection_t supported_node_vectors = { + {{0}, true}, + {{1}, true}, + {{2}, true}, + {{3}, true}, + {{4}, true}, + }; + + for (auto supported_node_vector: supported_node_vectors) { + auto subgraph_idx = supported_node_vector.first[0]; + std::unique_ptr sub_graph = GetSubGraph(supported_node_vector, graph, TRTGenerateId(graph), subgraph_idx); + + // Print supported_node_vector of std::pair, bool> + std::ostringstream oss; + oss << "First: ["; + for (auto item: supported_node_vector.first) { + oss << item << ", "; + } + LOGS_DEFAULT(VERBOSE) << "*#* supported_node_vector=" << oss.str() << "], Second (bool): " << supported_node_vector.second; + oss.str(""); + // Print sub_graph nodes of type std::unique_ptr -> Nodes() + const auto nodes = sub_graph->Nodes(); //std::vector Nodes() + for (const auto node: nodes) { + oss << node << ", "; + } + // Log the formatted string + LOGS_DEFAULT(VERBOSE) << "*#* sub_graph: Nodes=[" << oss.str() << "]"; + + result.push_back(ComputeCapability::Create(std::move(sub_graph))); + } + return result; + } + } // Generate unique kernel name for TRT graph @@ -2506,6 +2551,13 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, */ for (const auto& index : nodes_vector) { const auto& node = graph.GetNode(node_index[index]); + if (!node) { + LOGS_DEFAULT(VERBOSE) << "*#* Node at index " << node_index[index] << " is null!"; + continue; + } + if (!node->Name().empty()) { + LOGS_DEFAULT(VERBOSE) << "*#* node->Name()=" << node->Name() << " node->OpType()=" << node->OpType(); + } bool supported_node = true; /* If current node is control flow op, we take different approach based on following four cases: @@ -2541,9 +2593,20 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, if (exclude_set.find(node->OpType()) != exclude_set.end()) { supported_node = false; } + // supported_node = supported_node || GraphHasCtxNode(node); + LOGS_DEFAULT(VERBOSE) << "*#* supported_node=" << supported_node; + bool is_context_node = node->OpType() == "EPContext"; + LOGS_DEFAULT(VERBOSE) << "*#* Node " << index << " is_context_node=" << is_context_node; + + // std::vector, bool> > if (supported_node) { - if (new_subgraph) { + if (is_context_node) { + parser_nodes_vector.emplace_back(); + // Mark Context node is supported + parser_nodes_vector.back().second = true; + new_subgraph = true; + } else if (new_subgraph) { parser_nodes_vector.emplace_back(); // Mark all new graphs as "UnKnown" which will later be parsed by TRT parser parser_nodes_vector.back().second = false; @@ -2556,7 +2619,42 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, } bool early_termination = false; + // print parser_nodes_vector + std::ostringstream oss; + for (const auto& pair : parser_nodes_vector) { + oss << "[Vector: ["; + for (size_t i = 0; i < pair.first.size(); ++i) { + oss << pair.first[i]; + if (i < pair.first.size() - 1) { + oss << ", "; + } + } + oss << "]"; + + // Print the second element (bool) + oss << ", " << (pair.second ? "true" : "false") << "]"; + } + LOGS_DEFAULT(VERBOSE) << "*#* print parser_nodes_vector" << oss.str(); // std::vector, bool> > + supported_nodes_vector = GetSupportedList(parser_nodes_vector, 0, max_partition_iterations_, graph, &early_termination); + + // Print supported_nodes_vector + for (const auto& pair : supported_nodes_vector) { + // Print the first element (std::vector) + oss << "[Vector: ["; + for (size_t i = 0; i < pair.first.size(); ++i) { + oss << pair.first[i]; + if (i < pair.first.size() - 1) { + oss << ", "; + } + } + oss << "]"; + + // Print the second element (bool) + oss << ", " << (pair.second ? "true" : "false") << "]"; + } + LOGS_DEFAULT(VERBOSE) << "*#* print supported_nodes_vector" << oss.str(); // std::vector, bool> > + if (early_termination) { supported_nodes_vector.clear(); } @@ -2589,6 +2687,23 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, } } + // Print supported_nodes_vector + for (const auto& pair : supported_nodes_vector) { + // Print the first element (std::vector) + oss << "[Vector: ["; + for (size_t i = 0; i < pair.first.size(); ++i) { + oss << pair.first[i]; + if (i < pair.first.size() - 1) { + oss << ", "; + } + } + oss << "]"; + + // Print the second element (bool) + oss << ", " << (pair.second ? "true" : "false") << "]"; + } + LOGS_DEFAULT(VERBOSE) << "*#* After consolidation, print supported_nodes_vector" << oss.str(); // std::vector, bool> > + // Handle the case where the graph is subgraph of control flow op. // The purpose is to make control flow op as well as its subgraphs run on TRT. // Here we need to check whether subgraph is fully supported by TRT and don't fuse the nodes of the subgraph until control flow op level. @@ -2767,6 +2882,102 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil common::Status TensorrtExecutionProvider::Compile(const std::vector& fused_nodes_and_graphs, std::vector& node_compute_funcs) { + LOGS_DEFAULT(VERBOSE) << "*#* TensorrtExecutionProvider::Compile"; + // bool is_trt_ctx_model = GraphHasCtxNode(fused_nodes_and_graphs); + // // verify where to get context configs from + // bool context_cache_enabled_ = true; + // onnxruntime::PathString context_cache_path = ""; + // bool is_ctx_file_exist = false; + // bool share_ep_contexts_ = true; + // std::string context_cache_path_cfg_ = ""; + + // if (is_trt_ctx_model || context_cache_enabled_) { + // const onnxruntime::GraphViewer& graph_viewer_0(fused_nodes_and_graphs[0].filtered_graph); + // is_ctx_file_exist = qnn::ValidateContextCacheFilePath(is_trt_ctx_model, + // context_cache_path_cfg_, + // graph_viewer_0.ModelPath().native(), + // context_cache_path); + // } + + // ORT_RETURN_IF(is_ctx_file_exist && !is_trt_ctx_model && context_cache_enabled_, + // "The inference session is created from normal ONNX model. And an EP context model file is provided and existed. ", + // "Please remove the EP context model manually if you want to re-generate it."); + + // if (is_trt_ctx_model) { + // // Get TensorRT Model from EP shared contexts + // // if (share_ep_contexts_ && SharedContext::GetInstance().HasSharedTrtModels()) { + // // if (EpSharedContextsHasAllGraphs(fused_nodes_and_graphs, logger)) { + // // for (auto fused_node_and_graph : fused_nodes_and_graphs) { + // // const onnxruntime::GraphViewer& graph_viewer(fused_node_and_graph.filtered_graph); + // // const auto& ep_context_node = graph_viewer.Nodes().begin(); + // // const Node& fused_node = fused_node_and_graph.fused_node; + // // const std::string& graph_meta_id = fused_node.Name(); + // // std::string key = ep_context_node->Name(); + // // auto trt_model_shared = SharedContext::GetInstance().GetSharedTrtModel(key); + // // ORT_RETURN_IF(nullptr == trt_model_shared, "Graph: " + key + " not found from shared EP contexts."); + // // // ORT_RETURN_IF_ERROR(qnn_model_shared->SetGraphInputOutputInfo(graph_viewer, fused_node, logger)); + // // // ORT_RETURN_IF_ERROR(qnn_model_shared->SetupQnnInputOutput(logger)); + // // // qnn_models_.emplace(graph_meta_id, std::move(qnn_model_shared)); + + // // // ORT_RETURN_IF_ERROR(SetGraphInputOutputInfo(trt_model_shared, graph_viewer, fused_node, logger)); + // // // ORT_RETURN_IF_ERROR(SetupTrtInputOutput(trt_model_shared, logger)); + // // // trt_model_shared.emplace(graph_meta_id, std::move(trt_model_shared)); + // // // ORT_RETURN_IF_ERROR(CreateComputeFunc(node_compute_funcs, logger)); + // // } + // // return Status::OK(); + // // } + // // } + + // std::unordered_map> trt_models; + + // std::vector main_context_pos_list; + // ORT_RETURN_IF_ERROR(qnn::GetMainContextNode(fused_nodes_and_graphs, main_context_pos_list)); + + // for (auto main_context_pos : main_context_pos_list) { + // const onnxruntime::GraphViewer& main_ctx_graph_viewer(fused_nodes_and_graphs[main_context_pos].filtered_graph); + // // Create QNN context from the cached binary, deserialize the QNN graph from the binary + // ORT_RETURN_IF_ERROR(LoadTrtCtxFromOnnxGraph(main_ctx_graph_viewer, + // context_cache_path, + // trt_models, + // logger)); + // } + + // // for (auto fused_node_and_graph : fused_nodes_and_graphs) { + // // const onnxruntime::GraphViewer& graph_viewer(fused_node_and_graph.filtered_graph); + // // const auto& ep_context_node = graph_viewer.Nodes().begin(); + // // const Node& fused_node = fused_node_and_graph.fused_node; + // // const std::string& graph_meta_id = fused_node.Name(); + // // std::string key = ep_context_node->Name(); + // // ORT_RETURN_IF(trt_models.find(key) == trt_models.end(), key + " key name not exist in table trt_models."); + // // auto trt_models = std::move(trt_models[key]); + // // ORT_RETURN_IF_ERROR(trt_models->SetGraphInputOutputInfo(graph_viewer, fused_node, logger)); + // // ORT_RETURN_IF_ERROR(trt_models->SetupQnnInputOutput(logger)); + + // // // fused node name is QNNExecutionProvider_QNN_[hash_id]_[id] + // // // the name here must be same with context->node_name in compute_info + // // trt_models.emplace(graph_meta_id, std::move(trt_models)); + // // trt_models.erase(key); + + // // ORT_RETURN_IF_ERROR(CreateComputeFunc(node_compute_funcs, logger)); + // // } + + // // if (share_ep_contexts_ && trt_models.size() > 0) { + // // // std::vector> shared_qnn_models; + // // std::unordered_map> shared_trt_models; + + // // for (auto& [key, value] : trt_models) { + // // shared_trt_models.push_back(std::move(trt_models[key])); + // // } + // // std::string duplicate_graph_names; + // // bool has_duplicate_graph = SharedContext::GetInstance().SetSharedTrtModel(std::move(shared_trt_models), + // // duplicate_graph_names); + // // ORT_RETURN_IF(has_duplicate_graph, "Duplicate graph names detect across sessions: " + duplicate_graph_names); + // // } + + // } // end of is_trt_ctx_model + + LOGS_DEFAULT(VERBOSE) << "*#* fused_nodes_and_graphs.size()=" << fused_nodes_and_graphs.size(); + for (auto& fused_node_graph : fused_nodes_and_graphs) { const GraphViewer& graph_body_viewer = fused_node_graph.filtered_graph; const Node& fused_node = fused_node_graph.fused_node; @@ -2786,20 +2997,73 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorName()] = i; } + // Print current nodes in subgraph + // std::ostringstream oss; + // oss << "*#* Checking nodes of graph_body_viewer. graph_body_viewer.MaxNodeIndex()=" << graph_body_viewer.MaxNodeIndex() << "\n Nodes: ["; + // for (int i = 0; i < graph_body_viewer.MaxNodeIndex(); ++i) { + // const auto node = graph_body_viewer.GetNode(i); + // if (!node) { + // oss << " null \n"; + // continue; + // } + // if (!node->Name().empty() && !node->OpType().empty()) { + // oss << " node->Name()=" << node->Name() << " node->OpType()=" << node->OpType() << "\n"; + // } else { + // oss << " Either Name() or OpType() is empty\n"; + // } + // } + // LOGS_DEFAULT(VERBOSE) << oss.str() << "]"; + LOGS_DEFAULT(VERBOSE) << "*#* fused_node.Index()=" << fused_node.Index() << " Address of graph_body_viewer: " << &graph_body_viewer; + Status status; - if (GraphHasCtxNode(graph_body_viewer)) { + // if (GraphHasCtxNode(graph_body_viewer)) { + int ctx_node_idx = FindCtxNodeInGraph(graph_body_viewer); + if (ctx_node_idx >= 0) { + LOGS_DEFAULT(VERBOSE) << "*#* Calling CreateNodeComputeInfoFromPrecompiledEngine"; status = CreateNodeComputeInfoFromPrecompiledEngine(graph_body_viewer, fused_node, + ctx_node_idx, input_map, output_map, - node_compute_funcs); + node_compute_funcs + ); + // Should have embed info + // LOGS_DEFAULT(VERBOSE) << "*#* Create a new model from graph_body_viewer and add to trt_ep_context_models"; + // LOGS_DEFAULT(VERBOSE) << "*#* Checking nodes of graph_body_viewer after CreateNodeComputeInfoFromPrecompiledEngine"; + // LOGS_DEFAULT(VERBOSE) << "*#* graph_body_viewer.MaxNodeIndex()=" << graph_body_viewer.MaxNodeIndex(); + // LOGS_DEFAULT(VERBOSE) << "*#* Checking nodes of ctx_model_ptr"; + // auto ctx_model_ptr = graph_body_viewer.CreateModel(*GetLogger()); // Creates model without nodes + // auto& graph = ctx_model_ptr->MainGraph(); + // // LOGS_DEFAULT(VERBOSE) << "*#* graph.MaxNodeIndex()=" << graph.MaxNodeIndex(); + // for (int i = 0; i < graph_body_viewer.MaxNodeIndex(); ++i) { + // auto node = graph_body_viewer.GetNode(i); + // if (!node) { + // LOGS_DEFAULT(WARNING) << "Node at index " << i << " is null!"; + // continue; + // } + // if (node != nullptr && node->OpType() == EPCONTEXT_OP) { + // if (!node->Name().empty()) { + // LOGS_DEFAULT(VERBOSE) << "*#* node->Name()=" << node->Name() << " node->OpType()=" << node->OpType(); + // } + // // graph.AddNode(*node); + // } + // } + // LOGS_DEFAULT(VERBOSE) << "*#* graph.MaxNodeIndex()=" << graph.MaxNodeIndex(); + // Segmentation fault here? + // LOGS_DEFAULT(VERBOSE) << "*#* before trt_ep_context_models.emplace_back"; + // trt_ep_context_models.emplace_back(std::move(ctx_model_ptr)); + // LOGS_DEFAULT(VERBOSE) << "*#* After adding ctx_model_ptr, trt_ep_context_models.size()=" << trt_ep_context_models.size(); } else { + LOGS_DEFAULT(VERBOSE) << "*#* Calling CreateNodeComputeInfoFromGraph"; status = CreateNodeComputeInfoFromGraph(graph_body_viewer, fused_node, input_map, output_map, node_compute_funcs); + // LOGS_DEFAULT(VERBOSE) << "*#* After CreateNodeComputeInfoFromPrecompiledEngine graph_body_viewer.MaxNodeIndex()=" << graph_body_viewer.MaxNodeIndex(); + } if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, status.ErrorMessage()); } } + return Status::OK(); } @@ -3280,6 +3544,7 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView 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(VERBOSE) << "*#* engine_cache_enable_=" << engine_cache_enable_; if (engine_cache_enable_) { // Serialize engine profile if it has explicit profiles if (has_explicit_profile) { @@ -3317,6 +3582,7 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized timing cache " + timing_cache_path; } } + LOGS_DEFAULT(VERBOSE) << "*#* for static shape, dump_ep_context_model_=" << dump_ep_context_model_; // dump EP context node model if (dump_ep_context_model_) { // "ep_cache_context" node attribute should be a relative path to context model directory @@ -3328,15 +3594,62 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView if (engine_cache_enable_ && engine_hw_compatible_) { compute_capability_hw_compat = "80+"; } - 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())}; - DumpCtxModel(model_proto.get(), ctx_model_path_); + LOGS_DEFAULT(VERBOSE) << "*#* static shape, bool dump_ep_context_model_=" << dump_ep_context_model_; + // Old Context model + // 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())}; + // DumpCtxModel(model_proto.get(), ctx_model_path_); + + // New context model + // 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())}; + // std::unique_ptr new_graph_viewer = CreateCtxModel(graph_body_viewer, + // trt_ep_context_nodes = CreateCtxModel(graph_body_viewer, // *#* TODO Need to understand what happens when there's multiple context nodes + auto trt_ep_context_model_ptr = CreateCtxModel2(graph_body_viewer, + 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()); + auto& graph = trt_ep_context_model_ptr->MainGraph(); + LOGS_DEFAULT(VERBOSE) << "*#* graph.Nodes().size()=" << graph.Nodes().size(); + trt_ep_context_models.emplace_back(std::move(trt_ep_context_model_ptr)); + // auto new_graph_viewer = trt_ep_context_nodes->CreateGraphViewer(); + // std::unique_ptr model = new_graph_viewer->CreateModel(*GetLogger()); + // 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); + // model_proto = model_proto.release(); + // std::unique_ptr model_proto_ptr{model_proto}; + // EPCTX TODO populate trt_ep_context_nodes for static shape + // auto& subgraph = fused_node.filtered_graph.get(); + // auto local_registries = IOnnxRuntimeOpSchemaRegistryList{subgraph.GetSchemaRegistry()}; + // // auto model = onnxruntime::Model::CreateFromProto(model_proto, GetLogger()); + // auto model = Model::Create(model_proto, subgraph.ModelPath(), &local_registries, GetLogger()) + // const auto& graph = model_proto.get()->graph(); + + // const auto& graph = model->MainGraph(); + // for (const auto& node : graph.Nodes()) { + // LOGS_DEFAULT(VERBOSE) << "*#* node->Name" << node->Name(); + // // trt_ep_context_nodes.push_back(graph.GetNode(node.Index())); + // trt_ep_context_nodes.push_back(node); + // } + // DumpCtxModel(model_proto.get(), ctx_model_path_); } } } @@ -3434,17 +3747,42 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView if (engine_cache_enable_ && engine_hw_compatible_) { compute_capability_hw_compat = "80+"; } - model_proto_.reset(CreateCtxModel(graph_body_viewer, - ep_cache_context_attr_, - nullptr, - 0, - ep_context_embed_mode_, - compute_capability_hw_compat, - model_path_, - GetLogger())); - if (ep_context_embed_mode_ == 0) { - DumpCtxModel(model_proto_.get(), ctx_model_path_); - } + LOGS_DEFAULT(VERBOSE) << "*#* dynamic shape"; + // model_proto_.reset(CreateCtxModel(graph_body_viewer, + // ep_cache_context_attr_, + // nullptr, + // 0, + // ep_context_embed_mode_, + // compute_capability_hw_compat, + // model_path_, + // GetLogger())); + auto trt_ep_context_model_ptr = CreateCtxModel2(graph_body_viewer, + fused_node.Name(), + ep_cache_context_attr_, + nullptr, + 0, + ep_context_embed_mode_, + compute_capability_hw_compat, + model_path_, + GetLogger()); + + trt_ep_context_models.emplace_back(std::move(trt_ep_context_model_ptr)); + // std::unique_ptr ctx_model = new_graph_viewer->CreateModel(*GetLogger()); + // auto ctx_model_proto_ = ctx_model->ToProto(); + // new_graph_viewer->ToProto(*model_proto_->mutable_graph(), true, true); + // ctx_model_proto_->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + // // model_proto_.reset(ctx_model_proto_); + // model_proto_ = std::move(ctx_model_proto_); + // const auto& graph = model->MainGraph(); + // for (const auto& node : graph.Nodes()) { + // LOGS_DEFAULT(VERBOSE) << "*#* node->Name" << node->Name(); + // // trt_ep_context_nodes.push_back(graph.GetNode(node.Index())); + // trt_ep_context_nodes.push_back(node); + // } + // if (ep_context_embed_mode_ == 0) { + // DumpCtxModel(model_proto_.get(), ctx_model_path_); + // } + // TODO populate trt_ep_context_nodes for dynamic shape } // Create function state @@ -4035,11 +4373,60 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView return Status::OK(); } +std::string PrintGraphNodes(const GraphViewer& graph_body_viewer) { + std::ostringstream oss; // Create a new ostringstream instance + + // Append the message and graph details + oss << "graph_body_viewer.MaxNodeIndex()=" << graph_body_viewer.MaxNodeIndex() << "\n"; + oss << "graph_body_viewer.NumberOfNodes()=" << graph_body_viewer.NumberOfNodes() << "\n"; + oss << "Nodes: ["; + + // Iterate over the nodes in the graph + for (int i = 0; i < graph_body_viewer.MaxNodeIndex(); ++i) { + const auto node = graph_body_viewer.GetNode(i); + if (!node) { + oss << "null,"; + continue; + } + if (!node->Name().empty() && !node->OpType().empty()) { + oss << "node->Name()=" << node->Name() << " node->OpType()=" << node->OpType() << ","; + } else { + oss << "[Either Name() or OpType() is empty],"; + } + } + + oss << "]"; + return oss.str(); +} + Status TensorrtExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(const GraphViewer& graph_body_viewer, const Node& fused_node, + const int ctx_node_idx, std::unordered_map& input_map, std::unordered_map& output_map, std::vector& node_compute_funcs) { + // Print subgraph + // graph_body_viewer is virtually empty + auto model = graph_body_viewer.CreateModel(*GetLogger()); + auto model_proto = model->ToProto(); + + if (dump_subgraphs_) { + // Dump TensorRT subgraphs + std::fstream dump(fused_node.Name() + ".onnx", std::ios::out | std::ios::trunc | std::ios::binary); + model_proto->SerializeToOstream(dump); + std::string subgraph_str; + model_proto->SerializeToString(subgraph_str); + // LOGS_DEFAULT(VERBOSE) << "*#* Subgraph " << fused_node.Name() << ": " << subgraph_str; + LOGS_DEFAULT(VERBOSE) << "*#* Subgraph " << fused_node.Name() << ": \n" + << "fused_node.GetAttributes().size()=" << fused_node.GetAttributes().size(); + } + // Print nodes in graph_body_viewer + LOGS_DEFAULT(VERBOSE) << "*#* graph_body_viewer before priority toposort: " << PrintGraphNodes(graph_body_viewer); + + graph_body_viewer.ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/); + + LOGS_DEFAULT(VERBOSE) << "*#* graph_body_viewer after priority toposort: " << PrintGraphNodes(graph_body_viewer); + std::unique_ptr trt_engine; std::unique_ptr trt_context; std::unordered_map input_indexes; // TRT engine input name -> ORT kernel context input index @@ -4056,7 +4443,7 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(con onnx_model_bytestream_, onnx_model_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, ctx_node_idx); if (status != Status::OK()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, status.ErrorMessage()); } @@ -4360,6 +4747,45 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(con return Status::OK(); } +// static bool EpSharedContextsHasAllGraphs(const std::vector& fused_nodes_and_graphs, +// const logging::Logger& logger) { +// for (auto fused_node_and_graph : fused_nodes_and_graphs) { +// const onnxruntime::GraphViewer& graph_viewer(fused_node_and_graph.filtered_graph); +// const auto& ep_context_node = graph_viewer.Nodes().begin(); +// NodeAttrHelper node_helper(*ep_context_node); +// std::string cache_source = node_helper.Get(qnn::SOURCE, ""); + + +// const std::string& graph_name = ep_context_node->Name(); +// bool has_shared_trt_model = SharedContext::GetInstance().HasTrtModel(graph_name); +// if (!has_shared_trt_model) { +// LOGS(logger, VERBOSE) << "Graph: " << graph_name << " from EpContext node not found from shared EP contexts."; +// return false; +// } +// } + +// return true; +// } + +const InlinedVector TensorrtExecutionProvider::GetEpContextNodes() const { + LOGS_DEFAULT(VERBOSE) << "*#* trt_ep_context_models.size()=" << trt_ep_context_models.size(); + InlinedVector ep_context_nodes; + if (!trt_ep_context_models.empty()) { + for (const auto& context_model: trt_ep_context_models) { + const auto& graph = context_model->MainGraph(); + LOGS_DEFAULT(VERBOSE) << "*#* graph.Nodes().size()=" << graph.Nodes().size(); + LOGS_DEFAULT(VERBOSE) << "*#* graph.MaxNodeIndex()=" << graph.MaxNodeIndex(); + for (const auto& node: graph.Nodes()) { + LOGS_DEFAULT(VERBOSE) << "*#* node->Name()=" << node->Name() << " node->OpType()=" << node->OpType(); + // if (node.IsEpContextNode()) { // Check if it's an EP context node + ep_context_nodes.push_back(node); + } + } + } + LOGS_DEFAULT(VERBOSE) << "*#* ep_context_nodes.size()=" << ep_context_nodes.size(); + return ep_context_nodes; +} + void TensorrtExecutionProvider::RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry, AllocatorMap& allocators) const { auto allocator = allocators[GetOrtDeviceByMemType(OrtMemTypeCPU)]; RegisterCudaStreamHandles(stream_handle_registry, diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 9d8af02ba1..fde858ba63 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -273,6 +273,7 @@ class TensorrtExecutionProvider : public IExecutionProvider { bool IsGraphCaptureEnabled() const override; bool IsGraphCaptured(int graph_annotation_id) const override; Status ReplayGraph(int graph_annotation_id) override; + const InlinedVector GetEpContextNodes() const override; static common::Status RefitEngine(std::string onnx_model_filename, std::string& onnx_model_folder_path, @@ -331,6 +332,12 @@ class TensorrtExecutionProvider : public IExecutionProvider { std::string cache_prefix_; bool engine_hw_compatible_ = false; std::string op_types_to_exclude_; + // std::unique_ptr trt_ep_context_model_; + // InlinedVector trt_ep_context_nodes; + // std::unique_ptr trt_ep_context_nodes; + // std::unordered_map>; + // std::unordered_map> trt_ep_context_models; + std::vector> trt_ep_context_models; // *#* TODO need ep context node name as key? // The format is as for TENSORRT_VERSION: (MAJOR * 100 + MINOR) * 100 + PATCH int32_t trt_version_; @@ -567,6 +574,7 @@ class TensorrtExecutionProvider : public IExecutionProvider { */ Status CreateNodeComputeInfoFromPrecompiledEngine(const GraphViewer& graph_body_viewer, const Node& fused_node, + const int ctx_node_idx, std::unordered_map& input_map, std::unordered_map& output_map, std::vector& node_compute_funcs); @@ -591,4 +599,66 @@ class TensorrtExecutionProvider : public IExecutionProvider { */ nvinfer1::IBuilder* GetBuilder(TensorrtLogger& trt_logger) const; }; + +// class SharedContext { +// public: +// static SharedContext& GetInstance() { +// static SharedContext instance_; +// return instance_; +// } + +// bool HasSharedTrtModels() { +// const std::lock_guard lock(mtx_); +// return !shared_trt_models_.empty(); +// } + +// bool HasTrtModel(const std::string& model_name) { +// auto it = find_if(shared_trt_models_.begin(), shared_trt_models_.end(), +// [&model_name](const std::unique_ptr& trt_model) { return trt_model->Name() == model_name; }); +// return it != shared_trt_models_.end(); +// } + +// std::unique_ptr GetSharedTrtModel(const std::string& model_name) { +// const std::lock_guard lock(mtx_); +// auto it = find_if(shared_trt_models_.begin(), shared_trt_models_.end(), +// [&model_name](const std::unique_ptr& trt_model) { return trt_model->Name() == model_name; }); +// if (it == shared_trt_models_.end()) { +// return nullptr; +// } +// auto trt_model = std::move(*it); +// shared_trt_models_.erase(it); +// return trt_model; +// } + +// bool SetSharedTrtModel(std::vector>&& shared_trt_models, +// std::string& duplicate_graph_names) { +// const std::lock_guard lock(mtx_); +// bool graph_exist = false; +// for (auto& shared_trt_model : shared_trt_models) { +// auto& model_name = shared_trt_model->Name(); +// auto it = find_if(shared_trt_models_.begin(), shared_trt_models_.end(), +// [&model_name](const std::unique_ptr& trt_model) { return trt_model->Name() == model_name; }); +// if (it == shared_trt_models_.end()) { +// shared_trt_models_.push_back(std::move(shared_trt_model)); +// } else { +// duplicate_graph_names.append(model_name + " "); +// graph_exist = true; +// } +// } + +// return graph_exist; +// } + +// private: +// SharedContext() = default; +// ~SharedContext() = default; +// SharedContext(const SharedContext&) = delete; +// SharedContext& operator=(const SharedContext&) = delete; + +// // std::vector> shared_trt_models_; +// std::vector<> shared_trt_models_; +// // Producer sessions can be in parallel +// // Consumer sessions have to be after producer sessions initialized +// std::mutex mtx_; +// }; //SharedContext class } // namespace onnxruntime