diff --git a/docs/AddingCustomOp.md b/docs/AddingCustomOp.md index 6ad742617b..25bbfdef14 100644 --- a/docs/AddingCustomOp.md +++ b/docs/AddingCustomOp.md @@ -2,16 +2,12 @@ Adding a new op =============== ## A new op can be written and registered with ONNXRuntime in the following 3 ways -### 1. Using a dynamic shared library -* First write the implementation of the op and schema (if required) and assemble them in a shared library. -See [this](../onnxruntime/test/custom_op_shared_lib) for an example. Currently -this is supported for Linux only. - -Example of creating a shared lib using g++ on Linux: -```g++ -std=c++14 -shared test_custom_op.cc -o test_custom_op.so -fPIC -I. -Iinclude/onnxruntime -L. -lonnxruntime -DONNX_ML -DONNX_NAMESPACE=onnx``` - -* Register the shared lib with ONNXRuntime. -See [this](../onnxruntime/test/shared_lib/test_inference.cc) for an example. +### 1. Using the experimental custom op API in the C API (onnxruntime_c_api.h) +Note: These APIs are experimental and will change in the next release. They're released now for feedback and experimentation. +* Create an OrtCustomOpDomain with the domain name used by the custom ops +* Create an OrtCustomOp structure for each op and add them to the OrtCustomOpDomain with OrtCustomOpDomain_Add +* Call OrtAddCustomOpDomain to add the custom domain of ops to the session options +See [this](../onnxruntime/test/custom_op_shared_lib/test_custom_op.cc) for an example. ### 2. Using RegisterCustomRegistry API * Implement your kernel and schema (if required) using the OpKernel and OpSchema APIs (headers are in the include folder). diff --git a/include/onnxruntime/core/framework/op_kernel.h b/include/onnxruntime/core/framework/op_kernel.h index 05393030f3..d41a0f6da5 100644 --- a/include/onnxruntime/core/framework/op_kernel.h +++ b/include/onnxruntime/core/framework/op_kernel.h @@ -151,6 +151,7 @@ class OpKernelContext { const MLValue* GetInputMLValue(int index) const; const MLValue* GetImplicitInputMLValue(int index) const; MLValue* GetOutputMLValue(int index); + MLValue* OutputMLValue(int index, const TensorShape& shape); // Creates the MLValue* based on the shape, if it does not exist private: ORT_DISALLOW_COPY_AND_ASSIGNMENT(OpKernelContext); diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index fe06045c73..b93451adb8 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -66,7 +66,6 @@ extern "C" { #define NO_EXCEPTION #endif - // Copied from TensorProto::DataType // Currently, Ort doesn't support complex64, complex128, bfloat16 types typedef enum ONNXTensorElementDataType { @@ -152,6 +151,7 @@ ORT_RUNTIME_CLASS(TypeInfo); ORT_RUNTIME_CLASS(TensorTypeAndShapeInfo); ORT_RUNTIME_CLASS(SessionOptions); ORT_RUNTIME_CLASS(Callback); +ORT_RUNTIME_CLASS(CustomOpDomain); // When passing in an allocator to any ORT function, be sure that the allocator object // is not destroyed until the last allocated object using it is freed. @@ -511,6 +511,71 @@ ORT_API_STATUS(OrtGetValueCount, const OrtValue* value, size_t* out); ORT_API_STATUS(OrtCreateValue, OrtValue** const in, int num_values, enum ONNXType value_type, OrtValue** out); +/* + * EXPERIMENTAL APIS - Subject to change. Released as a preview to get feedback and enable early testing +*/ + +/* + * Steps to use a custom op: + * 1 Create an OrtCustomOpDomain with the domain name used by the custom ops + * 2 Create an OrtCustomOp structure for each op and add them to the domain + * 3 Call OrtAddCustomOpDomain to add the custom domain of ops to the session options +*/ +struct OrtKernelInfo; +typedef struct OrtKernelInfo OrtKernelInfo; + +/* + * These allow reading node attributes during kernel creation +*/ +ORT_API_STATUS(OrtKernelInfoGetAttribute_float, _In_ OrtKernelInfo* info, _In_ const char* name, _Out_ float* out); +ORT_API_STATUS(OrtKernelInfoGetAttribute_int64, _In_ OrtKernelInfo* info, _In_ const char* name, _Out_ int64_t* out); + +/* + * The OrtCustomOp structure defines a custom op's schema and its kernel callbacks. The callbacks are filled in by + * the implementor of the custom op. +*/ +struct OrtCustomOp { + uint32_t version; // Initialize to ORT_API_VERSION + + // This callback creates the kernel, which is a user defined parameter that is passed to the Kernel* callbacks below. + void(ORT_API_CALL* CreateKernel)(_In_ struct OrtCustomOp* op, _In_ OrtKernelInfo* info, _Out_ void** op_kernel); + + // Returns the name of the op + const char*(ORT_API_CALL* GetName)(_In_ struct OrtCustomOp* op); + + // Returns the count and types of the input & output tensors + ONNXTensorElementDataType(ORT_API_CALL* GetInputType)(_In_ struct OrtCustomOp* op, _In_ size_t index); + size_t(ORT_API_CALL* GetInputTypeCount)(_In_ struct OrtCustomOp* op); + ONNXTensorElementDataType(ORT_API_CALL* GetOutputType)(_In_ struct OrtCustomOp* op, _In_ size_t index); + size_t(ORT_API_CALL* GetOutputTypeCount)(_In_ struct OrtCustomOp* op); + + // Op kernel callbacks + void(ORT_API_CALL* KernelGetOutputShape)(_In_ void* op_kernel, _In_ OrtValue** inputs, _In_ size_t input_count, _In_ size_t output_index, _In_ OrtTensorTypeAndShapeInfo* output); + void(ORT_API_CALL* KernelCompute)(_In_ void* op_kernel, _In_ OrtValue** inputs, _In_ size_t input_count, _In_ OrtValue** outputs, _In_ size_t output_count); + void(ORT_API_CALL* KernelDestroy)(_In_ void* op_kernel); +}; +typedef struct OrtCustomOp OrtCustomOp; + +/* +* Create a custom op domain. After all sessions using it are released, call OrtReleaseCustomOpDomain +*/ +ORT_API(OrtCustomOpDomain*, OrtCreateCustomOpDomain, _In_ const char* domain, _In_ int op_version_start, _In_ int op_version_end); + +/* + * Add custom ops to the OrtCustomOpDomain + * Note: The OrtCustomOp* pointer must remain valid until the OrtCustomOpDomain using it is released +*/ +ORT_API_STATUS(OrtCustomOpDomain_Add, _In_ OrtCustomOpDomain* custom_op_domain, _In_ OrtCustomOp* op); + +/* + * Add a custom op domain to the OrtSessionOptions + * Note: The OrtCustomOpDomain* must not be deleted until the sessions using it are released +*/ +ORT_API_STATUS(OrtAddCustomOpDomain, _In_ OrtSessionOptions* options, OrtCustomOpDomain* custom_op_domain); +/* + * END EXPERIMENTAL +*/ + #ifdef __cplusplus } #endif diff --git a/onnxruntime/core/framework/op_kernel.cc b/onnxruntime/core/framework/op_kernel.cc index 6e26b70efb..dfd7489a2d 100644 --- a/onnxruntime/core/framework/op_kernel.cc +++ b/onnxruntime/core/framework/op_kernel.cc @@ -24,6 +24,11 @@ OpKernelContext::OpKernelContext(IExecutionFrame* frame, } Tensor* OpKernelContext::Output(int index, const TensorShape& shape) { + auto p_ml_value = OutputMLValue(index, shape); + return p_ml_value ? p_ml_value->GetMutable() : nullptr; +} + +MLValue* OpKernelContext::OutputMLValue(int index, const TensorShape& shape) { if (index < 0 || index >= OutputCount()) return nullptr; @@ -34,7 +39,7 @@ Tensor* OpKernelContext::Output(int index, const TensorShape& shape) { MLValue* p_ml_value = nullptr; Status status = execution_frame_->GetOrCreateNodeOutputMLValue(GetOutputArgIndex(index), &shape, p_ml_value); ORT_ENFORCE(status.IsOK(), status.ErrorMessage()); - return p_ml_value ? p_ml_value->GetMutable() : nullptr; + return p_ml_value; } int OpKernelContext::NumVariadicInputs(size_t arg_num) const { diff --git a/onnxruntime/core/framework/op_kernel_context_internal.h b/onnxruntime/core/framework/op_kernel_context_internal.h index 3ec850310d..5cc3a50a96 100644 --- a/onnxruntime/core/framework/op_kernel_context_internal.h +++ b/onnxruntime/core/framework/op_kernel_context_internal.h @@ -39,6 +39,10 @@ class OpKernelContextInternal : public OpKernelContext { return OpKernelContext::GetOutputMLValue(index); } + MLValue* OutputMLValue(int index, const TensorShape& shape) { + return OpKernelContext::OutputMLValue(index, shape); + } + std::unordered_map GetImplicitInputs() const { // we need to convert implicit_inputs_ to a name to MLValue map so it can be used in the ExecutionFrame // for a subgraph (the index numbers will be different there). diff --git a/onnxruntime/core/framework/tensor_type_and_shape.cc b/onnxruntime/core/framework/tensor_type_and_shape.cc index 3fd5323536..af3c52c51b 100644 --- a/onnxruntime/core/framework/tensor_type_and_shape.cc +++ b/onnxruntime/core/framework/tensor_type_and_shape.cc @@ -5,6 +5,7 @@ #include "core/framework/tensor_shape.h" #include "core/framework/ml_value.h" #include "core/framework/onnxruntime_typeinfo.h" +#include "core/framework/tensor_type_and_shape.h" #include #include @@ -15,16 +16,6 @@ using onnxruntime::DataTypeImpl; using onnxruntime::MLFloat16; using onnxruntime::Tensor; -struct OrtTensorTypeAndShapeInfo { - public: - ONNXTensorElementDataType type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; - onnxruntime::TensorShape shape; - - OrtTensorTypeAndShapeInfo() = default; - OrtTensorTypeAndShapeInfo(const OrtTensorTypeAndShapeInfo& other) = delete; - OrtTensorTypeAndShapeInfo& operator=(const OrtTensorTypeAndShapeInfo& other) = delete; -}; - #define API_IMPL_BEGIN try { #define API_IMPL_END \ } \ @@ -72,8 +63,7 @@ ORT_API(int64_t, OrtGetTensorShapeElementCount, _In_ const OrtTensorTypeAndShape struct OrtValue; -namespace { -inline ONNXTensorElementDataType MLDataTypeToOnnxRuntimeTensorElementDataType( +ONNXTensorElementDataType MLDataTypeToOnnxRuntimeTensorElementDataType( const onnxruntime::DataTypeImpl* cpp_type) { ONNXTensorElementDataType type; if (cpp_type == onnxruntime::DataTypeImpl::GetType()) { @@ -109,7 +99,41 @@ inline ONNXTensorElementDataType MLDataTypeToOnnxRuntimeTensorElementDataType( } return type; } -} // namespace + +const onnxruntime::DataTypeImpl* TensorElementDataTypeToMLDataType(ONNXTensorElementDataType type) { + switch (type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32: + return onnxruntime::DataTypeImpl::GetType(); + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64: + return onnxruntime::DataTypeImpl::GetType(); + default: + return nullptr; + } +} OrtStatus* GetTensorShapeAndType(const onnxruntime::TensorShape* shape, const onnxruntime::DataTypeImpl* tensor_data_type, OrtTensorTypeAndShapeInfo** out) { ONNXTensorElementDataType type = MLDataTypeToOnnxRuntimeTensorElementDataType(tensor_data_type); diff --git a/onnxruntime/core/framework/tensor_type_and_shape.h b/onnxruntime/core/framework/tensor_type_and_shape.h new file mode 100644 index 0000000000..9c829215b9 --- /dev/null +++ b/onnxruntime/core/framework/tensor_type_and_shape.h @@ -0,0 +1,13 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +struct OrtTensorTypeAndShapeInfo { + public: + ONNXTensorElementDataType type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; + onnxruntime::TensorShape shape; + + OrtTensorTypeAndShapeInfo() = default; + OrtTensorTypeAndShapeInfo(const OrtTensorTypeAndShapeInfo& other) = delete; + OrtTensorTypeAndShapeInfo& operator=(const OrtTensorTypeAndShapeInfo& other) = delete; +}; diff --git a/onnxruntime/core/providers/cpu/symbols.txt b/onnxruntime/core/providers/cpu/symbols.txt index 4f8aaaca4a..c8eac9cd59 100644 --- a/onnxruntime/core/providers/cpu/symbols.txt +++ b/onnxruntime/core/providers/cpu/symbols.txt @@ -1,3 +1,4 @@ +OrtAddCustomOpDomain OrtAllocatorAlloc OrtAllocatorFree OrtAllocatorGetInfo @@ -11,6 +12,7 @@ OrtCloneSessionOptions OrtCompareAllocatorInfo OrtCreateAllocatorInfo OrtCreateCpuAllocatorInfo +OrtCreateCustomOpDomain OrtCreateDefaultAllocator OrtCreateEnv OrtCreateEnvWithCustomLogger @@ -21,6 +23,7 @@ OrtCreateTensorAsOrtValue OrtCreateTensorTypeAndShapeInfo OrtCreateTensorWithDataAsOrtValue OrtCreateValue +OrtCustomOpDomain_Add OrtDisableCpuMemArena OrtDisableMemPattern OrtDisableProfiling @@ -48,6 +51,7 @@ OrtGetValueType OrtIsTensor OrtReleaseAllocator OrtReleaseAllocatorInfo +OrtReleaseCustomOpDomain OrtReleaseEnv OrtReleaseRunOptions OrtReleaseSession diff --git a/onnxruntime/core/session/abi_session_options_impl.h b/onnxruntime/core/session/abi_session_options_impl.h index 1af9e5f268..ca57a01b85 100644 --- a/onnxruntime/core/session/abi_session_options_impl.h +++ b/onnxruntime/core/session/abi_session_options_impl.h @@ -13,6 +13,7 @@ struct OrtSessionOptions { onnxruntime::SessionOptions value; std::vector custom_op_paths; + std::vector custom_op_domains_; std::vector> provider_factories; OrtSessionOptions() = default; ~OrtSessionOptions(); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 562e56fbbb..e70ff06e5c 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -22,6 +22,7 @@ #include "core/framework/allocatormgr.h" #include "core/framework/customregistry.h" #include "core/framework/environment.h" +#include "core/framework/error_code_helper.h" #include "core/framework/execution_frame.h" #include "core/framework/feeds_fetches_manager.h" #include "core/framework/graph_partitioner.h" @@ -31,11 +32,13 @@ #include "core/framework/mldata_type_utils.h" #include "core/framework/mlvalue_name_idx_map.h" #include "core/framework/sequential_executor.h" +#include "core/framework/op_kernel_context_internal.h" #include "core/framework/parallel_executor.h" #include "core/framework/path_lib.h" #include "core/framework/session_state.h" #include "core/framework/session_state_initializer.h" #include "core/framework/tensorprotoutils.h" +#include "core/framework/tensor_type_and_shape.h" #include "core/framework/utils.h" #include "core/optimizer/transformer_memcpy.h" #include "core/optimizer/graph_transformer.h" @@ -52,6 +55,77 @@ using namespace ONNX_NAMESPACE; +ONNXTensorElementDataType MLDataTypeToOnnxRuntimeTensorElementDataType(const onnxruntime::DataTypeImpl* cpp_type); +const onnxruntime::DataTypeImpl* TensorElementDataTypeToMLDataType(ONNXTensorElementDataType type); + +namespace onnxruntime { +const char* ElementTypeToString(MLDataType type) { + if (type == DataTypeImpl::GetType()) { + return "tensor(float)"; + } else if (type == DataTypeImpl::GetType()) { + return "tensor(bool)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(int32)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(double)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(string)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(uint8)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(uint16)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(int16)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(int64)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(uint32)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(uint64)"; + } + + else if (type == DataTypeImpl::GetType()) { + return "tensor(MLFloat16)"; + } else if (type == DataTypeImpl::GetType()) { + return "tensor(bfloat16)"; + } else { + return "unknown"; + } +} +} // namespace onnxruntime + +ORT_API_STATUS_IMPL(OrtKernelInfoGetAttribute_float, _In_ OrtKernelInfo* info, _In_ const char* name, _Out_ float* out) { + auto status = reinterpret_cast(info)->GetAttr(name, out); + if (status.IsOK()) + return nullptr; + return onnxruntime::ToOrtStatus(status); +} + +ORT_API_STATUS_IMPL(OrtKernelInfoGetAttribute_int64, _In_ OrtKernelInfo* info, _In_ const char* name, _Out_ int64_t* out) { + auto status = reinterpret_cast(info)->GetAttr(name, out); + if (status.IsOK()) + return nullptr; + return onnxruntime::ToOrtStatus(status); +} + namespace onnxruntime { namespace { template @@ -87,6 +161,41 @@ inline std::basic_string GetCurrentTimeString() { return std::basic_string(time_str); } } // namespace +struct CustomOpKernel : OpKernel { + CustomOpKernel(const OpKernelInfo& info, OrtCustomOp& op) : OpKernel(info), op_(op) { + op_.CreateKernel(&op_, reinterpret_cast(const_cast(&info)), &op_kernel_); + } + + ~CustomOpKernel() { + op_.KernelDestroy(op_kernel_); + } + + Status Compute(OpKernelContext* ctx) const override { + auto* ictx = static_cast(ctx); + std::vector input_tensors; + auto input_count = ictx->InputCount(); + for (int i = 0; i < input_count; i++) + input_tensors.emplace_back(const_cast(reinterpret_cast(ictx->GetInputMLValue(i)))); + + std::vector output_tensors; + auto output_count = ictx->OutputCount(); + for (int i = 0; i < output_count; i++) { + OrtTensorTypeAndShapeInfo info; + op_.KernelGetOutputShape(op_kernel_, input_tensors.data(), input_tensors.size(), i, &info); + output_tensors.emplace_back(reinterpret_cast(ictx->OutputMLValue(0, info.shape))); + } + + op_.KernelCompute(op_kernel_, input_tensors.data(), input_tensors.size(), output_tensors.data(), output_tensors.size()); + return Status::OK(); + } + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CustomOpKernel); + + OrtCustomOp& op_; + void* op_kernel_; +}; + class InferenceSession::Impl { public: Impl(const SessionOptions& session_options, logging::LoggingManager* logging_manager) @@ -156,6 +265,58 @@ class InferenceSession::Impl { return Status::OK(); } + common::Status AddCustomOpDomains(const std::vector& op_domains) { + auto custom_registry = std::make_shared(); + + for (auto& domain : op_domains) { + SchemasContainer schemas_container; + + schemas_container.domain = domain->domain_; + schemas_container.baseline_opset_version = domain->op_version_start_; + schemas_container.opset_version = domain->op_version_end_; + + for (auto& op : domain->custom_ops_) { + ONNX_NAMESPACE::OpSchema schema(op->GetName(op), "unknown", 0); + + auto input_count = op->GetInputTypeCount(op); + for (size_t i = 0; i < input_count; i++) { + auto type = op->GetInputType(op, i); + + schema.Input(i, "A", "Description", ElementTypeToString(TensorElementDataTypeToMLDataType(type))); + } + + auto output_count = op->GetOutputTypeCount(op); + for (size_t i = 0; i < output_count; i++) { + auto type = op->GetOutputType(op, i); + + schema.Output(i, "A", "Description", ElementTypeToString(TensorElementDataTypeToMLDataType(type))); + } + + schema.SinceVersion(domain->op_version_start_); + schema.AllowUncheckedAttributes(); + + schemas_container.schemas_list.push_back(schema); + + KernelDefBuilder def_builder; + def_builder.SetName(op->GetName(op)) + .SetDomain(onnxruntime::kOnnxDomain) + .SinceVersion(domain->op_version_start_) + .Provider(onnxruntime::kCpuExecutionProvider); + KernelCreateFn kernel_create_fn = [&op](const OpKernelInfo& info) -> OpKernel* { return new CustomOpKernel(info, *op); }; + KernelCreateInfo create_info(def_builder.Build(), kernel_create_fn); + + custom_registry->RegisterCustomKernel(create_info); + } + + ORT_RETURN_IF_ERROR(custom_registry->RegisterOpSet(schemas_container.schemas_list, + schemas_container.domain, + schemas_container.baseline_opset_version, + schemas_container.opset_version)); + } + RegisterCustomRegistry(custom_registry); + return Status::OK(); + } + common::Status RegisterCustomRegistry(std::shared_ptr& custom_registry) { if (custom_registry == nullptr) { return Status(common::ONNXRUNTIME, common::FAIL, "Received nullptr for custom registry"); @@ -163,6 +324,8 @@ class InferenceSession::Impl { // Insert session-level customized kernel registry. kernel_registry_manager_.RegisterKernelRegistry(custom_registry); + // if (custom_schema_registries_.empty()) + // custom_schema_registries_.push_back(); custom_schema_registries_.push_back(custom_registry); return Status::OK(); } @@ -1041,4 +1204,8 @@ common::Status InferenceSession::Run(IOBinding& io_binding) { common::Status InferenceSession::LoadCustomOps(const std::vector& dso_list) { return impl_->LoadCustomOps(dso_list); } + +common::Status InferenceSession::AddCustomOpDomains(const std::vector& ops) { + return impl_->AddCustomOpDomains(ops); +} } // namespace onnxruntime diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index aa43b4bbe3..3d88997935 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -20,6 +20,13 @@ namespace ONNX_NAMESPACE { class ModelProto; } // namespace ONNX_NAMESPACE +struct OrtCustomOpDomain { + std::string domain_; + int op_version_start_{}; + int op_version_end_{}; + std::vector custom_ops_; +}; + namespace onnxruntime { class IExecutionProvider; // forward decl class IOBinding; @@ -132,6 +139,8 @@ class InferenceSession { */ common::Status LoadCustomOps(const std::vector& dso_list); + common::Status AddCustomOpDomains(const std::vector& ops); + /** * Register a custom registry for operator schema and kernels. If you've one to register, * call this before invoking Initialize(). diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index f5fc9f545d..f359f5b932 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -337,18 +337,50 @@ ORT_API_STATUS_IMPL(OrtCreateTensorAsOrtValue, _Inout_ OrtAllocator* allocator, API_IMPL_END } -template -static OrtStatus* CreateSessionImpl(_In_ OrtEnv* env, _In_ T model_path, - _In_ const OrtSessionOptions* options, - _Out_ OrtSession** out) { +ORT_API(OrtCustomOpDomain*, OrtCreateCustomOpDomain, _In_ const char* domain, int op_version_start, int op_version_end) { + auto custom_op_domain = std::make_unique(); + custom_op_domain->domain_ = domain; + custom_op_domain->op_version_start_ = op_version_start; + custom_op_domain->op_version_end_ = op_version_end; + return custom_op_domain.release(); +} + +ORT_API(void, OrtReleaseCustomOpDomain, OrtCustomOpDomain* ptr) { + delete ptr; +} + +ORT_API_STATUS_IMPL(OrtCustomOpDomain_Add, _In_ OrtCustomOpDomain* custom_op_domain, OrtCustomOp* op) { + API_IMPL_BEGIN + custom_op_domain->custom_ops_.emplace_back(op); + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtAddCustomOpDomain, _In_ OrtSessionOptions* options, OrtCustomOpDomain* custom_op_domain) { + API_IMPL_BEGIN + options->custom_op_domains_.emplace_back(custom_op_domain); + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtCreateSession, _In_ OrtEnv* env, _In_ const ORTCHAR_T* model_path, + _In_ const OrtSessionOptions* options, _Out_ OrtSession** out) { API_IMPL_BEGIN auto sess = std::make_unique<::onnxruntime::InferenceSession>(options == nullptr ? onnxruntime::SessionOptions() : options->value, env->loggingManager); Status status; - if (options != nullptr && !options->custom_op_paths.empty()) { - status = sess->LoadCustomOps(options->custom_op_paths); - if (!status.IsOK()) - return ToOrtStatus(status); + if (options != nullptr) { + if (!options->custom_op_paths.empty()) { + status = sess->LoadCustomOps(options->custom_op_paths); + if (!status.IsOK()) + return ToOrtStatus(status); + } + if (!options->custom_op_domains_.empty()) { + status = sess->AddCustomOpDomains(options->custom_op_domains_); + if (!status.IsOK()) + return ToOrtStatus(status); + } } + if (options != nullptr) for (auto& factory : options->provider_factories) { auto provider = factory->CreateProvider(); @@ -366,22 +398,6 @@ static OrtStatus* CreateSessionImpl(_In_ OrtEnv* env, _In_ T model_path, API_IMPL_END } -#ifdef _WIN32 -ORT_API_STATUS_IMPL(OrtCreateSession, _In_ OrtEnv* env, _In_ const wchar_t* model_path, - _In_ const OrtSessionOptions* options, _Out_ OrtSession** out) { - API_IMPL_BEGIN - return CreateSessionImpl(env, model_path, options, out); - API_IMPL_END -} -#else -ORT_API_STATUS_IMPL(OrtCreateSession, _In_ OrtEnv* env, _In_ const char* model_path, - _In_ const OrtSessionOptions* options, _Out_ OrtSession** out) { - API_IMPL_BEGIN - return CreateSessionImpl(env, model_path, options, out); - API_IMPL_END -} -#endif - ORT_API_STATUS_IMPL(OrtRun, _In_ OrtSession* sess, _In_ OrtRunOptions* run_options, _In_ const char* const* input_names, _In_ const OrtValue* const* input, size_t input_len, diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 7737edb686..ad1129ec2a 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -63,7 +63,7 @@ void TestInference(OrtEnv* env, T model_uri, const std::vector& values_x, const std::vector& expected_dims_y, const std::vector& expected_values_y, - int provider_type, bool custom_op) { + int provider_type, bool custom_op, OrtCustomOpDomain* custom_op_domain_ptr = nullptr) { SessionOptionsWrapper sf(env); if (provider_type == 1) { @@ -93,6 +93,10 @@ void TestInference(OrtEnv* env, T model_uri, if (custom_op) { sf.AppendCustomOpLibPath("libonnxruntime_custom_op_shared_lib_test.so"); } + if (custom_op_domain_ptr) { + ORT_THROW_ON_ERROR(OrtAddCustomOpDomain(sf, custom_op_domain_ptr)); + } + std::unique_ptr inference_session(sf.OrtCreateSession(model_uri), OrtReleaseSession); std::unique_ptr default_allocator(std::make_unique()); @@ -169,6 +173,88 @@ TEST_F(CApiTest, DISABLED_custom_op) { } #endif +struct OrtTensorDimensions : std::vector { + OrtTensorDimensions(OrtValue* value) { + OrtTensorTypeAndShapeInfo* info; + ORT_THROW_ON_ERROR(OrtGetTensorShapeAndType(value, &info)); + auto dimensionCount = OrtGetNumOfDimensions(info); + resize(dimensionCount); + OrtGetDimensions(info, data(), dimensionCount); + OrtReleaseTensorTypeAndShapeInfo(info); + } + + size_t ElementCount() const { + int64_t count = 1; + for (int i = 0; i < size(); i++) + count *= (*this)[i]; + return count; + } +}; + +template +constexpr size_t countof(T (&)[N]) { return N; } + +struct MyCustomKernel { + MyCustomKernel(OrtKernelInfo& /*info*/) { + } + + void GetOutputShape(OrtValue** inputs, size_t /*input_count*/, size_t /*output_index*/, OrtTensorTypeAndShapeInfo* info) { + OrtTensorDimensions dimensions(inputs[0]); + ORT_THROW_ON_ERROR(OrtSetDims(info, dimensions.data(), dimensions.size())); + } + + void Compute(OrtValue** inputs, size_t /*input_count*/, OrtValue** outputs, size_t /*output_count*/) { + const float* X; + const float* Y; + ORT_THROW_ON_ERROR(OrtGetTensorMutableData(inputs[0], reinterpret_cast(const_cast(&X)))); + ORT_THROW_ON_ERROR(OrtGetTensorMutableData(inputs[1], reinterpret_cast(const_cast(&Y)))); + + float* out; + ORT_THROW_ON_ERROR(OrtGetTensorMutableData(outputs[0], reinterpret_cast(&out))); + + int64_t size = OrtTensorDimensions(inputs[0]).ElementCount(); + for (int64_t i = 0; i < size; i++) { + out[i] = X[i] + Y[i]; + } + } +}; + +struct MyCustomOp : OrtCustomOp { + MyCustomOp() { + OrtCustomOp::version = ORT_API_VERSION; + OrtCustomOp::CreateKernel = [](OrtCustomOp* /*this_*/, OrtKernelInfo* info, void** output) { *output = new MyCustomKernel(*info); }; + OrtCustomOp::GetName = [](OrtCustomOp* /*this_*/) { return "Foo"; }; + + static const ONNXTensorElementDataType c_inputTypes[] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; + OrtCustomOp::GetInputTypeCount = [](OrtCustomOp* /*this_*/) { return countof(c_inputTypes); }; + OrtCustomOp::GetInputType = [](OrtCustomOp* /*this_*/, size_t index) { return c_inputTypes[index]; }; + + static const ONNXTensorElementDataType c_outputTypes[] = {ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; + OrtCustomOp::GetOutputTypeCount = [](OrtCustomOp* /*this_*/) { return countof(c_outputTypes); }; + OrtCustomOp::GetOutputType = [](OrtCustomOp* /*this_*/, size_t index) { return c_outputTypes[index]; }; + + OrtCustomOp::KernelGetOutputShape = [](void* op_kernel, OrtValue** inputs, size_t input_count, size_t output_index, OrtTensorTypeAndShapeInfo* output) { static_cast(op_kernel)->GetOutputShape(inputs, input_count, output_index, output); }; + OrtCustomOp::KernelCompute = [](void* op_kernel, OrtValue** inputs, size_t input_count, OrtValue** outputs, size_t output_count) { static_cast(op_kernel)->Compute(inputs, input_count, outputs, output_count); }; + OrtCustomOp::KernelDestroy = [](void* op_kernel) { delete static_cast(op_kernel); }; + } +}; + +TEST_F(CApiTest, custom_op_handler) { + std::cout << "Running custom op inference" << std::endl; + std::vector dims_x = {3, 2}; + std::vector values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + + // prepare expected inputs and outputs + std::vector expected_dims_y = {3, 2}; + std::vector expected_values_y = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f}; + + MyCustomOp custom_op; + OrtCustomOpDomain* custom_op_domain = OrtCreateCustomOpDomain("", 5, 7); + ORT_THROW_ON_ERROR(OrtCustomOpDomain_Add(custom_op_domain, &custom_op)); + + TestInference(env, CUSTOM_OP_MODEL_URI, dims_x, values_x, expected_dims_y, expected_values_y, false, false, custom_op_domain); +} + #ifdef ORT_RUN_EXTERNAL_ONNX_TESTS TEST_F(CApiTest, create_session_without_session_option) { constexpr PATH_TYPE model_uri = TSTR("../models/opset8/test_squeezenet/model.onnx");