Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix TensorRT kernel conflict issue for subgraphs of control flow operators #6115

Merged
merged 3 commits into from
Dec 16, 2020
Merged

Fix TensorRT kernel conflict issue for subgraphs of control flow operators #6115

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
476f3f4
add static subgraph kernel index
stevenlix Dec 11, 2020
48523ab
change kernel naming to avoid conflicts
stevenlix Dec 15, 2020
86e7609
Merge branch 'master' into stevenlix/subgraph
stevenlix Dec 16, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,7 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
}
return true;
}

} // namespace

namespace google {
Expand Down Expand Up @@ -500,7 +501,7 @@ void ToGraphProtoInternal(const onnxruntime::Provider_GraphViewer& graph, Provid
}
}

std::unique_ptr<Provider_IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph_t graph_nodes_index, int& kernels_index, const onnxruntime::Provider_GraphViewer& graph) const {
std::unique_ptr<Provider_IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph_t graph_nodes_index, const onnxruntime::Provider_GraphViewer& graph) const {
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
std::unordered_set<size_t> node_set;
node_set.reserve(graph_nodes_index.first.size());
Expand Down Expand Up @@ -605,7 +606,7 @@ std::unique_ptr<Provider_IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph
// Assign inputs and outputs to subgraph's meta_def
auto meta_def = Provider_IndexedSubGraph_MetaDef::Create();
const std::string graph_type = graph.IsSubgraph() ? "subgraph" : "graph";
meta_def->name() = "TRTKernel_" + graph_type + "_" + graph.Name() + "_" + std::to_string(kernels_index++);
meta_def->name() = "TRTKernel_" + graph_type + "_" + graph.Name() + "_" + std::to_string(subgraph_id_++);
meta_def->domain() = kMSDomain;

for (const auto& input : inputs) {
Expand Down Expand Up @@ -771,11 +772,11 @@ void TensorrtExecutionProvider::RemoveTensorRTGraphCycles(SubGraphCollection_t&
std::unordered_map<int, std::string> index_to_node_map;
std::unordered_map<std::string, std::unordered_set<std::string>> input_to_nodes_map, node_to_outputs_map;
std::unordered_set<int> non_trt_node_index(node_index.begin(), node_index.end());
int counter = 0, id = 0;
int id = 0;
for (const auto& group : supported_nodes_vector) {
if (!group.first.empty()) {
// Construct subgraph from node list
std::unique_ptr<Provider_IndexedSubGraph> sub_graph = GetSubGraph(group, counter, graph);
std::unique_ptr<Provider_IndexedSubGraph> sub_graph = GetSubGraph(group, graph);

// Create node to inputs/outputs/index maps
const auto& meta_def = sub_graph->GetMetaDef();
Expand Down Expand Up @@ -901,10 +902,10 @@ TensorrtExecutionProvider::Provider_GetCapability(const onnxruntime::Provider_Gr

// Construct subgraph capability from node list
std::vector<std::unique_ptr<Provider_ComputeCapability>> result;
int counter = 0, number_of_trt_nodes = 0;
int number_of_trt_nodes = 0;
for (const auto& group : supported_nodes_vector) {
if (!group.first.empty()) {
std::unique_ptr<Provider_IndexedSubGraph> sub_graph = GetSubGraph(group, counter, graph);
std::unique_ptr<Provider_IndexedSubGraph> sub_graph = GetSubGraph(group, graph);
result.push_back(Provider_ComputeCapability::Create(std::move(sub_graph)));
number_of_trt_nodes += group.first.size();
}
Expand Down
8 changes: 4 additions & 4 deletions onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,9 +126,11 @@ class TensorrtExecutionProvider : public Provider_IExecutionProvider {
bool engine_cache_enable_ = false;
std::string cache_path_;
nvinfer1::IRuntime* runtime_ = nullptr;

OrtMutex tensorrt_mu_;
int device_id_;
AllocatorPtr allocator_;
mutable int subgraph_id_ = 0;

std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvonnxparser::IParser>> parsers_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>> engines_;
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext>> contexts_;
Expand All @@ -139,7 +141,7 @@ class TensorrtExecutionProvider : public Provider_IExecutionProvider {
std::unordered_map<std::string, std::unordered_map<std::string, std::unordered_map<int, std::pair<int64_t, int64_t>>>> input_shape_ranges_;

/**Get IndexedSubGraph based on node list of the subgraph*/
std::unique_ptr<Provider_IndexedSubGraph> GetSubGraph(SubGraph_t graph_nodes_index, int& kernels_index,
std::unique_ptr<Provider_IndexedSubGraph> GetSubGraph(SubGraph_t graph_nodes_index,
const onnxruntime::Provider_GraphViewer& graph) const;

/**
Expand All @@ -153,7 +155,5 @@ class TensorrtExecutionProvider : public Provider_IExecutionProvider {
const onnxruntime::Provider_GraphViewer& graph, bool* early_termination) const;

void RemoveTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const onnxruntime::Provider_GraphViewer& graph) const;

AllocatorPtr allocator_;
};
} // namespace onnxruntime