From b7fde84341f5e7e4fc8b202e9aabad4d087ec15c Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Wed, 1 Mar 2023 11:22:54 +1000 Subject: [PATCH] Changes to support standalone custom ops in a minimal build. (#14497) ### Description Changes to support standalone custom ops in a minimal build. Also incorporates changes from #14492 (needed to test builds prior to that being checked in). We first need to save the schema info from the operators used by the standalone op invoker in the ORT format model. Add mechanism for that. Merge the kernel lookup logic so the same is used in full and minimal build. NOTE: the version matching is now consistent with all other kernel lookups, and the call to CreateOp MUST use the exact version for the operator. Previously matching wasn't as strict, but this can lead to the incorrect kernel being chosen. Add tests. NOTE: There is currently no way to detect the ops/types/opsets used inside these custom ops as they don't exist until we create kernels, which is after model loading completes (which is the point the ORT format model is saved). Due to that they have to be manually added to the configuration used to do the reduced ops build. That shouldn't be too hard for the custom op author to add given the custom op implementation is specifying the op, opset and type constraints (i.e. they have the info and it's just a case of capturing/formatting it correctly). ### Motivation and Context Enable usage of the standalone op invoker by custom ops in a minimal build. --------- Co-authored-by: Edward Chen <18449977+edgchen1@users.noreply.github.com> --- .../core/framework/kernel_registry.h | 39 ++++-- .../onnxruntime/core/framework/op_kernel.h | 4 +- .../core/framework/op_kernel_info.h | 4 +- include/onnxruntime/core/graph/graph.h | 6 +- onnxruntime/core/framework/kernel_registry.cc | 120 +++++++++--------- onnxruntime/core/graph/graph.cc | 8 +- onnxruntime/core/session/custom_ops.cc | 12 +- onnxruntime/core/session/custom_ops.h | 25 +++- onnxruntime/core/session/inference_session.cc | 2 + onnxruntime/core/session/inference_session.h | 53 ++++---- .../core/session/standalone_op_invoker.cc | 87 +++++++++---- onnxruntime/core/util/thread_utils.cc | 3 +- .../framework/local_kernel_registry_test.cc | 16 +-- .../test/shared_lib/custom_op_utils.cc | 28 ++-- onnxruntime/test/shared_lib/custom_op_utils.h | 46 +++++-- onnxruntime/test/shared_lib/test_inference.cc | 25 +++- .../test/shared_lib/test_ort_format_models.cc | 20 +++ onnxruntime/test/testdata/foo_1.onnx.ort | Bin 1608 -> 2040 bytes .../required_ops.standalone_invoker.config | 3 + .../build_full_ort_and_create_ort_files.sh | 7 + 20 files changed, 328 insertions(+), 180 deletions(-) create mode 100644 onnxruntime/test/testdata/ort_minimal_e2e_test_data/required_ops.standalone_invoker.config diff --git a/include/onnxruntime/core/framework/kernel_registry.h b/include/onnxruntime/core/framework/kernel_registry.h index 68b610ac4b..dc5499d0a7 100644 --- a/include/onnxruntime/core/framework/kernel_registry.h +++ b/include/onnxruntime/core/framework/kernel_registry.h @@ -29,11 +29,21 @@ class KernelRegistry { // TODO(edgchen1) for TryFindKernel(), consider using `out` != nullptr as indicator of whether kernel was found and // Status as an indication of failure - // Check if an execution provider can create kernel for a node and return the kernel if so + // Check if an execution provider can create kernel for a node and return the kernel if so. + // Kernel matching uses the types from the node and the kernel_type_str_resolver. Status TryFindKernel(const Node& node, ProviderType exec_provider, const IKernelTypeStrResolver& kernel_type_str_resolver, const KernelCreateInfo** out) const; + // map of type constraint name to required type + using TypeConstraintMap = InlinedHashMap; + + // Check if an execution provider can create kernel for a node and return the kernel if so. + // Kernel matching uses the explicit type constraint name to required type map in type_constraints. + Status TryFindKernel(const Node& node, ProviderType exec_provider, + const TypeConstraintMap& type_constraints, + const KernelCreateInfo** out) const; + static bool HasImplementationOf(const KernelRegistry& r, const Node& node, ProviderType exec_provider, const IKernelTypeStrResolver& kernel_type_str_resolver) { @@ -42,13 +52,6 @@ class KernelRegistry { return st.IsOK(); } -#if !defined(ORT_MINIMAL_BUILD) - // Find KernelCreateInfo in instant mode - Status TryFindKernel(const std::string& op_name, const std::string& domain, const int& version, - const std::unordered_map& type_constraints, - ProviderType exec_provider, const KernelCreateInfo** out) const; -#endif // !defined(ORT_MINIMAL_BUILD) - bool IsEmpty() const { return kernel_creator_fn_map_.empty(); } #ifdef onnxruntime_PYBIND_EXPORT_OPSCHEMA @@ -59,6 +62,12 @@ class KernelRegistry { #endif private: + // TryFindKernel implementation. Either kernel_type_str_resolver or type_constraints is provided. + Status TryFindKernelImpl(const Node& node, ProviderType exec_provider, + const IKernelTypeStrResolver* kernel_type_str_resolver, + const TypeConstraintMap* type_constraints, + const KernelCreateInfo** out) const; + // 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 @@ -69,9 +78,17 @@ class KernelRegistry { // // Note that this is not intended for type-checking the node against the ONNX // type specification of the corresponding op, which is done before this check. - static bool VerifyKernelDef(const Node& node, - const KernelDef& kernel_def, - const IKernelTypeStrResolver& kernel_type_str_resolver, + // + // In typical usage kernel_type_str_resolver is provided and type information from the node is used with + // kernel_type_str_resolver. + // + // There is also usage from a node dynamically created within a custom op via OrtApi CreateOp where an explicit + // type value for each type constraint is provided in type_constraints. + // + // Either kernel_type_str_resolver or type_constraints is provided and not both. + static bool VerifyKernelDef(const Node& node, const KernelDef& kernel_def, + const IKernelTypeStrResolver* kernel_type_str_resolver, + const TypeConstraintMap* type_constraints, std::string& error_str); static std::string GetMapKey(std::string_view op_name, std::string_view domain, std::string_view provider) { diff --git a/include/onnxruntime/core/framework/op_kernel.h b/include/onnxruntime/core/framework/op_kernel.h index 392057cf03..b942c06e95 100644 --- a/include/onnxruntime/core/framework/op_kernel.h +++ b/include/onnxruntime/core/framework/op_kernel.h @@ -142,7 +142,9 @@ struct KernelCreateInfo { KernelCreateInfo(std::unique_ptr definition, KernelCreateFn create_func) : kernel_def(std::move(definition)), - kernel_create_func(create_func) {} + kernel_create_func(create_func) { + assert(kernel_def != nullptr); + } KernelCreateInfo(KernelCreateInfo&& other) noexcept : kernel_def(std::move(other.kernel_def)), diff --git a/include/onnxruntime/core/framework/op_kernel_info.h b/include/onnxruntime/core/framework/op_kernel_info.h index 4695c0a1c9..335e2635f1 100644 --- a/include/onnxruntime/core/framework/op_kernel_info.h +++ b/include/onnxruntime/core/framework/op_kernel_info.h @@ -12,9 +12,9 @@ namespace onnxruntime { -class OrtValueNameIdxMap; -class FuncManager; class DataTransferManager; +class FuncManager; +class OrtValueNameIdxMap; struct AllocPlanPerValue; // A very light-weight class, which works as an aggregated diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 0fc2762ca0..280f4e1a60 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -84,7 +84,7 @@ class Node { explicit Node() = default; -#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) Node(std::string_view name, std::string_view op_type, std::string_view description, @@ -565,7 +565,7 @@ class Node { private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Node); -#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) void Init(const std::string& name, const std::string& op_type, const std::string& description, @@ -573,7 +573,9 @@ class Node { const std::vector& output_args, const NodeAttributes* attributes, const std::string& domain); +#endif +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) // internal only method to allow selected classes to directly alter the input/output definitions and arg counts Definitions& MutableDefinitions() noexcept; diff --git a/onnxruntime/core/framework/kernel_registry.cc b/onnxruntime/core/framework/kernel_registry.cc index e2bc7c3e3c..b3efa53c77 100644 --- a/onnxruntime/core/framework/kernel_registry.cc +++ b/onnxruntime/core/framework/kernel_registry.cc @@ -39,8 +39,9 @@ bool IsTypeProtoCompatible(gsl::span enabled_types, const ONNX return true; } +// match the kernel using type info from the Node's args bool MatchKernelDefTypes(const Node& node, - const KernelDef& kernel_def, + const std::unordered_map>& kernel_type_constraints, const IKernelTypeStrResolver& kernel_type_str_resolver, std::string& mismatch_reason) { const auto actual_inputs = node.InputDefs(); @@ -63,11 +64,9 @@ bool MatchKernelDefTypes(const Node& node, // for each type constraint // map type constraint to arg // check arg type against type constraint enabled types - const auto& kernel_type_constraints = kernel_def.TypeConstraints(); for (const auto& [kernel_type_str, enabled_types] : kernel_type_constraints) { gsl::span constraint_args{}; - ORT_THROW_IF_ERROR(kernel_type_str_resolver.ResolveKernelTypeStr(node, kernel_type_str, - constraint_args)); + ORT_THROW_IF_ERROR(kernel_type_str_resolver.ResolveKernelTypeStr(node, kernel_type_str, constraint_args)); for (const auto& [arg_type, formal_arg_idx] : constraint_args) { const NodeArg* arg; @@ -100,58 +99,75 @@ bool MatchKernelDefTypes(const Node& node, return true; } + +bool MatchKernelDefTypes(const std::unordered_map>& kernel_type_constraints, + const KernelRegistry::TypeConstraintMap& type_constraints) { + bool match = true; + for (auto& constraint : type_constraints) { + auto iter = kernel_type_constraints.find(constraint.first); + if (iter == kernel_type_constraints.end() || + find(iter->second.begin(), iter->second.end(), constraint.second) == iter->second.end()) { + match = false; + break; + } + } + + return match; +} } // namespace bool KernelRegistry::VerifyKernelDef(const Node& node, const KernelDef& kernel_def, - const IKernelTypeStrResolver& kernel_type_str_resolver, + const IKernelTypeStrResolver* kernel_type_str_resolver, + const TypeConstraintMap* type_constraint_values, std::string& error_str) { // check if version matches + int node_version = node.SinceVersion(); int kernel_start_version; int kernel_end_version; kernel_def.SinceVersion(&kernel_start_version, &kernel_end_version); - int node_since_version = node.SinceVersion(); - // Ideal case is, if schema is Since(5), current opset version is opset 7, - // kernel_def Since(8) Invalid - // kernel_def Since(6) Valid - // kernel_def Since(5) Valid - // kernel_def Since(4) Invalid - // kernel_def Since(4, 6) Valid + bool valid_version = + // exact match. typical usage. + kernel_start_version == node_version || + // allow match if the kernel def has an end version. if it does not, all we know is that the kernel supported + // the start version when it was created, and not whether a new version of the operator was added since then + // that the kernel doesn't support. + (kernel_end_version != INT_MAX && + kernel_start_version <= node_version && kernel_end_version >= node_version); - // Right now there is no "until version" on schema, it is difficult to get opset version here.(require a lot of interface change.) - // As a trade off, we will temporary require kernel definition to have the same since version as schema definition. - // so kernel_def Since(6) will become invalid now. - // After ONNX add "until version" on the schema object, we will update this place - bool valid_version = kernel_start_version == node_since_version // the idea case this branch should be kernel_start_version >= node_version && kernel_start_version <= until_version - || (kernel_start_version < node_since_version && kernel_end_version != INT_MAX && kernel_end_version >= node_since_version); if (!valid_version) { std::ostringstream ostr; ostr << "Op with name (" << node.Name() << ")" << " and type (" << node.OpType() << ")" << " Version mismatch." - << " node_version: " << node_since_version + << " node_version: " << node_version << " kernel start version: " << kernel_start_version << " kernel_end_version: " << kernel_end_version; error_str = ostr.str(); return false; } - if (std::string mismatch_reason; - !MatchKernelDefTypes(node, kernel_def, kernel_type_str_resolver, mismatch_reason)) { + std::string mismatch_reason; + const auto& kernel_type_constraints = kernel_def.TypeConstraints(); + + bool matched = type_constraint_values ? MatchKernelDefTypes(kernel_type_constraints, *type_constraint_values) + : MatchKernelDefTypes(node, kernel_type_constraints, *kernel_type_str_resolver, + mismatch_reason); + + if (!matched) { std::ostringstream ostr; ostr << "Found kernel for Op with name (" << node.Name() << ")" << " and type (" << node.OpType() << ")" << " in the supported version range" - << " (node_version: " << node_since_version + << " (node_version: " << node_version << " kernel start version: " << kernel_start_version << " kernel_end_version: " << kernel_end_version << ")." << " However the types are incompatible. " << mismatch_reason; error_str = ostr.str(); - return false; } - return true; + return matched; } // It's often this function returns a failed status, but it is totally expected. @@ -159,10 +175,11 @@ bool KernelRegistry::VerifyKernelDef(const Node& node, // 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 Node& node, - ProviderType exec_provider, - const IKernelTypeStrResolver& kernel_type_str_resolver, - const KernelCreateInfo** out) const { +Status KernelRegistry::TryFindKernelImpl(const Node& node, + ProviderType exec_provider, + const IKernelTypeStrResolver* kernel_type_str_resolver, + const TypeConstraintMap* type_constraints, + const KernelCreateInfo** out) const { const auto& node_provider = node.GetExecutionProviderType(); const auto& expected_provider = (node_provider.empty() ? exec_provider : node_provider); @@ -173,10 +190,13 @@ Status KernelRegistry::TryFindKernel(const Node& node, for (auto i = range.first; i != range.second; ++i) { std::string error_str; - if (VerifyKernelDef(node, *i->second.kernel_def, kernel_type_str_resolver, error_str)) { - if (out) *out = &i->second; + if (VerifyKernelDef(node, *i->second.kernel_def, kernel_type_str_resolver, type_constraints, error_str)) { + if (out) { + *out = &i->second; + } return Status::OK(); } + verify_kernel_def_error_strs.push_back(error_str); } @@ -197,37 +217,17 @@ Status KernelRegistry::TryFindKernel(const Node& node, return Status(common::ONNXRUNTIME, common::FAIL, "Kernel not found"); } -#if !defined(ORT_MINIMAL_BUILD) -Status KernelRegistry::TryFindKernel(const std::string& op_name, const std::string& domain, const int& version, - const std::unordered_map& type_constraints, - ProviderType exec_provider, const KernelCreateInfo** kernel_out) const { - const KernelCreateInfo* kernel = nullptr; - auto range = kernel_creator_fn_map_.equal_range(GetMapKey(op_name, domain, exec_provider)); - for (auto i = range.first; i != range.second; ++i) { // loop through all kernels - const KernelCreateInfo& kci = i->second; - int start_ver{}; - int end_ver{}; - kci.kernel_def->SinceVersion(&start_ver, &end_ver); - if (start_ver <= version && end_ver >= version) { // try match the version - auto& kci_constraints = kci.kernel_def->TypeConstraints(); - bool match = true; - for (auto& constraint : type_constraints) { // try match type constraints - auto iter = kci_constraints.find(constraint.first); - if (iter == kci_constraints.end() || find(iter->second.begin(), iter->second.end(), constraint.second) == iter->second.end()) { - match = false; - break; - } - } // for - if (match) { - kernel = &kci; // found match, exit loop - break; - } - } // if - } // for - if (kernel_out) *kernel_out = kernel; - return kernel == nullptr ? Status(common::ONNXRUNTIME, common::FAIL, "Kernel not found") : Status::OK(); +Status KernelRegistry::TryFindKernel(const Node& node, ProviderType exec_provider, + const IKernelTypeStrResolver& kernel_type_str_resolver, + const KernelCreateInfo** out) const { + return TryFindKernelImpl(node, exec_provider, &kernel_type_str_resolver, nullptr, out); +} + +Status KernelRegistry::TryFindKernel(const Node& node, ProviderType exec_provider, + const TypeConstraintMap& type_constraints, + const KernelCreateInfo** out) const { + return TryFindKernelImpl(node, exec_provider, nullptr, &type_constraints, out); } -#endif // !defined(ORT_MINIMAL_BUILD) Status KernelRegistry::Register(KernelDefBuilder& kernel_builder, const KernelCreateFn& kernel_creator) { diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 81309ceb05..c47569b5ca 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -178,7 +178,7 @@ static std::string GenerateSchemaKey(const IndexedSubGraph& subgraph_ptr) { } #endif // !defined(ORT_MINIMAL_BUILD) -#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) NodeArg::NodeArg(const std::string& name, const TypeProto* p_node_arg_type) { node_arg_info_.set_name(name); // If the name is empty, it means the arg does not exist. @@ -195,7 +195,7 @@ NodeArg::NodeArg(const std::string& name, const TypeProto* p_node_arg_type) { type_ = nullptr; } } -#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +#endif // #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) NodeArg::NodeArg(NodeArgInfo&& node_arg_info) { node_arg_info_ = std::move(node_arg_info); @@ -821,7 +821,7 @@ Status Node::LoadEdgesFromOrtFormat(const onnxruntime::fbs::NodeEdge& fbs_node_e return Status::OK(); } -#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) void Node::Init(const std::string& name, const std::string& op_type, const std::string& description, @@ -859,7 +859,9 @@ void Node::Init(const std::string& name, } } } +#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) +#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) Node::Definitions& Node::MutableDefinitions() noexcept { // someone fetching these is going to change something graph_->SetGraphResolveNeeded(); diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index dae00b96ea..78509a8580 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -5,19 +5,21 @@ #pragma warning(disable : 4267) #endif +#include + #include "core/framework/data_types.h" -#include "core/framework/op_kernel_info.h" -#include "core/framework/op_kernel_context_internal.h" #include "core/framework/error_code_helper.h" -#include "core/framework/tensor_type_and_shape.h" #include "core/framework/onnxruntime_typeinfo.h" +#include "core/framework/op_kernel_context_internal.h" +#include "core/framework/op_kernel_info.h" +#include "core/framework/tensor_type_and_shape.h" #include "core/framework/tensorprotoutils.h" #include "core/graph/onnx_protobuf.h" #include "core/session/allocator_adapters.h" +#include "core/session/api_utils.h" +#include "core/session/custom_ops.h" #include "core/session/inference_session.h" #include "core/session/ort_apis.h" -#include -#include "api_utils.h" ORT_API_STATUS_IMPL(OrtApis::KernelInfoGetAttribute_float, _In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ float* out) { API_IMPL_BEGIN diff --git a/onnxruntime/core/session/custom_ops.h b/onnxruntime/core/session/custom_ops.h index 3420f54ccb..35b2abf1e5 100644 --- a/onnxruntime/core/session/custom_ops.h +++ b/onnxruntime/core/session/custom_ops.h @@ -1,8 +1,29 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#pragma once +#include +#include + +#include "core/common/status.h" + +struct OrtCustomOpDomain; namespace onnxruntime { +class CustomRegistry; -common::Status CreateCustomRegistry(gsl::span op_domains, std::shared_ptr& output); +common::Status CreateCustomRegistry(gsl::span op_domains, + std::shared_ptr& output); -} +#if !defined(ORT_MINIMAL_BUILD) +class Graph; +class KernelTypeStrResolver; + +namespace standalone { +// Register the schemas from any custom ops using the standalone invoker to call ORT kernels via OrtApi CreateOp. +// This is required so they can be captured when saving to an ORT format model. +// Implemented in standalone_op_invoker.cc +common::Status RegisterCustomOpNodeSchemas(KernelTypeStrResolver& kernel_type_str_resolver, Graph& graph); +} // namespace standalone +#endif + +} // namespace onnxruntime diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 79766bb7b8..11995da5be 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -593,6 +593,8 @@ common::Status InferenceSession::SaveToOrtFormat(const PathString& filepath) con flatbuffers::Offset fbs_kernel_type_str_resolver; KernelTypeStrResolver kernel_type_str_resolver{}; ORT_RETURN_IF_ERROR(kernel_type_str_resolver.RegisterGraphNodeOpSchemas(model_->MainGraph())); + ORT_RETURN_IF_ERROR(standalone::RegisterCustomOpNodeSchemas(kernel_type_str_resolver, model_->MainGraph())); + for (const auto* op_schema : saved_runtime_optimization_produced_node_op_schemas_) { ORT_RETURN_IF_ERROR(kernel_type_str_resolver.RegisterOpSchema(*op_schema)); } diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 95b0dde281..f0199a343b 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -32,11 +32,6 @@ #include #endif -namespace onnxruntime { // forward declarations -class GraphTransformer; -class Environment; -} // namespace onnxruntime - namespace ONNX_NAMESPACE { class ModelProto; } // namespace ONNX_NAMESPACE @@ -46,15 +41,18 @@ struct OrtCustomOpDomain { std::vector custom_ops_; }; -namespace onnxruntime { -class IExecutionProvider; // forward decl -class IOBinding; +namespace onnxruntime { // forward declarations class CustomRegistry; +class Environment; +class GraphTransformer; +class IExecutionProvider; +class IOBinding; struct Notification; + #ifdef ENABLE_TRAINING struct PartialGraphExecutionState; -typedef InlinedHashMap OrtValueCache; -typedef std::shared_ptr OrtValueCachePtr; +using OrtValueCache = InlinedHashMap; +using OrtValueCachePtr = std::shared_ptr; #endif namespace logging { @@ -214,7 +212,7 @@ class InferenceSession { * @return OK if success. */ [[nodiscard]] common::Status RegisterGraphTransformer(std::unique_ptr p_graph_transformer, - TransformerLevel level = TransformerLevel::Level2); + TransformerLevel level = TransformerLevel::Level2); #endif // !defined(ORT_MINIMAL_BUILD) @@ -307,9 +305,9 @@ class InferenceSession { [[nodiscard]] common::Status Initialize(); [[nodiscard]] common::Status Run(const RunOptions& run_options, gsl::span feed_names, - gsl::span feeds, gsl::span output_names, - std::vector* p_fetches, - const std::vector* p_fetches_device_info = nullptr); + gsl::span feeds, gsl::span output_names, + std::vector* p_fetches, + const std::vector* p_fetches_device_info = nullptr); /** * Run a pre-loaded and pre-intialized model. @@ -322,7 +320,7 @@ class InferenceSession { * @return OK if success. */ [[nodiscard]] common::Status Run(const NameMLValMap& feeds, gsl::span output_names, - std::vector* p_fetches); + std::vector* p_fetches); /** * See Run(const NameMLValMap& feeds, const std::vector& output_names, std::vector* p_fetches) @@ -330,8 +328,8 @@ class InferenceSession { * @param run_options use this to tune the Run call to your needs. */ [[nodiscard]] common::Status Run(const RunOptions& run_options, const NameMLValMap& feeds, - gsl::span output_names, - std::vector* p_fetches); + gsl::span output_names, + std::vector* p_fetches); /** * Creates a new binding object for binding inputs and outputs. @@ -465,7 +463,6 @@ class InferenceSession { Status SetTuningResults(const std::vector& trs, bool error_on_invalid = false); #endif - #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) MemoryProfiler& GetMemoryProfiler() { return memory_profiler_; @@ -515,7 +512,7 @@ class InferenceSession { [[nodiscard]] common::Status LoadOnnxModel(std::unique_ptr p_model_proto); [[nodiscard]] common::Status LoadWithLoader(std::function&)> loader, - const std::string& event_name); + const std::string& event_name); [[nodiscard]] common::Status DoPostLoadProcessing(onnxruntime::Model& model); @@ -618,13 +615,13 @@ class InferenceSession { void InitLogger(logging::LoggingManager* logging_manager); [[nodiscard]] common::Status CheckShapes(const std::string& input_name, const TensorShape& input_shape, - const TensorShape& expected_shape) const; + const TensorShape& expected_shape) const; [[nodiscard]] common::Status ValidateInputs(gsl::span feed_names, - gsl::span feeds) const; + gsl::span feeds) const; [[nodiscard]] common::Status ValidateOutputs(gsl::span output_names, - const std::vector* p_fetches) const; + const std::vector* p_fetches) const; [[nodiscard]] common::Status WaitForNotification(Notification* p_executor_done, int64_t timeout_in_ms); @@ -642,7 +639,7 @@ class InferenceSession { */ [[nodiscard]] common::Status ValidateAndParseShrinkArenaString(const std::string& ort_device_list, - /*out*/ InlinedVector& arenas_to_shrink) const; + /*out*/ InlinedVector& arenas_to_shrink) const; /* * Performs the shrinkage of arenas requested to be shrunk by the user @@ -658,11 +655,11 @@ class InferenceSession { RecordRuntimeOptimizationProducedNodeOpSchemaFn record_runtime_optimization_produced_op_schema_fn) const; [[nodiscard]] common::Status TransformGraph(onnxruntime::Graph& graph, - const onnxruntime::GraphTransformerManager& graph_transformer_mgr, - const ExecutionProviders& providers, KernelRegistryManager& kernel_registry_manager, - const InsertCastTransformer& insert_cast_transformer, - SessionState& session_state, - bool saving_model_in_ort_format); + const onnxruntime::GraphTransformerManager& graph_transformer_mgr, + const ExecutionProviders& providers, KernelRegistryManager& kernel_registry_manager, + const InsertCastTransformer& insert_cast_transformer, + SessionState& session_state, + bool saving_model_in_ort_format); onnxruntime::GraphTransformerManager graph_transformation_mgr_; diff --git a/onnxruntime/core/session/standalone_op_invoker.cc b/onnxruntime/core/session/standalone_op_invoker.cc index 48cfdd60d2..b14ca21e6d 100644 --- a/onnxruntime/core/session/standalone_op_invoker.cc +++ b/onnxruntime/core/session/standalone_op_invoker.cc @@ -8,12 +8,16 @@ #include "core/session/ort_apis.h" #include +#if !defined(ORT_MINIMAL_BUILD) +#include "core/graph/schema_registry.h" +#endif + #if defined(_MSC_VER) && !defined(__clang__) // disabling warning on calling of raw "delete" operator #pragma warning(disable : 26400) #endif -#ifdef ORT_MINIMAL_BUILD +#if defined(ORT_MINIMAL_BUILD) && !defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) ORT_API_STATUS_IMPL(OrtApis::CreateOpAttr, _In_ const char*, @@ -79,7 +83,6 @@ namespace onnxruntime { namespace standalone { using NodePtr = std::unique_ptr; - using ArgPtr = std::unique_ptr; using ArgPtrs = onnxruntime::InlinedVector; @@ -93,15 +96,45 @@ class NodeRepo { return node_repo; } + // create the kernel using the FuncManager NodeRepo owns for consistency + onnxruntime::Status CreateKernel(const KernelCreateInfo& kernel_create_info, + const OpKernelInfo& kernel_info, + std::unique_ptr& op_kernel) { + std::lock_guard guard(mutex_); + return kernel_create_info.kernel_create_func(func_mgr_, kernel_info, op_kernel); + } + onnxruntime::Status AddNode(const onnxruntime::OpKernel* kernel, NodePtr&& node_ptr, ArgPtrs&& args) { std::lock_guard guard(mutex_); auto ret = resource_map_.try_emplace(kernel, NodeResource{std::move(node_ptr), std::move(args)}); if (!ret.second) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "kernel already mapped to existing node"); } + return Status::OK(); } +#if !defined(ORT_MINIMAL_BUILD) + common::Status RegisterCustomOpNodeSchemas(KernelTypeStrResolver& kernel_type_str_resolver, Graph& graph) { + std::lock_guard guard(mutex_); + + for (auto cur = resource_map_.begin(), end = resource_map_.end(); cur != end; ++cur) { + // Lookup the schema for the operator so we include it in the ORT format model and can match the kernel + // in a minimal build. + // The opset version will not necessarily match the model, so we need to call GetSchema directly to plug that in. + // In theory this should never fail if the kernel lookup earlier was successful. + const Node& node = *cur->second.first; + auto* schema = graph.GetSchemaRegistry()->GetSchema(node.OpType(), node.SinceVersion(), node.Domain()); + + ORT_RETURN_IF_NOT(schema, "Unable to find schema for node. Domain:'", node.Domain(), + "' op_type:", node.OpType()); + ORT_RETURN_IF_ERROR(kernel_type_str_resolver.RegisterOpSchema(*schema)); + } + + return Status::OK(); + } +#endif + onnxruntime::Status ValidateInputOutputCounts(const onnxruntime::OpKernel* kernel, int input_count, int output_count) { @@ -141,8 +174,15 @@ class NodeRepo { std::mutex mutex_; NodeResourceMap resource_map_; + FuncManager func_mgr_; }; +#if !defined(ORT_MINIMAL_BUILD) +common::Status RegisterCustomOpNodeSchemas(KernelTypeStrResolver& kernel_type_str_resolver, Graph& graph) { + return NodeRepo::GetInstance().RegisterCustomOpNodeSchemas(kernel_type_str_resolver, graph); +} +#endif + // For invoking kernels without a graph class StandAloneKernelContext : public OpKernelContext { public: @@ -338,19 +378,13 @@ onnxruntime::Status CreateOp(const OrtKernelInfo* info, auto ep = reinterpret_cast(kernel_info->GetExecutionProvider()); auto kernel_registry = ep->GetKernelRegistry(); const KernelCreateInfo* kernel_create_info{}; - std::unordered_map type_constraint_map; + InlinedHashMap type_constraint_map; + for (int i = 0; i < type_constraint_count; ++i) { ONNX_NAMESPACE::TypeProto proto; proto.mutable_tensor_type()->set_elem_type(type_constraint_values[i]); type_constraint_map[type_constraint_names[i]] = DataTypeImpl::TypeFromProto(proto); } - auto status = kernel_registry->TryFindKernel(op_name, - domain, - version, - type_constraint_map, - ep->Type(), - &kernel_create_info); - ORT_RETURN_IF_ERROR(status); ArgPtrs arg_ptrs; std::vector input_args; @@ -366,29 +400,34 @@ onnxruntime::Status CreateOp(const OrtKernelInfo* info, output_args.push_back(arg_ptrs.back().get()); } - NodePtr node_ptr = std::make_unique(std::string("standalone_") + op_name, op_name, "", input_args, output_args, nullptr, domain); + NodePtr node_ptr = std::make_unique(std::string("standalone_") + op_name, op_name, "", + input_args, output_args, nullptr, domain); + for (int i = 0; i < attr_count; ++i) { auto attr_proto = reinterpret_cast(attr_values[i]); node_ptr->AddAttributeProto(*attr_proto); } - auto kernel_def_builder = KernelDefBuilder::Create(); - kernel_def_builder->SetName(op_name); - kernel_def_builder->SetDomain(domain); - kernel_def_builder->SinceVersion(version); - auto kernel_def = kernel_def_builder->Build(); + node_ptr->SetSinceVersion(version); - static std::unordered_map kEmptyValueMap; - static OrtValueNameIdxMap kEmptyNameMap; + auto status = kernel_registry->TryFindKernel(*node_ptr, ep->Type(), type_constraint_map, &kernel_create_info); + ORT_RETURN_IF_ERROR(status); - OpKernelInfo tmp_kernel_info(*node_ptr.get(), *kernel_def, *ep, kEmptyValueMap, kEmptyNameMap, kernel_info->GetDataTransferManager()); + auto& kernel_def = kernel_create_info->kernel_def; + ORT_RETURN_IF_NOT(kernel_def, "Kernel definition was not found for node Domain:'", + node_ptr->Domain(), "' op_type:", node_ptr->OpType()); + + static const std::unordered_map kEmptyValueMap; + static const OrtValueNameIdxMap kEmptyNameMap; + + OpKernelInfo tmp_kernel_info(*node_ptr.get(), *kernel_def, *ep, kEmptyValueMap, kEmptyNameMap, + kernel_info->GetDataTransferManager()); std::unique_ptr op_kernel; - static FuncManager kFuncMgr; - status = kernel_create_info->kernel_create_func(kFuncMgr, tmp_kernel_info, op_kernel); - ORT_RETURN_IF_ERROR(status); - status = NodeRepo::GetInstance().AddNode(op_kernel.get(), std::move(node_ptr), std::move(arg_ptrs)); - ORT_RETURN_IF_ERROR(status); + auto& node_repo = NodeRepo::GetInstance(); + ORT_RETURN_IF_ERROR(node_repo.CreateKernel(*kernel_create_info, tmp_kernel_info, op_kernel)); + ORT_RETURN_IF_ERROR(node_repo.AddNode(op_kernel.get(), std::move(node_ptr), std::move(arg_ptrs))); + *op = reinterpret_cast(op_kernel.release()); return status; } diff --git a/onnxruntime/core/util/thread_utils.cc b/onnxruntime/core/util/thread_utils.cc index 1535b52aa5..54602e70a0 100644 --- a/onnxruntime/core/util/thread_utils.cc +++ b/onnxruntime/core/util/thread_utils.cc @@ -91,7 +91,6 @@ CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options) { if (!options.affinity_str.empty()) { #if defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) ORT_THROW("Setting thread affinity is not implemented in this build."); - return nullptr; #else to.affinities = ReadThreadAffinityConfig(options.affinity_str); // Limiting the number of affinities to be of thread_pool_size - 1, @@ -233,7 +232,7 @@ ORT_API_STATUS_IMPL(SetGlobalIntraOpThreadAffinity, _Inout_ OrtThreadingOptions* } tp_options->intra_op_thread_pool_params.affinity_str = affinity_string; return nullptr; - #endif +#endif } } // namespace OrtApis diff --git a/onnxruntime/test/framework/local_kernel_registry_test.cc b/onnxruntime/test/framework/local_kernel_registry_test.cc index cb120f8222..a4adbadae8 100644 --- a/onnxruntime/test/framework/local_kernel_registry_test.cc +++ b/onnxruntime/test/framework/local_kernel_registry_test.cc @@ -317,18 +317,14 @@ TEST(CustomKernelTests, CustomKernelWithOptionalOutput) { RunSession(session_object, dims_x, values_x, expected_dims_y, expected_values_y); } -// Regression test for OnnxRuntimeOpSchemaRegistry::GetSchemaAndHistory +// Regression test for OnnxRuntimeOpSchemaRegistry::GetSchemaAndHistory needing to reset `version` before +// falling through to the ONNX schema lookup. // -// When there was a custom registry that matched the ONNX domain but not the current op, the earliest version from that -// was being used in the fallthrough search of the ONNX schemas. This results in invalid matching. -// e.g. -// Model with ONNX opset of 12. -// Custom op using ONNX domain with version range of 1..1000 (which is what the custom op registration code -// available via the ORT API uses by default) and custom op (called 'Foo' in our test). +// If there is a custom registry that matches the ONNX domain but not the current op, we fall though but need to +// use the original opset version and ignore any version values found in the custom registry. // -// If the model has another ONNX op (Clip in our test), GetSchemaAndHistory was using the custom registry first, -// 'version' gets set to 1 and the ONNX schema lookup matches that opset and not the model's opset. This leads to an -// invalid schema being selected. +// If we regress we will match Clip(1) which only had one input. The model uses Clip(11) and has two inputs. The ONNX +// checker will fail if this happens. TEST(CustomKernelTests, CustomOnnxKernelSchemaLookup) { SessionOptions so; so.session_logid = "CustomOnnxKernelSchemaLookup"; diff --git a/onnxruntime/test/shared_lib/custom_op_utils.cc b/onnxruntime/test/shared_lib/custom_op_utils.cc index ebb9a6c1cc..34ceb9e9e1 100644 --- a/onnxruntime/test/shared_lib/custom_op_utils.cc +++ b/onnxruntime/test/shared_lib/custom_op_utils.cc @@ -72,12 +72,11 @@ void MyCustomKernelSecondInputOnCpu::Compute(OrtKernelContext* context) { ASSERT_EQ(y_mem_type, OrtMemType::OrtMemTypeCPUInput); // copy the second input to GPU - const int64_t y_size = input_Y.GetTensorTypeAndShapeInfo().GetElementCount(); - float* Y_cuda {}; + const int64_t y_size = input_Y.GetTensorTypeAndShapeInfo().GetElementCount(); + float* Y_cuda{}; cudaMalloc(&Y_cuda, y_size * sizeof(float)); cudaMemcpy(Y_cuda, Y, y_size * sizeof(float), cudaMemcpyHostToDevice); - // Setup output auto dimensions = input_X.GetTensorTypeAndShapeInfo().GetShape(); auto output = ctx.GetOutput(0, dimensions); @@ -356,18 +355,22 @@ StandaloneCustomKernel::StandaloneCustomKernel(const OrtKernelInfo* k_info) { const char* add_type_constraint_names[1] = {"T"}; ONNXTensorElementDataType add_type_constraint_values[1] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; - op_add_ = Ort::Op::Create(info_copy_, "Add", "", 14, + op_add_ = Ort::Op::Create(info_copy_, "Add", "", /* must match onnx version number exactly */ 14, add_type_constraint_names, add_type_constraint_values, 1, nullptr, 0, 2, 1); +#if !defined(REDUCED_OPS_BUILD) InitTopK(); InitGru(); +#endif } +#if !defined(REDUCED_OPS_BUILD) void StandaloneCustomKernel::InitTopK() { const char* type_constraint_names[2] = {"T", "I"}; - ONNXTensorElementDataType type_constraint_values[2] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}; + ONNXTensorElementDataType type_constraint_values[2] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, + ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}; constexpr int64_t axis_value = -1; auto axis = Ort::OpAttr("axis", &axis_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT); @@ -379,7 +382,7 @@ void StandaloneCustomKernel::InitTopK() { auto sorted = Ort::OpAttr("sorted", &sorted_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT); Ort::OpAttr top_attrs[3] = {std::move(axis), std::move(largest), std::move(sorted)}; - op_topk_ = Ort::Op::Create(info_copy_, "TopK", "", 14, + op_topk_ = Ort::Op::Create(info_copy_, "TopK", "", /* must match onnx version number exactly */ 11, type_constraint_names, type_constraint_values, 2, top_attrs, 3, 2, 2); @@ -418,7 +421,8 @@ void StandaloneCustomKernel::InvokeTopK(OrtKernelContext* context) { void StandaloneCustomKernel::InitGru() { const char* type_constraint_names[2] = {"T", "T1"}; - ONNXTensorElementDataType type_constraint_values[2] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32}; + ONNXTensorElementDataType type_constraint_values[2] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, + ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32}; const char* activition_names[4] = {"LeakyRelu", "Tanh", "Sigmoid", "ScaledTanh"}; Ort::OpAttr activations = Ort::OpAttr("activations", activition_names, 4, OrtOpAttrType::ORT_OP_ATTR_STRINGS); @@ -433,7 +437,8 @@ void StandaloneCustomKernel::InitGru() { Ort::OpAttr direction = Ort::OpAttr("direction", direction_string, 1, OrtOpAttrType::ORT_OP_ATTR_STRING); int64_t linear_before_reset_value = 0; - Ort::OpAttr linear_before_reset = Ort::OpAttr("linear_before_reset", &linear_before_reset_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT); + Ort::OpAttr linear_before_reset = Ort::OpAttr("linear_before_reset", &linear_before_reset_value, 1, + OrtOpAttrType::ORT_OP_ATTR_INT); int64_t hidden_size_value = 2; Ort::OpAttr hidden_size = Ort::OpAttr("hidden_size", &hidden_size_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT); @@ -443,7 +448,7 @@ void StandaloneCustomKernel::InitGru() { std::move(activation_beta), std::move(direction), std::move(linear_before_reset), std::move(hidden_size)}; - op_gru_ = Ort::Op::Create(info_copy_, "GRU", "", 14, + op_gru_ = Ort::Op::Create(info_copy_, "GRU", "", /* must match onnx version number exactly */ 14, type_constraint_names, type_constraint_values, 2, gru_attrs, 6, 6, 2); @@ -593,6 +598,8 @@ void StandaloneCustomKernel::InitInvokeConv(OrtKernelContext* context) { } } +#endif // !defined(REDUCED_OPS_BUILD) + void StandaloneCustomKernel::Compute(OrtKernelContext* context) { Ort::KernelContext ctx(context); auto input_X = ctx.GetInput(0); @@ -605,7 +612,8 @@ void StandaloneCustomKernel::Compute(OrtKernelContext* context) { OrtValue* outputs[1] = {output}; op_add_.Invoke(context, inputs, 2, outputs, 1); -#ifndef USE_CUDA + +#if !defined(USE_CUDA) && !defined(REDUCED_OPS_BUILD) InvokeTopK(context); InvokeGru(context); InitInvokeConv(context); diff --git a/onnxruntime/test/shared_lib/custom_op_utils.h b/onnxruntime/test/shared_lib/custom_op_utils.h index 873af327de..74975b53d7 100644 --- a/onnxruntime/test/shared_lib/custom_op_utils.h +++ b/onnxruntime/test/shared_lib/custom_op_utils.h @@ -56,9 +56,13 @@ struct MyCustomOp : Ort::CustomOpBase { }; struct MyCustomOpSecondInputOnCpu : Ort::CustomOpBase { - explicit MyCustomOpSecondInputOnCpu(const char* provider, void* compute_stream) : provider_(provider), compute_stream_(compute_stream) {} + explicit MyCustomOpSecondInputOnCpu(const char* provider, void* compute_stream) + : provider_(provider), compute_stream_(compute_stream) {} + + void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { + return new MyCustomKernelSecondInputOnCpu(info, compute_stream_); + }; - void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { return new MyCustomKernelSecondInputOnCpu(info, compute_stream_); }; const char* GetName() const { return "Foo"; }; const char* GetExecutionProviderType() const { return provider_; }; @@ -69,7 +73,9 @@ struct MyCustomOpSecondInputOnCpu : Ort::CustomOpBase { +struct MyCustomOpMultipleDynamicInputs : Ort::CustomOpBase { explicit MyCustomOpMultipleDynamicInputs(const char* provider) : provider_(provider) {} void* CreateKernel(const OrtApi& api, const OrtKernelInfo* info) const { return new MyCustomKernelMultipleDynamicInputs(api, info); @@ -123,7 +130,10 @@ struct MyCustomKernelWithOptionalInput { struct MyCustomOpWithOptionalInput : Ort::CustomOpBase { explicit MyCustomOpWithOptionalInput(const char* provider) : provider_(provider) {} - void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { return new MyCustomKernelWithOptionalInput(info); }; + void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { + return new MyCustomKernelWithOptionalInput(info); + }; + const char* GetName() const { return "FooBar"; }; const char* GetExecutionProviderType() const { return provider_; }; @@ -179,11 +189,15 @@ struct TemplatedCustomOp : Ort::CustomOpBase, T> { bool input_homogeneity, std::vector output_types, std::vector output_characs, int output_min_arity, bool output_homogeneity) - : op_name_(op_name), input_types_(std::move(input_types)), - input_characs_(std::move(input_characs)), input_min_arity_(input_min_arity), - input_homogeneity_(input_homogeneity), output_types_(std::move(output_types)), - output_characs_(std::move(output_characs)), output_min_arity_(output_min_arity), - output_homogeneity_(output_homogeneity) {} + : op_name_(op_name), + input_types_(std::move(input_types)), + input_characs_(std::move(input_characs)), + input_min_arity_(input_min_arity), + input_homogeneity_(input_homogeneity), + output_types_(std::move(output_types)), + output_characs_(std::move(output_characs)), + output_min_arity_(output_min_arity), + output_homogeneity_(output_homogeneity) {} void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { return new T(info); @@ -268,7 +282,10 @@ struct MyCustomKernelWithAttributes { struct MyCustomOpWithAttributes : Ort::CustomOpBase { explicit MyCustomOpWithAttributes(const char* provider) : provider_(provider) {} - void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MyCustomKernelWithAttributes(info); }; + void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { + return new MyCustomKernelWithAttributes(info); + }; + const char* GetName() const { return "FooBar_Attr"; }; const char* GetExecutionProviderType() const { return provider_; }; @@ -325,6 +342,7 @@ struct StandaloneCustomKernel { void Compute(OrtKernelContext* context); private: +#if !defined(REDUCED_OPS_BUILD) void InitTopK(); void InvokeTopK(OrtKernelContext* context); @@ -333,10 +351,12 @@ struct StandaloneCustomKernel { void InitInvokeConv(OrtKernelContext* context); // create Conv and invoke in Compute(...) - Ort::KernelInfo info_copy_{nullptr}; - Ort::Op op_add_{nullptr}; Ort::Op op_topk_{nullptr}; Ort::Op op_gru_{nullptr}; +#endif + + Ort::KernelInfo info_copy_{nullptr}; + Ort::Op op_add_{nullptr}; }; struct StandaloneCustomOp : Ort::CustomOpBase { diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 6a0f5b39e2..b984c90164 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -92,8 +92,11 @@ static void TestInference(Ort::Env& env, const std::basic_string& mod OrtCustomOpDomain* custom_op_domain_ptr, const ORTCHAR_T* custom_op_library_filename, bool test_session_creation_only = false, - void* cuda_compute_stream = nullptr) { - Ort::SessionOptions session_options; + void* cuda_compute_stream = nullptr, + Ort::SessionOptions* predefined_session_options = nullptr) { + Ort::SessionOptions default_session_options; + Ort::SessionOptions& session_options = predefined_session_options ? *predefined_session_options + : default_session_options; if (provider_type == 1) { #ifdef USE_CUDA @@ -444,9 +447,8 @@ TEST(CApiTest, custom_op_set_input_memory_type) { } #endif -#if !defined(ORT_MINIMAL_BUILD) && !defined(REDUCED_OPS_BUILD) -// disable test in reduced-op-build since TOPK and GRU are excluded there -TEST(CApiTest, standalone_op_handler) { +#if !defined(ORT_MINIMAL_BUILD) +TEST(CApiTest, StandaloneOpHandler) { std::vector inputs(1); Input& input = inputs[0]; input.name = "X"; @@ -466,9 +468,18 @@ TEST(CApiTest, standalone_op_handler) { custom_op_domain.Add(&standalone_op); #ifdef USE_CUDA - TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 1, custom_op_domain, nullptr); + TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 1, + custom_op_domain, nullptr); #else - TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 0, custom_op_domain, nullptr); + Ort::SessionOptions session_options; + const std::basic_string ort_file = ORT_TSTR("testdata/foo_1.onnx.test_output.ort"); + session_options.SetOptimizedModelFilePath(ort_file.c_str()); + + TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 0, + custom_op_domain, nullptr, false, nullptr, &session_options); + + TestInference(*ort_env, ort_file, inputs, "Y", expected_dims_y, expected_values_y, 0, + custom_op_domain, nullptr); #endif } #endif diff --git a/onnxruntime/test/shared_lib/test_ort_format_models.cc b/onnxruntime/test/shared_lib/test_ort_format_models.cc index c46b46f27d..d67c5a3048 100644 --- a/onnxruntime/test/shared_lib/test_ort_format_models.cc +++ b/onnxruntime/test/shared_lib/test_ort_format_models.cc @@ -141,6 +141,26 @@ TEST(OrtFormatCustomOpTests, LoadOrtModel) { TestInference(*ort_env, ort_file, inputs, "Y", expected_dims_y, expected_values_y, custom_op_domain); } + +TEST(OrtFormatCustomOpTests, LoadOrtModelStandaloneCustomOpImplementation) { + const std::basic_string ort_file = ORT_TSTR("testdata/foo_1.onnx.ort"); + + StandaloneCustomOp standalone_op{onnxruntime::kCpuExecutionProvider}; + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&standalone_op); + + // load the ORT format model and execute it + std::vector inputs(1); + Input& input = inputs[0]; + input.name = "X"; + input.dims = {3, 2}; + input.values = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + + std::vector expected_dims_y = {3, 2}; + std::vector expected_values_y = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f}; + + TestInference(*ort_env, ort_file, inputs, "Y", expected_dims_y, expected_values_y, custom_op_domain); +} #endif #endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS) diff --git a/onnxruntime/test/testdata/foo_1.onnx.ort b/onnxruntime/test/testdata/foo_1.onnx.ort index 11a5c42f2f8b55929eee85313218caf2f11b35e9..3c2ee4c4a1820de788d3e590c0b3bc28c8d2e753 100644 GIT binary patch delta 543 zcmZ9Iy-EW?6ov1`EJ?9o3JVKC3qgw@30TN>LO>)5#xORL!Xj9SrnC^^D=b8?u(Ys9 zBlrYXDO2Pzd;p0*<99Y$A{n@I?wz^koS8d{K3Dr|W_Hqc&B#g?S&=FaS{j3#*|=c# zKJ`AI-k4nkv%5FyA^Gl?_tMZk?RG$v2xt}`wa?<}zBTBPvUJtss%B8d>09G>Cf=i| z!l{HrZV@#?f@NZh$c*m2_mqJ|&fzN`fFa2fxB@|+F|NahMsDoN^x|xi&bLF4paxds4 K4?C|Po$wFT{B+&` delta 112 zcmeyte}X4af`Ng-KPbc(NHOp*2mmn$ki`O|L3{>=2S6;t#=uba=l_2eAe#}0LG%S6 sAEXz=W(Hy(AT9u6k;#lq@{> \ + /home/onnxruntimedev/.test_data/required_ops.ort_models.config +cat /onnxruntime_src/onnxruntime/test/testdata/ort_minimal_e2e_test_data/required_ops.standalone_invoker.config >> \ + /home/onnxruntimedev/.test_data/required_ops_and_types.ort_models.config + # Test that we can convert an ONNX model with custom ops to ORT format mkdir /home/onnxruntimedev/.test_data/custom_ops_model cp /onnxruntime_src/onnxruntime/test/testdata/custom_op_library/*.onnx /home/onnxruntimedev/.test_data/custom_ops_model/