mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-11 17:48:34 +00:00
Add kernel def hash logic for minimal build (#4891)
* Add hash based lookup of kernels
This commit is contained in:
parent
db7669b225
commit
728e886bba
7 changed files with 182 additions and 82 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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_);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 "));
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue