[TensorRT EP] Make TRT EP use priority-based topo sort (#20512)

This PR is needed for
https://github.com/microsoft/onnxruntime/pull/20411 to make sure TRT EP
use priority-based topo sort for consistency across TRT EP.
This commit is contained in:
Chi Lo 2024-04-29 16:00:43 -07:00 committed by GitHub
parent 8c31f27dd1
commit a1558fe117
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 12 additions and 10 deletions

View file

@ -848,7 +848,7 @@ struct ProviderHost {
virtual bool GraphViewer__GetInitializedTensor(const GraphViewer* p, const std::string& tensor_name, const ONNX_NAMESPACE::TensorProto*& value) = 0;
virtual const std::unordered_map<std::string, int>& GraphViewer__DomainToVersionMap(const GraphViewer* p) = 0;
virtual const std::vector<NodeIndex>& GraphViewer__GetNodesInTopologicalOrder(const GraphViewer* p) = 0;
virtual const std::vector<NodeIndex>& GraphViewer__GetNodesInTopologicalOrder(const GraphViewer* p, int execution_order) = 0;
virtual const std::vector<const NodeArg*>& GraphViewer__GetInputsIncludingInitializers(const GraphViewer* p) noexcept = 0;
virtual void GraphViewer__ToProto(const GraphViewer* p,

View file

@ -884,7 +884,7 @@ class GraphViewer final {
const std::unordered_map<std::string, int>& DomainToVersionMap() const noexcept { return g_host->GraphViewer__DomainToVersionMap(this); }
const std::vector<NodeIndex>& GetNodesInTopologicalOrder() const { return g_host->GraphViewer__GetNodesInTopologicalOrder(this); }
const std::vector<NodeIndex>& GetNodesInTopologicalOrder(int execution_order = 0) const { return g_host->GraphViewer__GetNodesInTopologicalOrder(this, execution_order); }
const std::vector<const NodeArg*>& GetInputsIncludingInitializers() const noexcept { return g_host->GraphViewer__GetInputsIncludingInitializers(this); }
void ToProto(ONNX_NAMESPACE::GraphProto& graph_proto,

View file

@ -1783,7 +1783,7 @@ bool TensorrtExecutionProvider::AllNodesAssignedToSpecificEP(const GraphViewer&
const int number_of_ort_nodes = graph.NumberOfNodes();
std::vector<size_t> nodes_vector(number_of_ort_nodes);
std::iota(std::begin(nodes_vector), std::end(nodes_vector), 0);
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
for (const auto& index : nodes_vector) {
const auto& node = graph.GetNode(node_index[index]);
if (node->GetExecutionProviderType() != provider_type) {
@ -1807,7 +1807,7 @@ bool TensorrtExecutionProvider::IsSubGraphFullySupported(SubGraphCollection_t su
}
std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph_t graph_nodes_index, const GraphViewer& graph, const HashValue& model_hash, int subgraph_index) const {
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
std::unordered_set<size_t> node_set;
node_set.reserve(graph_nodes_index.first.size());
for (const auto& index : graph_nodes_index.first) {
@ -1971,7 +1971,7 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect
}
iterations++;
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
for (const auto& group : nodes_vector_input) {
// Construct subgraph
if (!group.first.empty()) {
@ -2135,7 +2135,7 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect
trt_parser->supportsModel(string_buf.data(), string_buf.size(), parser_nodes_list, model_path_);
SubGraphCollection_t next_nodes_list;
const std::vector<NodeIndex>& subgraph_node_index = graph_viewer->GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& subgraph_node_index = graph_viewer->GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
next_nodes_list = GetSupportedList(parser_nodes_list, iterations, max_iterations, *graph_viewer, early_termination);
for (size_t i = 0, end = next_nodes_list.size(); i < end; ++i) {
for (size_t j = 0, end = next_nodes_list[i].first.size(); j < end; ++j) {
@ -2151,7 +2151,7 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect
// Detect and remove cycles from supported node list
bool TensorrtExecutionProvider::DetectTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const GraphViewer& graph, const HashValue& model_hash, bool remove_cycles) const {
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
bool trt_cycle = true, cycle_detected = false;
while (trt_cycle) {
trt_cycle = false;
@ -2299,7 +2299,7 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph,
std::iota(std::begin(nodes_vector), std::end(nodes_vector), 0);
std::vector<size_t> filtered_nodes_vector;
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
for (const auto& index : nodes_vector) {
const auto& node = graph.GetNode(node_index[index]);
@ -2374,7 +2374,7 @@ TensorrtExecutionProvider::GetCapability(const GraphViewer& graph,
// 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.
if (IsSubGraphOfControlFlowOp(graph) && IsSubGraphFullySupported(supported_nodes_vector, number_of_ort_nodes)) {
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder();
const std::vector<NodeIndex>& node_index = graph.GetNodesInTopologicalOrder(1 /*priority-based topological sort*/);
bool all_subgraphs_are_supported = true;
// "If" control flow op has two subgraph bodies, "then" body and "else" body respectively.

View file

@ -1086,7 +1086,9 @@ struct ProviderHostImpl : ProviderHost {
const std::unordered_map<std::string, int>& GraphViewer__DomainToVersionMap(const GraphViewer* p) override { return p->DomainToVersionMap(); }
const std::vector<NodeIndex>& GraphViewer__GetNodesInTopologicalOrder(const GraphViewer* p) override { return p->GetNodesInTopologicalOrder(); }
const std::vector<NodeIndex>& GraphViewer__GetNodesInTopologicalOrder(const GraphViewer* p, int execution_order) override {
return p->GetNodesInTopologicalOrder(static_cast<ExecutionOrder>(execution_order));
}
const std::vector<const NodeArg*>& GraphViewer__GetInputsIncludingInitializers(const GraphViewer* p) noexcept override { return p->GetInputsIncludingInitializers(); }
void GraphViewer__ToProto(const GraphViewer* p,
ONNX_NAMESPACE::GraphProto& graph_proto,