Add handling for variadic inputs/outputs in a function. (#13140)

### Description
<!-- Describe your changes. -->
Add handling for variadic inputs/outputs in a function.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

#13121

onnxruntime/core/graph/function_utils.cc

@@ -53,7 +53,7 @@ std::unique_ptr<ONNX_NAMESPACE::OpSchema> CreateSchema(
 
   int i = 0;
   for (const auto& input : meta_def->inputs) {
-    const auto *input_arg = graph.GetNodeArg(input);
+    const auto* input_arg = graph.GetNodeArg(input);
     // inputs must have a type. can be inferred for outputs.
     ORT_ENFORCE(input_arg->Type() != nullptr);
     op_schema->Input(i, input, "",
@@ -114,17 +114,29 @@ static void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto& onnx_fun
 
   std::function<void(const ONNX_NAMESPACE::NodeProto&)> process_node = [&](const ONNX_NAMESPACE::NodeProto& node) {
     auto it = opset_imports.find(node.domain());
-    ORT_ENFORCE(it != opset_imports.end(), "No opset registered for domain " + node.domain() + " in function opset imports.");
+    ORT_ENFORCE(it != opset_imports.end(),
+                "No opset registered for domain " + node.domain() + " in function opset imports.");
     int domain_version = it->second;
-    ORT_ENFORCE(domain_version != -1, "No opset registered for domain " + node.domain() + " in function opset imports.");
-    const auto node_op_schema =
-        schema_registry->GetSchema(node.op_type(), domain_version, node.domain());
+    ORT_ENFORCE(domain_version != -1,
+                "No opset registered for domain " + node.domain() + " in function opset imports.");
+
+    const auto* node_op_schema = schema_registry->GetSchema(node.op_type(), domain_version, node.domain());
+    int variadic_arg_idx = -1;
     for (int i = 0; i < node.input_size(); ++i) {
       auto& in_name = node.input().Get(i);
       auto iter = input_name_idx_map.find(in_name);
       if (iter != input_name_idx_map.end()) {
         int idx = iter->second;
-        std::string type_str = node_op_schema ? node_op_schema->inputs().at(i).GetTypeStr() + "in" + std::to_string(idx) : "Tin" + std::to_string(idx);
+        // if we have hit a variadic arg it is the last input in the schema, so we need to use that index not i.
+        auto schema_idx = variadic_arg_idx != -1 ? variadic_arg_idx : i;
+
+        std::string type_str;
+        if (node_op_schema) {
+          type_str = node_op_schema->inputs().at(schema_idx).GetTypeStr() + "in" + std::to_string(idx);
+        } else {
+          type_str = "Tin" + std::to_string(idx);
+        }
+
         input_types_list[idx] = std::make_pair(in_name, type_str);
         if (!type_constraint_map.count(type_str)) {
           // If schema is available for the node then get the allowed types from the schema
@@ -133,7 +145,7 @@ static void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto& onnx_fun
           // the requested types.
           auto& dest_types = type_constraint_map[type_str];
           if (node_op_schema) {
-            const auto& types = node_op_schema->inputs().at(i).GetTypes();
+            const auto& types = node_op_schema->inputs().at(schema_idx).GetTypes();
             dest_types.reserve(dest_types.size() + types.size());
             for (const auto* s : types) {
               dest_types.emplace_back(*s);
@@ -145,14 +157,31 @@ static void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto& onnx_fun
             }
           }
         }
+
+        // if this is a variadic input there are no more inputs in the schema
+        if (node_op_schema && variadic_arg_idx == -1 &&
+            node_op_schema->inputs().at(schema_idx).GetOption() == OpSchema::FormalParameterOption::Variadic) {
+          variadic_arg_idx = i;
+        }
       }
     }
+
+    variadic_arg_idx = -1;
     for (int i = 0; i < node.output_size(); ++i) {
       auto& out_name = node.output().Get(i);
       auto iter = output_name_idx_map.find(out_name);
       if (iter != output_name_idx_map.end()) {
         int idx = iter->second;
-        std::string type_str = node_op_schema ? node_op_schema->outputs().at(i).GetTypeStr() + "out" + std::to_string(i) : "Tout" + std::to_string(i);
+        // if we have hit a variadic arg it is the last output in the schema, so we need to use that index.
+        auto schema_idx = variadic_arg_idx != -1 ? variadic_arg_idx : i;
+
+        std::string type_str;
+        if (node_op_schema) {
+          type_str = node_op_schema->outputs().at(schema_idx).GetTypeStr() + "out" + std::to_string(idx);
+        } else {
+          type_str = "Tout" + std::to_string(idx);
+        }
+
         output_types_list[idx] = std::make_pair(out_name, type_str);
         if (!type_constraint_map.count(type_str)) {
           // If schema is available for the node then get the allowed types from the schema
@@ -161,7 +190,7 @@ static void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto& onnx_fun
           // the requested types.
           auto& dest_types = type_constraint_map[type_str];
           if (node_op_schema) {
-            const auto& types = node_op_schema->outputs().at(i).GetTypes();
+            const auto& types = node_op_schema->outputs().at(schema_idx).GetTypes();
             dest_types.reserve(dest_types.size() + types.size());
             for (auto* data_type : types) {
               dest_types.emplace_back(*data_type);
@@ -173,6 +202,12 @@ static void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto& onnx_fun
             }
           }
         }
+
+        // if this is a variadic output there are no more outputs in the schema
+        if (node_op_schema && variadic_arg_idx == -1 &&
+            node_op_schema->outputs().at(schema_idx).GetOption() == OpSchema::FormalParameterOption::Variadic) {
+          variadic_arg_idx = i;
+        }
       }
     }
 

onnxruntime/test/framework/function_test.cc

@@ -13,7 +13,7 @@
 #include "test/test_environment.h"
 #include "test/framework/test_utils.h"
 #include "test/common/tensor_op_test_utils.h"
-
+#include "test/util/include/asserts.h"
 
 // Unit tests to check the implementation of functions, model-local functions,
 // function-inlining etc.
@@ -339,5 +339,18 @@ TEST(FunctionTest, NestedConstant) {
   Check(code, "x", {1.0, 2.0, 3.0}, "y", {1.0, 2.0, 3.0, 0.0, 0.0, 0.0});
 }
 
+// GH13121. Model with function body that has variadic inputs (or outputs) was not loading.
+// Add handling for variadics to IOTypeConstraintHelper. Test model has a Concat and Split to test both variadic
+// inputs and outputs.
+TEST(FunctionTest, Variadics) {
+  Status status;
+  auto model_uri = ORT_TSTR("testdata/function_with_variadics.onnx");
+
+  SessionOptions so;
+  so.session_logid = "FunctionTest.Variadics";
+  InferenceSession session_object{so, GetEnvironment()};
+  ASSERT_STATUS_OK(session_object.Load(model_uri));
+  ASSERT_STATUS_OK(session_object.Initialize());
+}
 }  // namespace test
 }  // namespace onnxruntime