From 3bbcc2799fe19bfb26962851d7b3b6002a91d323 Mon Sep 17 00:00:00 2001 From: Adrian Lizarraga Date: Fri, 23 Dec 2022 11:41:15 -0800 Subject: [PATCH] Support for custom op variadic inputs/outputs (#13946) ### Description Adds support for variadic inputs and outputs to custom operators. ### Motivation and Context Needed for custom ops that wrap external runtimes/models and maybe TensorRT plugins. --- .../core/session/onnxruntime_c_api.h | 28 +- .../core/session/onnxruntime_cxx_api.h | 29 ++ onnxruntime/core/session/custom_ops.cc | 64 ++- onnxruntime/core/session/onnxruntime_c_api.cc | 6 +- .../test/shared_lib/custom_op_utils.cc | 85 ++++ onnxruntime/test/shared_lib/custom_op_utils.h | 94 +++++ onnxruntime/test/shared_lib/test_inference.cc | 385 ++++++++++++++++++ .../test/testdata/custom_op_variadic_io.onnx | Bin 0 -> 261 bytes .../testdata/custom_op_variadic_undef_io.onnx | 30 ++ 9 files changed, 699 insertions(+), 22 deletions(-) create mode 100644 onnxruntime/test/testdata/custom_op_variadic_io.onnx create mode 100644 onnxruntime/test/testdata/custom_op_variadic_undef_io.onnx diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index dbf3236f81..e6e960287c 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -30,7 +30,7 @@ * * This value is used by some API functions to behave as this version of the header expects. */ -#define ORT_API_VERSION 13 +#define ORT_API_VERSION 14 #ifdef __cplusplus extern "C" { @@ -3651,10 +3651,14 @@ struct OrtApi { // Specify if the inputs/outputs are one of: // 1) Non-optional (input/output must be present in the node) // 2) Optional (input/output may be absent in the node) +// 3) Variadic: A variadic input or output specifies N (i.e., the minimum arity) or more operands. +// Only the last input or output of a custom op may be marked as variadic. +// The homogeneity of the variadic input or output determines whether all operands must be of the same +// tensor element type. typedef enum OrtCustomOpInputOutputCharacteristic { - // TODO: Support 'Variadic' inputs/outputs INPUT_OUTPUT_REQUIRED = 0, INPUT_OUTPUT_OPTIONAL, + INPUT_OUTPUT_VARIADIC, } OrtCustomOpInputOutputCharacteristic; /* @@ -3692,8 +3696,26 @@ struct OrtCustomOp { // to place the inputs on specific devices. By default, it returns // OrtMemTypeDefault, which means the input is placed on the default device for // the execution provider. If the inputs need to be with different memory tyeps, - // this function can be overriden to return the specific memory types. + // this function can be overridden to return the specific memory types. OrtMemType(ORT_API_CALL* GetInputMemoryType)(_In_ const struct OrtCustomOp* op, _In_ size_t index); + + // Returns the minimum number of input arguments expected for the variadic input. + // Applicable only for custom ops that have a variadic input. + int(ORT_API_CALL* GetVariadicInputMinArity)(_In_ const struct OrtCustomOp* op); + + // Returns true (non-zero) if all arguments of a variadic input have to be of the same type (homogeneous), + // and false (zero) otherwise. + // Applicable only for custom ops that have a variadic input. + int(ORT_API_CALL* GetVariadicInputHomogeneity)(_In_ const struct OrtCustomOp* op); + + // Returns the minimum number of output values expected for the variadic output. + // Applicable only for custom ops that have a variadic output. + int(ORT_API_CALL* GetVariadicOutputMinArity)(_In_ const struct OrtCustomOp* op); + + // Returns true (non-zero) if all outputs values of a variadic output have to be of the same type (homogeneous), + // and false (zero) otherwise. + // Applicable only for custom ops that have a variadic output. + int(ORT_API_CALL* GetVariadicOutputHomogeneity)(_In_ const struct OrtCustomOp* op); }; /* diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index c5ee9b28aa..85dc2ab515 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1668,6 +1668,11 @@ struct CustomOpBase : OrtCustomOp { #endif OrtCustomOp::GetInputCharacteristic = [](const OrtCustomOp* this_, size_t index) { return static_cast(this_)->GetInputCharacteristic(index); }; OrtCustomOp::GetOutputCharacteristic = [](const OrtCustomOp* this_, size_t index) { return static_cast(this_)->GetOutputCharacteristic(index); }; + + OrtCustomOp::GetVariadicInputMinArity = [](const OrtCustomOp* this_) { return static_cast(this_)->GetVariadicInputMinArity(); }; + OrtCustomOp::GetVariadicInputHomogeneity = [](const OrtCustomOp* this_) { return static_cast(static_cast(this_)->GetVariadicInputHomogeneity()); }; + OrtCustomOp::GetVariadicOutputMinArity = [](const OrtCustomOp* this_) { return static_cast(this_)->GetVariadicOutputMinArity(); }; + OrtCustomOp::GetVariadicOutputHomogeneity = [](const OrtCustomOp* this_) { return static_cast(static_cast(this_)->GetVariadicOutputHomogeneity()); }; } // Default implementation of GetExecutionProviderType that returns nullptr to default to the CPU provider @@ -1687,6 +1692,30 @@ struct CustomOpBase : OrtCustomOp { OrtMemType GetInputMemoryType(size_t /*index*/) const { return OrtMemTypeDefault; } + + // Default implementation of GetVariadicInputMinArity() returns 1 to specify that a variadic input + // should expect at least 1 argument. + int GetVariadicInputMinArity() const { + return 1; + } + + // Default implementation of GetVariadicInputHomegeneity() returns true to specify that all arguments + // to a variadic input should be of the same type. + bool GetVariadicInputHomogeneity() const { + return true; + } + + // Default implementation of GetVariadicOutputMinArity() returns 1 to specify that a variadic output + // should produce at least 1 output value. + int GetVariadicOutputMinArity() const { + return 1; + } + + // Default implementation of GetVariadicOutputHomegeneity() returns true to specify that all output values + // produced by a variadic output should be of the same type. + bool GetVariadicOutputHomogeneity() const { + return true; + } }; } // namespace Ort diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index 2f5e808430..489066a65d 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -176,43 +176,74 @@ 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; - auto input_count = op->GetInputTypeCount(op); + 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; - // Only since the ORT API version 8 and onwards does the OrtCustomOp interface have the relevant methods exposed to query - // if an input/output is required/optional. So, query the relevant methods ONLY from API version 8 onwards. - if (op->version >= 8 && op->GetInputCharacteristic(op, i) == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) { - option = onnx::OpSchema::FormalParameterOption::Optional; + // 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)); + } } - auto type = op->GetInputType(op, i); + 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); + 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); + DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type)), option, + is_homogeneous, min_arity); } } - auto output_count = op->GetOutputTypeCount(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; - // Only since the ORT API version 8 and onwards does the OrtCustomOp interface have the relevant methods exposed to query - // if an input/output is required/optional. So, query the relevant methods ONLY from API version 8 onwards. - if (op->version >= 8 && op->GetOutputCharacteristic(op, i) == OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL) { - option = onnx::OpSchema::FormalParameterOption::Optional; + // 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)); + } } - auto type = op->GetOutputType(op, i); + const auto type = op->GetOutputType(op, i); if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) { // Dynamic typed output ORT_ENFORCE(type_id_counter == 1, "There must be one (and only one) dynamic typed input to the custom op. " @@ -221,10 +252,11 @@ common::Status CreateCustomRegistry(gsl::span op_domai "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); + 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); + DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type)), option, + is_homogeneous, min_arity); } } diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index d691cff0f7..9a3411bb15 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -2335,7 +2335,7 @@ Second example, if we wanted to add and remove some members, we'd do this: In GetApi we now make it return ort_api_3 for version 3. */ -static constexpr OrtApi ort_api_1_to_12 = { +static constexpr OrtApi ort_api_1_to_14 = { // NOTE: The ordering of these fields MUST not change after that version has shipped since existing binaries depend on this ordering. // Shipped as version 1 - DO NOT MODIFY (see above text for more information) @@ -2634,13 +2634,13 @@ static_assert(offsetof(OrtApi, ReleaseCANNProviderOptions) / sizeof(void*) == 22 static_assert(std::string_view(ORT_VERSION) == "1.14.0", "ORT_Version change detected, please follow below steps to ensure OrtApi is updated properly"); // 1. Update the hardcoded version string in above static_assert to silence it -// 2. If there were any APIs added to ort_api_1_to_13 above: +// 2. If there were any APIs added to ort_api_1_to_14 above: // a. Add the 'End of version #' markers (pattern above should be obvious) // b. Add a static_assert in the directly above list of version sizes to ensure nobody adds any more functions to the just shipped API version ORT_API(const OrtApi*, OrtApis::GetApi, uint32_t version) { if (version >= 1 && version <= ORT_API_VERSION) - return &ort_api_1_to_12; + return &ort_api_1_to_14; fprintf(stderr, "The given version [%u] is not supported, only version 1 to %u is supported in this build.\n", version, ORT_API_VERSION); diff --git a/onnxruntime/test/shared_lib/custom_op_utils.cc b/onnxruntime/test/shared_lib/custom_op_utils.cc index e11a834afb..ebb9a6c1cc 100644 --- a/onnxruntime/test/shared_lib/custom_op_utils.cc +++ b/onnxruntime/test/shared_lib/custom_op_utils.cc @@ -168,6 +168,91 @@ void MyCustomKernelWithOptionalInput::Compute(OrtKernelContext* context) { } } +void MyCustomStringLengthsKernel::Compute(OrtKernelContext* context) { + Ort::KernelContext kcontext(context); + constexpr std::array output_shape = {1}; + const size_t num_inputs = kcontext.GetInputCount(); + + // Each output is set to the length of the corresponding input string. + for (size_t i = 0; i < num_inputs; ++i) { + auto input = kcontext.GetInput(i); + auto output = kcontext.GetOutput(i, output_shape.data(), output_shape.size()); + int64_t* str_len_ptr = output.GetTensorMutableData(); + + *str_len_ptr = input.GetStringTensorElementLength(0); + } +} + +void AddInputForCustomStringLengthsKernel(std::string input_str, OrtAllocator* allocator, + std::vector& ort_inputs, std::vector& input_names, + std::vector& output_names, + std::vector>& expected_dims, + std::vector>& expected_outputs) { + const size_t input_index = ort_inputs.size(); + constexpr std::array input_dims = {1}; + Ort::Value& ort_value = ort_inputs.emplace_back( + Ort::Value::CreateTensor(allocator, input_dims.data(), input_dims.size(), + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING)); + std::ostringstream oss(std::ostringstream::ate); + + oss.str("input_"); + oss << input_index; + input_names.emplace_back(oss.str()); + + oss.str("output_"); + oss << input_index; + output_names.emplace_back(oss.str()); + + expected_dims.push_back({1}); + expected_outputs.push_back({static_cast(input_str.size())}); + ort_value.FillStringTensorElement(input_str.data(), 0); +} + +void MyCustomEchoReversedArgsKernel::Compute(OrtKernelContext* context) { + Ort::KernelContext kcontext(context); + constexpr std::array output_shape = {1}; + const size_t num_inputs = kcontext.GetInputCount(); + + for (size_t i = 0; i < num_inputs; ++i) { + const size_t out_index = num_inputs - i - 1; + auto input = kcontext.GetInput(i); + auto output = kcontext.GetOutput(out_index, output_shape.data(), output_shape.size()); + + auto type_shape_info = input.GetTensorTypeAndShapeInfo(); + auto elem_type = type_shape_info.GetElementType(); + + // Only support STRING, INT64_T, and FLOAT + switch (elem_type) { + case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING: { + const size_t str_len = input.GetStringTensorElementLength(0); + std::string str; + + str.resize(str_len); + input.GetStringTensorElement(str.size(), 0, str.data()); + output.FillStringTensorElement(str.c_str(), 0); + break; + } + case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: { + int64_t* out_ptr = output.GetTensorMutableData(); + const int64_t* inp_ptr = input.GetTensorData(); + + out_ptr[0] = inp_ptr[0]; + break; + } + case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: { + float* out_ptr = output.GetTensorMutableData(); + const float* inp_ptr = input.GetTensorData(); + + out_ptr[0] = inp_ptr[0]; + break; + } + default: + ORT_CXX_API_THROW("MyCustomEchoReversedArgsKernel only supports tensor inputs of type STRING, INT64_T, and FLOAT", + OrtErrorCode::ORT_INVALID_GRAPH); + } + } +} + void MyCustomKernelWithAttributes::Compute(OrtKernelContext* context) { // Setup inputs Ort::KernelContext ctx(context); diff --git a/onnxruntime/test/shared_lib/custom_op_utils.h b/onnxruntime/test/shared_lib/custom_op_utils.h index bc7dfee0ad..873af327de 100644 --- a/onnxruntime/test/shared_lib/custom_op_utils.h +++ b/onnxruntime/test/shared_lib/custom_op_utils.h @@ -147,6 +147,100 @@ struct MyCustomOpWithOptionalInput : Ort::CustomOpBase& ort_inputs, std::vector& input_names, + std::vector& output_names, + std::vector>& expected_dims, + std::vector>& expected_outputs); + +// Custom kernel that echos input arguments (shape [1]) in reversed order. +// Used to test variadic custom ops with heterogenous input types. +struct MyCustomEchoReversedArgsKernel { + explicit MyCustomEchoReversedArgsKernel(const OrtKernelInfo* /* info */) {} + void Compute(OrtKernelContext* context); +}; + +// Utility custom op class that can be configured with a kernel class (T) and input/output +// configurations. +template +struct TemplatedCustomOp : Ort::CustomOpBase, T> { + TemplatedCustomOp(const char* op_name, std::vector input_types, + std::vector input_characs, int input_min_arity, + bool input_homogeneity, std::vector output_types, + std::vector output_characs, int output_min_arity, + bool output_homogeneity) + : op_name_(op_name), input_types_(std::move(input_types)), + input_characs_(std::move(input_characs)), input_min_arity_(input_min_arity), + input_homogeneity_(input_homogeneity), output_types_(std::move(output_types)), + output_characs_(std::move(output_characs)), output_min_arity_(output_min_arity), + output_homogeneity_(output_homogeneity) {} + + void* CreateKernel(const OrtApi& /* api */, const OrtKernelInfo* info) const { + return new T(info); + } + + const char* GetName() const noexcept { return op_name_; } + + size_t GetInputTypeCount() const noexcept { return input_types_.size(); } + + ONNXTensorElementDataType GetInputType(size_t index) const noexcept { + return input_types_[index]; + } + + OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t index) const noexcept { + return input_characs_[index]; + } + + int GetVariadicInputMinArity() const noexcept { + return input_min_arity_; + } + + bool GetVariadicInputHomogeneity() const noexcept { + return input_homogeneity_; + } + + size_t GetOutputTypeCount() const noexcept { return output_types_.size(); } + + ONNXTensorElementDataType GetOutputType(size_t index) const noexcept { + return output_types_[index]; + } + + OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t index) const noexcept { + return output_characs_[index]; + } + + int GetVariadicOutputMinArity() const noexcept { + return output_min_arity_; + } + + bool GetVariadicOutputHomogeneity() const noexcept { + return output_homogeneity_; + } + + private: + const char* op_name_; + + std::vector input_types_; + std::vector input_characs_; + int input_min_arity_; + bool input_homogeneity_; + + std::vector output_types_; + std::vector output_characs_; + int output_min_arity_; + bool output_homogeneity_; +}; + struct MyCustomKernelWithAttributes { MyCustomKernelWithAttributes(const OrtKernelInfo* kernel_info) { Ort::ConstKernelInfo info{kernel_info}; diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 8c4eb5ced7..5eb8bed2b7 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -173,6 +173,9 @@ static constexpr PATH_TYPE VARIED_INPUT_CUSTOM_OP_MODEL_URI = TSTR("testdata/Var static constexpr PATH_TYPE VARIED_INPUT_CUSTOM_OP_MODEL_URI_2 = TSTR("testdata/foo_3.onnx"); static constexpr PATH_TYPE OPTIONAL_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI = TSTR("testdata/foo_bar_1.onnx"); static constexpr PATH_TYPE OPTIONAL_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI_2 = TSTR("testdata/foo_bar_2.onnx"); +static constexpr PATH_TYPE VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI = TSTR("testdata/custom_op_variadic_io.onnx"); +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"); #if !defined(DISABLE_SPARSE_TENSORS) static constexpr PATH_TYPE SPARSE_OUTPUT_MODEL_URI = TSTR("testdata/sparse_initializer_as_output.onnx"); @@ -650,6 +653,388 @@ TEST(CApiTest, multiple_varied_input_custom_op_handler) { #endif } +TEST(CApiTest, variadic_input_output_custom_op) { + // Create a custom op with 1 variadic input and 1 variadic output. + // The model passes in 3 string inputs and expects 3 int64_t outputs. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + std::vector ort_inputs; + Ort::AllocatorWithDefaultOptions allocator; + std::vector> expected_dims; + std::vector> expected_lens; + std::vector input_names; + std::vector output_names; + + // Create inputs. + AddInputForCustomStringLengthsKernel("hello", allocator, ort_inputs, input_names, output_names, expected_dims, + expected_lens); + AddInputForCustomStringLengthsKernel("", allocator, ort_inputs, input_names, output_names, expected_dims, + expected_lens); + AddInputForCustomStringLengthsKernel("123", allocator, ort_inputs, input_names, output_names, expected_dims, + expected_lens); + + // Create arrays of c-strings for input and output names. + auto get_c_str = [](const std::string& str) { return str.c_str(); }; + std::vector input_name_cstrs(input_names.size()); + std::transform(input_names.begin(), input_names.end(), input_name_cstrs.begin(), get_c_str); + std::vector output_name_cstrs(output_names.size()); + std::transform(output_names.begin(), output_names.end(), output_name_cstrs.begin(), get_c_str); + + Ort::Session session(*ort_env, VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + auto ort_outputs = session.Run(Ort::RunOptions{}, input_name_cstrs.data(), ort_inputs.data(), ort_inputs.size(), + output_name_cstrs.data(), output_name_cstrs.size()); + ASSERT_EQ(ort_outputs.size(), 3u); + + // Validate outputs. + for (size_t i = 0; i < ort_outputs.size(); ++i) { + auto type_info = ort_outputs[i].GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), expected_dims[i]); + ASSERT_EQ(type_info.GetElementCount(), 1u); + + int64_t* lens_data = ort_outputs[i].GetTensorMutableData(); + ASSERT_EQ(lens_data[0], expected_lens[i][0]); + } +} + +TEST(CApiTest, mixed_variadic_input_output_custom_op) { + // Create a custom op with 2 inputs (required, variadic) and 2 outputs (required, variadic). + // The model passes in 3 string inputs and expects 3 int64_t outputs. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING, + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED, + OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED, + OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + std::vector ort_inputs; + Ort::AllocatorWithDefaultOptions allocator; + std::vector> expected_dims; + std::vector> expected_lens; + std::vector input_names; + std::vector output_names; + + // Create inputs. + AddInputForCustomStringLengthsKernel("mixed variadic", allocator, ort_inputs, input_names, output_names, + expected_dims, expected_lens); + AddInputForCustomStringLengthsKernel("", allocator, ort_inputs, input_names, output_names, expected_dims, + expected_lens); + AddInputForCustomStringLengthsKernel("abcd", allocator, ort_inputs, input_names, output_names, expected_dims, + expected_lens); + + // Create arrays of c-strings for input and output names. + auto get_c_str = [](const std::string& str) { return str.c_str(); }; + std::vector input_name_cstrs(input_names.size()); + std::transform(input_names.begin(), input_names.end(), input_name_cstrs.begin(), get_c_str); + std::vector output_name_cstrs(output_names.size()); + std::transform(output_names.begin(), output_names.end(), output_name_cstrs.begin(), get_c_str); + + Ort::Session session(*ort_env, VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + auto ort_outputs = session.Run(Ort::RunOptions{}, input_name_cstrs.data(), ort_inputs.data(), ort_inputs.size(), + output_name_cstrs.data(), output_name_cstrs.size()); + ASSERT_EQ(ort_outputs.size(), 3u); + + // Validate outputs. + for (size_t i = 0; i < ort_outputs.size(); ++i) { + auto type_info = ort_outputs[i].GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), expected_dims[i]); + ASSERT_EQ(type_info.GetElementCount(), 1u); + + int64_t* lens_data = ort_outputs[i].GetTensorMutableData(); + ASSERT_EQ(lens_data[0], expected_lens[i][0]); + } +} + +TEST(CApiTest, variadic_undef_input_output_custom_op) { + // Create a custom op with 1 variadic input and 1 variadic output. + // Both the input and output are of undefined element type and allowed to differ in type (hetergeneous). + // The model passes in inputs (string, int64_t, and float) which are then echoed in + // reversed order (float, int64_t, string). + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + false, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + false); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + std::vector ort_inputs; + Ort::AllocatorWithDefaultOptions allocator; + Ort::ConstMemoryInfo mem_info = allocator.GetInfo(); + std::vector input_dims = {1}; + + // Set string input. + std::string str_input("hello_ort"); + Ort::Value& str_input_val = ort_inputs.emplace_back( + Ort::Value::CreateTensor(allocator, input_dims.data(), input_dims.size(), + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING)); + str_input_val.FillStringTensorElement(str_input.c_str(), 0); + + // Set int64_t input. + std::array int_inps = {23}; + ort_inputs.emplace_back(Ort::Value::CreateTensor(mem_info, int_inps.data(), int_inps.size(), + input_dims.data(), input_dims.size())); + + // Set float input. + std::array float_inps = {10.0f}; + ort_inputs.emplace_back(Ort::Value::CreateTensor(mem_info, float_inps.data(), float_inps.size(), + input_dims.data(), input_dims.size())); + + constexpr std::array input_names = {"input_0", "input_1", "input_2"}; + constexpr std::array output_names = {"output_0", "output_1", "output_2"}; + + Ort::Session session(*ort_env, VARIADIC_UNDEF_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + auto ort_outputs = session.Run(Ort::RunOptions{}, input_names.data(), ort_inputs.data(), ort_inputs.size(), + output_names.data(), output_names.size()); + ASSERT_EQ(ort_outputs.size(), 3u); + + // Validate outputs. + + // First output should be a float. + { + auto& ort_output = ort_outputs[0]; + auto type_info = ort_output.GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), input_dims); + ASSERT_EQ(type_info.GetElementCount(), 1u); + ASSERT_EQ(type_info.GetElementType(), ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + + const float* out_ptr = ort_output.GetTensorData(); + ASSERT_EQ(out_ptr[0], float_inps[0]); + } + + // Second output should be a int64_t. + { + auto& ort_output = ort_outputs[1]; + auto type_info = ort_output.GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), input_dims); + ASSERT_EQ(type_info.GetElementCount(), 1u); + ASSERT_EQ(type_info.GetElementType(), ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + + const int64_t* out_ptr = ort_output.GetTensorData(); + ASSERT_EQ(out_ptr[0], int_inps[0]); + } + + // Last output should be a string. + { + auto& ort_output = ort_outputs[2]; + auto type_info = ort_output.GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), input_dims); + ASSERT_EQ(type_info.GetElementCount(), 1u); + ASSERT_EQ(type_info.GetElementType(), ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING); + + const size_t str_len = ort_output.GetStringTensorElementLength(0); + ASSERT_EQ(str_len, str_input.size()); + + std::string str; + str.resize(str_len); + + ort_output.GetStringTensorElement(str_len, 0, str.data()); + ASSERT_EQ(str, str_input); + } +} + +TEST(CApiTest, invalid_variadic_input_not_last_custom_op) { + // Create an invalid custom op with 2 inputs. The first input is variadic and the last is not. + // Expect an error because only the last input may be marked as variadic. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING, + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC, + OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED}, + 1, + true, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED, + OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + try { + Ort::Session session(*ort_env, VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + FAIL(); + } catch (const Ort::Exception& excpt) { + ASSERT_THAT(excpt.what(), testing::HasSubstr("Only the last input to a custom op may be marked variadic.")); + } +} + +TEST(CApiTest, invalid_variadic_output_not_last_custom_op) { + // Create an invalid custom op with 2 outputs. The first output is variadic and the last is not. + // Expect an error because only the last output may be marked as variadic. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING, + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED, + OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, + ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC, + OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED}, + 1, + true); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + try { + Ort::Session session(*ort_env, VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + FAIL(); + } catch (const Ort::Exception& excpt) { + ASSERT_THAT(excpt.what(), testing::HasSubstr("Only the last output to a custom op may be marked variadic.")); + } +} + +TEST(CApiTest, invalid_variadic_input_min_arity_custom_op) { + // Create a custom op with a variadic input with a minimum arity of 4. + // Expect an error because the model passes in less than 4 inputs to the op. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 4, + true, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + try { + Ort::Session session(*ort_env, VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + FAIL(); + } catch (const Ort::Exception& excpt) { + ASSERT_THAT(excpt.what(), testing::HasSubstr("Error Node (VariadicNode0) has input size 3 not in range [min=4")); + } +} + +TEST(CApiTest, invalid_variadic_output_min_arity_custom_op) { + // Create a custom op with a variadic output with a minimum arity of 4. + // Expect an error because the model instantiates the op with less than 4 outputs. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true, + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 4, + true); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + try { + Ort::Session session(*ort_env, VARIADIC_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI, session_options); + FAIL(); + } catch (const Ort::Exception& excpt) { + ASSERT_THAT(excpt.what(), testing::HasSubstr("Error Node (VariadicNode0) has output size 3 not in range [min=4")); + } +} + +TEST(CApiTest, invalid_variadic_input_homogeneity_custom_op) { + // Create a custom op with a homogeneous variadic input. The model has heterogeneous inputs, + // so we expect an error. + TemplatedCustomOp custom_op( + "VariadicNode", + // Input config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + true, // Input homogeneity requirement will cause error! + // Output config + {ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED}, + {OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC}, + 1, + false); + + Ort::CustomOpDomain custom_op_domain("test"); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + try { + 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 " + "to different types")); + } +} + TEST(CApiTest, optional_input_output_custom_op_handler) { MyCustomOpWithOptionalInput custom_op{onnxruntime::kCpuExecutionProvider}; diff --git a/onnxruntime/test/testdata/custom_op_variadic_io.onnx b/onnxruntime/test/testdata/custom_op_variadic_io.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3274a9e6d58cbd029a11c577d993df84691250b7 GIT binary patch literal 261 zcmdWx6~o1znO9I+5^n%w7{VAvLLB*}C15!?%Mi{olHv_ZEXqtw$xQakPf0aU z;z8zGv6Q41mk5a_mll`g=f>w3#Fs%8$7kk83B#-~5aQrs(^b literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/custom_op_variadic_undef_io.onnx b/onnxruntime/test/testdata/custom_op_variadic_undef_io.onnx new file mode 100644 index 0000000000..8bb13eae23 --- /dev/null +++ b/onnxruntime/test/testdata/custom_op_variadic_undef_io.onnx @@ -0,0 +1,30 @@ +:ˆ +\ +input_0 +input_1 +input_2output_0output_1output_2 VariadicNode0" VariadicNode:testcustom_op_variadic_undef_ioZ +input_0 + + +Z +input_1 + + +Z +input_2 + + +b +output_0 + + +b +output_1 + + +b +output_2 + + +B +test \ No newline at end of file