mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Add enhanced partitioning utils for use by compiling EPs (#7991)
* 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.
This commit is contained in:
Scott McKay
GitHub
parent
35ca3c99d1
commit
00d48d9c30
8 changed files with 1024 additions and 109 deletions
336
onnxruntime/core/providers/partitioning_utils.cc
Normal file
336
onnxruntime/core/providers/partitioning_utils.cc
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/providers/partitioning_utils.h"
|
||||
|
||||
#include <queue>
|
||||
|
||||
#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<const Node*> PartitionAwareTopoSort(const GraphViewer& graph_viewer,
|
||||
const std::unordered_set<const Node*>& supported_nodes) {
|
||||
std::queue<const Node*> supported_to_visit, unsupported_to_visit;
|
||||
std::unordered_map<NodeIndex, size_t> in_degree;
|
||||
std::vector<const Node*> 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<int>(topo_order.size()),
|
||||
"Partition aware topological sort has produced invalid output.");
|
||||
|
||||
return topo_order;
|
||||
}
|
||||
|
||||
std::unordered_set<const Node*> CreateExcludedNodeSet(const GraphViewer& graph_viewer,
|
||||
const std::unordered_set<std::string>& stop_ops) {
|
||||
std::unordered_set<const Node*> 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<const Node*> 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<ComputeCapability> MakeComputeCapability(const GraphViewer& graph_viewer,
|
||||
const std::vector<const Node*>& group,
|
||||
const std::function<std::string()>& generate_metadef_name,
|
||||
const std::string& execution_provider_name) {
|
||||
std::unordered_set<const Node*> node_set;
|
||||
node_set.reserve(group.size());
|
||||
node_set.insert(group.cbegin(), group.cend());
|
||||
|
||||
std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
|
||||
|
||||
std::unordered_set<const NodeArg*> node_outputs;
|
||||
std::unordered_set<const NodeArg*> subgraph_inputs;
|
||||
std::unordered_set<const NodeArg*> subgraph_outputs;
|
||||
std::vector<const NodeArg*> ordered_subgraph_inputs;
|
||||
std::vector<const NodeArg*> ordered_subgraph_outputs;
|
||||
|
||||
const auto& graph_output_list = graph_viewer.GetOutputs();
|
||||
std::unordered_set<const NodeArg*> 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<ComputeCapability>(std::move(sub_graph));
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<ComputeCapability>>
|
||||
CreateSupportedPartitions(const GraphViewer& graph_viewer,
|
||||
const std::unordered_set<const Node*>& supported_nodes,
|
||||
const std::unordered_set<std::string>& stop_ops,
|
||||
const std::function<std::string()>& generate_metadef_name,
|
||||
const std::string& execution_provider_name,
|
||||
bool debug_output) {
|
||||
// find any nodes we need to exclude
|
||||
std::unordered_set<const Node*> 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<const Node*>& 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<const Node*> 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<const Node*> nodes_to_process; // supported nodes to process
|
||||
std::unordered_set<const Node*> processed_nodes; // supported nodes we have processed
|
||||
std::map<NodeIndex, std::vector<const Node*>> node_groups;
|
||||
std::vector<const Node*> 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<std::unique_ptr<ComputeCapability>> 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
|
||||
67
onnxruntime/core/providers/partitioning_utils.h
Normal file
67
onnxruntime/core/providers/partitioning_utils.h
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#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<std::unique_ptr<ComputeCapability>> CreateSupportedPartitions(
|
||||
const GraphViewer& graph_viewer,
|
||||
const std::unordered_set<const Node*>& supported_nodes,
|
||||
const std::unordered_set<std::string>& stop_ops,
|
||||
const std::function<std::string()>& 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<ComputeCapability> MakeComputeCapability(const GraphViewer& graph_viewer,
|
||||
const std::vector<const Node*>& group,
|
||||
const std::function<std::string()>& generate_metadef_name,
|
||||
const std::string& execution_provider_name);
|
||||
} // namespace utils
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -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 <queue>
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
constexpr const char* INTERNAL_TESTING_EP = "InternalTestingEP";
|
||||
|
||||
InternalTestingExecutionProvider::InternalTestingExecutionProvider(const std::unordered_set<std::string>& ops)
|
||||
InternalTestingExecutionProvider::InternalTestingExecutionProvider(const std::unordered_set<std::string>& ops,
|
||||
const std::unordered_set<std::string>& 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<CPUAllocator>(OrtMemoryInfo(INTERNAL_TESTING_EP,
|
||||
OrtAllocatorType::OrtDeviceAllocator));
|
||||
OrtAllocatorType::OrtDeviceAllocator));
|
||||
});
|
||||
|
||||
InsertAllocator(CreateAllocator(device_info));
|
||||
|
|
@ -37,120 +46,50 @@ InternalTestingExecutionProvider::~InternalTestingExecutionProvider() {}
|
|||
|
||||
std::vector<std::unique_ptr<ComputeCapability>>
|
||||
InternalTestingExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_viewer,
|
||||
const std::vector<const KernelRegistry*>& /*kernel_registries*/) const {
|
||||
std::vector<std::unique_ptr<ComputeCapability>> result;
|
||||
const std::vector<const KernelRegistry*>& /*registries*/) const {
|
||||
// find all supported nodes
|
||||
std::unordered_set<const Node*> 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<std::vector<NodeIndex>> node_groups;
|
||||
std::vector<NodeIndex> 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<const NodeArg*> graph_outputs(graph_output_list.cbegin(), graph_output_list.cend());
|
||||
|
||||
for (const auto& group : node_groups) {
|
||||
std::unordered_set<NodeIndex> node_set;
|
||||
node_set.reserve(group.size());
|
||||
for (const auto& index : group) {
|
||||
node_set.insert(index);
|
||||
}
|
||||
|
||||
std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
|
||||
|
||||
std::unordered_set<const NodeArg*> node_outputs;
|
||||
std::unordered_set<const NodeArg*> subgraph_inputs;
|
||||
std::unordered_set<const NodeArg*> subgraph_outputs;
|
||||
std::vector<const NodeArg*> ordered_subgraph_inputs;
|
||||
std::vector<const NodeArg*> 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<ComputeCapability>(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<FusedNodeAndGraph>& fused_nodes,
|
||||
|
|
|
|||
|
|
@ -8,7 +8,9 @@
|
|||
namespace onnxruntime {
|
||||
class InternalTestingExecutionProvider : public IExecutionProvider {
|
||||
public:
|
||||
InternalTestingExecutionProvider(const std::unordered_set<std::string>& ops);
|
||||
InternalTestingExecutionProvider(const std::unordered_set<std::string>& ops,
|
||||
const std::unordered_set<std::string>& stop_ops = {},
|
||||
bool debug_output = false);
|
||||
virtual ~InternalTestingExecutionProvider();
|
||||
|
||||
std::vector<std::unique_ptr<ComputeCapability>>
|
||||
|
|
@ -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<std::string> 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<std::string> stop_ops_;
|
||||
|
||||
const bool debug_output_;
|
||||
};
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -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 <queue>
|
||||
|
||||
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<InferenceSessionWrapper>(so, GetEnvironment());
|
||||
|
||||
const std::unordered_set<std::string> supported_ops{op};
|
||||
|
||||
ASSERT_STATUS_OK(session->RegisterExecutionProvider(
|
||||
std::make_unique<InternalTestingExecutionProvider>(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<InferenceSessionWrapper>(so, GetEnvironment());
|
||||
|
||||
const std::unordered_set<std::string> supported_ops{"Add"};
|
||||
|
||||
ASSERT_STATUS_OK(session->RegisterExecutionProvider(
|
||||
std::make_unique<InternalTestingExecutionProvider>(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<std::string> GetNnapiSupportedOps() {
|
||||
return std::unordered_set<std::string>{
|
||||
"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<std::string>& stop_ops,
|
||||
const std::vector<std::string>& additional_supported_ops,
|
||||
PartitionStats& stats) {
|
||||
SessionOptions so;
|
||||
so.graph_optimization_level = optimize ? TransformerLevel::Level3 : TransformerLevel::Level1;
|
||||
|
||||
auto session = std::make_unique<InferenceSessionWrapper>(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<InternalTestingExecutionProvider>(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<std::string, int> 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<const Node*> 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<const Node*> 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<std::string> 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
|
||||
43
onnxruntime/test/testdata/ep_partitioning_test_1.onnx
vendored
Normal file
43
onnxruntime/test/testdata/ep_partitioning_test_1.onnx
vendored
Normal file
|
|
@ -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
|
||||
44
onnxruntime/test/testdata/ep_partitioning_test_2.onnx
vendored
Normal file
44
onnxruntime/test/testdata/ep_partitioning_test_2.onnx
vendored
Normal file
|
|
@ -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
|
||||
128
onnxruntime/test/testdata/ep_partitioning_tests.py
vendored
Normal file
128
onnxruntime/test/testdata/ep_partitioning_tests.py
vendored
Normal file
|
|
@ -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')
|
||||
Loading…
Reference in a new issue