Add kernel def hash logic for minimal build (#4891)

* Add hash based lookup of kernels
This commit is contained in:
Scott McKay 2020-08-23 14:39:07 +10:00 committed by GitHub
parent db7669b225
commit 728e886bba
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 182 additions and 82 deletions

View file

@ -242,4 +242,11 @@ inline std::wstring ToWideString(const std::wstring& s) { return s; }
inline std::string ToWideString(const std::string& s) { return s; }
#endif
// from http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2014/n3876.pdf
template <class T>
inline void HashCombine(std::uint64_t& seed, const T& v) {
std::hash<T> hasher;
seed ^= hasher(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
}
} // namespace onnxruntime

View file

@ -88,9 +88,37 @@ class KernelDef {
bool IsConflict(const KernelDef& other) const;
uint64_t GetHash() const noexcept {
// if we need to support different hash versions we can update CalculateHash to take a version number
// and calculate any non-default versions dynamically. we only use this during kernel lookup so
// it's not performance critical
return hash_;
}
private:
friend class KernelDefBuilder;
// call once the KernelDef has been built
void CalculateHash() {
// use name, start/end, domain, provider and the type constraints.
// we wouldn't have two kernels that only differed by the inplace or alias info or memory types.
// currently nothing sets exec_queue_id either (and would assumably be a runtime thing and not part of the base
// kernel definition)
hash_ = 0; // reset in case this is called multiple times
HashCombine(hash_, op_name_);
HashCombine(hash_, op_since_version_start_);
HashCombine(hash_, op_since_version_end_);
HashCombine(hash_, op_domain_);
HashCombine(hash_, provider_type_);
for (const auto& key_value : type_constraints_) {
HashCombine(hash_, key_value.first);
for (const auto& data_type : key_value.second) {
// need to construct a std::string so it doesn't hash the address of a const char*
HashCombine(hash_, std::string(DataTypeImpl::ToString(data_type)));
}
}
}
// The operator name supported by <*this> kernel..
std::string op_name_;
@ -128,6 +156,9 @@ class KernelDef {
OrtMemType default_inputs_mem_type_{OrtMemTypeDefault};
// Default memory type for all outputs
OrtMemType default_outputs_mem_type_{OrtMemTypeDefault};
// hash of kernel definition for lookup in minimal build
uint64_t hash_ = 0;
};
class KernelDefBuilder {
@ -283,6 +314,7 @@ class KernelDefBuilder {
Return the kernel definition, passing ownership of the KernelDef to the caller
*/
std::unique_ptr<KernelDef> Build() {
kernel_def_->CalculateHash();
return std::move(kernel_def_);
}

View file

@ -19,6 +19,14 @@ class KernelRegistry {
Status Register(KernelCreateInfo&& create_info) ORT_MUST_USE_RESULT;
#if !defined(ORT_MINIMAL_BUILD)
static bool HasImplementationOf(const KernelRegistry& r, const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider) {
const KernelCreateInfo* info;
Status st = r.TryFindKernel(node, exec_provider, &info);
return st.IsOK();
}
// factory functions should always return a unique_ptr for maximum flexibility
// for its clients unless the factory is managing the lifecycle of the pointer
// itself.
@ -30,15 +38,16 @@ class KernelRegistry {
std::unique_ptr<OpKernel>& op_kernel) const ORT_MUST_USE_RESULT;
// Check if an execution provider can create kernel for a node and return the kernel if so
Status TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider, const KernelCreateInfo** out) const;
Status TryFindKernel(const onnxruntime::Node& node, onnxruntime::ProviderType exec_provider,
const KernelCreateInfo** out) const;
static bool HasImplementationOf(const KernelRegistry& r, const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider) {
const KernelCreateInfo* info;
Status st = r.TryFindKernel(node, exec_provider, &info);
return st.IsOK();
}
#endif
// Check if an execution provider can create kernel for a node and return the kernel if so.
// Kernel matching is via kernel_def_hash.
Status TryFindKernel(const onnxruntime::Node& node, onnxruntime::ProviderType exec_provider,
uint64_t kernel_def_hash,
const KernelCreateInfo** out) const;
bool IsEmpty() const { return kernel_creator_fn_map_.empty(); }
@ -50,6 +59,7 @@ class KernelRegistry {
#endif
private:
#if !defined(ORT_MINIMAL_BUILD)
// Check whether the types of inputs/outputs of the given node match the extra
// type-constraints of the given kernel. This serves two purposes: first, to
// select the right kernel implementation based on the types of the arguments
@ -67,6 +77,7 @@ class KernelRegistry {
static bool VerifyKernelDef(const onnxruntime::Node& node,
const KernelDef& kernel_def,
std::string& error_str);
#endif
static std::string GetMapKey(const std::string& op_name, const std::string& domain, const std::string& provider) {
std::string key(op_name);

View file

@ -29,6 +29,7 @@ std::vector<std::unique_ptr<ComputeCapability>>
IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
const std::vector<const KernelRegistry*>& kernel_registries) const {
std::vector<std::unique_ptr<ComputeCapability>> result;
#if !defined(ORT_MINIMAL_BUILD)
for (auto& node : graph.Nodes()) {
for (auto registry : kernel_registries) {
if (KernelRegistry::HasImplementationOf(*registry, node, Type())) {
@ -39,6 +40,9 @@ IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
}
}
}
#else
ORT_NOT_IMPLEMENTED("IExecutionProvider::GetCapability is not supported in this build.");
#endif
return result;
}

View file

@ -200,8 +200,102 @@ bool KernelRegistry::VerifyKernelDef(const onnxruntime::Node& node,
return true;
}
Status KernelRegistry::TryCreateKernel(const onnxruntime::Node& node,
const IExecutionProvider& execution_provider,
const std::unordered_map<int, OrtValue>& constant_initialized_tensors,
const OrtValueNameIdxMap& ort_value_name_idx_map,
const FuncManager& funcs_mgr,
const DataTransferManager& data_transfer_mgr,
/*out*/ std::unique_ptr<OpKernel>& op_kernel) const {
const KernelCreateInfo* kernel_create_info = nullptr;
ORT_RETURN_IF_ERROR(TryFindKernel(node, execution_provider.Type(), &kernel_create_info));
OpKernelInfo kernel_info(node,
*kernel_create_info->kernel_def,
execution_provider,
constant_initialized_tensors,
ort_value_name_idx_map,
funcs_mgr,
data_transfer_mgr);
op_kernel.reset(kernel_create_info->kernel_create_func(kernel_info));
return Status::OK();
}
static std::string ToString(const std::vector<std::string>& error_strs) {
std::ostringstream ostr;
std::for_each(std::begin(error_strs), std::end(error_strs),
[&ostr](const std::string& str) { ostr << str << "\n"; });
return ostr.str();
}
Status KernelRegistry::TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider,
const KernelCreateInfo** out) const {
return TryFindKernel(node, exec_provider, uint64_t(0), out);
}
#endif // !defined(ORT_MINIMAL_BUILD)
// It's often this function returns a failed status, but it is totally expected.
// It just means this registry doesn't have such a kernel, please search it elsewhere.
// if this function is called before graph partition, then node.provider is not set.
// In this case, the kernel's provider must equal to exec_provider
// otherwise, kernel_def.provider must equal to node.provider. exec_provider is ignored.
Status KernelRegistry::TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider,
uint64_t kernel_def_hash,
const KernelCreateInfo** out) const {
const auto& node_provider = node.GetExecutionProviderType();
const auto& expected_provider = (node_provider.empty() ? exec_provider : node_provider);
auto range = kernel_creator_fn_map_.equal_range(GetMapKey(node.OpType(), node.Domain(), expected_provider));
*out = nullptr;
// if we have a hash (ORT format model) use only that.
if (kernel_def_hash != 0) {
for (auto i = range.first; i != range.second; ++i) {
if (i->second.kernel_def->GetHash() == kernel_def_hash) {
*out = &i->second;
return Status::OK();
}
}
std::ostringstream oss;
oss << "Op with name (" << node.Name() << ")"
<< " and type (" << node.OpType() << ")"
<< " kernel not found in " << expected_provider << "."
<< " No matching hash for " << kernel_def_hash;
return Status(ONNXRUNTIME, FAIL, oss.str());
}
#if !defined(ORT_MINIMAL_BUILD)
else {
std::vector<std::string> verify_kernel_def_error_strs;
for (auto i = range.first; i != range.second; ++i) {
std::string error_str;
if (VerifyKernelDef(node, *i->second.kernel_def, error_str)) {
*out = &i->second;
return Status::OK();
}
verify_kernel_def_error_strs.push_back(error_str);
}
if (!verify_kernel_def_error_strs.empty()) {
std::ostringstream oss;
oss << "Op with name (" << node.Name() << ")"
<< " and type (" << node.OpType() << ")"
<< " kernel is not supported in " << expected_provider << "."
<< " Encountered following errors: (" << ToString(verify_kernel_def_error_strs) << ")";
return Status(ONNXRUNTIME, FAIL, oss.str());
}
}
return Status(ONNXRUNTIME, FAIL, "Kernel not found");
#else
ORT_THROW("Kernel hash must be provided in minimal build.")
#endif
}
Status KernelRegistry::Register(KernelDefBuilder& kernel_builder,
const KernelCreateFn& kernel_creator) {
return Register(KernelCreateInfo(kernel_builder.Build(), kernel_creator));
@ -229,71 +323,4 @@ Status KernelRegistry::Register(KernelCreateInfo&& create_info) {
return Status::OK();
}
Status KernelRegistry::TryCreateKernel(const onnxruntime::Node& node,
const IExecutionProvider& execution_provider,
const std::unordered_map<int, OrtValue>& constant_initialized_tensors,
const OrtValueNameIdxMap& ort_value_name_idx_map,
const FuncManager& funcs_mgr,
const DataTransferManager& data_transfer_mgr,
/*out*/ std::unique_ptr<OpKernel>& op_kernel) const {
const KernelCreateInfo* kernel_create_info = nullptr;
ORT_RETURN_IF_ERROR(TryFindKernel(node, execution_provider.Type(), &kernel_create_info));
OpKernelInfo kernel_info(node,
*kernel_create_info->kernel_def,
execution_provider,
constant_initialized_tensors,
ort_value_name_idx_map,
funcs_mgr,
data_transfer_mgr);
op_kernel.reset(kernel_create_info->kernel_create_func(kernel_info));
return Status::OK();
}
static std::string ToString(const std::vector<std::string>& error_strs) {
std::ostringstream ostr;
std::for_each(std::begin(error_strs), std::end(error_strs),
[&ostr](const std::string& str) { ostr << str << "\n"; });
return ostr.str();
}
// It's often this function returns a failed status, but it is totally expected.
// It just means this registry doesn't have such a kernel, please search it elsewhere.
// if this function is called before graph partition, then node.provider is not set.
// In this case, the kernel's provider must equal to exec_provider
// otherwise, kernel_def.provider must equal to node.provider. exec_provider is ignored.
Status KernelRegistry::TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider, const KernelCreateInfo** out) const {
const auto& node_provider = node.GetExecutionProviderType();
const auto& expected_provider = (node_provider.empty() ? exec_provider : node_provider);
auto range = kernel_creator_fn_map_.equal_range(GetMapKey(node.OpType(), node.Domain(), expected_provider));
std::vector<std::string> verify_kernel_def_error_strs;
for (auto i = range.first; i != range.second; ++i) {
#if !defined(ORT_MINIMAL_BUILD)
std::string error_str;
if (VerifyKernelDef(node, *i->second.kernel_def, error_str)) {
*out = &i->second;
return Status::OK();
}
verify_kernel_def_error_strs.push_back(error_str);
#else
// TODO: Add hash based lookup
#endif
}
if (!verify_kernel_def_error_strs.empty()) {
std::ostringstream oss;
oss << "Op with name (" << node.Name() << ")"
<< " and type (" << node.OpType() << ")"
<< " kernel is not supported in " << expected_provider << "."
<< " Encountered following errors: (" << ToString(verify_kernel_def_error_strs) << ")";
return Status(ONNXRUNTIME, FAIL, oss.str());
}
*out = nullptr;
return Status(ONNXRUNTIME, FAIL, "Kernel not found");
}
} // namespace onnxruntime

View file

@ -45,6 +45,7 @@ Status KernelRegistryManager::RegisterKernels(const ExecutionProviders& executio
return Status::OK();
}
#if !defined(ORT_MINIMAL_BUILD)
void KernelRegistryManager::RegisterKernelRegistry(std::shared_ptr<KernelRegistry> kernel_registry) {
if (nullptr == kernel_registry) {
return;
@ -61,6 +62,12 @@ bool KernelRegistryManager::HasImplementationOf(const KernelRegistryManager& r,
Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node,
/*out*/ const KernelCreateInfo** kernel_create_info) const {
return SearchKernelRegistry(node, uint64_t(0), kernel_create_info);
}
#endif
Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node, uint64_t kernel_def_hash,
/*out*/ const KernelCreateInfo** kernel_create_info) const {
Status status;
auto create_error_message = [&node, &status](const std::string& prefix) {
@ -77,10 +84,12 @@ Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node
return Status(ONNXRUNTIME, FAIL, create_error_message("The node is not placed on any Execution Provider. "));
}
#if !defined(ORT_MINIMAL_BUILD)
for (auto& registry : custom_kernel_registries_) {
status = registry->TryFindKernel(node, std::string(), kernel_create_info);
status = registry->TryFindKernel(node, std::string(), kernel_def_hash, kernel_create_info);
if (status.IsOK()) return status;
}
#endif
KernelRegistry* p = nullptr;
auto iter = provider_type_to_registry_.find(ptype);
@ -89,8 +98,10 @@ Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node
}
if (p != nullptr) {
status = p->TryFindKernel(node, std::string(), kernel_create_info);
if (status.IsOK()) return status;
status = p->TryFindKernel(node, std::string(), kernel_def_hash, kernel_create_info);
if (status.IsOK()) {
return status;
}
}
return Status(ONNXRUNTIME, NOT_IMPLEMENTED, create_error_message("Failed to find kernel for "));

View file

@ -32,6 +32,8 @@ class KernelRegistryManager {
// Register kernels from providers
Status RegisterKernels(const ExecutionProviders& execution_providers) ORT_MUST_USE_RESULT;
#if !defined(ORT_MINIMAL_BUILD)
// The registry passed in this function has highest priority than anything already in this KernelRegistryManager,
// and anything registered from RegisterKernels
// For example, if you do:
@ -46,10 +48,6 @@ class KernelRegistryManager {
Status SearchKernelRegistry(const onnxruntime::Node& node,
/*out*/ const KernelCreateInfo** kernel_create_info) const;
std::unique_ptr<OpKernel> CreateKernel(const onnxruntime::Node& node, const IExecutionProvider& execution_provider,
const SessionState& session_state,
const KernelCreateInfo& kernel_create_info) const ORT_MUST_USE_RESULT;
/**
* Whether this node can be run on this provider
*/
@ -70,6 +68,16 @@ class KernelRegistryManager {
if (iter != provider_type_to_registry_.end()) result.push_back(iter->second.get());
return result;
}
#endif
Status SearchKernelRegistry(const onnxruntime::Node& node,
uint64_t kernel_def_hash,
/*out*/ const KernelCreateInfo** kernel_create_info) const;
std::unique_ptr<OpKernel> CreateKernel(const onnxruntime::Node& node,
const IExecutionProvider& execution_provider,
const SessionState& session_state,
const KernelCreateInfo& kernel_create_info) const ORT_MUST_USE_RESULT;
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(KernelRegistryManager);