diff --git a/onnxruntime/core/graph/function_utils.cc b/onnxruntime/core/graph/function_utils.cc index eb7655cb65..c16d10851f 100644 --- a/onnxruntime/core/graph/function_utils.cc +++ b/onnxruntime/core/graph/function_utils.cc @@ -53,7 +53,7 @@ std::unique_ptr 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 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; + } } } diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index e65aa01b11..2fdacd46d3 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/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 diff --git a/onnxruntime/test/testdata/function_with_variadics.onnx b/onnxruntime/test/testdata/function_with_variadics.onnx new file mode 100644 index 0000000000..4a1b7056f1 Binary files /dev/null and b/onnxruntime/test/testdata/function_with_variadics.onnx differ