mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-20 19:12:24 +00:00
Ryanunderhill/custom op (#550)
* Prototype version that demonstrates it can work * Switched to OrtValue and removed the OrtCustomOpTensor code. * Support multiple outputs and reading of attributes * Add custom domain handling to custom ops * Update documentation * more wording changes
This commit is contained in:
Ryan Hill
GitHub
parent
0b143d0703
commit
af9c554dd3
13 changed files with 441 additions and 50 deletions
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Tensor>() : 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<Tensor>() : nullptr;
|
||||
return p_ml_value;
|
||||
}
|
||||
|
||||
int OpKernelContext::NumVariadicInputs(size_t arg_num) const {
|
||||
|
|
|
|||
|
|
@ -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<std::string, const MLValue*> 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).
|
||||
|
|
|
|||
|
|
@ -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 <assert.h>
|
||||
#include <stdexcept>
|
||||
|
|
@ -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<float>()) {
|
||||
|
|
@ -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<float>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
|
||||
return onnxruntime::DataTypeImpl::GetType<uint8_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
|
||||
return onnxruntime::DataTypeImpl::GetType<int8_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
|
||||
return onnxruntime::DataTypeImpl::GetType<uint16_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16:
|
||||
return onnxruntime::DataTypeImpl::GetType<int16_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
|
||||
return onnxruntime::DataTypeImpl::GetType<int32_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
|
||||
return onnxruntime::DataTypeImpl::GetType<int64_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
|
||||
return onnxruntime::DataTypeImpl::GetType<std::string>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL:
|
||||
return onnxruntime::DataTypeImpl::GetType<bool>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
|
||||
return onnxruntime::DataTypeImpl::GetType<MLFloat16>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16:
|
||||
return onnxruntime::DataTypeImpl::GetType<BFloat16>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
|
||||
return onnxruntime::DataTypeImpl::GetType<double>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
|
||||
return onnxruntime::DataTypeImpl::GetType<uint32_t>();
|
||||
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
|
||||
return onnxruntime::DataTypeImpl::GetType<uint64_t>();
|
||||
default:
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
OrtStatus* GetTensorShapeAndType(const onnxruntime::TensorShape* shape, const onnxruntime::DataTypeImpl* tensor_data_type, OrtTensorTypeAndShapeInfo** out) {
|
||||
ONNXTensorElementDataType type = MLDataTypeToOnnxRuntimeTensorElementDataType(tensor_data_type);
|
||||
|
|
|
|||
13
onnxruntime/core/framework/tensor_type_and_shape.h
Normal file
13
onnxruntime/core/framework/tensor_type_and_shape.h
Normal file
|
|
@ -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;
|
||||
};
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@
|
|||
struct OrtSessionOptions {
|
||||
onnxruntime::SessionOptions value;
|
||||
std::vector<std::string> custom_op_paths;
|
||||
std::vector<OrtCustomOpDomain*> custom_op_domains_;
|
||||
std::vector<std::shared_ptr<onnxruntime::IExecutionProviderFactory>> provider_factories;
|
||||
OrtSessionOptions() = default;
|
||||
~OrtSessionOptions();
|
||||
|
|
|
|||
|
|
@ -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<float>()) {
|
||||
return "tensor(float)";
|
||||
} else if (type == DataTypeImpl::GetType<bool>()) {
|
||||
return "tensor(bool)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<int32_t>()) {
|
||||
return "tensor(int32)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<double>()) {
|
||||
return "tensor(double)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<std::string>()) {
|
||||
return "tensor(string)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<uint8_t>()) {
|
||||
return "tensor(uint8)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<uint16_t>()) {
|
||||
return "tensor(uint16)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<int16_t>()) {
|
||||
return "tensor(int16)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<int64_t>()) {
|
||||
return "tensor(int64)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<uint32_t>()) {
|
||||
return "tensor(uint32)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<uint64_t>()) {
|
||||
return "tensor(uint64)";
|
||||
}
|
||||
|
||||
else if (type == DataTypeImpl::GetType<MLFloat16>()) {
|
||||
return "tensor(MLFloat16)";
|
||||
} else if (type == DataTypeImpl::GetType<BFloat16>()) {
|
||||
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<onnxruntime::OpKernelInfo*>(info)->GetAttr<float>(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<onnxruntime::OpKernelInfo*>(info)->GetAttr<int64_t>(name, out);
|
||||
if (status.IsOK())
|
||||
return nullptr;
|
||||
return onnxruntime::ToOrtStatus(status);
|
||||
}
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace {
|
||||
template <typename T>
|
||||
|
|
@ -87,6 +161,41 @@ inline std::basic_string<T> GetCurrentTimeString() {
|
|||
return std::basic_string<T>(time_str);
|
||||
}
|
||||
} // namespace
|
||||
struct CustomOpKernel : OpKernel {
|
||||
CustomOpKernel(const OpKernelInfo& info, OrtCustomOp& op) : OpKernel(info), op_(op) {
|
||||
op_.CreateKernel(&op_, reinterpret_cast<OrtKernelInfo*>(const_cast<OpKernelInfo*>(&info)), &op_kernel_);
|
||||
}
|
||||
|
||||
~CustomOpKernel() {
|
||||
op_.KernelDestroy(op_kernel_);
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* ctx) const override {
|
||||
auto* ictx = static_cast<OpKernelContextInternal*>(ctx);
|
||||
std::vector<OrtValue*> input_tensors;
|
||||
auto input_count = ictx->InputCount();
|
||||
for (int i = 0; i < input_count; i++)
|
||||
input_tensors.emplace_back(const_cast<OrtValue*>(reinterpret_cast<const OrtValue*>(ictx->GetInputMLValue(i))));
|
||||
|
||||
std::vector<OrtValue*> 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<OrtValue*>(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<OrtCustomOpDomain*>& op_domains) {
|
||||
auto custom_registry = std::make_shared<CustomRegistry>();
|
||||
|
||||
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<CustomRegistry>& 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<std::string>& dso_list) {
|
||||
return impl_->LoadCustomOps(dso_list);
|
||||
}
|
||||
|
||||
common::Status InferenceSession::AddCustomOpDomains(const std::vector<OrtCustomOpDomain*>& ops) {
|
||||
return impl_->AddCustomOpDomains(ops);
|
||||
}
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -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<OrtCustomOp*> custom_ops_;
|
||||
};
|
||||
|
||||
namespace onnxruntime {
|
||||
class IExecutionProvider; // forward decl
|
||||
class IOBinding;
|
||||
|
|
@ -132,6 +139,8 @@ class InferenceSession {
|
|||
*/
|
||||
common::Status LoadCustomOps(const std::vector<std::string>& dso_list);
|
||||
|
||||
common::Status AddCustomOpDomains(const std::vector<OrtCustomOpDomain*>& ops);
|
||||
|
||||
/**
|
||||
* Register a custom registry for operator schema and kernels. If you've one to register,
|
||||
* call this before invoking Initialize().
|
||||
|
|
|
|||
|
|
@ -337,18 +337,50 @@ ORT_API_STATUS_IMPL(OrtCreateTensorAsOrtValue, _Inout_ OrtAllocator* allocator,
|
|||
API_IMPL_END
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
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<OrtCustomOpDomain>();
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -63,7 +63,7 @@ void TestInference(OrtEnv* env, T model_uri,
|
|||
const std::vector<float>& values_x,
|
||||
const std::vector<int64_t>& expected_dims_y,
|
||||
const std::vector<float>& 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<OrtSession, decltype(&OrtReleaseSession)>
|
||||
inference_session(sf.OrtCreateSession(model_uri), OrtReleaseSession);
|
||||
std::unique_ptr<MockedOrtAllocator> default_allocator(std::make_unique<MockedOrtAllocator>());
|
||||
|
|
@ -169,6 +173,88 @@ TEST_F(CApiTest, DISABLED_custom_op) {
|
|||
}
|
||||
#endif
|
||||
|
||||
struct OrtTensorDimensions : std::vector<int64_t> {
|
||||
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 <typename T, size_t N>
|
||||
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<void**>(const_cast<float**>(&X))));
|
||||
ORT_THROW_ON_ERROR(OrtGetTensorMutableData(inputs[1], reinterpret_cast<void**>(const_cast<float**>(&Y))));
|
||||
|
||||
float* out;
|
||||
ORT_THROW_ON_ERROR(OrtGetTensorMutableData(outputs[0], reinterpret_cast<void**>(&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<MyCustomKernel*>(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<MyCustomKernel*>(op_kernel)->Compute(inputs, input_count, outputs, output_count); };
|
||||
OrtCustomOp::KernelDestroy = [](void* op_kernel) { delete static_cast<MyCustomKernel*>(op_kernel); };
|
||||
}
|
||||
};
|
||||
|
||||
TEST_F(CApiTest, custom_op_handler) {
|
||||
std::cout << "Running custom op inference" << std::endl;
|
||||
std::vector<size_t> dims_x = {3, 2};
|
||||
std::vector<float> values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
|
||||
|
||||
// prepare expected inputs and outputs
|
||||
std::vector<int64_t> expected_dims_y = {3, 2};
|
||||
std::vector<float> 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<PATH_TYPE>(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");
|
||||
|
|
|
|||
Loading…
Reference in a new issue