APIs for custom op to invoke ort operator directly (#10713)

* draft kernel creation

* setup eager context

* call into kernel in eager mode

* redefine test case

* refact eager context

* add comment

* remove header

* rename argument

* redefine API definition with types

* list outputs as argument

* switch to int to represent length

* fix compile err

* create attribute API

* add test case for topk

* remove bool from c api

* add gru test case

* remove var

* fix compile warnings

* rename status

* fix compile err

* exclude sparse tensor

* fix comments

* fix comments

* fix build err

* rename file and move location

* format code

* move file to session folder

* fix comments

Co-authored-by: Randy <Randy@randysmac.attlocal.net>
This commit is contained in:
RandySheriffH 2022-05-03 14:16:30 -07:00 committed by GitHub
parent a3e38d7c90
commit 8d69b9398b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 949 additions and 29 deletions

View file

@ -45,6 +45,11 @@ class KernelRegistry {
Status TryFindKernel(const Node& node, ProviderType exec_provider,
const KernelCreateInfo** out) const;
// Find KernelCreateInfo in instant mode
Status TryFindKernel(const std::string& op_name, const std::string& domain, const int& version,
const std::unordered_map<std::string, MLDataType>& type_constraints,
ProviderType exec_provider, const KernelCreateInfo** out) const;
#endif
// Try to find the kernel given a kernel def hash.

View file

@ -21,10 +21,10 @@ class OpKernelContext {
@param arg_num The operator argument number.
@returns Number of inputs the argument has.
*/
int NumVariadicInputs(size_t arg_num) const;
virtual int NumVariadicInputs(size_t arg_num) const;
MLDataType InputType(int index) const;
MLDataType OutputType(int index) const;
virtual MLDataType InputType(int index) const;
virtual MLDataType OutputType(int index) const;
const OrtValue* GetInputOrtValue(int index) const {
return GetInputMLValue(index);
@ -98,29 +98,29 @@ class OpKernelContext {
// Retrieve indexed shape obtained from memory planning before actual
// computation. If the indexed shape cannot be inferred, this function returns
// false.
bool TryGetInferredInputShape(int index, TensorShape& shape) const;
virtual bool TryGetInferredInputShape(int index, TensorShape& shape) const;
// Retrieve indexed shape obtained from memory planning before actual
// computation. If the indexed shape cannot be inferred, this function returns
// false.
bool TryGetInferredOutputShape(int index, TensorShape& shape) const;
virtual bool TryGetInferredOutputShape(int index, TensorShape& shape) const;
const logging::Logger& Logger() const {
return *logger_;
}
// always >= 0
int InputCount() const {
virtual int InputCount() const {
return static_cast<int>(kernel_->Node().InputDefs().size());
}
// always >= 0
int ImplicitInputCount() const {
virtual int ImplicitInputCount() const {
return static_cast<int>(kernel_->Node().ImplicitInputDefs().size());
}
// always >= 0
int OutputCount() const {
virtual int OutputCount() const {
return static_cast<int>(kernel_->Node().OutputDefs().size());
}
@ -128,7 +128,7 @@ class OpKernelContext {
Return an allocator on device 0, with memtype of OrtMemTypeDefault.
@remarks Use SafeInt when calculating the size of memory to allocate using AllocatorPtr->Alloc.
*/
Status GetTempSpaceAllocator(AllocatorPtr* output) const ORT_MUST_USE_RESULT;
virtual Status GetTempSpaceAllocator(AllocatorPtr* output) const ORT_MUST_USE_RESULT;
/**
Return the allocator associated with the CPU EP with memtype of OrtMemTypeDefault.
@ -142,7 +142,7 @@ class OpKernelContext {
@returns Point to the Fence of the input OrtValue.
It is null if the input OrtValue doesn't have fence or the input is optional.
*/
Fence_t InputFence(int index) const;
virtual Fence_t InputFence(int index) const;
/**
Return the fence of current node's implicit input.
@ -150,7 +150,7 @@ class OpKernelContext {
@returns Point to the Fence of the implicit input OrtValue.
It is null if the input OrtValue doesn't have fence or the input is optional.
*/
Fence_t ImplicitInputFence(int index) const;
virtual Fence_t ImplicitInputFence(int index) const;
/**
Return the fence of current node's output identifed by index.
@ -158,12 +158,12 @@ class OpKernelContext {
@returns Point to the Fence of the output OrtValue.
It is null if the output OrtValue doesn't have fence or the output is optional.
*/
Fence_t OutputFence(int index) const;
virtual Fence_t OutputFence(int index) const;
/**
Return the device id that current kernel runs on.
*/
int GetDeviceId() const {
virtual int GetDeviceId() const {
return kernel_->Info().GetExecutionProvider()->GetDeviceId();
}
@ -171,7 +171,7 @@ class OpKernelContext {
Return the compute stream associated with the EP that the kernel is partitioned to.
For EPs that do not have a compute stream (e.g. CPU EP), a nullptr is returned.
*/
void* GetComputeStream() const {
virtual void* GetComputeStream() const {
return kernel_->Info().GetExecutionProvider()->GetComputeStream();
}
@ -203,9 +203,12 @@ class OpKernelContext {
}
protected:
OpKernelContext(concurrency::ThreadPool* threadpool, const logging::Logger& logger);
onnxruntime::NodeIndex GetNodeIndex() const;
const OrtValue* GetInputMLValue(int index) const;
virtual const OrtValue* GetInputMLValue(int index) const;
const OrtValue* GetImplicitInputMLValue(int index) const;
OrtValue* GetOutputMLValue(int index);
@ -214,21 +217,20 @@ class OpKernelContext {
#endif
// Creates the OrtValue* based on the shape, if it does not exist
OrtValue* OutputMLValue(int index, const TensorShape& shape);
virtual OrtValue* OutputMLValue(int index, const TensorShape& shape);
virtual OrtValue* GetOrCreateOutputMLValue(int index);
private:
ORT_DISALLOW_COPY_AND_ASSIGNMENT(OpKernelContext);
OrtValue* GetOrCreateOutputMLValue(int index);
int GetInputArgIndex(int index) const;
int GetImplicitInputArgIndex(int index) const;
int GetOutputArgIndex(int index) const;
IExecutionFrame* const execution_frame_;
const OpKernel* const kernel_;
concurrency::ThreadPool* const threadpool_;
const logging::Logger* const logger_;
IExecutionFrame* const execution_frame_{};
const OpKernel* const kernel_{};
concurrency::ThreadPool* const threadpool_{};
const logging::Logger* const logger_{};
// The argument starting index in ExecutionFrame.
int node_input_start_index_{-1};

View file

@ -79,6 +79,7 @@ class Node {
Fused = 1, ///< The node refers to a function.
};
explicit Node() = default;
~Node() = default;
/**
@ -620,7 +621,7 @@ class Node {
NodeAttributes attributes_;
// Graph that contains this Node
Graph* graph_;
Graph* graph_ = nullptr;
// Map of attribute name to the Graph instance created from the GraphProto attribute
std::unordered_map<std::string, gsl::not_null<Graph*>> attr_to_subgraph_map_;

View file

@ -227,6 +227,16 @@ typedef enum OrtErrorCode {
ORT_EP_FAIL,
} OrtErrorCode;
typedef enum OrtOpAttrType {
ORT_OP_ATTR_UNDEFINED = 0,
ORT_OP_ATTR_INT,
ORT_OP_ATTR_INTS,
ORT_OP_ATTR_FLOAT,
ORT_OP_ATTR_FLOATS,
ORT_OP_ATTR_STRING,
ORT_OP_ATTR_STRINGS,
} OrtOpAttrType;
//! @}
#define ORT_RUNTIME_CLASS(X) \
struct Ort##X; \
@ -257,6 +267,8 @@ ORT_RUNTIME_CLASS(ArenaCfg);
ORT_RUNTIME_CLASS(PrepackedWeightsContainer);
ORT_RUNTIME_CLASS(TensorRTProviderOptionsV2);
ORT_RUNTIME_CLASS(CUDAProviderOptionsV2);
ORT_RUNTIME_CLASS(Op);
ORT_RUNTIME_CLASS(OpAttr);
#ifdef _WIN32
typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr;
@ -3329,6 +3341,86 @@ struct OrtApi {
ORT_API2_STATUS(AddExternalInitializers, _In_ OrtSessionOptions* options,
_In_reads_(input_len) const char* const* initializer_names,
_In_reads_(input_len) const OrtValue* const* initializers, size_t initializers_num);
/** \brief: Create attribute of onnxruntime operator
*
* \param[in] name of the attribute
* \param[in] data of the attribute
* \param[in] data length
* \param[in] data type
* \param[out] attribute that has been created, which must be released by OrtApi::ReleaseOpAttr
*
* \since Version 1.12.
*/
ORT_API2_STATUS(CreateOpAttr,
_In_ const char* name,
_In_ const void* data,
_In_ int len,
_In_ OrtOpAttrType type,
_Outptr_ OrtOpAttr** op_attr);
/* \brief: Release op attribute
*
* \param[in] attribute created by OrtApi::CreateOpAttr
*
* \since Version 1.12.
*/
ORT_CLASS_RELEASE(OpAttr);
/** \brief: Create onnxruntime native operator
*
* \param[in] kernel info
* \param[in] operator name
* \param[in] operator domain
* \param[in] operator opset
* \param[in] name of the type contraints, such as "T" or "T1"
* \param[in] type of each contraints
* \param[in] number of contraints
* \param[in] attributes used to initialize the operator
* \param[in] number of the attributes
* \param[out] operator that has been created
*
* \since Version 1.12.
*/
ORT_API2_STATUS(CreateOp,
_In_ const OrtKernelInfo* info,
_In_ const char* op_name,
_In_ const char* domain,
_In_ int version,
_In_opt_ const char** type_constraint_names,
_In_opt_ const ONNXTensorElementDataType* type_constraint_values,
_In_opt_ int type_constraint_count,
_In_opt_ const OrtOpAttr* const* attr_values,
_In_opt_ int attr_count,
_Outptr_ OrtOp** ort_op);
/** \brief: Invoke the operator created by OrtApi::CreateOp
* The inputs must follow the order as specified in onnx specification
*
* \param[in] kernel context
* \param[in] operator that has been created
* \param[in] inputs
* \param[in] number of inputs
* \param[in] outputs
* \param[in] number of outputs
*
* \since Version 1.12.
*/
ORT_API2_STATUS(InvokeOp,
_In_ const OrtKernelContext* context,
_In_ const OrtOp* ort_op,
_In_ const OrtValue* const* input_values,
_In_ int input_count,
_Inout_ OrtValue* const* output_values,
_In_ int output_count);
/* \brief: Release an onnxruntime operator
*
* \param[in] operator created by OrtApi::CreateOp
*
* \since Version 1.12.
*/
ORT_CLASS_RELEASE(Op);
};
/*
@ -3398,6 +3490,7 @@ ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_CUDA, _In_ OrtSessionOpt
* \param device_id HIP device id, starts from zero.
*/
ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_MIGraphX, _In_ OrtSessionOptions* options, int device_id);
#ifdef __cplusplus
}
#endif

View file

@ -973,6 +973,34 @@ struct CustomOpApi {
void ThrowOnError(OrtStatus* result);
void CreateOpAttr(_In_ const char* name,
_In_ const void* data,
_In_ int len,
_In_ OrtOpAttrType type,
_Outptr_ OrtOpAttr** op_attr);
void ReleaseOpAttr(_In_ OrtOpAttr* op_attr);
void CreateOp(_In_ const OrtKernelInfo* info,
_In_ const char* op_name,
_In_ const char* domain,
_In_ int version,
_In_opt_ const char** type_constraint_names,
_In_opt_ const ONNXTensorElementDataType* type_constraint_values,
_In_opt_ int type_constraint_count,
_In_opt_ const OrtOpAttr* const* attr_values,
_In_opt_ int attr_count,
_Outptr_ OrtOp** ort_op);
void InvokeOp(_In_ const OrtKernelContext* context,
_In_ const OrtOp* ort_op,
_In_ const OrtValue* const* input_values,
_In_ int input_count,
_Inout_ OrtValue* const* output_values,
_In_ int output_count);
void ReleaseOp(_In_ OrtOp* ort_op);
private:
const OrtApi& api_;
};

View file

@ -1206,6 +1206,44 @@ inline void* CustomOpApi::KernelContext_GetGPUComputeStream(const OrtKernelConte
return out;
}
inline void CustomOpApi::CreateOpAttr(_In_ const char* name,
_In_ const void* data,
_In_ int len,
_In_ OrtOpAttrType type,
_Outptr_ OrtOpAttr** op_attr) {
ThrowOnError(api_.CreateOpAttr(name, data, len, type, op_attr));
}
inline void CustomOpApi::ReleaseOpAttr(_Frees_ptr_opt_ OrtOpAttr* op_attr) {
api_.ReleaseOpAttr(op_attr);
}
inline void CustomOpApi::CreateOp(_In_ const OrtKernelInfo* info,
_In_ const char* op_name,
_In_ const char* domain,
_In_ int version,
_In_opt_ const char** type_constraint_names,
_In_opt_ const ONNXTensorElementDataType* type_constraint_values,
_In_opt_ int type_constraint_count,
_In_opt_ const OrtOpAttr* const* attr_values,
_In_opt_ int attr_count,
_Outptr_ OrtOp** ort_op) {
ThrowOnError(api_.CreateOp(info, op_name, domain, version, type_constraint_names, type_constraint_values, type_constraint_count, attr_values, attr_count, ort_op));
}
inline void CustomOpApi::InvokeOp(_In_ const OrtKernelContext* context,
_In_ const OrtOp* ort_op,
_In_ const OrtValue* const* input_values,
_In_ int input_count,
_Inout_ OrtValue* const* output_values,
_In_ int output_count) {
ThrowOnError(api_.InvokeOp(context, ort_op, input_values, input_count, output_values, output_count));
}
inline void CustomOpApi::ReleaseOp(_Frees_ptr_opt_ OrtOp* ort_op) {
api_.ReleaseOp(ort_op);
}
inline SessionOptions& SessionOptions::DisablePerSessionThreads() {
ThrowOnError(GetApi().DisablePerSessionThreads(p_));
return *this;

View file

@ -315,6 +315,35 @@ Status KernelRegistry::TryFindKernel(const Node& node,
return Status(common::ONNXRUNTIME, common::FAIL, "Kernel not found");
}
Status KernelRegistry::TryFindKernel(const std::string& op_name, const std::string& domain, const int& version,
const std::unordered_map<std::string, MLDataType>& type_constraints,
ProviderType exec_provider, const KernelCreateInfo** out) const {
*out = 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) {
*out = &kci; //found match, exit loop
break;
}
} //if
} //for
return *out == nullptr ? Status(common::ONNXRUNTIME, common::FAIL, "Kernel not found") : Status::OK();
}
#endif // !defined(ORT_MINIMAL_BUILD)
bool KernelRegistry::TryFindKernelByHash(HashValue kernel_def_hash, const KernelCreateInfo** out) const {

View file

@ -36,6 +36,9 @@ OpKernelContext::OpKernelContext(_Inout_ IExecutionFrame* frame, _In_ const OpKe
node_output_start_index_ = node_implicit_input_start_index_ + ImplicitInputCount();
}
OpKernelContext::OpKernelContext(concurrency::ThreadPool* threadpool,
const logging::Logger& logger) : threadpool_(threadpool), logger_(&logger) {}
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;

View file

@ -45,7 +45,7 @@ class OpKernelContextInternal : public OpKernelContext {
return session_state_.GetSubgraphSessionState(GetNodeIndex(), attribute_name);
}
const OrtValue* GetInputMLValue(int index) const {
const OrtValue* GetInputMLValue(int index) const override {
return OpKernelContext::GetInputMLValue(index);
}
@ -59,7 +59,7 @@ class OpKernelContextInternal : public OpKernelContext {
}
#endif
OrtValue* OutputMLValue(int index, const TensorShape& shape) {
OrtValue* OutputMLValue(int index, const TensorShape& shape) override {
return OpKernelContext::OutputMLValue(index, shape);
}

View file

@ -896,9 +896,10 @@ void Node::CreateSubgraph(const std::string& attr_name) {
void Node::AddAttributeProto(AttributeProto value) {
utils::SetNodeAttribute(std::move(value), attributes_);
graph_->SetGraphResolveNeeded();
graph_->SetGraphProtoSyncNeeded();
if (graph_) {
graph_->SetGraphResolveNeeded();
graph_->SetGraphProtoSyncNeeded();
}
}
#define ADD_ATTR_SINGLE_IMPL(Type) \

View file

@ -2522,6 +2522,11 @@ static constexpr OrtApi ort_api_1_to_12 = {
&OrtApis::SessionOptionsAppendExecutionProvider_MIGraphX,
// End of Version 11 - DO NOT MODIFY ABOVE (see above text for more information)
&OrtApis::AddExternalInitializers,
&OrtApis::CreateOpAttr,
&OrtApis::ReleaseOpAttr,
&OrtApis::CreateOp,
&OrtApis::InvokeOp,
&OrtApis::ReleaseOp,
};
// Asserts to do a some checks to ensure older Versions of the OrtApi never change (will detect an addition or deletion but not if they cancel out each other)

View file

@ -339,8 +339,40 @@ ORT_API_STATUS_IMPL(UpdateCUDAProviderOptions, _Inout_ OrtCUDAProviderOptionsV2*
size_t num_keys);
ORT_API_STATUS_IMPL(GetCUDAProviderOptionsAsString, _In_ const OrtCUDAProviderOptionsV2* cuda_options, _Inout_ OrtAllocator* allocator, _Outptr_ char** ptr);
ORT_API(void, ReleaseCUDAProviderOptions, _Frees_ptr_opt_ OrtCUDAProviderOptionsV2*);
ORT_API_STATUS_IMPL(AddExternalInitializers, _In_ OrtSessionOptions* options,
_In_reads_(initializers_num) const char* const* initializer_names,
_In_reads_(initializers_num) const OrtValue* const* initializers, size_t initializers_num);
ORT_API_STATUS_IMPL(CreateOpAttr,
_In_ const char* name,
_In_ const void* data,
_In_ int len,
_In_ OrtOpAttrType type,
_Outptr_ OrtOpAttr** op_attr);
ORT_API(void, ReleaseOpAttr, _Frees_ptr_opt_ OrtOpAttr* op_attr);
ORT_API_STATUS_IMPL(CreateOp,
_In_ const OrtKernelInfo* info,
_In_ const char* op_name,
_In_ const char* domain,
_In_ int version,
_In_opt_ const char** type_constraint_names,
_In_opt_ const ONNXTensorElementDataType* type_constraint_values,
_In_opt_ int type_constraint_count,
_In_opt_ const OrtOpAttr* const* attr_values,
_In_opt_ int attr_count,
_Outptr_ OrtOp** ort_op);
ORT_API_STATUS_IMPL(InvokeOp,
_In_ const OrtKernelContext* context,
_In_ const OrtOp* ort_op,
_In_ const OrtValue* const* input_values,
_In_ int input_count,
_Inout_ OrtValue* const* output_values,
_In_ int output_count);
ORT_API(void, ReleaseOp, _Frees_ptr_opt_ OrtOp* op);
} // namespace OrtApis

View file

@ -0,0 +1,407 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/session/inference_session.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/error_code_helper.h"
#include "core/session/ort_apis.h"
#include <unordered_map>
#ifdef ORT_MINIMAL_BUILD
ORT_API_STATUS_IMPL(OrtApis::CreateOpAttr,
_In_ const char*,
_In_ const void*,
_In_ int,
_In_ OrtOpAttrType,
_Out_ OrtOpAttr**) {
API_IMPL_BEGIN
return CreateStatus(ORT_NOT_IMPLEMENTED, "CreateOpAttr is not implemented for minimal build.");
API_IMPL_END
}
ORT_API(void, OrtApis::ReleaseOpAttr, _Frees_ptr_opt_ OrtOpAttr*) {
}
ORT_API_STATUS_IMPL(OrtApis::CreateOp,
_In_ const OrtKernelInfo*,
_In_ const char*,
_In_ const char*,
_In_ int,
_In_ const char**,
_In_ const ONNXTensorElementDataType*,
_In_ int,
_In_ const OrtOpAttr* const*,
_In_ int,
_Out_ OrtOp**) {
API_IMPL_BEGIN
return CreateStatus(ORT_NOT_IMPLEMENTED, "CreateOp is not implemented for minimal build.");
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::InvokeOp,
_In_ const OrtKernelContext*,
_In_ const OrtOp*,
_In_ const OrtValue* const*,
_In_ int,
_Inout_ OrtValue* const*,
_In_ int) {
API_IMPL_BEGIN
return CreateStatus(ORT_NOT_IMPLEMENTED, "InvokeOp is not implemented for minimal build.");
API_IMPL_END
}
ORT_API(void, OrtApis::ReleaseOp, _Frees_ptr_opt_ OrtOp*) {
}
#else
namespace onnxruntime {
namespace standalone {
// For invoking kernels without a graph
class StandAloneKernelContext : public OpKernelContext {
public:
StandAloneKernelContext(const OrtValue* const* input_values,
int input_count,
OrtValue* const* output_values,
int output_count,
AllocatorPtr allocator,
onnxruntime::concurrency::ThreadPool* threadpool,
const logging::Logger& logger) : OpKernelContext(threadpool, logger),
input_values_(input_values),
input_count_(input_count),
output_values_(output_values),
output_count_(output_count),
allocator_(allocator) {}
int NumVariadicInputs(size_t arg_num) const override {
ORT_ENFORCE(arg_num < static_cast<size_t>(input_count_), "invalid arg_num.");
auto ort_value = input_values_[arg_num];
if (ort_value->IsTensor()) {
return static_cast<int>(ort_value->Get<Tensor>().Shape().Size());
} else if (ort_value->IsTensorSequence()) {
return static_cast<int>(ort_value->Get<TensorSeq>().Size());
} else if (ort_value->IsSparseTensor()) {
#ifdef DISABLE_SPARSE_TENSORS
ORT_THROW("sparse tensor is not supported in this build.");
#else
return static_cast<int>(ort_value->Get<SparseTensor>().Values().Shape().Size());
#endif
} else {
return 0;
}
}
MLDataType InputType(int index) const override {
if (index >= input_count_) {
return nullptr;
} else {
return input_values_[index]->Type();
}
}
MLDataType OutputType(int index) const override {
if (index >= output_count_) {
return nullptr;
} else {
return output_values_[index]->Type();
}
}
bool TryGetInferredInputShape(int, TensorShape&) const override {
return false;
}
bool TryGetInferredOutputShape(int, TensorShape&) const override {
return false;
}
int InputCount() const override {
return input_count_;
}
int ImplicitInputCount() const override {
return 0;
}
int OutputCount() const override {
return static_cast<int>(output_count_);
}
Status GetTempSpaceAllocator(AllocatorPtr* output) const override ORT_MUST_USE_RESULT {
*output = allocator_;
return Status::OK();
}
Fence_t InputFence(int index) const override {
if (index >= input_count_) {
return nullptr;
} else {
return input_values_[index]->Fence();
}
}
Fence_t ImplicitInputFence(int) const override {
return nullptr;
}
Fence_t OutputFence(int index) const override {
if (index >= output_count_) {
return nullptr;
} else {
return output_values_[index]->Fence();
}
}
int GetDeviceId() const override {
return 0;
}
void* GetComputeStream() const override {
return nullptr;
}
protected:
const OrtValue* GetInputMLValue(int index) const override {
if (index >= input_count_) {
return nullptr;
} else {
return input_values_[index];
}
}
OrtValue* OutputMLValue(int index, const TensorShape& shape) override {
if (index >= output_count_) {
return nullptr;
}
OrtValue& ort_value = *output_values_[index];
if (!ort_value.IsAllocated()) {
if (ort_value.IsTensor()) {
Tensor::InitOrtValue(ort_value.Type(), shape, allocator_, ort_value);
} else if (ort_value.IsTensorSequence()) {
auto ml_type = ort_value.Type();
auto element_type = ml_type->AsSequenceTensorType()->GetElementType();
auto p_sequence = std::make_unique<TensorSeq>(element_type);
auto ml_tensor_sequence = DataTypeImpl::GetType<TensorSeq>();
ort_value.Init(p_sequence.release(), ml_tensor_sequence, ml_tensor_sequence->GetDeleteFunc());
} else if (ort_value.IsSparseTensor()) {
#ifdef DISABLE_SPARSE_TENSORS
ORT_THROW("sparse tensor is not supported in this build.");
#else
auto ml_type = ort_value.Type();
auto element_type = ml_type->AsSparseTensorType()->GetElementType();
SparseTensor::InitOrtValue(element_type, shape, allocator_, ort_value);
#endif
}
}
return &ort_value;
}
OrtValue* GetOrCreateOutputMLValue(int index) override {
if (index >= output_count_) {
return nullptr;
} else {
return output_values_[index];
}
}
const OrtValue* const* input_values_;
const int input_count_;
OrtValue* const* output_values_;
const int output_count_;
AllocatorPtr allocator_;
}; // StandAloneKernelContext
onnxruntime::Status CreateOpAttr(const char* name, const void* data, int len, OrtOpAttrType type, OrtOpAttr** op_attr) {
auto attr = std::make_unique<ONNX_NAMESPACE::AttributeProto>();
onnxruntime::Status status = onnxruntime::Status::OK();
attr->set_name(std::string{name});
const int* ints = reinterpret_cast<const int*>(data);
const float* floats = reinterpret_cast<const float*>(data);
auto str = reinterpret_cast<const char*>(data);
auto strs = reinterpret_cast<const char* const*>(data);
switch (type) {
case OrtOpAttrType::ORT_OP_ATTR_INT:
attr->set_i(ints[0]);
attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INT);
break;
case OrtOpAttrType::ORT_OP_ATTR_INTS:
for (int j = 0; j < len; ++j) {
attr->add_ints(ints[j]);
}
attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_INTS);
break;
case OrtOpAttrType::ORT_OP_ATTR_FLOAT:
attr->set_f(floats[0]);
attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT);
break;
case OrtOpAttrType::ORT_OP_ATTR_FLOATS:
for (int j = 0; j < len; ++j) {
attr->add_floats(floats[j]);
}
attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_FLOATS);
break;
case OrtOpAttrType::ORT_OP_ATTR_STRING:
attr->set_s(std::string{str});
attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRING);
break;
case OrtOpAttrType::ORT_OP_ATTR_STRINGS:
for (int j = 0; j < len; ++j) {
attr->add_strings(std::string{strs[j]});
}
attr->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType::AttributeProto_AttributeType_STRINGS);
break;
default:
status = Status(common::ONNXRUNTIME, common::FAIL, "Attribute type not supported yet.");
break;
}
*op_attr = reinterpret_cast<OrtOpAttr*>(attr.release());
return status;
}
onnxruntime::Status CreateOp(const OrtKernelInfo* info,
const char* op_name,
const char* domain,
int version,
const char** type_constraint_names,
const ONNXTensorElementDataType* type_constraint_values,
int type_constraint_count,
const OrtOpAttr* const* attr_values,
int attr_count,
OrtOp** op) {
*op = nullptr;
auto kernel_info = reinterpret_cast<const OpKernelInfo*>(info);
auto ep = reinterpret_cast<const IExecutionProvider*>(kernel_info->GetExecutionProvider());
auto kernel_registry = ep->GetKernelRegistry();
const KernelCreateInfo* kernel_create_info{};
std::unordered_map<std::string, MLDataType> 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);
onnxruntime::Node node;
for (int i = 0; i < attr_count; ++i) {
auto attr_proto = reinterpret_cast<const ONNX_NAMESPACE::AttributeProto*>(attr_values[i]);
node.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);
OpKernelInfo instant_kernel_info(node, *kernel_def_builder->Build(), *ep, {}, {}, {});
std::unique_ptr<onnxruntime::OpKernel> op_kernel;
FuncManager func_mgr;
status = kernel_create_info->kernel_create_func(func_mgr, instant_kernel_info, op_kernel);
ORT_RETURN_IF_ERROR(status);
*op = reinterpret_cast<OrtOp*>(op_kernel.release());
return status;
}
onnxruntime::Status InvokeOp(_In_ const OrtKernelContext* context,
_In_ const OrtOp* ort_op,
_In_ const OrtValue* const* input_values,
_In_ int input_count,
_Inout_ OrtValue* const* output_values,
_In_ int output_count) {
auto ctx = reinterpret_cast<const OpKernelContext*>(context);
AllocatorPtr allocator{};
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&allocator));
StandAloneKernelContext standalone_kernel_ctx(input_values,
input_count,
output_values,
output_count,
allocator,
ctx->GetOperatorThreadPool(),
ctx->Logger());
auto kernel = reinterpret_cast<const OpKernel*>(ort_op);
return kernel->Compute(&standalone_kernel_ctx);
}
} // namespace standalone
} // namespace onnxruntime
ORT_API_STATUS_IMPL(OrtApis::CreateOpAttr,
_In_ const char* name,
_In_ const void* data,
_In_ int len,
_In_ OrtOpAttrType type,
_Outptr_ OrtOpAttr** op_attr) {
API_IMPL_BEGIN
auto status = onnxruntime::standalone::CreateOpAttr(name, data, len, type, op_attr);
if (status.IsOK()) {
return nullptr;
} else {
return CreateStatus(static_cast<OrtErrorCode>(status.Code()), status.ErrorMessage().c_str());
}
API_IMPL_END
}
ORT_API(void, OrtApis::ReleaseOpAttr, _Frees_ptr_opt_ OrtOpAttr* op_attr) {
if (op_attr) {
delete reinterpret_cast<ONNX_NAMESPACE::AttributeProto*>(op_attr);
}
}
ORT_API_STATUS_IMPL(OrtApis::CreateOp,
_In_ const OrtKernelInfo* info,
_In_ const char* op_name,
_In_ const char* domain,
_In_ int version,
_In_opt_ const char** type_constraint_names,
_In_opt_ const ONNXTensorElementDataType* type_constraint_values,
_In_opt_ int type_constraint_count,
_In_opt_ const OrtOpAttr* const* attr_values,
_In_opt_ int attr_count,
_Outptr_ OrtOp** ort_op) {
API_IMPL_BEGIN
auto status = onnxruntime::standalone::CreateOp(info,
op_name,
domain,
version,
type_constraint_names,
type_constraint_values,
type_constraint_count,
attr_values,
attr_count,
ort_op);
if (status.IsOK()) {
return nullptr;
} else {
return CreateStatus(static_cast<OrtErrorCode>(status.Code()), status.ErrorMessage().c_str());
}
API_IMPL_END
}
ORT_API_STATUS_IMPL(OrtApis::InvokeOp,
_In_ const OrtKernelContext* context,
_In_ const OrtOp* ort_op,
_In_ const OrtValue* const* input_values,
_In_ int input_count,
_Inout_ OrtValue* const* output_values,
_In_ int output_count) {
API_IMPL_BEGIN
auto status = onnxruntime::standalone::InvokeOp(context, ort_op, input_values, input_count, output_values, output_count);
if (status.IsOK()) {
return nullptr;
} else {
return CreateStatus(static_cast<OrtErrorCode>(status.Code()), status.ErrorMessage().c_str());
}
API_IMPL_END
}
ORT_API(void, OrtApis::ReleaseOp, _Frees_ptr_opt_ OrtOp* op) {
if (op) {
delete reinterpret_cast<onnxruntime::OpKernel*>(op);
}
}
#endif

View file

@ -206,3 +206,220 @@ void SliceCustomOpKernel::Compute(OrtKernelContext* context) {
ORT_THROW("Unsupported input type");
}
}
InstantCustomKernel::InstantCustomKernel(Ort::CustomOpApi ort, const OrtKernelInfo* info, void*) : ort_(ort) {
const char* add_type_constrait_names[1] = {"T"};
int add_type_constrait_values[1] = {1};
ort.CreateOp(info, "Add", "", 14,
(const char**)add_type_constrait_names,
(const ONNXTensorElementDataType*)add_type_constrait_values,
1, nullptr, 0, &op_add);
ORT_ENFORCE(op_add, "op_add not initialzied");
InitTopK(ort, info);
ORT_ENFORCE(op_topk, "op_add not initialzied");
InitGru(ort, info);
ORT_ENFORCE(op_gru, "op_add not initialzied");
}
void InstantCustomKernel::InitTopK(Ort::CustomOpApi ort, const OrtKernelInfo* info) {
const char* type_constrait_names[2] = {"T", "I"};
int type_constrait_values[2] = {1, 7};
int axis_value = -1;
OrtOpAttr* axis{};
ort.CreateOpAttr("axis", &axis_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT, &axis);
int largest_value = 0; // return in ascending order
OrtOpAttr* largest{};
ort.CreateOpAttr("largest", &largest_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT, &largest);
int sorted_value = 1;
OrtOpAttr* sorted{};
ort.CreateOpAttr("sorted", &sorted_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT, &sorted);
if (!axis || !largest || !sorted) {
ORT_THROW("Failed to create attributes for topk.");
}
OrtOpAttr* top_attrs[3] = {axis, largest, sorted};
ort.CreateOp(info, "TopK", "", 14,
(const char**)type_constrait_names,
(const ONNXTensorElementDataType*)type_constrait_values,
2, top_attrs, 3, &op_topk);
ort.ReleaseOpAttr(axis);
ort.ReleaseOpAttr(largest);
ort.ReleaseOpAttr(sorted);
}
void InstantCustomKernel::InvokeTopK(OrtKernelContext* context) {
auto mem_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeCPU);
float raw_x[10] = {6., 3., 4., 8., 7., 1., 9., 0., 5., 2.};
int64_t raw_x_shape[1] = {10};
auto topk_x = Ort::Value::CreateTensor(mem_info, raw_x, 10, raw_x_shape, 1);
int64_t raw_k[1] = {2};
int64_t raw_k_shape[1] = {1};
auto topk_k = Ort::Value::CreateTensor(mem_info, raw_k, 1, raw_k_shape, 1);
float raw_values[2] = {};
int64_t raw_values_shape[1] = {2};
auto topk_values = Ort::Value::CreateTensor(mem_info, raw_values, 2, raw_values_shape, 1);
int64_t raw_indices[2] = {};
int64_t raw_indices_shape[1] = {2};
auto topk_indices = Ort::Value::CreateTensor(mem_info, raw_indices, 2, raw_indices_shape, 1);
const OrtValue* topk_inputs[2] = {(OrtValue*)topk_x, (OrtValue*)topk_k};
OrtValue* topk_outputs[2] = {(OrtValue*)topk_values, (OrtValue*)topk_indices};
ort_.InvokeOp(context, op_topk, topk_inputs, 2, topk_outputs, 2);
if (std::abs(raw_values[0] - 0.) > 1e-6 || std::abs(raw_values[1] - 1.) > 1e-6) {
ORT_THROW("topk instant operator returns wrong values");
}
if (raw_indices[0] != 7 || raw_indices[1] != 5) {
ORT_THROW("topk instant operator returns wrong indices");
}
}
void InstantCustomKernel::InitGru(Ort::CustomOpApi ort, const OrtKernelInfo* info) {
const char* type_constrait_names[2] = {"T", "T1"};
int type_constrait_values[2] = {1, 6};
const char* activition_names[4] = {"LeakyRelu", "Tanh", "Sigmoid", "ScaledTanh"};
OrtOpAttr* activations{};
ort.CreateOpAttr("activations", activition_names, 4, OrtOpAttrType::ORT_OP_ATTR_STRINGS, &activations);
float alphas[2] = {0.5f, 2.f};
OrtOpAttr* activation_alpha{};
ort.CreateOpAttr("activation_alpha ", alphas, 2, OrtOpAttrType::ORT_OP_ATTR_FLOATS, &activation_alpha);
float betas[1] = {2.f};
OrtOpAttr* activation_beta{};
ort.CreateOpAttr("activation_beta ", betas, 1, OrtOpAttrType::ORT_OP_ATTR_FLOATS, &activation_beta);
const char* direction_string = "bidirectional";
OrtOpAttr* direction{};
ort.CreateOpAttr("direction", direction_string, 1, OrtOpAttrType::ORT_OP_ATTR_STRING, &direction);
int linear_before_reset_value = 0;
OrtOpAttr* linear_before_reset{};
ort.CreateOpAttr("linear_before_reset", &linear_before_reset_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT, &linear_before_reset);
int hidden_size_value = 2;
OrtOpAttr* hidden_size{};
ort.CreateOpAttr("hidden_size", &hidden_size_value, 1, OrtOpAttrType::ORT_OP_ATTR_INT, &hidden_size);
if (!activations || !activation_alpha || !activation_beta || !direction || !linear_before_reset || !hidden_size) {
ORT_THROW("failed to create attributes for gru.");
}
OrtOpAttr* gru_attrs[6] = {activations, activation_alpha, activation_beta, direction, linear_before_reset, hidden_size};
ort.CreateOp(info, "GRU", "", 14,
(const char**)type_constrait_names,
(const ONNXTensorElementDataType*)type_constrait_values,
2, gru_attrs, 6, &op_gru);
ort.ReleaseOpAttr(activations);
ort.ReleaseOpAttr(activation_alpha);
ort.ReleaseOpAttr(activation_beta);
ort.ReleaseOpAttr(direction);
ort.ReleaseOpAttr(linear_before_reset);
ort.ReleaseOpAttr(hidden_size);
}
void InstantCustomKernel::InvokeGru(OrtKernelContext* context) {
auto mem_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeCPU);
float raw_x[2] = {1.0f, 2.0f};
int64_t raw_x_shape[3] = {1, 1, 2};
auto X = Ort::Value::CreateTensor(mem_info, raw_x, 2, raw_x_shape, 3);
float raw_w[24] = {
-0.494659f, 0.0453352f, -0.487793f, 0.417264f, // Wz
-0.0091708f, -0.255364f, -0.106952f, -0.266717f, // Wr
-0.0888852f, -0.428709f, -0.283349f, 0.208792f, // Wh
-0.494659f, 0.0453352f, -0.487793f, 0.417264f, // WBz
-0.0091708f, -0.255364f, -0.106952f, -0.266717f, // WBr
-0.0888852f, -0.428709f, -0.283349f, 0.208792f // WBh
};
int64_t raw_w_shape[3] = {2, 6, 2};
auto W = Ort::Value::CreateTensor(mem_info, raw_w, 24, raw_w_shape, 3);
float raw_r[24] = {
0.146626f, -0.0620289f, -0.0815302f, 0.100482f, // Rz
-0.228172f, 0.405972f, 0.31576f, 0.281487f, // Rr
-0.394864f, 0.42111f, -0.386624f, -0.390225f, // Rh
0.146626f, -0.0620289f, -0.0815302f, 0.100482f, // RBz
-0.228172f, 0.405972f, 0.31576f, 0.281487f, // RBr
-0.394864f, 0.42111f, -0.386624f, -0.390225f}; // RBh
int64_t raw_r_shape[3] = {2, 6, 2};
auto R = Ort::Value::CreateTensor(mem_info, raw_r, 24, raw_r_shape, 3);
float raw_b[24] = {
0.381619f, 0.0323954f, // Wbz
-0.258721f, 0.45056f, // Wbr
-0.250755f, 0.0967895f, // Wbh
0.0f, 0.0f, // Rbz
-0.0f, 0.0f, // Rbr
-0.0f, 0.0f, // Rbh
0.381619f, 0.0323954f, // WBbz
-0.258721f, 0.45056f, // WBbr
-0.250755f, 0.0967895f, // WBbh
0.0f, 0.0f, // RBbz
-0.0f, 0.0f, // RBbr
-0.0f, 0.0f}; // RBbh
int64_t raw_b_shape[2] = {2, 12};
auto B = Ort::Value::CreateTensor(mem_info, raw_b, 24, raw_b_shape, 2);
int32_t raw_seq_lens = 1;
int64_t seq_lens_shape[1] = {1};
auto sequence_lens = Ort::Value::CreateTensor(mem_info, &raw_seq_lens, 1, seq_lens_shape, 1);
std::vector<float> raw_initial_h(4, 0.25f);
int64_t initial_h_shape[3] = {2, 1, 2};
auto initial_h = Ort::Value::CreateTensor(mem_info, raw_initial_h.data(), 4, initial_h_shape, 3);
float raw_y[4] = {};
int64_t raw_y_shape[64] = {1, 2, 1, 2};
auto Y = Ort::Value::CreateTensor(mem_info, raw_y, 4, raw_y_shape, 4);
float raw_yh[4] = {};
int64_t raw_yh_shape[64] = {2, 1, 2};
auto YH = Ort::Value::CreateTensor(mem_info, raw_yh, 4, raw_yh_shape, 3);
const OrtValue* inputs[6] = {(OrtValue*)X, (OrtValue*)W, (OrtValue*)R, (OrtValue*)B, (OrtValue*)sequence_lens, (OrtValue*)initial_h};
OrtValue* outputs[2] = {(OrtValue*)Y, (OrtValue*)YH};
const float expected_y[4] = {-0.832559f,
0.236267f,
0.124924f,
0.148701f};
ort_.InvokeOp(context, op_gru, inputs, 6, outputs, 2);
for (int i = 0; i < 4; ++i) {
if (std::abs(raw_y[i] - expected_y[i]) > 1e-6) {
ORT_THROW("GRU op give unexpected output.");
}
}
}
void InstantCustomKernel::Compute(OrtKernelContext* context) {
const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0);
const OrtValue* input_Y = ort_.KernelContext_GetInput(context, 1);
OrtTensorDimensions dimensions(ort_, input_X);
OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions.data(), dimensions.size());
const OrtValue* inputs[2] = {input_X, input_Y};
OrtValue* outputs[1] = {output};
ort_.InvokeOp(context, op_add, inputs, 2, outputs, 1);
InvokeTopK(context);
InvokeGru(context);
}
InstantCustomKernel::~InstantCustomKernel() {
ort_.ReleaseOp(op_add);
ort_.ReleaseOp(op_topk);
ort_.ReleaseOp(op_gru);
}

View file

@ -207,3 +207,42 @@ struct SliceCustomOp : Ort::CustomOpBase<SliceCustomOp, SliceCustomOpKernel> {
private:
const char* provider_;
};
struct InstantCustomKernel {
InstantCustomKernel(Ort::CustomOpApi ort, const OrtKernelInfo* info, void*);
~InstantCustomKernel();
void Compute(OrtKernelContext* context);
private:
void InitTopK(Ort::CustomOpApi ort, const OrtKernelInfo* info);
void InvokeTopK(OrtKernelContext* context);
void InitGru(Ort::CustomOpApi ort, const OrtKernelInfo* info);
void InvokeGru(OrtKernelContext* context);
Ort::CustomOpApi ort_;
OrtOp* op_add{};
OrtOp* op_topk{};
OrtOp* op_gru{};
};
struct InstantCustomOp : Ort::CustomOpBase<InstantCustomOp, InstantCustomKernel> {
explicit InstantCustomOp(const char* provider, void* compute_stream) : provider_(provider), compute_stream_(compute_stream) {}
void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new InstantCustomKernel(api, info, compute_stream_); };
const char* GetName() const { return "Foo"; };
const char* GetExecutionProviderType() const { return provider_; };
size_t GetInputTypeCount() const { return 2; };
ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
};
size_t GetOutputTypeCount() const { return 1; };
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
private:
const char* provider_;
void* compute_stream_;
};

View file

@ -405,6 +405,26 @@ TEST(CApiTest, custom_op_handler) {
#endif
}
#if !defined(ORT_MINIMAL_BUILD)
TEST(CApiTest, instant_op_handler) {
std::vector<Input> 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<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};
InstantCustomOp instant_op{onnxruntime::kCpuExecutionProvider, nullptr};
Ort::CustomOpDomain custom_op_domain("");
custom_op_domain.Add(&instant_op);
TestInference<float>(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 0,
custom_op_domain, nullptr);
}
#endif
#ifdef ENABLE_EXTENSION_CUSTOM_OPS
// test enabled ort-customops negpos
TEST(CApiTest, test_enable_ort_customops_negpos) {