diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 7ec87fb363..5ca4b9b1e0 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -7,7 +7,7 @@ #include // This value is used in structures passed to ORT so that a newer version of ORT will still work with them -#define ORT_API_VERSION 7 +#define ORT_API_VERSION 8 #ifdef __cplusplus extern "C" { @@ -1189,6 +1189,16 @@ struct OrtApi { */ #define OrtCustomOpApi OrtApi +// Specifies some characteristics of inputs/outputs of custom ops: +// 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) +typedef enum OrtCustomOpInputOutputCharacteristic { + // TODO: Support 'Variadic' inputs/outputs + INPUT_OUTPUT_REQUIRED = 0, + INPUT_OUTPUT_OPTIONAL, +} OrtCustomOpInputOutputCharacteristic; + /* * The OrtCustomOp structure defines a custom op's schema and its kernel callbacks. The callbacks are filled in by * the implementor of the custom op. @@ -1215,11 +1225,11 @@ struct OrtCustomOp { // Op kernel callbacks void(ORT_API_CALL* KernelCompute)(_In_ void* op_kernel, _In_ OrtKernelContext* context); void(ORT_API_CALL* KernelDestroy)(_In_ void* op_kernel); -}; -/* - * END EXPERIMENTAL -*/ + // Returns the characteristics of the input & output tensors + OrtCustomOpInputOutputCharacteristic(ORT_API_CALL* GetInputCharacteristic)(_In_ const struct OrtCustomOp* op, _In_ size_t index); + OrtCustomOpInputOutputCharacteristic(ORT_API_CALL* GetOutputCharacteristic)(_In_ const struct OrtCustomOp* op, _In_ size_t index); +}; #ifdef __cplusplus } diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index be43d9cd21..f9c87c0a0c 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -623,10 +623,23 @@ struct CustomOpBase : OrtCustomOp { OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) { static_cast(op_kernel)->Compute(context); }; OrtCustomOp::KernelDestroy = [](void* op_kernel) { delete static_cast(op_kernel); }; + + 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); }; } // Default implementation of GetExecutionProviderType that returns nullptr to default to the CPU provider const char* GetExecutionProviderType() const { return nullptr; } + + // Default implementations of GetInputCharacteristic() and GetOutputCharacteristic() below + // (inputs and outputs are required by default) + OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t /*index*/) const { + return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED; + } + + OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t /*index*/) const { + return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED; + } }; } // namespace Ort diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index bc4feafa58..25d97a6b00 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/onnxruntime/core/session/custom_ops.cc @@ -125,18 +125,35 @@ common::Status CreateCustomRegistry(const std::vector& op_do size_t type_id_counter = 0; auto input_count = op->GetInputTypeCount(op); for (size_t i = 0; i < input_count; i++) { + onnx::OpSchema::FormalParameterOption option = onnx::OpSchema::FormalParameterOption::Single; + + // 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; + } + 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)); + schema.Input(i, "Input" + std::to_string(i), "", "T" + std::to_string(type_id_counter), option); 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))); + schema.Input(i, "Input" + std::to_string(i), "", + DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type)), option); } } auto output_count = op->GetOutputTypeCount(op); for (size_t i = 0; i < output_count; i++) { + onnx::OpSchema::FormalParameterOption option = onnx::OpSchema::FormalParameterOption::Single; + + // 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; + } + auto type = op->GetOutputType(op, i); if (ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED == type) { // Dynamic typed output ORT_ENFORCE(type_id_counter == 1, @@ -146,9 +163,10 @@ common::Status CreateCustomRegistry(const std::vector& op_do "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"); + schema.Output(i, "Output" + std::to_string(i), "", "T0", option); } else { - schema.Output(i, "Output" + std::to_string(i), "", DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type))); + schema.Output(i, "Output" + std::to_string(i), "", + DataTypeImpl::ToString(onnxruntime::DataTypeImpl::TensorTypeFromONNXEnum(type)), option); } } diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index ecafbc4ce5..24d94bc625 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -1916,7 +1916,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_7 = { +static constexpr OrtApi ort_api_1_to_8 = { // 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) @@ -2109,6 +2109,8 @@ static constexpr OrtApi ort_api_1_to_7 = { &OrtApis::SetCurrentGpuDeviceId, &OrtApis::GetCurrentGpuDeviceId, // End of Version 7 - DO NOT MODIFY ABOVE (see above text for more information) + + // Version 8 - In development, feel free to add/remove/rearrange here }; // Assert to do a limited check to ensure Version 1 of OrtApi never changes (will detect an addition or deletion but not if they cancel out each other) @@ -2117,7 +2119,7 @@ static_assert(offsetof(OrtApi, ReleaseCustomOpDomain) / sizeof(void*) == 101, "S ORT_API(const OrtApi*, OrtApis::GetApi, uint32_t version) { if (version >= 1 && version <= ORT_API_VERSION) - return &ort_api_1_to_7; + return &ort_api_1_to_8; 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/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index cdb755ff7b..4f4a03c818 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -165,6 +165,8 @@ static constexpr PATH_TYPE NAMED_AND_ANON_DIM_PARAM_URI = TSTR("testdata/capi_sy static constexpr PATH_TYPE MODEL_WITH_CUSTOM_MODEL_METADATA = TSTR("testdata/model_with_valid_ort_config_json.onnx"); static constexpr PATH_TYPE VARIED_INPUT_CUSTOM_OP_MODEL_URI = TSTR("testdata/VariedInputCustomOp.onnx"); 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"); #ifdef ENABLE_LANGUAGE_INTEROP_OPS static constexpr PATH_TYPE PYOP_FLOAT_MODEL_URI = TSTR("testdata/pyop_1.onnx"); @@ -347,7 +349,7 @@ struct SliceCustomOp : Ort::CustomOpBase { default: ORT_THROW("Invalid output index: ", index); } - }; + } private: const char* provider_; @@ -443,6 +445,91 @@ TEST(CApiTest, multiple_varied_input_custom_op_handler) { } } +TEST(CApiTest, optional_input_output_custom_op_handler) { + MyCustomOpWithOptionalInput custom_op{onnxruntime::kCpuExecutionProvider}; + + // `MyCustomOpFooBar` defines a custom op with atmost 3 inputs and the second input is optional. + // In this test, we are going to try and run 2 models - one with the optional input and one without + // the optional input. + Ort::CustomOpDomain custom_op_domain(""); + custom_op_domain.Add(&custom_op); + + Ort::SessionOptions session_options; + session_options.Add(custom_op_domain); + + std::vector ort_inputs; + + Ort::MemoryInfo info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault); + + // input 0 + std::vector input_0_data = {1.f}; + std::vector input_0_dims = {1}; + ort_inputs.emplace_back( + Ort::Value::CreateTensor(info, const_cast(input_0_data.data()), + input_0_data.size(), input_0_dims.data(), input_0_dims.size())); + + // input 1 + std::vector input_1_data = {1.f}; + std::vector input_1_dims = {1}; + ort_inputs.emplace_back( + Ort::Value::CreateTensor(info, const_cast(input_1_data.data()), + input_1_data.size(), input_1_dims.data(), input_1_dims.size())); + + // input 2 + std::vector input_2_data = {1.f}; + std::vector input_2_dims = {1}; + ort_inputs.emplace_back( + Ort::Value::CreateTensor(info, const_cast(input_2_data.data()), + input_2_data.size(), input_2_dims.data(), input_2_dims.size())); + + const char* output_name = "Y"; + + // Part 1: Model with optional input present + { + std::vector input_names = {"X1", "X2", "X3"}; + Ort::Session session(*ort_env, OPTIONAL_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_name, 1); + ASSERT_EQ(ort_outputs.size(), 1u); + + // Validate results + std::vector y_dims = {1}; + std::vector values_y = {3.f}; + auto type_info = ort_outputs[0].GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), y_dims); + size_t total_len = type_info.GetElementCount(); + ASSERT_EQ(values_y.size(), total_len); + + float* f = ort_outputs[0].GetTensorMutableData(); + for (size_t i = 0; i != total_len; ++i) { + ASSERT_EQ(values_y[i], f[i]); + } + } + + // Part 2: Model with optional input absent + { + std::vector input_names = {"X1", "X2"}; + ort_inputs.erase(ort_inputs.begin() + 2); // remove the last input in the container + Ort::Session session(*ort_env, OPTIONAL_INPUT_OUTPUT_CUSTOM_OP_MODEL_URI_2, session_options); + auto ort_outputs = session.Run(Ort::RunOptions{}, input_names.data(), ort_inputs.data(), ort_inputs.size(), + &output_name, 1); + ASSERT_EQ(ort_outputs.size(), 1u); + + // Validate results + std::vector y_dims = {1}; + std::vector values_y = {2.f}; + auto type_info = ort_outputs[0].GetTensorTypeAndShapeInfo(); + ASSERT_EQ(type_info.GetShape(), y_dims); + size_t total_len = type_info.GetElementCount(); + ASSERT_EQ(values_y.size(), total_len); + + float* f = ort_outputs[0].GetTensorMutableData(); + for (size_t i = 0; i != total_len; ++i) { + ASSERT_EQ(values_y[i], f[i]); + } + } +} + // Tests registration of a custom op of the same name for both CPU and CUDA EPs #ifdef USE_CUDA TEST(CApiTest, RegisterCustomOpForCPUAndCUDA) { diff --git a/onnxruntime/test/shared_lib/utils.cc b/onnxruntime/test/shared_lib/utils.cc index 97cda33b1a..6267cc8d0d 100644 --- a/onnxruntime/test/shared_lib/utils.cc +++ b/onnxruntime/test/shared_lib/utils.cc @@ -64,4 +64,26 @@ void MyCustomKernelMultipleDynamicInputs::Compute(OrtKernelContext* context) { out[i] = static_cast(X[i] + Y[i]); } #endif -} \ No newline at end of file +} + +void MyCustomKernelWithOptionalInput::Compute(OrtKernelContext* context) { + // Setup inputs + const OrtValue* input_X1 = ort_.KernelContext_GetInput(context, 0); + const OrtValue* input_X2 = ort_.KernelContext_GetInput(context, 1); + const OrtValue* input_X3 = ort_.KernelContext_GetInput(context, 2); + + const float* X1 = ort_.GetTensorData(input_X1); + // The second input may or may not be present + const float* X2 = (input_X2 != nullptr) ? ort_.GetTensorData(input_X2) : nullptr; + const float* X3 = ort_.GetTensorData(input_X3); + + // Setup output + int64_t output_dim_value = 1; + OrtValue* output = ort_.KernelContext_GetOutput(context, 0, &output_dim_value, 1); + float* out = ort_.GetTensorMutableData(output); + + // Only CPU EP is supported in this kernel + for (int64_t i = 0; i < output_dim_value; i++) { + out[i] = X1[i] + (X2 != nullptr ? X2[i] : 0) + X3[i]; + } +} diff --git a/onnxruntime/test/shared_lib/utils.h b/onnxruntime/test/shared_lib/utils.h index aa934b69b6..19b8467f3f 100644 --- a/onnxruntime/test/shared_lib/utils.h +++ b/onnxruntime/test/shared_lib/utils.h @@ -74,6 +74,43 @@ struct MyCustomOpMultipleDynamicInputs : Ort::CustomOpBase { + explicit MyCustomOpWithOptionalInput(const char* provider) : provider_(provider) {} + + void* CreateKernel(Ort::CustomOpApi api, const OrtKernelInfo* info) const { return new MyCustomKernelMultipleDynamicInputs(api, info); }; + const char* GetName() const { return "FooBar"; }; + const char* GetExecutionProviderType() const { return provider_; }; + + size_t GetInputTypeCount() const { return 3; }; + ONNXTensorElementDataType GetInputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }; + OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t index) const { + // The second input (index == 1) is optional + if (index == 1) + return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_OPTIONAL; + + return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED; + } + + size_t GetOutputTypeCount() const { return 1; }; + ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }; + OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t /*index*/) const { + return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_REQUIRED; + } + private: const char* provider_; }; \ No newline at end of file diff --git a/onnxruntime/test/testdata/foo_bar_1.onnx b/onnxruntime/test/testdata/foo_bar_1.onnx new file mode 100644 index 0000000000..8b0cbc5da2 --- /dev/null +++ b/onnxruntime/test/testdata/foo_bar_1.onnx @@ -0,0 +1,22 @@ + onnx-example:l + +X1 +X2 +X3Y"FooBar +test-modelZ +X1 + + +Z +X2 + + +Z +X3 + + +b +Y + + +B \ No newline at end of file diff --git a/onnxruntime/test/testdata/foo_bar_2.onnx b/onnxruntime/test/testdata/foo_bar_2.onnx new file mode 100644 index 0000000000..449997e6eb Binary files /dev/null and b/onnxruntime/test/testdata/foo_bar_2.onnx differ