diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index ad2a0d8bb6..6cd4044b19 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -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(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(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& out) -> Status { + out = std::make_unique(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(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(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& 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(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 op_domains, std::shared_ptr& output) { output = std::make_shared(); for (const auto& domain : op_domains) { - // Create an OpSchema for each op and register them - - // Container to hold type template parameters - std::unordered_map> type_constraint_ids; - #if !defined(ORT_MINIMAL_BUILD) + std::unordered_map schema_map; + std::unordered_map> 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 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 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(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(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 schemas; + for (auto schema_iter : schema_map) { + schemas.push_back(schema_iter.second); + InlinedVector 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 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 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(); } diff --git a/onnxruntime/test/shared_lib/custom_op_utils.h b/onnxruntime/test/shared_lib/custom_op_utils.h index 74975b53d7..8ead4056b1 100644 --- a/onnxruntime/test/shared_lib/custom_op_utils.h +++ b/onnxruntime/test/shared_lib/custom_op_utils.h @@ -377,3 +377,91 @@ struct StandaloneCustomOp : Ort::CustomOpBase { + 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 { + 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 { + 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 { + 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 { + 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; + } +}; \ No newline at end of file diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 07ce0f773d..0b79b29b20 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -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(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() == 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 diff --git a/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc b/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc index 16ce9e49c8..fbc520cefa 100644 --- a/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc +++ b/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc @@ -101,6 +101,65 @@ struct CustomOpTwo : Ort::CustomOpBase { ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; }; }; +//////////////////////////////////////////////// + +template +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(); + int64_t output_shape = 1; + auto tensor_out = ctx.GetOutput(0, &output_shape, 1); + auto float_out = tensor_out.GetTensorMutableData(); + *float_out = MulTopCompute(float_in[0], float_in[1]); + } +}; + +struct MulTopOpFloat : Ort::CustomOpBase { + 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(); + int64_t output_shape = 1; + auto tensor_out = ctx.GetOutput(0, &output_shape, 1); + auto int_out = tensor_out.GetTensorMutableData(); + *int_out = MulTopCompute(int_in[0], int_in[1]); + } +}; + +struct MulTopOpInt32 : Ort::CustomOpBase { + 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_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([&]() { diff --git a/onnxruntime/test/testdata/custom_op_single_schema_multi_kernel.onnx b/onnxruntime/test/testdata/custom_op_single_schema_multi_kernel.onnx new file mode 100644 index 0000000000..7032437d88 --- /dev/null +++ b/onnxruntime/test/testdata/custom_op_single_schema_multi_kernel.onnx @@ -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 \ No newline at end of file