saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve VerifyKernelDef() performance when op has many inputs/outputs/type constraints. (#347)

Browse source 
* Improve VerifyKernelDef() performance when op has many inputs/outputs/type constraints.

* Added two modes for resolving type binding.

* Updated TypeBindingResolver to avoid heap allocation.

* Tweaked TypeBindingResolver for performance.
This commit is contained in:
edgchen1 2019-01-24 17:34:03 -08:00 committed by Ke Zhang
parent 644c13050b
commit f922d427d3
1 changed files with 120 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

160
onnxruntime/core/framework/kernel_registry.cc
View file

@ -1,48 +1,117 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <memory>
#include <unordered_map>
#include "core/framework/kernel_registry.h"
using namespace ::onnxruntime::common;
namespace onnxruntime {
// Find the type that name is bound to in the given node.
// "name" can represent either a type parameter or an input/output parameter.
// Returns null if a match is not found.
const ::ONNX_NAMESPACE::TypeProto* FindTypeBinding(const onnxruntime::Node& node, const std::string& name) {
namespace {
// Traverses the node's formal parameters and calls TraverseFn with the formal
// parameter and its associated TypeProto.
// node - the node to traverse
// param_filter_fn - called to determine whether to consider a given formal parameter:
// bool ParamFilterFn(const ONNX_NAMESPACE::OpSchema::FormalParameter& param)
// param - the formal parameter
// returns true if the formal parameter should be considered, false otherwise
// traverse_fn - called to process the formal parameter and its associated TypeProto:
// bool TraverseFn(const ONNX_NAMESPACE::OpSchema::FormalParameter& param,
// const ONNX_NAMESPACE::TypeProto* type)
// param - the formal paremeter
// type - the associated TypeProto
// returns true if traversal should continue, false otherwise
template <typename ParamFilterFn, typename TraverseFn>
void TraverseFormalParametersWithTypeProto(const Node& node,
ParamFilterFn param_filter_fn,
TraverseFn traverse_fn) {
const ONNX_NAMESPACE::OpSchema& op_schema = *node.Op();
// search inputs:
// process inputs:
const size_t len = node.InputArgCount().size();
ORT_ENFORCE(len <= op_schema.inputs().size());
int actual_index = 0;
for (size_t formal_index = 0; formal_index != len; ++formal_index) {
auto& param = op_schema.inputs()[formal_index];
if ((param.GetTypeStr() == name) || (param.GetName() == name)) {
// return type of any corresponding actual parameter, if present
const auto& param = op_schema.inputs()[formal_index];
if (param_filter_fn(param)) {
// get type of any corresponding actual parameter, if present
for (int i = 0, end = node.InputArgCount()[formal_index]; i < end; ++i) {
const onnxruntime::NodeArg* arg = node.InputDefs()[actual_index + i];
const NodeArg* arg = node.InputDefs()[actual_index + i];
if (!arg->Exists()) continue; // a missing optional argument
return arg->TypeAsProto();
if (!traverse_fn(param, arg->TypeAsProto())) return;
}
}
actual_index += node.InputArgCount()[formal_index];
}
// search outputs:
// process outputs:
auto& actual_outputs = node.OutputDefs();
auto num_actual_outputs = actual_outputs.size();
auto last_formal = op_schema.outputs().size() - 1;
const auto num_actual_outputs = actual_outputs.size();
const auto last_formal = op_schema.outputs().size() - 1;
for (size_t i = 0; i != num_actual_outputs; ++i) {
const onnxruntime::NodeArg* arg = actual_outputs[i];
const auto& formal = op_schema.outputs()[std::min(i, last_formal)];
if (!param_filter_fn(formal)) continue;
const NodeArg* arg = actual_outputs[i];
if (!arg->Exists()) continue;
auto& formal = op_schema.outputs()[std::min(i, last_formal)];
const auto& formal_typestr = formal.GetTypeStr(); // for easier debugging
const auto& formal_name = formal.GetName(); // for easier debugging
if ((formal_typestr == name) || (formal_name == name)) {
return arg->TypeAsProto();
if (!traverse_fn(formal, arg->TypeAsProto())) return;
}
}
class TypeBindingResolver {
public:
TypeBindingResolver(const Node& node, bool use_lookup_map)
: node_{node},
type_binding_map_{} {
if (use_lookup_map) {
type_binding_map_ = std::make_unique<TypeBindingMap>();
TraverseFormalParametersWithTypeProto(
node_,
[](const ONNX_NAMESPACE::OpSchema::FormalParameter&) -> bool { return true; },
[this](const ONNX_NAMESPACE::OpSchema::FormalParameter& param,
const ONNX_NAMESPACE::TypeProto* type) -> bool {
type_binding_map_->emplace(param.GetName(), type);
type_binding_map_->emplace(param.GetTypeStr(), type);
return true;
});
}
}
return nullptr;
}
// Resolves a name to a TypeProto* for a given node.
// The name can represent either a type parameter or an input/output parameter.
// Returns the resolved TypeProto* or nullptr if unable to resolve.
const ONNX_NAMESPACE::TypeProto* Resolve(const std::string& name_or_type_str) const {
// lookup if available
if (type_binding_map_) {
auto found_it = type_binding_map_->find(name_or_type_str);
if (found_it == type_binding_map_->end()) return nullptr;
return found_it->second;
}
// fall back to node parameter traversal
const ONNX_NAMESPACE::TypeProto* result{};
TraverseFormalParametersWithTypeProto(
node_,
[&name_or_type_str](const ONNX_NAMESPACE::OpSchema::FormalParameter& param) -> bool {
return param.GetName() == name_or_type_str || param.GetTypeStr() == name_or_type_str;
},
[&result](const ONNX_NAMESPACE::OpSchema::FormalParameter&,
const ONNX_NAMESPACE::TypeProto* type) -> bool {
result = type;
return false;
});
return result;
}
private:
// map from input/output name or type string to TypeProto pointer
using TypeBindingMap = std::unordered_map<std::string, const ONNX_NAMESPACE::TypeProto*>;
const Node& node_;
std::unique_ptr<TypeBindingMap> type_binding_map_;
};
}; // namespace
std::vector<std::string> KernelRegistry::GetAllRegisteredOpNames() const {
std::vector<std::string> ret(kernel_creator_fn_map_.size());
@ -86,7 +155,7 @@ bool KernelRegistry::VerifyKernelDef(const onnxruntime::Node& node,
ostr << "Op: " << node.OpType()
<< " Execution provider mismatch."
<< " Expected: " << expected_provider
<< " Acutal: " << kernel_def.Provider();
<< " Actual: " << kernel_def.Provider();
error_str = ostr.str();
return false;
}
@ -122,10 +191,21 @@ bool KernelRegistry::VerifyKernelDef(const onnxruntime::Node& node,
// check if type matches
auto& kernel_type_constraints = kernel_def.TypeConstraints();
// Note: The number of formal input/output parameters is N and the number of
// type constraints is M. We select between an O(N*M) and an O(N+M) approach.
// The O(N*M) approach has lower initial overhead.
// kTypeBindingResolverComplexityThreshold is the value of N*M above which we
// will use the O(N+M) approach.
constexpr int kTypeBindingResolverComplexityThreshold = 50 * 50;
const bool use_lookup_map = (kernel_type_constraints.size() * (node.Op()->inputs().size() + node.Op()->outputs().size()) >
kTypeBindingResolverComplexityThreshold);
TypeBindingResolver type_binding_resolver{node, use_lookup_map};
for (auto& constraint : kernel_type_constraints) {
const std::string& name = constraint.first;
const std::vector<MLDataType>& allowed_types = constraint.second;
const ::ONNX_NAMESPACE::TypeProto* actual_type = FindTypeBinding(node, name);
const ONNX_NAMESPACE::TypeProto* actual_type = type_binding_resolver.Resolve(name);
// If actual_type is null, this represents a type-constraint on a
// missing optional parameter, which can be skipped.
@ -202,23 +282,23 @@ static std::string ToString(const std::vector<std::string>& error_strs) {
const KernelCreateInfo* KernelRegistry::TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider) const {
auto range = kernel_creator_fn_map_.equal_range(node.OpType());
std::vector<std::string> error_strs;
for (auto i = range.first; i != range.second; ++i) {
if (!i->second.status.IsOK()) {
LOGS_DEFAULT(ERROR) << "Failed to create kernel for op: " << node.OpType()
<< " since it was ill-formed during registration";
continue;
}
std::string error_str;
if (VerifyKernelDef(node, *i->second.kernel_def, error_str, exec_provider)) {
return &i->second;
}
error_strs.push_back(error_str);
}
LOGS_DEFAULT(INFO) << node.OpType() << " kernel is not supported in " << exec_provider
<< " Encountered following errors: " << ToString(error_strs);
return nullptr;
auto range = kernel_creator_fn_map_.equal_range(node.OpType());
std::vector<std::string> error_strs;
for (auto i = range.first; i != range.second; ++i) {
if (!i->second.status.IsOK()) {
LOGS_DEFAULT(ERROR) << "Failed to create kernel for op: " << node.OpType()
<< " since it was ill-formed during registration";
continue;
}
std::string error_str;
if (VerifyKernelDef(node, *i->second.kernel_def, error_str, exec_provider)) {
return &i->second;
}
error_strs.push_back(error_str);
}
LOGS_DEFAULT(INFO) << node.OpType() << " kernel is not supported in " << exec_provider
<< " Encountered following errors: " << ToString(error_strs);
return nullptr;
}
} // namespace onnxruntime