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')