diff --git a/include/onnxruntime/core/common/common.h b/include/onnxruntime/core/common/common.h index f6ff7d3fdf..eda0f429c6 100644 --- a/include/onnxruntime/core/common/common.h +++ b/include/onnxruntime/core/common/common.h @@ -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 +inline void HashCombine(std::uint64_t& seed, const T& v) { + std::hash hasher; + seed ^= hasher(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2); +} + } // namespace onnxruntime diff --git a/include/onnxruntime/core/framework/kernel_def_builder.h b/include/onnxruntime/core/framework/kernel_def_builder.h index 1b7027668b..0b47e6743c 100644 --- a/include/onnxruntime/core/framework/kernel_def_builder.h +++ b/include/onnxruntime/core/framework/kernel_def_builder.h @@ -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 Build() { + kernel_def_->CalculateHash(); return std::move(kernel_def_); } diff --git a/include/onnxruntime/core/framework/kernel_registry.h b/include/onnxruntime/core/framework/kernel_registry.h index ca92da91c6..0542183d6d 100644 --- a/include/onnxruntime/core/framework/kernel_registry.h +++ b/include/onnxruntime/core/framework/kernel_registry.h @@ -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& 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); diff --git a/onnxruntime/core/framework/execution_provider.cc b/onnxruntime/core/framework/execution_provider.cc index 4c34291315..f771f33f00 100644 --- a/onnxruntime/core/framework/execution_provider.cc +++ b/onnxruntime/core/framework/execution_provider.cc @@ -29,6 +29,7 @@ std::vector> IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, const std::vector& kernel_registries) const { std::vector> 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; } diff --git a/onnxruntime/core/framework/kernel_registry.cc b/onnxruntime/core/framework/kernel_registry.cc index da5096e430..97dc6cbd47 100644 --- a/onnxruntime/core/framework/kernel_registry.cc +++ b/onnxruntime/core/framework/kernel_registry.cc @@ -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& constant_initialized_tensors, + const OrtValueNameIdxMap& ort_value_name_idx_map, + const FuncManager& funcs_mgr, + const DataTransferManager& data_transfer_mgr, + /*out*/ std::unique_ptr& 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& 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 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& constant_initialized_tensors, - const OrtValueNameIdxMap& ort_value_name_idx_map, - const FuncManager& funcs_mgr, - const DataTransferManager& data_transfer_mgr, - /*out*/ std::unique_ptr& 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& 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 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 diff --git a/onnxruntime/core/framework/kernel_registry_manager.cc b/onnxruntime/core/framework/kernel_registry_manager.cc index 737746cd09..b7e772701f 100644 --- a/onnxruntime/core/framework/kernel_registry_manager.cc +++ b/onnxruntime/core/framework/kernel_registry_manager.cc @@ -45,6 +45,7 @@ Status KernelRegistryManager::RegisterKernels(const ExecutionProviders& executio return Status::OK(); } +#if !defined(ORT_MINIMAL_BUILD) void KernelRegistryManager::RegisterKernelRegistry(std::shared_ptr 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 ")); diff --git a/onnxruntime/core/framework/kernel_registry_manager.h b/onnxruntime/core/framework/kernel_registry_manager.h index 205ddae975..0e34e1e57d 100644 --- a/onnxruntime/core/framework/kernel_registry_manager.h +++ b/onnxruntime/core/framework/kernel_registry_manager.h @@ -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 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 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);