Support optional inputs/outputs in custom op development (#6727)

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -7,7 +7,7 @@
 #include <string.h>
 
 // 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
 }

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -623,10 +623,23 @@ struct CustomOpBase : OrtCustomOp {
 
     OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) { static_cast<TKernel*>(op_kernel)->Compute(context); };
     OrtCustomOp::KernelDestroy = [](void* op_kernel) { delete static_cast<TKernel*>(op_kernel); };
+
+    OrtCustomOp::GetInputCharacteristic = [](const OrtCustomOp* this_, size_t index) { return static_cast<const TOp*>(this_)->GetInputCharacteristic(index); };
+    OrtCustomOp::GetOutputCharacteristic = [](const OrtCustomOp* this_, size_t index) { return static_cast<const TOp*>(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

onnxruntime/core/session/custom_ops.cc

@@ -125,18 +125,35 @@ common::Status CreateCustomRegistry(const std::vector<OrtCustomOpDomain*>& 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<OrtCustomOpDomain*>& 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);
         }
       }
 

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);

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<SliceCustomOp, SliceCustomOpKernel> {
       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::Value> ort_inputs;
+
+  Ort::MemoryInfo info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault);
+
+  // input 0
+  std::vector<float> input_0_data = {1.f};
+  std::vector<int64_t> input_0_dims = {1};
+  ort_inputs.emplace_back(
+      Ort::Value::CreateTensor<float>(info, const_cast<float*>(input_0_data.data()),
+                                      input_0_data.size(), input_0_dims.data(), input_0_dims.size()));
+
+  // input 1
+  std::vector<float> input_1_data = {1.f};
+  std::vector<int64_t> input_1_dims = {1};
+  ort_inputs.emplace_back(
+      Ort::Value::CreateTensor<float>(info, const_cast<float*>(input_1_data.data()),
+                                      input_1_data.size(), input_1_dims.data(), input_1_dims.size()));
+
+  // input 2
+  std::vector<float> input_2_data = {1.f};
+  std::vector<int64_t> input_2_dims = {1};
+  ort_inputs.emplace_back(
+      Ort::Value::CreateTensor<float>(info, const_cast<float*>(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<const char*> 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<int64_t> y_dims = {1};
+    std::vector<float> 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<float>();
+    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<const char*> 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<int64_t> y_dims = {1};
+    std::vector<float> 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<float>();
+    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) {

onnxruntime/test/shared_lib/utils.cc

@@ -64,4 +64,26 @@ void MyCustomKernelMultipleDynamicInputs::Compute(OrtKernelContext* context) {
     out[i] = static_cast<float>(X[i] + Y[i]);
   }
 #endif
-}
+}
+
+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<float>(input_X1);
+  // The second input may or may not be present
+  const float* X2 = (input_X2 != nullptr) ? ort_.GetTensorData<float>(input_X2) : nullptr;
+  const float* X3 = ort_.GetTensorData<float>(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<float>(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];
+  }
+}

onnxruntime/test/shared_lib/utils.h

@@ -74,6 +74,43 @@ struct MyCustomOpMultipleDynamicInputs : Ort::CustomOpBase<MyCustomOpMultipleDyn
   size_t GetOutputTypeCount() const { return 1; };
   ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; };
 
+ private:
+  const char* provider_;
+};
+
+struct MyCustomKernelWithOptionalInput {
+  MyCustomKernelWithOptionalInput(Ort::CustomOpApi ort, const OrtKernelInfo* /*info*/) : ort_(ort) {
+  }
+
+  void Compute(OrtKernelContext* context);
+
+ private:
+  Ort::CustomOpApi ort_;
+};
+
+struct MyCustomOpWithOptionalInput : Ort::CustomOpBase<MyCustomOpWithOptionalInput, MyCustomKernelWithOptionalInput> {
+  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_;
 };

onnxruntime/test/testdata/foo_bar_1.onnx

@@ -0,0 +1,22 @@
+onnx-example:l
+
+X1
+X2
+X3
Y"FooBar
+test-modelZ
+X1
+
+

+
Z
+X2
+
+

+
Z
+X3
+
+

+
b
+
Y
+
+

+
B