From 00d48d9c30051f279a99f70e38191adb782f2bbf Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Fri, 11 Jun 2021 15:23:21 +1000 Subject: [PATCH] Add enhanced partitioning utils for use by compiling EPs (#7991) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add enhanced partitioning utils and convert internal testing EP to use it. Will convert NNAPI EP once checked in. Background: Currently most EPs do their partitioning by iterating the model in the topologically sorted order. Whilst this works, it doesn’t ensure that all nodes which could possibly be added to the current group are, as the group is closed as when the first unsupported node is seen. Changes: - Ask EP for all nodes it supports first - Do partitioning aware topological sort - Groups nodes and flips between processing supported and unsupported nodes to maximise inputs that will be available for each partition - Create groups of nodes for the partition using the new order of nodes - Create ComputeCapability for each group There’s also an additional ability to specify operators to stop at. The processing will find all downstream nodes from ‘stop at’ operators and exclude them. If NonMaxSuppression is specified we can prevent the post-processing from SSD Mobilenet and MobileDet attempting to use NNAPI (so easy way to have parity with the TF Lite behavior). I don’t think there’s an automated way to determine what if any ‘stop at’ operators are required for a model, so this will need to be a configuration parameter for the EP and we’ll need to document recommended values for popular models. --- .../core/providers/partitioning_utils.cc | 336 +++++++++++++++++ .../core/providers/partitioning_utils.h | 67 ++++ .../internal_testing_execution_provider.cc | 155 +++----- .../internal_testing_execution_provider.h | 16 +- .../internal_testing_partitioning_tests.cc | 344 ++++++++++++++++++ .../test/testdata/ep_partitioning_test_1.onnx | 43 +++ .../test/testdata/ep_partitioning_test_2.onnx | 44 +++ .../test/testdata/ep_partitioning_tests.py | 128 +++++++ 8 files changed, 1024 insertions(+), 109 deletions(-) create mode 100644 onnxruntime/core/providers/partitioning_utils.cc create mode 100644 onnxruntime/core/providers/partitioning_utils.h create mode 100644 onnxruntime/test/providers/internal_testing/internal_testing_partitioning_tests.cc create mode 100644 onnxruntime/test/testdata/ep_partitioning_test_1.onnx create mode 100644 onnxruntime/test/testdata/ep_partitioning_test_2.onnx create mode 100644 onnxruntime/test/testdata/ep_partitioning_tests.py diff --git a/onnxruntime/core/providers/partitioning_utils.cc b/onnxruntime/core/providers/partitioning_utils.cc new file mode 100644 index 0000000000..bd85850aff --- /dev/null +++ b/onnxruntime/core/providers/partitioning_utils.cc @@ -0,0 +1,336 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/partitioning_utils.h" + +#include + +#include "core/framework/compute_capability.h" +#include "core/framework/execution_provider.h" +#include "core/graph/model.h" + +namespace onnxruntime { +namespace utils { + +// internal helpers +namespace { +// Kahn's topological sort with awareness of whether a node should be in a supported or unsupported partition +std::vector PartitionAwareTopoSort(const GraphViewer& graph_viewer, + const std::unordered_set& supported_nodes) { + std::queue supported_to_visit, unsupported_to_visit; + std::unordered_map in_degree; + std::vector topo_order; + + auto num_nodes = graph_viewer.NumberOfNodes(); + topo_order.reserve(num_nodes); + in_degree.reserve(num_nodes); + + auto add_to_visit = [&](const Node& node) { + if (supported_nodes.count(&node)) { + supported_to_visit.push(&node); + } else { + unsupported_to_visit.push(&node); + } + }; + + // find root nodes + for (auto& node : graph_viewer.Nodes()) { + size_t input_edge_count = node.GetInputEdgesCount(); + in_degree.insert({node.Index(), input_edge_count}); + if (input_edge_count == 0) { + add_to_visit(node); + } + } + + // prefer unsupported nodes first. this will increase the number of inputs potentially available to the first + // partition handled by this EP. + bool processing_supported_nodes = false; + + while (!supported_to_visit.empty() || !unsupported_to_visit.empty()) { + const Node* current = nullptr; + + // see if we need to flip + if ((processing_supported_nodes && supported_to_visit.empty()) || + (!processing_supported_nodes && unsupported_to_visit.empty())) { + processing_supported_nodes = !processing_supported_nodes; + continue; + } + + // get next node from same partition + if (processing_supported_nodes) { + current = supported_to_visit.front(); + supported_to_visit.pop(); + } else { + current = unsupported_to_visit.front(); + unsupported_to_visit.pop(); + } + + // when in_degree is zero all the inputs to the node are available + for (auto node_it = current->OutputNodesBegin(), end = current->OutputNodesEnd(); node_it != end; ++node_it) { + in_degree[node_it->Index()]--; + + if (in_degree[node_it->Index()] == 0) { + add_to_visit(*node_it); + } + } + + topo_order.push_back(&*current); + } + + // check we didn't break something + ORT_ENFORCE(graph_viewer.NumberOfNodes() == static_cast(topo_order.size()), + "Partition aware topological sort has produced invalid output."); + + return topo_order; +} + +std::unordered_set CreateExcludedNodeSet(const GraphViewer& graph_viewer, + const std::unordered_set& stop_ops) { + std::unordered_set excluded_nodes; + const auto end_stop_ops = stop_ops.cend(); + + for (const NodeIndex node_index : graph_viewer.GetNodesInTopologicalOrder()) { + const Node& node = *graph_viewer.GetNode(node_index); + + if (excluded_nodes.find(&node) == excluded_nodes.cend() && + stop_ops.find(node.OpType()) != end_stop_ops) { + excluded_nodes.insert(&node); + + // add all the downstream nodes + std::queue nodes_to_process; + nodes_to_process.push(&node); + + while (!nodes_to_process.empty()) { + const Node* cur_node = nodes_to_process.front(); + nodes_to_process.pop(); + + std::for_each(cur_node->OutputNodesBegin(), cur_node->OutputNodesEnd(), + [&nodes_to_process, &excluded_nodes](const Node& output_node) { + nodes_to_process.push(&output_node); + excluded_nodes.insert(&output_node); + }); + } + } + } + + return excluded_nodes; +} + +} // namespace + +std::unique_ptr MakeComputeCapability(const GraphViewer& graph_viewer, + const std::vector& group, + const std::function& generate_metadef_name, + const std::string& execution_provider_name) { + std::unordered_set node_set; + node_set.reserve(group.size()); + node_set.insert(group.cbegin(), group.cend()); + + std::unique_ptr sub_graph = std::make_unique(); + + std::unordered_set node_outputs; + std::unordered_set subgraph_inputs; + std::unordered_set subgraph_outputs; + std::vector ordered_subgraph_inputs; + std::vector ordered_subgraph_outputs; + + const auto& graph_output_list = graph_viewer.GetOutputs(); + std::unordered_set graph_outputs(graph_output_list.cbegin(), graph_output_list.cend()); + + for (const Node* node : group) { + sub_graph->nodes.push_back(node->Index()); + + for (const auto* input : node->InputDefs()) { + // if the node input was not produced by this subgraph, add it to the subgraph inputs. + if (node_outputs.count(input) == 0) { + if (subgraph_inputs.count(input) == 0) { + subgraph_inputs.insert(input); + ordered_subgraph_inputs.push_back(input); + } + } + } + + const auto& output_defs = node->OutputDefs(); + for (const auto* output_def : output_defs) { + node_outputs.insert(output_def); + // if output is overall graph output we need to produce it. + if (graph_outputs.count(output_def) != 0) { + ordered_subgraph_outputs.push_back(output_def); + } + } + + // if output connects to a node not in this subgraph we need to add it + // unless it was already added as an overall graph output, + for (auto it = node->OutputEdgesBegin(), end = node->OutputEdgesEnd(); it != end; ++it) { + if (node_set.count(&it->GetNode()) == 0) { + const auto* output_def = output_defs[it->GetSrcArgIndex()]; + if (subgraph_outputs.count(output_def) == 0 && graph_outputs.count(output_def) == 0) { + subgraph_outputs.insert(output_def); + ordered_subgraph_outputs.push_back(output_def); + } + } + } + } + + // Assign inputs and outputs to subgraph's meta_def + auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>(); + meta_def->name = generate_metadef_name(); + meta_def->domain = execution_provider_name; + meta_def->since_version = 1; + meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL; + + for (const auto& input : ordered_subgraph_inputs) { + meta_def->inputs.push_back(input->Name()); + } + + for (const auto& output : ordered_subgraph_outputs) { + meta_def->outputs.push_back(output->Name()); + } + + sub_graph->SetMetaDef(std::move(meta_def)); + + return std::make_unique(std::move(sub_graph)); +} + +std::vector> +CreateSupportedPartitions(const GraphViewer& graph_viewer, + const std::unordered_set& supported_nodes, + const std::unordered_set& stop_ops, + const std::function& generate_metadef_name, + const std::string& execution_provider_name, + bool debug_output) { + // find any nodes we need to exclude + std::unordered_set excluded_nodes = CreateExcludedNodeSet(graph_viewer, stop_ops); + +#ifndef NDEBUG + auto node_str = [](const Node& node) { + std::ostringstream oss; + oss << node.Index() << " '" << node.Name() << "'(" << node.OpType() << ")"; + return oss.str(); + }; + + auto group_str = [&node_str](const std::vector& group) { + const Node& start_node = *group.front(); + const Node& end_node = *group.back(); + std::ostringstream oss; + oss << node_str(start_node) << " to " << node_str(end_node) << "\n"; + return oss.str(); + }; +#endif + + // partition aware sort. this groups all the nodes we can and can't handle + const std::vector new_order = PartitionAwareTopoSort(graph_viewer, supported_nodes); + + // create groups using the new sort order + auto cur_topo_node = new_order.cbegin(); + auto end_topo_nodes = new_order.cend(); + + std::queue nodes_to_process; // supported nodes to process + std::unordered_set processed_nodes; // supported nodes we have processed + std::map> node_groups; + std::vector cur_group; + + bool check_excluded_nodes = !excluded_nodes.empty(); + const auto excluded_nodes_end = excluded_nodes.cend(); + + while (cur_topo_node != end_topo_nodes) { + const Node* node = *cur_topo_node; + ++cur_topo_node; + + if (processed_nodes.find(node) != processed_nodes.cend()) { + continue; + } + + if (check_excluded_nodes && excluded_nodes.find(node) != excluded_nodes_end) { + processed_nodes.insert(node); + continue; + } + + bool supported = supported_nodes.count(node) != 0; + bool in_partition = !cur_group.empty(); + + // check if end of a partition. + if (in_partition && !supported) { +#ifndef NDEBUG + if (debug_output) { + LOGS_DEFAULT(VERBOSE) << "New partition due to " << node_str(*node) + << ". Nodes in old partition: " << cur_group.size() << "\n"; + LOGS_DEFAULT(VERBOSE) << group_str(cur_group) << "\n"; + } +#else + ORT_UNUSED_PARAMETER(debug_output); +#endif + node_groups.insert({cur_group.front()->Index(), std::move(cur_group)}); + } + + // add the node and any connected downstream nodes that we can handle if supported. + // if not mark as processed so we know its inputs are available + if (supported) { + nodes_to_process.push(node); + } else { + processed_nodes.insert(node); + } + + while (!nodes_to_process.empty()) { + node = nodes_to_process.front(); + nodes_to_process.pop(); + + if (processed_nodes.find(node) == processed_nodes.cend()) { + // add to partition if all inputs available + bool inputs_available = true; + for (auto cur = node->InputNodesBegin(), end = node->InputNodesEnd(); cur != end; ++cur) { + if (processed_nodes.find(&*cur) == processed_nodes.cend()) { + inputs_available = false; + break; + } + } + + if (inputs_available) { + cur_group.push_back(node); + processed_nodes.insert(node); + + for (auto cur = node->OutputNodesBegin(), end = node->OutputNodesEnd(); cur != end; ++cur) { + const Node& downstream_node = *cur; + + // nodes will get added to the queue once per input from a supported node. + // we need this to happen as they can't be added to the group until all inputs are known to be available. + if (supported_nodes.count(&downstream_node) != 0) { + nodes_to_process.push(&downstream_node); + } + } + } else { + // we need all other nodes providing input to this node to have been processed + // before it can be added to cur_group. + // + // e.g. given A B with a topological order of A, B, C. + // \ / + // C + // + // When we process A we add C via the output edge to nodes_to_process. After we finish with A we look at C + // as the next node in nodes_to_process, but the input from B is missing. + // There are no more entries in nodes_to_process so we move to the next node in the topological order and + // process B. Again C is added to nodes_to_process via the output edge. After we finish with B we look at C + // again as the next node in nodes_to_process. + // Now all the inputs are available and C is added to the current group. + } + } + } + } + + if (!cur_group.empty()) { + node_groups.insert({cur_group.front()->Index(), std::move(cur_group)}); + } + + // create ComputeCapability instances + std::vector> results; + results.reserve(node_groups.size()); + + for (const auto& idx_to_group : node_groups) { + results.push_back( + MakeComputeCapability(graph_viewer, idx_to_group.second, generate_metadef_name, execution_provider_name)); + } + + return results; +} +} // namespace utils +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/partitioning_utils.h b/onnxruntime/core/providers/partitioning_utils.h new file mode 100644 index 0000000000..e739fafe2a --- /dev/null +++ b/onnxruntime/core/providers/partitioning_utils.h @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "core/graph/basic_types.h" + +namespace onnxruntime { +struct ComputeCapability; +class GraphViewer; +class NodeArg; +class Node; + +namespace utils { + +/** +Create the supported partitions for the execution provider. + +@param graph_viewer GraphViewer that IExecutionProvider::GetCapability is called with. +@param supported_nodes Set of nodes that the execution provider wants to handle. +@param stop_ops Set of operator names at which we stop considering nodes for assignment to this execution provider. +@param generate_metadef_name Functor to create the name for the MetaDef. + Most likely should call IExecutionProvider::GenerateMetaDefId. + See onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.cc + for example usage. +@param execution_provider_name Name of execution provider creating the ComputeCapability instance. +@param debug_output Print diagnostic output about the partitions and reasons for partition breaks. + No-op in a release build. + +@returns ComputeCapability instances for all partitions assigned to the execution provider. +*/ +std::vector> CreateSupportedPartitions( + const GraphViewer& graph_viewer, + const std::unordered_set& supported_nodes, + const std::unordered_set& stop_ops, + const std::function& generate_metadef_name, + const std::string& execution_provider_name, + bool debug_output = false); + +/** +Create a ComputeCapability instance from the group of nodes. +Will automatically determine the inputs and outputs required. + +@param graph_viewer GraphViewer that IExecutionProvider::GetCapability is called with. +@param group Group of nodes to include in the ComputeCapability instance. +@param generate_metadef_name Functor to create the name for the MetaDef. + Most likely should call IExecutionProvider::GenerateMetaDefId. + See onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.cc for + example usage. +@param execution_provider_name Name of execution provider creating the ComputeCapability instance. + +@returns New ComputeCapability instance. + +@remarks Prefer using CreateSupportedPartitions where possible, but if you need custom handling this provides a + convenient way to correctly create the ComputeCapability instance. +*/ +std::unique_ptr MakeComputeCapability(const GraphViewer& graph_viewer, + const std::vector& group, + const std::function& generate_metadef_name, + const std::string& execution_provider_name); +} // namespace utils +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.cc b/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.cc index 0ef998faf4..ccc552b581 100644 --- a/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.cc +++ b/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.cc @@ -11,15 +11,24 @@ #include "core/framework/tensorprotoutils.h" #include "core/framework/utils.h" #include "core/graph/model.h" +#include "core/providers/partitioning_utils.h" #include "core/session/onnxruntime_cxx_api.h" +#include + namespace onnxruntime { constexpr const char* INTERNAL_TESTING_EP = "InternalTestingEP"; -InternalTestingExecutionProvider::InternalTestingExecutionProvider(const std::unordered_set& ops) +InternalTestingExecutionProvider::InternalTestingExecutionProvider(const std::unordered_set& ops, + const std::unordered_set& stop_ops, + bool debug_output) : IExecutionProvider{utils::kInternalTestingExecutionProvider, true}, - ops_{ops} { + ep_name_{INTERNAL_TESTING_EP}, + ops_{ops}, + stop_ops_{stop_ops}, + debug_output_{debug_output} { + // // TODO: Allocation planner calls GetAllocator for the individual EP. It would be better if it goes through // the session state to get the allocator so it's per-device (or for the allocation planner to try the EP first // and fall back to using session state next by passing in a functor it can use to call SessionState::GetAllocator). @@ -27,7 +36,7 @@ InternalTestingExecutionProvider::InternalTestingExecutionProvider(const std::un AllocatorCreationInfo device_info( [](int) { return std::make_unique(OrtMemoryInfo(INTERNAL_TESTING_EP, - OrtAllocatorType::OrtDeviceAllocator)); + OrtAllocatorType::OrtDeviceAllocator)); }); InsertAllocator(CreateAllocator(device_info)); @@ -37,120 +46,50 @@ InternalTestingExecutionProvider::~InternalTestingExecutionProvider() {} std::vector> InternalTestingExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_viewer, - const std::vector& /*kernel_registries*/) const { - std::vector> result; + const std::vector& /*registries*/) const { + // find all supported nodes + std::unordered_set supported_nodes; - /* - Very basic search for groups of nodes that can be handled by the EP. - This doesn't work perfectly if you have a scenario like the following where A and D could be handled by the EP - but B is between them in the topological sort as you'll get two single node capabilities. However if can also - be advantageous if C and E could be handled by the EP as they would be combined with D even though not connected. - Not sure how often each of these scenarios happens. + const auto& topo_nodes = graph_viewer.GetNodesInTopologicalOrder(); + std::for_each(topo_nodes.cbegin(), topo_nodes.cend(), + [this, &supported_nodes, &graph_viewer](NodeIndex node_index) { + const Node* node = graph_viewer.GetNode(node_index); + bool supported = ops_.count(node->OpType()) != 0; + if (supported) { + supported_nodes.insert(node); + } + }); - A B C - | / | - D E - | | - - Would probably be better to walk the edges for each node the EP can handle as they are iterated in topological order, - accumulating nodes (and saving which ones have been taken) until you run out. This would guarantee all - connected nodes that can be handled are grouped together. - */ - - std::vector> node_groups; - std::vector cur_group; - for (NodeIndex node_index : graph_viewer.GetNodesInTopologicalOrder()) { - if (ops_.count(graph_viewer.GetNode(node_index)->OpType())) { - cur_group.push_back(node_index); - } else if (!cur_group.empty()) { - node_groups.push_back(std::move(cur_group)); - } + if (supported_nodes.empty()) { + return {}; } - if (!cur_group.empty()) { - node_groups.push_back(std::move(cur_group)); - } + // NOTE: GetCapability is called for all subgraphs from the bottom up, for one execution provider at a time. + // i.e. each execution provider will see the entire model individually. + // If your execution provider may selectively handle a control flow node (Scan/Loop/If) if it can process all nodes + // in the subgraph, here would be the place to check if: + // - you're processing a subgraph (graph_viewer.IsSubgraph() returns true) + // - and all nodes are handled (supported_nodes.size == graph_viewer.NumberOfNodes()) + // + // If that is the case and you wish to take the control flow node containing the subgraph: + // - return an empty vector so the nodes are left as is + // - note the node containing the subgraph (graph_viewer.ParentNode()) so that when GetCapability is called for + // the graph containing the parent node you can either: + // - include that node in supported_nodes if your Compile implementation can handle it potentially + // being part of a larger partition; or + // - create a ComputeCapability instance for just the control flow node by calling utils::MakeComputeCapability + // and adding the instance to the partitions returned by CreateSupportedPartitions - if (node_groups.empty()) { - return result; - } - - const auto& graph_output_list = graph_viewer.GetOutputs(); - std::unordered_set graph_outputs(graph_output_list.cbegin(), graph_output_list.cend()); - - for (const auto& group : node_groups) { - std::unordered_set node_set; - node_set.reserve(group.size()); - for (const auto& index : group) { - node_set.insert(index); - } - - std::unique_ptr sub_graph = std::make_unique(); - - std::unordered_set node_outputs; - std::unordered_set subgraph_inputs; - std::unordered_set subgraph_outputs; - std::vector ordered_subgraph_inputs; - std::vector ordered_subgraph_outputs; - - for (const auto& index : group) { - sub_graph->nodes.push_back(index); - const auto* node = graph_viewer.GetNode(index); - - for (const auto* input : node->InputDefs()) { - // if the node input was not produced by this subgraph, add it to the subgraph inputs. - if (node_outputs.count(input) == 0) { - if (subgraph_inputs.count(input) == 0) { - subgraph_inputs.insert(input); - ordered_subgraph_inputs.push_back(input); - } - } - } - - const auto& output_defs = node->OutputDefs(); - for (const auto* output_def : output_defs) { - node_outputs.insert(output_def); - // if output is overall graph output we need to produce it. - if (graph_outputs.count(output_def) != 0) { - ordered_subgraph_outputs.push_back(output_def); - } - } - - // if output connects to a node not in this subgraph we need to produce it - for (auto it = node->OutputEdgesBegin(), end = node->OutputEdgesEnd(); it != end; ++it) { - if (node_set.count(it->GetNode().Index()) == 0) { - const auto* output_def = output_defs[it->GetSrcArgIndex()]; - if (subgraph_outputs.count(output_def) == 0) { - subgraph_outputs.insert(output_def); - ordered_subgraph_outputs.push_back(output_def); - } - } - } - } - - // Assign inputs and outputs to subgraph's meta_def + // create functor to generate a guaranteed unique metadef id + auto generate_metadef_name = [this, &graph_viewer]() { uint64_t model_hash; int metadef_id = GenerateMetaDefId(graph_viewer, model_hash); auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>(); - meta_def->name = "InternalTestingEP_" + std::to_string(model_hash) + "_" + std::to_string(metadef_id); - meta_def->domain = "InternalTesting"; - meta_def->since_version = 1; - meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL; + return ep_name_ + "_" + std::to_string(model_hash) + "_" + std::to_string(metadef_id); + }; - for (const auto& input : ordered_subgraph_inputs) { - meta_def->inputs.push_back(input->Name()); - } - - for (const auto& output : ordered_subgraph_outputs) { - meta_def->outputs.push_back(output->Name()); - } - - sub_graph->SetMetaDef(std::move(meta_def)); - - result.push_back(std::make_unique(std::move(sub_graph))); - } - - return result; + return utils::CreateSupportedPartitions(graph_viewer, supported_nodes, stop_ops_, + generate_metadef_name, ep_name_, debug_output_); } common::Status InternalTestingExecutionProvider::Compile(const std::vector& fused_nodes, diff --git a/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.h b/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.h index 938a91ffc3..82df260094 100644 --- a/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.h +++ b/onnxruntime/test/providers/internal_testing/internal_testing_execution_provider.h @@ -8,7 +8,9 @@ namespace onnxruntime { class InternalTestingExecutionProvider : public IExecutionProvider { public: - InternalTestingExecutionProvider(const std::unordered_set& ops); + InternalTestingExecutionProvider(const std::unordered_set& ops, + const std::unordered_set& stop_ops = {}, + bool debug_output = false); virtual ~InternalTestingExecutionProvider(); std::vector> @@ -23,6 +25,18 @@ class InternalTestingExecutionProvider : public IExecutionProvider { } private: + const std::string ep_name_; + + // List of operators that the EP will claim nodes for const std::unordered_set ops_; + + // operators that we stop processing at. + // all nodes of an operator in this list and all their downstream nodes will be skipped + // e.g. NonMaxSuppression is the beginning of post-processing in an SSD model. It's unsupported for NNAPI, + // so from the NMS node on we want to use the CPU EP as the remaining work to do is far cheaper than + // the cost of going back to NNAPI. + const std::unordered_set stop_ops_; + + const bool debug_output_; }; } // namespace onnxruntime diff --git a/onnxruntime/test/providers/internal_testing/internal_testing_partitioning_tests.cc b/onnxruntime/test/providers/internal_testing/internal_testing_partitioning_tests.cc new file mode 100644 index 0000000000..953c1bfa64 --- /dev/null +++ b/onnxruntime/test/providers/internal_testing/internal_testing_partitioning_tests.cc @@ -0,0 +1,344 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/logging/logging.h" +#include "core/framework/compute_capability.h" +#include "core/framework/utils.h" +#include "core/session/inference_session.h" + +#include "test/framework/test_utils.h" +#include "test/test_environment.h" +#include "test/providers/internal_testing/internal_testing_execution_provider.h" +#include "test/util/include/asserts.h" +#include "test/util/include/inference_session_wrapper.h" +#include "test/util/include/test_utils.h" + +#include "gtest/gtest.h" +#include "gmock/gmock.h" + +#include + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::logging; + +namespace onnxruntime { + +namespace test { + +// tests use onnx format models currently so exclude them from a minimal build. +// it would be possible to use ORT format models but the same partitioning code would run either way +#if !defined(ORT_MINIMAL_BUILD) + +// model has an unsupported node between the supported nodes after the initial topo sort. +// the partition aware topo sort should result in the unsupported node moving to earlier in the order, +// and allow a single partition of supported nodes to be created. +TEST(InternalTestingEP, TestSortResultsInSinglePartition) { + auto run_test = [](const std::string& op) { + SessionOptions so; + auto session = std::make_unique(so, GetEnvironment()); + + const std::unordered_set supported_ops{op}; + + ASSERT_STATUS_OK(session->RegisterExecutionProvider( + std::make_unique(supported_ops))); + + const ORTCHAR_T* model_path = ORT_TSTR("testdata/ep_partitioning_test_1.onnx"); + ASSERT_STATUS_OK(session->Load(model_path)); + const auto& graph = session->GetGraph(); + GraphViewer viewer{graph}; + + ASSERT_STATUS_OK(session->Initialize()); + + const auto& func_mgr = session->GetSessionState().GetFuncMgr(); + NodeComputeInfo* compute_func = nullptr; + + int num_partitions{0}, num_other_nodes{0}; + + for (const auto& node : graph.Nodes()) { + EXPECT_EQ(supported_ops.count(node.OpType()), size_t(0)) + << "Nodes with supported op types should have been replaced. Node with type " << node.OpType() << " was not."; + if (node.GetExecutionProviderType() == utils::kInternalTestingExecutionProvider) { + EXPECT_STATUS_OK(func_mgr.GetFuncs(node.Name(), compute_func)); + EXPECT_NE(compute_func, nullptr); + ++num_partitions; + } else { + ++num_other_nodes; + } + } + + ASSERT_EQ(num_partitions, 1) << "Partition aware topological sort should have resulted in a single partition." + << " Op=" << op << " Partitions=" << num_partitions + << " OtherNodes=" << num_other_nodes; + }; + + // see testdata/ep_partitioning_tests.py for model description + // There should be only one partition, regardless of whether Add or Sub is supported by the EP. + run_test("Add"); + run_test("Sub"); +} + +// Test that when doing the partition aware sort and selecting groups that input dependencies are correctly handled +TEST(InternalTestingEP, TestDependenciesCorrectlyHandled) { + SessionOptions so; + auto session = std::make_unique(so, GetEnvironment()); + + const std::unordered_set supported_ops{"Add"}; + + ASSERT_STATUS_OK(session->RegisterExecutionProvider( + std::make_unique(supported_ops))); + + const ORTCHAR_T* model_path = ORT_TSTR("testdata/ep_partitioning_test_2.onnx"); + ASSERT_STATUS_OK(session->Load(model_path)); + const auto& graph = session->GetGraph(); + GraphViewer viewer{graph}; + + ASSERT_STATUS_OK(session->Initialize()); // this should fail if we don't process dependencies correctly + + const auto& func_mgr = session->GetSessionState().GetFuncMgr(); + NodeComputeInfo* compute_func = nullptr; + + int num_partitions{0}; + int num_other_nodes{0}; + + for (const auto& node : graph.Nodes()) { + EXPECT_EQ(supported_ops.count(node.OpType()), size_t(0)) + << "Nodes with supported op types should have been replaced. Node with type " << node.OpType() << " was not."; + if (node.GetExecutionProviderType() == utils::kInternalTestingExecutionProvider) { + EXPECT_STATUS_OK(func_mgr.GetFuncs(node.Name(), compute_func)); + EXPECT_NE(compute_func, nullptr); + ++num_partitions; + } else { + ++num_other_nodes; + } + } + + ASSERT_EQ(num_partitions, 2); + ASSERT_EQ(num_other_nodes, 2); +} + +// Infrastructure that was used to check NNAPI coverage. +// Ideally this could be updated to read the model paths, supported ops and stop ops from input files +// and provide info on the partitions so no code changes are required to investigate different scenarios. +static std::unordered_set GetNnapiSupportedOps() { + return std::unordered_set{ + "Add", + "Sub", + "Mul", + "Div", + "QLinearAdd", + "Pow", + "Relu", + "Transpose", + "Reshape", + "BatchNormalization", + "GlobalAveragePool", + "GlobalMaxPool", + "AveragePool", + "MaxPool", + "QLinearAveragePool", + "Conv", + "QLinearConv", + "Cast", + "Softmax", + "Identity", + "Gemm", + "MatMul", + "QLinearMatMul", + "Abs", + "Exp", + "Floor", + "Log", + "Sigmoid", + "Neg", + "Sin", + "Sqrt", + "Tanh", + "QLinearSigmoid", + "Concat", + "Squeeze", + "QuantizeLinear", + "DequantizeLinear", + "LRN", + "Clip", + "Resize", + "Flatten", + "Min", + "Max"}; +} + +struct PartitionStats { + int num_nodes_handled; + int num_nodes_not_handled; + int num_compiled_nodes; +}; + +static void TestNnapiPartitioning(const std::string& test_name, const std::string& model_uri, + bool optimize, bool debug_output, + const std::unordered_set& stop_ops, + const std::vector& additional_supported_ops, + PartitionStats& stats) { + SessionOptions so; + so.graph_optimization_level = optimize ? TransformerLevel::Level3 : TransformerLevel::Level1; + + auto session = std::make_unique(so, GetEnvironment()); + + // we disable NCHWc in mobile scenarios as it's not relevant to ARM + if (optimize) { + ASSERT_STATUS_OK(session->FilterEnabledOptimizers({"NchwcTransformer"})); + } + + auto ops = GetNnapiSupportedOps(); + for (const auto& op_type : additional_supported_ops) { + ops.insert(op_type); + } + + ASSERT_STATUS_OK(session->RegisterExecutionProvider( + std::make_unique(ops, stop_ops, debug_output))); + + ASSERT_STATUS_OK(session->Load(model_uri)); + const auto& graph = session->GetGraph(); + GraphViewer viewer{graph}; + + // save node count before optimization/partitioning. we lose some accuracy if optimizations replace nodes + const auto num_nodes = graph.NumberOfNodes(); + + ASSERT_STATUS_OK(session->Initialize()); + + // log the unsupported ops after initializer so that anything removed by constant folding etc. isn't listed. + std::unordered_map unsupported_ops; + std::ostringstream oss; + std::string unsupported_op_str; + + for (const Node& node : graph.Nodes()) { + if (node.GetExecutionProviderType() != utils::kInternalTestingExecutionProvider && + ops.count(node.OpType()) == 0) { + auto entry = unsupported_ops.find(node.OpType()); + if (entry != unsupported_ops.end()) { + ++entry->second; + } else { + unsupported_ops[node.OpType()] = 1; + } + } + } + + if (!unsupported_ops.empty()) { + bool add_comma = false; + for (const auto& pair : unsupported_ops) { + if (add_comma) { + oss << ","; + } + + oss << pair.first << "(" << pair.second << ")"; + add_comma = true; + } + + unsupported_op_str = oss.str(); + } + + const auto& func_mgr = session->GetSessionState().GetFuncMgr(); + NodeComputeInfo* compute_func = nullptr; + + stats.num_nodes_not_handled = 0; + stats.num_compiled_nodes = 0; + + // find the nodes downstream of the excluded nodes to check that they were assigned correctly + std::unordered_set excluded_nodes; + if (!stop_ops.empty()) { + for (const auto& node : graph.Nodes()) { + if (stop_ops.find(node.OpType()) != stop_ops.cend()) { + excluded_nodes.insert(&node); + + // add all the downstream nodes to the excluded list + std::queue nodes_to_add; + nodes_to_add.push(&node); + while (!nodes_to_add.empty()) { + const Node* cur_node = nodes_to_add.front(); + nodes_to_add.pop(); + + std::for_each(cur_node->OutputNodesBegin(), cur_node->OutputNodesEnd(), + [&nodes_to_add, &excluded_nodes](const Node& output_node) { + nodes_to_add.push(&output_node); + excluded_nodes.insert(&output_node); + }); + } + } + } + } + + for (const auto& node : graph.Nodes()) { + if (stop_ops.empty() || excluded_nodes.find(&node) == excluded_nodes.cend()) { + EXPECT_EQ(ops.count(node.OpType()), size_t(0)) + << "Nodes with supported op types should have been replaced. Node with type " + << node.OpType() << " was not."; + } else { + EXPECT_NE(node.GetExecutionProviderType(), utils::kInternalTestingExecutionProvider) + << "Node is downstream from a 'stop at' node and should not have been taken. Node:" + << node.Name(); + } + + if (node.GetExecutionProviderType() == utils::kInternalTestingExecutionProvider) { + EXPECT_STATUS_OK(func_mgr.GetFuncs(node.Name(), compute_func)); + EXPECT_NE(compute_func, nullptr); + ++stats.num_compiled_nodes; + } else { + ++stats.num_nodes_not_handled; + } + } + + stats.num_nodes_handled = num_nodes - stats.num_nodes_not_handled; + + auto pad_str = [](std::string const& str, size_t len = 10) { + return (str.size() < len) ? str + std::string(len - str.size(), ' ') : str; + }; + + std::cout << pad_str(test_name, 25) + << ": Compiled=" << stats.num_compiled_nodes + << " Handled=" << stats.num_nodes_handled + << " NotHandled=" << stats.num_nodes_not_handled + << " UnsupportedOps=" << unsupported_op_str + << "\n"; +} + +// DISABLED - manually update model_paths and enable to test coverage as needed +TEST(InternalTestingEP, DISABLED_TestNnapiPartitioningMlPerfModels) { + const auto supported_ops = GetNnapiSupportedOps(); + + // list the models you want to test here + std::vector model_paths = { + "deeplabv3_mnv2_ade20k_float.onnx", + "mobilebert.onnx", + "mobiledet.onnx", + }; + + for (const auto& model_uri : model_paths) { + auto run_tests = [&](bool optimize) { + std::cout << "\n================================\n"; + std::cout << "Model: " << model_uri; + if (optimize) { + std::cout << " (optimized)"; + } + std::cout << std::endl; + + const bool debug_output = false; + PartitionStats stats{}, stop_at_nms_stats{}, slice_stats{}; + + // arbitrary examples of running different combinations to test what partitioning results + TestNnapiPartitioning("Base", model_uri, optimize, debug_output, + {}, {}, stats); + + TestNnapiPartitioning("StopAt[NMS]", model_uri, optimize, debug_output, + {"NonMaxSuppression"}, {}, stop_at_nms_stats); + + TestNnapiPartitioning("ExtraOps[Slice]", model_uri, optimize, debug_output, + {}, {"Slice"}, slice_stats); + }; + + run_tests(false); + // run_tests(true); // optimized - if models have already be optimized this isn't helpful + } +} + +#endif // !defined(ORT_MINIMAL_BUILD) + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/testdata/ep_partitioning_test_1.onnx b/onnxruntime/test/testdata/ep_partitioning_test_1.onnx new file mode 100644 index 0000000000..ff9ab68284 --- /dev/null +++ b/onnxruntime/test/testdata/ep_partitioning_test_1.onnx @@ -0,0 +1,43 @@ +:´ + +input0 +input11A1"Add + +input0 +input12S1"Sub + +2 +input13_outA2"Add + +2 +input14S2"Sub + +1 +45_outA3"Add + +input1 +input26_outA4"AddgraphZ +input0 + + +Z +input1 + + +Z +input2 + + +b +3_out + + +b +5_out + + +b +6_out + + +B \ No newline at end of file diff --git a/onnxruntime/test/testdata/ep_partitioning_test_2.onnx b/onnxruntime/test/testdata/ep_partitioning_test_2.onnx new file mode 100644 index 0000000000..d0c62971ef --- /dev/null +++ b/onnxruntime/test/testdata/ep_partitioning_test_2.onnx @@ -0,0 +1,44 @@ +:š +! +input0 +input1s1_outS1"Sub +! +s1_out +input2a1_outA1"Add +! +a1_out +input0a2_outA2"Add +! +a1_out +input1s2_outS2"Sub +! +a2_out +input2a3_outA3"Add +! +a2_out +s2_outa4_outA4"Add +! +a3_out +input0a5_outA5"Add +! +a4_out +input1a6_outA6"Add +! +a5_out +a6_outa7_outA7"AddgraphZ +input0 + + +Z +input1 + + +Z +input2 + + +b +a7_out + + +B \ No newline at end of file diff --git a/onnxruntime/test/testdata/ep_partitioning_tests.py b/onnxruntime/test/testdata/ep_partitioning_tests.py new file mode 100644 index 0000000000..23bad4cdcf --- /dev/null +++ b/onnxruntime/test/testdata/ep_partitioning_tests.py @@ -0,0 +1,128 @@ +import numpy as np +import onnx +from onnx import helper +from onnx import TensorProto + + +# Create graph with Add and Sub nodes that can be used to test partitioning when one of the operators +# can run using the test EP and the other cannot. +# As the operators take 2 inputs and produce one output we can easily create edges to test different scenarios +def create_model_1(): + # Assume the EP can handle either Add or Sub but not both, and we need to minimize the partitions for nodes + # the EP can handle (as going to/from the EP has significant performance cost). + # + # graph inputs + # / \ \ + # a1 s1 \ + # | / \ \ + # | a2 s2 \ + # \ / | + # \ / | + # a3 a4 + # + # Assuming the initial topological sort is top down, left to right, we get a1, s1, a2, s2, a3, a4 + # + # Naively creating groups based on iterating this order and whether a node is supported gives the following groups + # (a1), (s1), (a2), (s2), (a3, a4). This is similar to what most EPs do currently. + # + # If we also consider downstream nodes with all inputs available when adding via the topological sort we get two + # less groups as s2 gets added with s1. + # (a1), (s1, s2), (a2, a3, a4) + # + # If the EP handles Sub that's fine. If the EP handles Add that's not. + # + # Finally, if we do a partition aware sort, and prefer unhandled nodes first to maximize the inputs that would be + # available each time we go to the EP, we can choose either of the root nodes (a1 or s1) to start at. + # + # If the EP is handling Sub we would start with a1 and get the same groups as above - which is perfectly fine as + # there's a single partition with (s1, s2) run on the EP. + # + # If the EP is handling Add we would start with s1 (due to preferring unhandled nodes first) and get the following + # groups, which also achieves a single partition of the handled nodes. + # (s2, s2), (a1, a2, a3, a4) + # + # So if this model is loaded in a partitioning test, there should only be one partition running on the EP regardless + # of whether Add or Sub is supported by it. + graph = helper.make_graph( + nodes= + [ + helper.make_node("Add", ['input0', 'input1'], ['1'], "A1"), + helper.make_node("Sub", ['input0', 'input1'], ["2"], "S1"), + helper.make_node("Add", ['2', 'input1'], ['3_out'], "A2"), + helper.make_node("Sub", ['2', 'input1'], ['4'], "S2"), + helper.make_node("Add", ['1', '4'], ['5_out'], "A3"), + helper.make_node("Add", ['input1', 'input2'], ['6_out'], "A4"), + ], + name="graph", + inputs= + [ + helper.make_tensor_value_info('input0', TensorProto.INT64, [1]), + helper.make_tensor_value_info('input1', TensorProto.INT64, [1]), + helper.make_tensor_value_info('input2', TensorProto.INT64, [1]), + ], + outputs= + [ + helper.make_tensor_value_info('3_out', TensorProto.INT64, [1]), + helper.make_tensor_value_info('5_out', TensorProto.INT64, [1]), + helper.make_tensor_value_info('6_out', TensorProto.INT64, [1]), + ], + initializer=[] + ) + + model = helper.make_model(graph) + return model + + +def create_model_2(): + # Create a model where there's a node that can't be run breaking up the partition. + # Partition aware topo sort should give us + # s1, [a1, a2, a3, a5], s2, [a4, a6, a7] + # + # graph inputs + # s1 + # | + # a1 + # / \ + # a2 s2 + # / \ / + # a3 a4 + # | | + # a5 a6 + # \ / + # a7 + graph = helper.make_graph( + nodes= + [ + helper.make_node("Sub", ['input0', 'input1'], ['s1_out'], "S1"), + helper.make_node("Add", ['s1_out', 'input2'], ['a1_out'], "A1"), + helper.make_node("Add", ['a1_out', 'input0'], ['a2_out'], "A2"), + helper.make_node("Sub", ['a1_out', 'input1'], ['s2_out'], "S2"), + helper.make_node("Add", ['a2_out', 'input2'], ['a3_out'], "A3"), + helper.make_node("Add", ['a2_out', 's2_out'], ['a4_out'], "A4"), + helper.make_node("Add", ['a3_out', 'input0'], ['a5_out'], "A5"), + helper.make_node("Add", ['a4_out', 'input1'], ['a6_out'], "A6"), + helper.make_node("Add", ['a5_out', 'a6_out'], ['a7_out'], "A7"), + ], + name="graph", + inputs= + [ + helper.make_tensor_value_info('input0', TensorProto.INT64, [1]), + helper.make_tensor_value_info('input1', TensorProto.INT64, [1]), + helper.make_tensor_value_info('input2', TensorProto.INT64, [1]), + ], + outputs= + [ + helper.make_tensor_value_info('a7_out', TensorProto.INT64, [1]), + ], + initializer=[] + ) + + model = helper.make_model(graph) + return model + + +if __name__ == '__main__': + model = create_model_1() + onnx.save(model, 'ep_partitioning_test_1.onnx') + model = create_model_2() + onnx.save(model, 'ep_partitioning_test_2.onnx')