saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Single-schema-multi-kernel (#15184)

Browse source 
The PR is to allow custom op of different input types to have same op
name in a graph.
The idea to go over all ops of same name and merge their input/output
types into a type-inference function.
With the enhancement, custom op node inside a graph can have same
op-type given that the input/output types are different.

---------

Co-authored-by: Randy Shuai <rashuai@microsoft.com>
This commit is contained in:
RandySheriffH 2023-04-27 13:39:59 -07:00 committed by GitHub
parent d3d232b047
commit 9773e76c44
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 582 additions and 107 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

398
onnxruntime/core/session/custom_ops.cc
View file

@ -414,17 +414,274 @@ struct CustomOpKernel : OpKernel {
void* op_kernel_;
};
#if !defined(ORT_MINIMAL_BUILD)
KernelCreateInfo CreateKernelCreateInfo(const std::string& domain, const OrtCustomOp* op) {
const size_t input_count = op->GetInputTypeCount(op);
const size_t output_count = op->GetOutputTypeCount(op);
KernelDefBuilder def_builder;
def_builder.SetName(op->GetName(op))
.SetDomain(domain)
.SinceVersion(1);
// GetInputMemoryType was introduced in ver 13. This check allows custom ops compiled using older versions
// to work with newer versions (> 12) of the ORT binary.
if (op->version > 12) {
for (size_t i = 0; i < input_count; i++) {
def_builder.InputMemoryType(op->GetInputMemoryType(op, i), i);
}
}
for (size_t i = 0; i < input_count; i++) {
const auto input_type = op->GetInputType(op, i);
const auto input_name = "Input" + std::to_string(i);
if (input_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) {
def_builder.TypeConstraint(input_name, DataTypeImpl::AllTensorTypes());
} else {
def_builder.TypeConstraint(input_name, DataTypeImpl::TensorTypeFromONNXEnum(static_cast<int>(input_type))->AsTensorType());
}
}
for (size_t i = 0; i < output_count; i++) {
const auto output_type = op->GetOutputType(op, i);
const auto output_name = "Output" + std::to_string(i);
if (output_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) {
def_builder.TypeConstraint(output_name, DataTypeImpl::AllTensorTypes());
} else {
def_builder.TypeConstraint(output_name, DataTypeImpl::TensorTypeFromONNXEnum(static_cast<int>(output_type))->AsTensorType());
}
}
if (const char* provider_type = op->GetExecutionProviderType(op)) {
def_builder.Provider(provider_type);
} else {
def_builder.Provider(onnxruntime::kCpuExecutionProvider);
}
KernelCreateFn kernel_create_fn = [op](FuncManager&, const OpKernelInfo& info, std::unique_ptr<OpKernel>& out) -> Status {
out = std::make_unique<CustomOpKernel>(info, *op);
return Status::OK();
};
return KernelCreateInfo(def_builder.Build(), kernel_create_fn);
}
ONNX_NAMESPACE::OpSchema CreateSchema(const std::string& domain, const OrtCustomOp* op) {
constexpr uint32_t min_ort_version_with_optional_io_support = 8;
constexpr uint32_t min_ort_version_with_variadic_io_support = 14;
const size_t input_count = op->GetInputTypeCount(op);
const size_t output_count = op->GetOutputTypeCount(op);
int undefined = 0;
ONNX_NAMESPACE::OpSchema schema(op->GetName(op), "custom op registered at runtime", 0);
for (size_t i = 0; i < input_count; i++) {
onnx::OpSchema::FormalParameterOption option = onnx::OpSchema::FormalParameterOption::Single;
bool is_homogeneous = true;
int min_arity = 1;
// The OrtCustomOp interface did not support the methods to query input/output characteristics before
// ORT API version 8. So, query the relevant methods ONLY from API version 8 onwards.
if (op->version >= min_ort_version_with_optional_io_support) {
const auto characteristic = op->GetInputCharacteristic(op, i);
// Support for optional and variadic inputs/output was added in versions 8 and 14, respectively.
if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) {
option = onnx::OpSchema::FormalParameterOption::Optional;
} else if ((op->version >= min_ort_version_with_variadic_io_support) &&
(characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC)) {
ORT_ENFORCE(i == input_count - 1, "Only the last input to a custom op may be marked variadic.");
option = onnx::OpSchema::FormalParameterOption::Variadic;
min_arity = op->GetVariadicInputMinArity(op);
is_homogeneous = static_cast<bool>(op->GetVariadicInputHomogeneity(op));
}
}
const auto type = op->GetInputType(op, i);
if (type == ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) {
undefined++;
}
std::string input_name = "Input" + std::to_string(i);
schema.Input(i, input_name, "", input_name, option, is_homogeneous, min_arity);
// support all types as input here in schema, and handle the type inference in TypeShapeInference func
schema.TypeConstraint(input_name, DataTypeImpl::ToString(DataTypeImpl::AllTensorTypes()), "all types");
}
for (size_t i = 0; i < output_count; i++) {
onnx::OpSchema::FormalParameterOption option = onnx::OpSchema::FormalParameterOption::Single;
bool is_homogeneous = true;
int min_arity = 1;
// The OrtCustomOp interface did not support the methods to query input/output characteristics before
// ORT API version 8. So, query the relevant methods ONLY from API version 8 onwards.
if (op->version >= min_ort_version_with_optional_io_support) {
const auto characteristic = op->GetOutputCharacteristic(op, i);
// Support for optional and variadic inputs/output was added in versions 8 and 14, respectively.
if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) {
option = onnx::OpSchema::FormalParameterOption::Optional;
} else if ((op->version >= min_ort_version_with_variadic_io_support) &&
(characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC)) {
ORT_ENFORCE(i == output_count - 1, "Only the last output to a custom op may be marked variadic.");
option = onnx::OpSchema::FormalParameterOption::Variadic;
min_arity = op->GetVariadicOutputMinArity(op);
is_homogeneous = static_cast<bool>(op->GetVariadicOutputHomogeneity(op));
}
}
const auto type = op->GetOutputType(op, i);
if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) {
if (op->GetOutputCharacteristic(op, i) == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED) {
ORT_ENFORCE(1 == undefined,
"There must be one (and only one) dynamic typed input to the custom op. "
"Its type info at runtime will be used to infer the type info of this dynamic typed output "
"which is required for the success of the model loading step. "
"More than one dynamic typed inputs are currently not supported as differing types at runtime means the output type "
"cannot be inferred without which model loading cannot proceed.");
}
}
std::string output_name = "Output" + std::to_string(i);
schema.Output(i, output_name, "", output_name, option, is_homogeneous, min_arity);
// support all types as input here in schema, and handle the type inference in TypeShapeInference func
schema.TypeConstraint(output_name, DataTypeImpl::ToString(DataTypeImpl::AllTensorTypes()), "all types");
}
schema.SetDomain(domain);
schema.SinceVersion(1);
schema.AllowUncheckedAttributes();
return schema;
}
Status IsCompatible(const ONNX_NAMESPACE::OpSchema& schema, const OrtCustomOp* op) {
const size_t input_count = op->GetInputTypeCount(op);
const size_t output_count = op->GetOutputTypeCount(op);
// check inputs
const auto& input_parameters = schema.inputs();
ORT_RETURN_IF_NOT(input_parameters.size() == input_count, "input count does not match");
for (size_t i = 0; i < input_parameters.size(); ++i) {
const auto characteristic = op->GetInputCharacteristic(op, i);
const auto& formal_parameter = input_parameters[i];
if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) {
ORT_RETURN_IF_NOT(formal_parameter.GetOption() == onnx::OpSchema::FormalParameterOption::Optional,
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" input to be of optional type");
} else if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC) {
ORT_RETURN_IF_NOT(formal_parameter.GetOption() == onnx::OpSchema::FormalParameterOption::Variadic,
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" input to be of variadic type");
} else {
ORT_RETURN_IF_NOT(formal_parameter.GetOption() == onnx::OpSchema::FormalParameterOption::Single,
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" input to be of single type");
}
ORT_RETURN_IF_NOT(formal_parameter.GetIsHomogeneous() == (op->GetVariadicOutputHomogeneity(op) != 0),
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" input to keep same homogeneity");
ORT_RETURN_IF_NOT(formal_parameter.GetMinArity() == op->GetVariadicInputMinArity(op),
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" input to keep same arity");
}
// check outputs
const auto& output_parameters = schema.outputs();
ORT_RETURN_IF_NOT(output_parameters.size() == output_count, "output count does not match");
for (size_t i = 0; i < output_parameters.size(); ++i) {
const auto characteristic = op->GetOutputCharacteristic(op, i);
const auto& formal_parameter = output_parameters[i];
if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) {
ORT_RETURN_IF_NOT(formal_parameter.GetOption() == onnx::OpSchema::FormalParameterOption::Optional,
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" output to be of optional type");
} else if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC) {
ORT_RETURN_IF_NOT(formal_parameter.GetOption() == onnx::OpSchema::FormalParameterOption::Variadic,
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" output to be of variadic type");
} else {
ORT_RETURN_IF_NOT(formal_parameter.GetOption() == onnx::OpSchema::FormalParameterOption::Single,
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" output to be of single type");
}
ORT_RETURN_IF_NOT(formal_parameter.GetIsHomogeneous() == (op->GetVariadicOutputHomogeneity(op) != 0),
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" output to keep same homogeneity");
ORT_RETURN_IF_NOT(formal_parameter.GetMinArity() == op->GetVariadicInputMinArity(op),
"custom op schemas mismatch, expecting ", i + 1,
i == 0 ? "st" : (i == 1 ? "nd" : "th"),
" output to keep same arity");
}
return Status::OK();
}
void InferOutputTypes(const InlinedVector<const KernelDef*>& kernel_defs,
ONNX_NAMESPACE::InferenceContext& infer_ctx) {
for (const auto& kernel_def : kernel_defs) {
const auto& type_constraints = kernel_def->TypeConstraints();
auto num_inputs = infer_ctx.getNumInputs();
bool matched = true;
ONNXTensorElementDataType undef = ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
// first, make sure there is a constraint for every input
for (size_t i = 0; i < num_inputs && matched; ++i) {
auto input_name = "Input" + std::to_string(i);
auto input_type = infer_ctx.getInputType(i);
if (input_type) {
auto elem_type = static_cast<ONNXTensorElementDataType>(input_type->tensor_type().elem_type());
auto tc_iter = type_constraints.find(input_name);
if (tc_iter != type_constraints.end()) {
if (tc_iter->second.size() > 1) {
undef = elem_type;
} else if (tc_iter->second.size() != 1 || tc_iter->second[0] != DataTypeImpl::TensorTypeFromONNXEnum(elem_type)) {
matched = false;
}
} else {
matched = false;
}
} else {
matched = false;
}
} // for
// next, ensure that there is a constraint for every output
auto num_outputs = infer_ctx.getNumOutputs();
for (size_t i = 0; i < num_outputs && matched; i++) {
auto output_name = "Output" + std::to_string(i);
auto tc_iter = type_constraints.find(output_name);
if (tc_iter == type_constraints.end() || tc_iter->second.size() < 1) {
matched = false;
}
}
if (matched) {
for (size_t i = 0; i < num_outputs; i++) {
auto output_name = "Output" + std::to_string(i);
auto output_type = infer_ctx.getOutputType(i);
auto tc_iter = type_constraints.find(output_name);
if (tc_iter->second.size() > 1) {
output_type->mutable_tensor_type()->set_elem_type(undef);
} else {
output_type->mutable_tensor_type()->set_elem_type(tc_iter->second[0]->GetTypeProto()->tensor_type().elem_type());
}
}
break;
}
}
}
#endif
common::Status CreateCustomRegistry(gsl::span<OrtCustomOpDomain* const> op_domains,
std::shared_ptr<CustomRegistry>& output) {
output = std::make_shared<CustomRegistry>();
for (const auto& domain : op_domains) {
// Create an OpSchema for each op and register them
// Container to hold type template parameters
std::unordered_map<const OrtCustomOp*, std::vector<std::string>> type_constraint_ids;
#if !defined(ORT_MINIMAL_BUILD)
std::unordered_map<std::string, ONNX_NAMESPACE::OpSchema> schema_map;
std::unordered_map<std::string, InlinedVector<const KernelDef*>> kernel_def_map;
// Domain is not empty - add it to the DomainToVersion ONNX map
// If domain is empty, it is assumed to be part of the ONNX domain
if (!domain->domain_.empty()) {
@ -438,120 +695,48 @@ common::Status CreateCustomRegistry(gsl::span<OrtCustomOpDomain* const> op_domai
}
}
constexpr uint32_t min_ort_version_with_optional_io_support = 8;
constexpr uint32_t min_ort_version_with_variadic_io_support = 14;
std::vector<ONNX_NAMESPACE::OpSchema> schemas_list;
for (const auto* op : domain->custom_ops_) {
ONNX_NAMESPACE::OpSchema schema(op->GetName(op), "custom op registered at runtime", 0);
size_t type_id_counter = 0;
const size_t input_count = op->GetInputTypeCount(op);
for (size_t i = 0; i < input_count; i++) {
onnx::OpSchema::FormalParameterOption option = onnx::OpSchema::FormalParameterOption::Single;
bool is_homogeneous = true;
int min_arity = 1;
// The OrtCustomOp interface did not support the methods to query input/output characteristics before
// ORT API version 8. So, query the relevant methods ONLY from API version 8 onwards.
if (op->version >= min_ort_version_with_optional_io_support) {
const auto characteristic = op->GetInputCharacteristic(op, i);
// Support for optional and variadic inputs/output was added in versions 8 and 14, respectively.
if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) {
option = onnx::OpSchema::FormalParameterOption::Optional;
} else if ((op->version >= min_ort_version_with_variadic_io_support) &&
(characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC)) {
ORT_ENFORCE(i == input_count - 1, "Only the last input to a custom op may be marked variadic.");
option = onnx::OpSchema::FormalParameterOption::Variadic;
min_arity = op->GetVariadicInputMinArity(op);
is_homogeneous = static_cast<bool>(op->GetVariadicInputHomogeneity(op));
}
}
const auto type = op->GetInputType(op, i);
if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) { // Dynamic typed input
schema.Input(i, "Input" + std::to_string(i), "", "T" + std::to_string(type_id_counter), option,
is_homogeneous, min_arity);
schema.TypeConstraint("T" + std::to_string(type_id_counter), DataTypeImpl::ToString(DataTypeImpl::AllTensorTypes()), "all types");
type_constraint_ids[op].push_back("T" + std::to_string(type_id_counter++));
} else {
schema.Input(i, "Input" + std::to_string(i), "",
DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type)), option,
is_homogeneous, min_arity);
}
// define kernel
auto kernel_create_info = CreateKernelCreateInfo(domain->domain_, op);
kernel_def_map[op->GetName(op)].push_back(kernel_create_info.kernel_def.get());
ORT_RETURN_IF_ERROR(output->RegisterCustomKernel(kernel_create_info));
// define schema
auto schema_map_iter = schema_map.find(op->GetName(op));
if (schema_map_iter == schema_map.end()) {
auto schema = CreateSchema(domain->domain_, op);
schema_map.emplace(schema.Name(), schema);
} else {
ORT_RETURN_IF_ERROR(IsCompatible(schema_map_iter->second, op));
}
const size_t output_count = op->GetOutputTypeCount(op);
for (size_t i = 0; i < output_count; i++) {
onnx::OpSchema::FormalParameterOption option = onnx::OpSchema::FormalParameterOption::Single;
bool is_homogeneous = true;
int min_arity = 1;
// The OrtCustomOp interface did not support the methods to query input/output characteristics before
// ORT API version 8. So, query the relevant methods ONLY from API version 8 onwards.
if (op->version >= min_ort_version_with_optional_io_support) {
const auto characteristic = op->GetOutputCharacteristic(op, i);
// Support for optional and variadic inputs/output was added in versions 8 and 14, respectively.
if (characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) {
option = onnx::OpSchema::FormalParameterOption::Optional;
} else if ((op->version >= min_ort_version_with_variadic_io_support) &&
(characteristic == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC)) {
ORT_ENFORCE(i == output_count - 1, "Only the last output to a custom op may be marked variadic.");
option = onnx::OpSchema::FormalParameterOption::Variadic;
min_arity = op->GetVariadicOutputMinArity(op);
is_homogeneous = static_cast<bool>(op->GetVariadicOutputHomogeneity(op));
}
}
const auto type = op->GetOutputType(op, i);
if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) { // Dynamic typed output
if (op->GetOutputCharacteristic(op, i) == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED) {
ORT_ENFORCE(type_id_counter == 1,
"There must be one (and only one) dynamic typed input to the custom op. "
"Its type info at runtime will be used to infer the type info of this dynamic typed output "
"which is required for the success of the model loading step. "
"More than one dynamic typed inputs are currently not supported as differing types at runtime means the output type "
"cannot be inferred without which model loading cannot proceed.");
}
schema.Output(i, "Output" + std::to_string(i), "", "T0", option, is_homogeneous, min_arity);
} else {
schema.Output(i, "Output" + std::to_string(i), "",
DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type)), option,
is_homogeneous, min_arity);
}
}
schema.SetDomain(domain->domain_);
schema.SinceVersion(1);
schema.AllowUncheckedAttributes();
schemas_list.push_back(schema);
}
ORT_RETURN_IF_ERROR(output->RegisterOpSet(schemas_list,
std::vector<ONNX_NAMESPACE::OpSchema> schemas;
for (auto schema_iter : schema_map) {
schemas.push_back(schema_iter.second);
InlinedVector<const KernelDef*> kernel_defs = std::move(kernel_def_map[schema_iter.first]);
ONNX_NAMESPACE::InferenceFunction infer_fn = [kernel_defs](ONNX_NAMESPACE::InferenceContext& infer_ctx) {
InferOutputTypes(kernel_defs, infer_ctx);
};
schemas.back().TypeAndShapeInferenceFunction(infer_fn);
}
ORT_RETURN_IF_ERROR(output->RegisterOpSet(schemas,
domain->domain_,
1 /* baseline opset version */,
1000 /* opset version */));
#else
// For a minimal build, we may not need any of the ONNX schema stuff but we still need to track
// the type template parameters to be used during the kernel def building step below
for (const auto* op : domain->custom_ops_) {
size_t type_id_counter = 0;
auto input_count = op->GetInputTypeCount(op);
size_t undefined = 0;
size_t input_count = op->GetInputTypeCount(op);
for (size_t i = 0; i < input_count; i++) {
auto type = op->GetInputType(op, i);
if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) { // Dynamic typed input
type_constraint_ids[op].push_back("T" + std::to_string(type_id_counter++));
if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) {
undefined++;
}
}
}
#endif
// create the KernelDef for each op and register it
for (const auto* op : domain->custom_ops_) {
KernelDefBuilder def_builder;
def_builder.SetName(op->GetName(op))
.SetDomain(domain->domain_)
@ -566,8 +751,8 @@ common::Status CreateCustomRegistry(gsl::span<OrtCustomOpDomain* const> op_domai
}
}
for (auto& id : type_constraint_ids[op]) {
def_builder.TypeConstraint(id, DataTypeImpl::AllTensorTypes());
for (size_t i = 0; i < undefined; i++) {
def_builder.TypeConstraint("T" + std::to_string(i), DataTypeImpl::AllTensorTypes());
}
if (const char* provider_type = op->GetExecutionProviderType(op)) {
@ -584,7 +769,8 @@ common::Status CreateCustomRegistry(gsl::span<OrtCustomOpDomain* const> op_domai
KernelCreateInfo create_info(def_builder.Build(), kernel_create_fn);
ORT_RETURN_IF_ERROR(output->RegisterCustomKernel(create_info));
}
}
#endif
} // for each domain
return Status::OK();
}

88
onnxruntime/test/shared_lib/custom_op_utils.h
View file

@ -377,3 +377,91 @@ struct StandaloneCustomOp : Ort::CustomOpBase<StandaloneCustomOp, StandaloneCust
private:
const char* provider_;
};
/////////////// structures to test multi-kernls-single-schema ///////////////
struct MulTopKernelFloat {
MulTopKernelFloat(const OrtKernelInfo*){};
~MulTopKernelFloat() = default;
void Compute(OrtKernelContext*){};
};
struct MulTopOpFloat : Ort::CustomOpBase<MulTopOpFloat, MulTopKernelFloat> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelFloat(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }
size_t GetOutputTypeCount() const { return 1; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }
};
struct MulTopKernelInt32 {
MulTopKernelInt32(const OrtKernelInfo*){};
~MulTopKernelInt32() = default;
void Compute(OrtKernelContext*){};
};
struct MulTopOpInt32 : Ort::CustomOpBase<MulTopOpInt32, MulTopKernelInt32> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelInt32(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; }
size_t GetOutputTypeCount() const { return 1; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; }
};
struct MulTopKernelDouble {
MulTopKernelDouble(const OrtKernelInfo*){};
~MulTopKernelDouble() = default;
void Compute(OrtKernelContext*){};
};
// MulTopOpDouble and MulTopOpFloat has input count mismatch
struct MulTopOpDouble : Ort::CustomOpBase<MulTopOpDouble, MulTopKernelDouble> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelDouble(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 2; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; }
size_t GetOutputTypeCount() const { return 1; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; }
};
struct MulTopKernelInt16 {
MulTopKernelInt16(const OrtKernelInfo*){};
~MulTopKernelInt16() = default;
void Compute(OrtKernelContext*){};
};
// MulTopOpInt16 and MulTopOpFloat has output count mismatch
struct MulTopOpInt16 : Ort::CustomOpBase<MulTopOpInt16, MulTopKernelInt16> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelInt16(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; }
size_t GetOutputTypeCount() const { return 2; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; }
};
// MulTopKernelFloat16 and MulTopOpFloat has input characteristic mismatch
struct MulTopKernelFloat16 {
MulTopKernelFloat16(const OrtKernelInfo*){};
~MulTopKernelFloat16() = default;
void Compute(OrtKernelContext*){};
};
struct MulTopOpFloat16 : Ort::CustomOpBase<MulTopOpFloat16, MulTopKernelFloat16> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelFloat16(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; }
size_t GetOutputTypeCount() const { return 1; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; }
OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t) const {
return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL;
}
};

109
onnxruntime/test/shared_lib/test_inference.cc
View file

@ -187,6 +187,7 @@ static constexpr PATH_TYPE VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI = TSTR("tes
static constexpr PATH_TYPE VARIADIC_UNDEF_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI = TSTR(
"testdata/custom_op_variadic_undef_io.onnx");
static constexpr PATH_TYPE CUSTOM_OP_MODEL_WITH_ATTRIBUTES_URI = TSTR("testdata/foo_bar_3.onnx");
static constexpr PATH_TYPE CUSTOM_OP_SINGLE_SCHEMA_MULTI_KERNEL = TSTR("testdata/custom_op_single_schema_multi_kernel.onnx");
#if !defined(DISABLE_SPARSE_TENSORS)
static constexpr PATH_TYPE SPARSE_OUTPUT_MODEL_URI = TSTR("testdata/sparse_initializer_as_output.onnx");
#ifndef DISABLE_CONTRIB_OPS
@ -1047,7 +1048,7 @@ TEST(CApiTest, invalid_variadic_input_homogeneity_custom_op) {
Ort::Session session(*ort_env, VARIADIC_UNDEF_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options);
FAIL();
} catch (const Ort::Exception& excpt) {
ASSERT_THAT(excpt.what(), testing::HasSubstr("Type Error: Type parameter (T0) of Optype (VariadicNode) bound "
ASSERT_THAT(excpt.what(), testing::HasSubstr("Type Error: Type parameter (Input0) of Optype (VariadicNode) bound "
"to different types"));
}
}
@ -2825,3 +2826,109 @@ TEST(CApiTest, TestMultiStreamInferenceSimpleSSD) {
ASSERT_TRUE(output_dims == expected_output_dims);
}
#endif
#if !defined(ORT_MINIMAL_BUILD)
TEST(MultiKernelSingleSchemaTest, valid) {
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_options.SetLogSeverityLevel(0);
#if defined(_WIN32)
session_options.RegisterCustomOpsLibrary(ORT_TSTR("custom_op_library.dll"));
#elif defined(__APPLE__)
session_options.RegisterCustomOpsLibrary(ORT_TSTR("libcustom_op_library.dylib"));
#else
session_options.RegisterCustomOpsLibrary(ORT_TSTR("./libcustom_op_library.so"));
#endif
Ort::Session session(*ort_env, CUSTOM_OP_SINGLE_SCHEMA_MULTI_KERNEL, session_options);
const char* input_names[] = {"X"};
const char* output_names[] = {"Y", "Z"};
float x_value[] = {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f};
int64_t x_dim[] = {10};
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
Ort::Value input_tensors[1] = {
Ort::Value::CreateTensor<float>(memory_info, x_value, 10, x_dim, 1),
};
Ort::RunOptions run_optoins;
auto output_tensors = session.Run(run_optoins, input_names, input_tensors, 1, output_names, 2);
ASSERT_TRUE(*output_tensors[1].GetTensorData<int32_t>() == 72);
}
// expect input count mismatch exception
TEST(MultiKernelSingleSchemaTest, InputCountMismatch) {
Ort::CustomOpDomain v2_domain("v2");
MulTopOpFloat mul_top_f32;
MulTopOpDouble mul_top_double;
v2_domain.Add(&mul_top_f32);
v2_domain.Add(&mul_top_double);
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_options.SetLogSeverityLevel(0);
session_options.Add(v2_domain);
EXPECT_THROW(Ort::Session session(*ort_env, CUSTOM_OP_SINGLE_SCHEMA_MULTI_KERNEL, session_options), std::exception);
}
// expect output count mismatch exception
TEST(MultiKernelSingleSchemaTest, OutputMismatch) {
Ort::CustomOpDomain v2_domain("v2");
MulTopOpFloat mul_top_f32;
MulTopOpInt16 mul_top_int64;
v2_domain.Add(&mul_top_f32);
v2_domain.Add(&mul_top_int64);
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_options.SetLogSeverityLevel(0);
session_options.Add(v2_domain);
EXPECT_THROW(Ort::Session session(*ort_env, CUSTOM_OP_SINGLE_SCHEMA_MULTI_KERNEL, session_options), std::exception);
}
// expect characteristic mismatch exception
TEST(MultiKernelSingleSchemaTest, CharacterMismatch) {
Ort::CustomOpDomain v2_domain("v2");
MulTopOpFloat mul_top_f32;
MulTopOpFloat16 mul_top_f16;
v2_domain.Add(&mul_top_f32);
v2_domain.Add(&mul_top_f16);
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_options.SetLogSeverityLevel(0);
session_options.Add(v2_domain);
EXPECT_THROW(Ort::Session session(*ort_env, CUSTOM_OP_SINGLE_SCHEMA_MULTI_KERNEL, session_options), std::exception);
}
TEST(MultiKernelSingleSchemaTest, DuplicateKernel) {
Ort::CustomOpDomain v2_domain("v2");
MulTopOpFloat mul_top_f32_1;
MulTopOpFloat mul_top_f32_2;
MulTopOpInt32 mul_top_i32;
v2_domain.Add(&mul_top_f32_1);
v2_domain.Add(&mul_top_f32_2);
v2_domain.Add(&mul_top_i32);
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_options.SetLogSeverityLevel(0);
session_options.Add(v2_domain);
EXPECT_NO_THROW(Ort::Session session(*ort_env, CUSTOM_OP_SINGLE_SCHEMA_MULTI_KERNEL, session_options));
}
#endif

68
onnxruntime/test/testdata/custom_op_library/custom_op_library.cc vendored
View file

@ -101,6 +101,65 @@ struct CustomOpTwo : Ort::CustomOpBase<CustomOpTwo, KernelTwo> {
ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; };
};
////////////////////////////////////////////////
template <typename T>
T MulTopCompute(const T& input_0, const T& input_1) {
return input_0 * input_1;
}
struct MulTopKernelFloat {
MulTopKernelFloat(const OrtKernelInfo*){};
~MulTopKernelFloat() = default;
void Compute(OrtKernelContext* context) {
Ort::KernelContext ctx(context);
auto tensor_in = ctx.GetInput(0);
const float* float_in = tensor_in.GetTensorData<float>();
int64_t output_shape = 1;
auto tensor_out = ctx.GetOutput(0, &output_shape, 1);
auto float_out = tensor_out.GetTensorMutableData<float>();
*float_out = MulTopCompute(float_in[0], float_in[1]);
}
};
struct MulTopOpFloat : Ort::CustomOpBase<MulTopOpFloat, MulTopKernelFloat> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelFloat(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }
size_t GetOutputTypeCount() const { return 1; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }
};
////////////////////////////////////////////////
struct MulTopKernelInt32 {
MulTopKernelInt32(const OrtKernelInfo*){};
~MulTopKernelInt32() = default;
void Compute(OrtKernelContext* context) {
Ort::KernelContext ctx(context);
auto tensor_in = ctx.GetInput(0);
const int32_t* int_in = tensor_in.GetTensorData<int32_t>();
int64_t output_shape = 1;
auto tensor_out = ctx.GetOutput(0, &output_shape, 1);
auto int_out = tensor_out.GetTensorMutableData<int32_t>();
*int_out = MulTopCompute(int_in[0], int_in[1]);
}
};
struct MulTopOpInt32 : Ort::CustomOpBase<MulTopOpInt32, MulTopKernelInt32> {
void* CreateKernel(const OrtApi&, const OrtKernelInfo* info) const { return new MulTopKernelInt32(info); }
const char* GetName() const { return "MulTop"; }
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; }
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; }
size_t GetOutputTypeCount() const { return 1; }
ONNXTensorElementDataType GetOutputType(size_t) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; }
};
////////////////////////////////////////////////
static void AddOrtCustomOpDomainToContainer(Ort::CustomOpDomain&& domain) {
static std::vector<Ort::CustomOpDomain> ort_custom_op_domain_container;
static std::mutex ort_custom_op_domain_mutex;
@ -114,6 +173,9 @@ OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtA
static const CustomOpOne c_CustomOpOne;
static const CustomOpTwo c_CustomOpTwo;
static const MulTopOpFloat c_MulTopOpFloat;
static const MulTopOpInt32 c_MulTopOpInt32;
OrtStatus* result = nullptr;
ORT_TRY {
@ -121,9 +183,15 @@ OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtA
domain.Add(&c_CustomOpOne);
domain.Add(&c_CustomOpTwo);
Ort::CustomOpDomain domain_v2{"v2"};
domain_v2.Add(&c_MulTopOpFloat);
domain_v2.Add(&c_MulTopOpInt32);
Ort::UnownedSessionOptions session_options(options);
session_options.Add(domain);
session_options.Add(domain_v2);
AddOrtCustomOpDomainToContainer(std::move(domain));
AddOrtCustomOpDomainToContainer(std::move(domain_v2));
}
ORT_CATCH(const std::exception& e) {
ORT_HANDLE_EXCEPTION([&]() {

26
onnxruntime/test/testdata/custom_op_single_schema_multi_kernel.onnx vendored Normal file
View file

@ -0,0 +1,26 @@
:Ç
5 const_two"Constant*
value*:B const_two 
$
X
const_twotop_f32top_i"TopK
#
top_f32top_int"Cast*
to 

top_f32mul_f32"MulTop:v2

top_intmul_int"MulTop:v2

mul_f32Y"Identity

mul_intZ"IdentitygraphZ
X

ÿÿÿÿÿÿÿÿÿb
Y

ÿÿÿÿÿÿÿÿÿb
Z

ÿÿÿÿÿÿÿÿÿB