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

[TensorRT EP] Fix wrong input order when generating IndexedSubGraph (#22857)

Browse source 
The input order of generated indexedSubGraph needs to be consistent with
the input order of original graph.

This PR will also fix the github issue
https://github.com/microsoft/onnxruntime/issues/22729
This commit is contained in:
Chi Lo 2024-12-02 01:45:29 -08:00 committed by GitHub
parent 49a80df77f
commit e234023d11
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 17 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1952,7 +1952,7 @@ std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph
// Find inputs and outputs of the subgraph
std::unique_ptr<IndexedSubGraph> sub_graph = onnxruntime::IndexedSubGraph::Create();
std::unordered_map<const NodeArg*, int> fused_inputs, fused_outputs, fused_outputs_to_add, graph_outputs_to_add;
std::unordered_map<const NodeArg*, int> original_inputs, fused_inputs, fused_outputs, fused_outputs_to_add, graph_outputs_to_add;
std::unordered_set<const NodeArg*> erased;
int input_order = 0;
int output_order = 0;
@ -2044,12 +2044,25 @@ std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph
fused_outputs.insert(fused_outputs_to_add.begin(), fused_outputs_to_add.end());
fused_outputs.insert(graph_outputs_to_add.begin(), graph_outputs_to_add.end());
// Sort inputs and outputs by the order they were added
std::multimap<int, const NodeArg*> inputs, outputs;
for (auto it = fused_inputs.begin(), end = fused_inputs.end(); it != end; ++it) {
inputs.insert(std::pair<int, const NodeArg*>(it->second, it->first));
// Get the input order of the original graph
int order = 0;
for (const auto* input : graph.GetInputs()) {
original_inputs[input] = order++;
}
// input order needs to be consistent with original graph's input order
for (auto it = fused_inputs.begin(), end = fused_inputs.end(); it != end; ++it) {
const auto& iter = original_inputs.find(it->first);
if (iter != original_inputs.end()) {
inputs.insert(std::pair<int, const NodeArg*>(iter->second, iter->first));
} else {
inputs.insert(std::pair<int, const NodeArg*>(it->second, it->first));
}
}
// Sort outputs by the order they were added
for (auto it = fused_outputs.begin(), end = fused_outputs.end(); it != end; ++it) {
outputs.insert(std::pair<int, const NodeArg*>(it->second, it->first));
}