saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-13 18:08:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add comments

Browse source
This commit is contained in:
Chi Lo 2025-01-28 09:05:07 -08:00
parent 309341e86c
commit d0cbc65382
4 changed files with 82 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -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<std::vector<std::unique_ptr<ComputeCapability>>(const GraphViewer&)> selection_func;
auto status = g_host->GetEPOptimizerByName("ConstantFoldingDQ", graph_transformer_mgr, selection_func);
std::vector<std::unique_ptr<ComputeCapability>> selection_cc;
@ -2666,14 +2677,15 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph,
selection_cc = selection_func(graph);
}
std::unordered_set<NodeIndex> trt_selection_node_set;
std::unordered_set<NodeIndex> trt_selection_node_set; // The qualified dq nodes selected by TRT EP
std::unordered_map<NodeIndex, NodeIndex> 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<IndexedSubGraph> 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<ComputeCapability> TensorrtExecutionProvider::CreateOptimizationComputeCapability(ComputeCapability* selection_cc,
std::unordered_set<NodeIndex>& trt_selection_node_set,
ComputeCapability* trt_cc) const {
auto sub_graph = onnxruntime::IndexedSubGraph::Create();
std::unordered_set<NodeIndex> 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
*/

25
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h
View file

@ -592,9 +592,26 @@ class TensorrtExecutionProvider : public IExecutionProvider {
*/
nvinfer1::IBuilder* GetBuilder(TensorrtLogger& trt_logger) const;
void CreateConsumerToDqMap(const GraphViewer& graph, std::unordered_set<NodeIndex>& selection_node_set, std::unordered_map<NodeIndex, NodeIndex>& consumer_to_dq) const;
std::unique_ptr<ComputeCapability> TensorrtExecutionProvider::CreateOptimizationComputeCapability(ComputeCapability* selection_cc,
std::unordered_set<NodeIndex>& 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<NodeIndex>& selection_node_set,
std::unordered_map<NodeIndex, NodeIndex>& 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<ComputeCapability> CreateOptimizationComputeCapability(ComputeCapability* selection_cc,
std::unordered_set<NodeIndex>& trt_selection_node_set,
ComputeCapability* trt_cc) const;
};
} // namespace onnxruntime

42
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_helper.cc
View file

@ -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<NodeIndex>& selection_node_set,
std::unordered_map<NodeIndex, NodeIndex>& 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<ComputeCapability> TensorrtExecutionProvider::CreateOptimizationComputeCapability(ComputeCapability* selection_cc,
std::unordered_set<NodeIndex>& trt_selection_node_set,
ComputeCapability* trt_cc) const {
auto sub_graph = onnxruntime::IndexedSubGraph::Create();
std::unordered_set<NodeIndex> 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

8
onnxruntime/core/session/provider_bridge_ort.cc
View file

@ -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::Logger*>(&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::unique_ptr<ComputeCapability>> {
std::vector<std::unique_ptr<ComputeCapability>> result;
std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
@ -356,7 +356,7 @@ struct ProviderHostImpl : ProviderHost {
return result;
};
// optimizer lookup table
// Optimizer lookup table
static std::unordered_map<std::string, std::function<std::vector<std::unique_ptr<ComputeCapability>>(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;