diff --git a/docs/reference/operators/add-custom-op.md b/docs/reference/operators/add-custom-op.md index 99439e7032..e331a9edd0 100644 --- a/docs/reference/operators/add-custom-op.md +++ b/docs/reference/operators/add-custom-op.md @@ -15,16 +15,180 @@ ONNX Runtime provides options to run custom operators that are not official ONNX * TOC placeholder {:toc} -## Register a custom operator -A new op can be registered with ONNX Runtime using the Custom Operator API in [onnxruntime_c_api](https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_c_api.h). +## Define and register a custom operator +A custom operator class inherits from `Ort::CustomOpBase` and provides implementations for member functions that define the operator's characteristics and functionality. For example, the following snippet shows the class definition for a basic custom operator named "MyCustomOp" with 2 inputs and 1 output. -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 OrtCustomOpDomain with OrtCustomOpDomain_Add. -3. Call OrtAddCustomOpDomain to add the custom domain of ops to the session options. +```c++ +struct MyCustomOp : Ort::CustomOpBase { + void* CreateKernel(const OrtApi& api, const OrtKernelInfo* info) const { + return std::make_unique(api, info).release(); + }; + + // Returns the name of the custom operator. + const char* GetName() const { return "MyCustomOp"; }; + // Returns the custom operator's execution provider. + const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }; -## Examples -{: .no_toc} + // Returns the number of inputs. + size_t GetInputTypeCount() const { return 2; }; + + // Returns the type of each input. Both inputs are tensor(float). + ONNXTensorElementDataType GetInputType(size_t index) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }; + + // Returns the number of outputs. + size_t GetOutputTypeCount() const { return 1; }; + + // Returns the type of each output. The single output is a tensor(float). + ONNXTensorElementDataType GetOutputType(size_t index) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }; +}; +``` + +Refer to the [OrtCustomOp struct](https://onnxruntime.ai/docs/api/c/struct_ort_custom_op.html) or the [Ort::CustomOpBase struct](https://onnxruntime.ai/docs/api/c/struct_ort_1_1_custom_op_base.html) definitions for a listing of all custom operator member functions. + +A custom operator returns a custom kernel via its `CreateKernel` method. A kernel exposes a `Compute` method that is called during model inference to compute the operator's outputs. For example, the following snippet shows the class definition for a basic custom kernel that adds two tensors. + +```c++ +struct MyCustomKernel { + MyCustomKernel(const OrtApi& api, const OrtKernelInfo* info) {} + + void Compute(OrtKernelContext* context) { + // Setup inputs + Ort::KernelContext ctx(context); + Ort::ConstValue input_X = ctx.GetInput(0); + Ort::ConstValue input_Y = ctx.GetInput(1); + const float* X = input_X.GetTensorData(); + const float* Y = input_Y.GetTensorData(); + + // Setup output, which is assumed to have the same dimensions as the inputs. + std::vector dimensions = input_X.GetTensorTypeAndShapeInfo().GetShape(); + + Ort::UnownedValue output = ctx.GetOutput(0, dimensions); + float* out = output.GetTensorMutableData(); + + const size_t size = output.GetTensorTypeAndShapeInfo().GetElementCount(); + + // Do computation + for (size_t i = 0; i < size; i++) { + out[i] = X[i] + Y[i]; + } + } +}; +``` + +Refer to the API documentation for information on all available custom operator kernel APIs: +- [C APIs for OrtKernelInfo and OrtKernelContext](https://onnxruntime.ai/docs/api/c/struct_ort_api.html) +- [C++ APIs for Ort::KernelInfo](https://onnxruntime.ai/docs/api/c/struct_ort_1_1_kernel_info.html) +- [C++ APIs for Ort::KernelContext](https://onnxruntime.ai/docs/api/c/struct_ort_1_1_kernel_context.html) + +The following snippet shows how to use an `Ort::CustomOpDomain` to register a custom operator with an ONNX Runtime session. + +```c++ +const MyCustomOp my_custom_op; + +Ort::Env env; +Ort::CustomOpDomain domain("my.customop.domain"); +domain.Add(&my_custom_op); // Add a custom op instance to the domain. + +Ort::SessionOptions session_options; +session_options.Add(domain); // Add the domain to the session options. + +// Create a session. +Ort::Session session(env, "my_model_with_custom_ops.onnx", session_options); +``` +## Create a library of custom operators +Custom operators can be defined in a separate shared library (e.g., a .dll on Windows or a .so on Linux). A custom operator library must export and implement a `RegisterCustomOps` function. The `RegisterCustomOps` function adds a `Ort::CustomOpDomain` containing the library's custom operators to the provided session options. + +The following code snippets show how to write a shared library with two custom operators. Refer to a [complete example](https://github.com/microsoft/onnxruntime/tree/main/onnxruntime/test/testdata/custom_op_library) for more details. + +```c++ +// custom_op_library.h + +#pragma once + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +ORT_EXPORT OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api_base); + +#ifdef __cplusplus +} +#endif +``` + +```c++ +// custom_op_library.cc + +#include "custom_op_library.h" + +// Custom operator libraries are not typically linked with ONNX Runtime. +// Therefore, must define ORT_API_MANUAL_INIT before including onnxruntime_cxx_api.h +// to indicate that the OrtApi object will be initialized manually. +#define ORT_API_MANUAL_INIT +#include "onnxruntime_cxx_api.h" +#undef ORT_API_MANUAL_INIT + +#include +#include +#include + +// Define custom operators and kernels ... +struct MyCustomOp : Ort::CustomOpBase { + // ... +}; + +struct MyOtherCustomOp : Ort::CustomOpBase { + // ... +}; + +// This function shows one way of keeping domains alive until the library is unloaded. +static void AddOrtCustomOpDomainToContainer(Ort::CustomOpDomain&& domain) { + static std::vector ort_custom_op_domain_container; + static std::mutex ort_custom_op_domain_mutex; + std::lock_guard lock(ort_custom_op_domain_mutex); + ort_custom_op_domain_container.push_back(std::move(domain)); +} + +// Called by ONNX Runtime to register the library's custom operators with the provided session options. +OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api) { + Ort::InitApi(api->GetApi(ORT_API_VERSION)); // Manually initialize the OrtApi to enable use of C++ API classes and functions. + + // Custom operators are static to ensure they remain valid until the library is unloaded. + static const MyCustomOp my_custom_op; + static const MyOtherCustomOp my_other_custom_op; + + OrtStatus* result = nullptr; + + try { + Ort::CustomOpDomain domain{c_OpDomain}; + domain.Add(&c_CustomOpOne); + domain.Add(&c_CustomOpTwo); + + Ort::UnownedSessionOptions session_options(options); + session_options.Add(domain); + AddOrtCustomOpDomainToContainer(std::move(domain)); + } catch (const std::exception& e) { + Ort::Status status{e}; + result = status.release(); + } + return result; +} +``` +Once compiled, the custom operator shared library can then be registered with an ONNX Runtime session. + +```c++ +Ort::Env env; +Ort::SessionOptions session_options; + +session_options.RegisterCustomOpsLibrary(L"my_custom_op.dll"); + +Ort::Session session(env, "my_model.onnx", session_options); +``` + +### Examples * [C++ helper API](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/test/shared_lib/test_inference.cc): custom ops `MyCustomOp` and `SliceCustomOp` use the [C++ helper API](https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_cxx_api.h). The test file also demonstrates an option to compile the custom ops into a shared library to be used to run a model via the C++ API. @@ -52,6 +216,94 @@ When using CUDA custom ops, to ensure synchronization between ORT's CUDA kernels For example, see how the afore-mentioned `MyCustomOp` is being launched and how the Session using this custom op is created. +## Wrapping an external inference runtime in a custom operator +A custom operator can wrap an entire model that is then inferenced with an external API or runtime. This can facilitate the integration of external inference engines or APIs with ONNX Runtime. + +As an example, consider the following ONNX model with a custom operator named "OpenVINO_Wrapper". The "OpenVINO_Wrapper" node encapsulates an entire MNIST model in OpenVINO's native model format (XML and BIN data). The model data is serialized into the node's attributes and later retrieved by the custom operator's kernel to build an in-memory representation of the model and run inference with OpenVINO C++ APIs. + +

ONNX model of a custom operator wrapping an OpenVINO MNIST model

+ +The following code snippet shows how the custom operator is defined. + +```c++ +struct CustomOpOpenVINO : Ort::CustomOpBase { + explicit CustomOpOpenVINO(Ort::ConstSessionOptions session_options); + + CustomOpOpenVINO(const CustomOpOpenVINO&) = delete; + CustomOpOpenVINO& operator=(const CustomOpOpenVINO&) = delete; + + void* CreateKernel(const OrtApi& api, const OrtKernelInfo* info) const; + + constexpr const char* GetName() const noexcept { + return "OpenVINO_Wrapper"; + } + + constexpr const char* GetExecutionProviderType() const noexcept { + return "CPUExecutionProvider"; + } + + // IMPORTANT: In order to wrap a generic runtime-specific model, the custom operator + // must have a single non-homogeneous variadic input and output. + + constexpr size_t GetInputTypeCount() const noexcept { + return 1; + } + + constexpr size_t GetOutputTypeCount() const noexcept { + return 1; + } + + constexpr ONNXTensorElementDataType GetInputType(size_t /* index */) const noexcept { + return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; + } + + constexpr ONNXTensorElementDataType GetOutputType(size_t /* index */) const noexcept { + return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; + } + + constexpr OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t /* index */) const noexcept { + return INPUT_OUTPUT_VARIADIC; + } + + constexpr OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t /* index */) const noexcept { + return INPUT_OUTPUT_VARIADIC; + } + + constexpr bool GetVariadicInputHomogeneity() const noexcept { + return false; // heterogenous + } + + constexpr bool GetVariadicOutputHomogeneity() const noexcept { + return false; // heterogeneous + } + + // The "device_type" is configurable at the session level. + std::vector GetSessionConfigKeys() const { return {"device_type"}; } + + private: + std::unordered_map session_configs_; +}; +``` + +Note that the custom operator is defined to have a single variadic/heterogenous input and a single variadic/heterogeneous output. This is necessary to enable wrapping OpenVINO models with varying input and output types and shapes (not just an MNIST model). For more information on input and output characteristics, refer to the [OrtCustomOp struct documentation](https://onnxruntime.ai/docs/api/c/struct_ort_custom_op.html). + +Additionally, the custom operator declares "device_type" as a session configuration that can be set by the application. The following code snippet shows how to register and configure a custom operator library containing the aforementioned custom operator. + +```c++ +Ort::Env env; +Ort::SessionOptions session_options; +Ort::CustomOpConfigs custom_op_configs; + +// Create local session config entries for the custom op. +custom_op_configs.AddConfig("OpenVINO_Wrapper", "device_type", "CPU"); + +// Register custom op library and pass in the custom op configs (optional). +session_options.RegisterCustomOpsLibrary("MyOpenVINOWrapper_Lib.so", custom_op_configs); + +Ort::Session session(env, "custom_op_mnist_ov_wrapper.onnx", session_options); +``` + +Refer to the [complete OpenVINO custom operator wrapper example](https://github.com/microsoft/onnxruntime/tree/main/onnxruntime/test/testdata/custom_op_openvino_wrapper_library) for more details. To create an ONNX model that wraps an external model or weights, refer to the [create_custom_op_wrapper.py tool](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/custom_op_wrapper/create_custom_op_wrapper.py). ## Contrib ops diff --git a/images/custom_op_wrapper.png b/images/custom_op_wrapper.png new file mode 100644 index 0000000000..ebbfb6ea97 Binary files /dev/null and b/images/custom_op_wrapper.png differ