From d0cbc653822d79fcd8b9ee403c809dc6b46a6129 Mon Sep 17 00:00:00 2001 From: Chi Lo Date: Tue, 28 Jan 2025 09:05:07 -0800 Subject: [PATCH] add comments --- .../tensorrt/tensorrt_execution_provider.cc | 46 +++++++------------ .../tensorrt/tensorrt_execution_provider.h | 25 ++++++++-- .../tensorrt_execution_provider_helper.cc | 42 ++++++++++++++++- .../core/session/provider_bridge_ort.cc | 8 ++-- 4 files changed, 82 insertions(+), 39 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 452fa545ae..c9499a10c3 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -2657,8 +2657,19 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, } } - // Enable EP related L2+ graph optimizations: - // 1. Dequantize INT32, UINT16, INT16 constant to FP32 (Apply constant folding on DQ nodes) + /** + * Enable EP related L2+ graph optimizations with steps: + * + * 1. call provider bridge API to lookup pre-defined optimizer by name and get selection function + * - Run selection function to get selection ComputeCapability + * - ComputeCapability.optimize_func would be set by the optimizer to the function that does the optimization + * + * + * + * Current available optimizations: + * - (ConstantFoldingDQ) constant folding on DQ nodes -> Dequantize INT32, UINT16, INT16 constant to FP32. + */ + std::function>(const GraphViewer&)> selection_func; auto status = g_host->GetEPOptimizerByName("ConstantFoldingDQ", graph_transformer_mgr, selection_func); std::vector> selection_cc; @@ -2666,14 +2677,15 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, selection_cc = selection_func(graph); } - std::unordered_set trt_selection_node_set; + std::unordered_set trt_selection_node_set; // The qualified dq nodes selected by TRT EP std::unordered_map consumer_to_dq; // consumer node -> dq node - CreateConsumerToDqMap(graph, trt_selection_node_set, consumer_to_dq); + SelectQualifiedDQNode(graph, trt_selection_node_set, consumer_to_dq); - // Create compute capability + // Create ComputeCapability int number_of_trt_nodes = 0, subgraph_index = 0; for (const auto& group : supported_nodes_vector) { if (!group.first.empty()) { + // TODO: Use consumer_to_dq table to include DQ node that is filtered out by TRT parser. std::unique_ptr sub_graph = GetSubGraph(group, graph, model_hash, subgraph_index); auto compute_capability = ComputeCapability::Create(std::move(sub_graph)); @@ -2703,30 +2715,6 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph, return result; } -std::unique_ptr TensorrtExecutionProvider::CreateOptimizationComputeCapability(ComputeCapability* selection_cc, - std::unordered_set& trt_selection_node_set, - ComputeCapability* trt_cc) const { - auto sub_graph = onnxruntime::IndexedSubGraph::Create(); - std::unordered_set selection_node_set; - - for (auto index : selection_cc->SubGraph()->Nodes()) { - selection_node_set.insert(index); - } - - for (auto index : trt_cc->SubGraph()->Nodes()) { - if (selection_node_set.find(index) == selection_node_set.end()) { - continue; - } - if (trt_selection_node_set.find(index) == trt_selection_node_set.end()) { - continue; - } - sub_graph->Nodes().push_back(index); - } - auto compute_capability = ComputeCapability::Create(std::move(sub_graph)); - compute_capability->copy_optimization_func(selection_cc); - return compute_capability; -} - /** * Refit the weight-stripped engine */ diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 8fc1827ea6..2908c3061f 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -592,9 +592,26 @@ class TensorrtExecutionProvider : public IExecutionProvider { */ nvinfer1::IBuilder* GetBuilder(TensorrtLogger& trt_logger) const; - void CreateConsumerToDqMap(const GraphViewer& graph, std::unordered_set& selection_node_set, std::unordered_map& consumer_to_dq) const; - std::unique_ptr TensorrtExecutionProvider::CreateOptimizationComputeCapability(ComputeCapability* selection_cc, - std::unordered_set& trt_selection_node_set, - ComputeCapability* trt_cc) const; + /** + * This is the helper function for ConstantFoldingDQ graph transformer. + * + * It selects the qualified/required DQ node to be optimized as well as provides a mapping table + * to help TRT EP later include the DQ node which is filtered out by TRT parser. + */ + void SelectQualifiedDQNode(const GraphViewer& graph, + std::unordered_set& selection_node_set, + std::unordered_map& consumer_to_dq) const; + + /** + * This function returns an optimization ComputeCapability that is limited to: + * 1. the DQ nodes in this individual TRT ComputeCapability + * 2. the DQ nodes that are qualified and selected by TRT EP + * + * It also needs to make sure the DQ nodes is a subset of the complete list of DQ nodes to optimize in original selection ComputeCapability. + * Finally, copy the optimization function from the original selection ComputeCapability. + */ + std::unique_ptr CreateOptimizationComputeCapability(ComputeCapability* selection_cc, + std::unordered_set& trt_selection_node_set, + ComputeCapability* trt_cc) const; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_helper.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_helper.cc index 84c2a5e376..b819711679 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_helper.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_helper.cc @@ -259,7 +259,13 @@ void TensorrtExecutionProvider::SetAllGraphInputs(Graph& graph) const { graph.SetInputs(graph_inputs_including_initializers); } -void TensorrtExecutionProvider::CreateConsumerToDqMap(const GraphViewer& graph, +/** + * This is the helper function for ConstantFoldingDQ graph transformer. + * + * It selects the qualified/required DQ node to be optimized as well as provides a mapping table + * to help TRT EP later include the DQ node which is filtered out by TRT parser. + */ +void TensorrtExecutionProvider::SelectQualifiedDQNode(const GraphViewer& graph, std::unordered_set& selection_node_set, std::unordered_map& consumer_to_dq) const { LOGS_DEFAULT(VERBOSE) << "Select qualified DQ nodes ..."; @@ -289,7 +295,39 @@ void TensorrtExecutionProvider::CreateConsumerToDqMap(const GraphViewer& graph, consumer_to_dq[consumer_node.Index()] = index; LOGS_DEFAULT(VERBOSE) << consumer_node.Name() << " <- " << node->Name(); } - LOGS_DEFAULT(VERBOSE) << "Total " << selection_node_set.size() << " DequantizeLinear nodes are selected."; } + LOGS_DEFAULT(VERBOSE) << "Total " << selection_node_set.size() << " DequantizeLinear node(s) are selected."; +} + +/** + * This function returns an optimization ComputeCapability that is limited to: + * 1. the DQ nodes in this individual TRT ComputeCapability + * 2. the DQ nodes that are qualified and selected by TRT EP + * + * It also needs to make sure the DQ nodes is a subset of the complete list of DQ nodes to optimize in original selection ComputeCapability. + * Finally, copy the optimization function from the original selection ComputeCapability. + */ +std::unique_ptr TensorrtExecutionProvider::CreateOptimizationComputeCapability(ComputeCapability* selection_cc, + std::unordered_set& trt_selection_node_set, + ComputeCapability* trt_cc) const { + auto sub_graph = onnxruntime::IndexedSubGraph::Create(); + std::unordered_set selection_node_set; + + for (auto index : selection_cc->SubGraph()->Nodes()) { + selection_node_set.insert(index); + } + + for (auto index : trt_cc->SubGraph()->Nodes()) { + if (selection_node_set.find(index) == selection_node_set.end()) { + continue; + } + if (trt_selection_node_set.find(index) == trt_selection_node_set.end()) { + continue; + } + sub_graph->Nodes().push_back(index); + } + auto compute_capability = ComputeCapability::Create(std::move(sub_graph)); + compute_capability->copy_optimization_func(selection_cc); + return compute_capability; } } // namespace onnxruntime diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 959c9743db..fbe299f68f 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -259,10 +259,10 @@ struct ProviderHostImpl : ProviderHost { static const GraphTransformerManager& graph_transformer_mgr = transformer_mgr; std::string optimizer_name(name); - // pre-defined graph transformers/optimizers + // Pre-defined graph transformers/optimizers static const std::string kEP_GRAPH_TRANSFORMER_CONSTANT_FOLDING_DQ = "ConstantFoldingDQ"; - // ConstantFoldingDQ's optimization function + // ConstantFoldingDQ optimization function auto constant_folding_dq_optimization = [&](Graph& graph, const ComputeCapability& optimization_cc, ComputeCapability& cc_to_update) -> Status { std::string optimizer_name = kEP_GRAPH_TRANSFORMER_CONSTANT_FOLDING_DQ; auto logger = const_cast(&logging::LoggingManager::DefaultLogger()); @@ -331,7 +331,7 @@ struct ProviderHostImpl : ProviderHost { return Status::OK(); }; - // ConstantFoldingDQ's selection function + // ConstantFoldingDQ selection function auto constant_folding_dq_selection = [&](const GraphViewer& graph_viewer) -> std::vector> { std::vector> result; std::unique_ptr sub_graph = std::make_unique(); @@ -356,7 +356,7 @@ struct ProviderHostImpl : ProviderHost { return result; }; - // optimizer lookup table + // Optimizer lookup table static std::unordered_map>(const GraphViewer&)>> optimizer_to_selection_function; if (optimizer_to_selection_function.find(kEP_GRAPH_TRANSFORMER_CONSTANT_FOLDING_DQ) == optimizer_to_selection_function.end()) { optimizer_to_selection_function[kEP_GRAPH_TRANSFORMER_CONSTANT_FOLDING_DQ] = constant_folding_dq_selection;