skottmckay/bugfix/SubgraphInput (#4004)

Description:
Fix 2 edge cases as described here: #3755 (comment)

Create a NodeArg for subgraph inputs even if they have no type. If they are only used as an implicit input to another level of nested subgraph we will not create a NodeArg via any other path

Allow an If output to have no shape. Obscure edge case where a loop carried dependency to a Loop node is passed through a nested If node subgraph (i.e. the Loop subgraph contains an If node with a nested subgraph for the else_branch/then_branch). We can't infer a shape for a loop carried dependency (they may change across iterations), which means we can't infer a shape for the nested If subgraph output either. We have delayed allocation support for If outputs so use that.

Motivation and Context
#3755

onnxruntime/core/graph/graph.cc

@@ -208,8 +208,8 @@ void NodeArg::ClearShape() {
   }
 }
 
-common::Status NodeArg::UpdateTypeAndShape(const ONNX_NAMESPACE::TypeProto& input_type, bool strict, bool override_types,
-                                           const logging::Logger& logger) {
+common::Status NodeArg::UpdateTypeAndShape(const ONNX_NAMESPACE::TypeProto& input_type, bool strict,
+                                           bool override_types, const logging::Logger& logger) {
   if (!utils::HasType(node_arg_info_)) {
     *node_arg_info_.mutable_type() = input_type;
     type_ = DataTypeUtils::ToType(node_arg_info_.type());
@@ -784,9 +784,17 @@ Graph::Graph(const Model& owning_model,
   // process graph inputs first as we want the type/shape from them to be preferred if a graph input
   // has a matching initializer
   for (auto& graph_input : graph_proto_->input()) {
-    if (utils::HasName(graph_input) && utils::HasType(graph_input)) {
-      name_to_type_map[graph_input.name()] = graph_input.type();
-      GetOrCreateNodeArg(graph_input.name(), &graph_input.type());
+    if (utils::HasName(graph_input)) {
+      if (utils::HasType(graph_input)) {
+        name_to_type_map[graph_input.name()] = graph_input.type();
+        GetOrCreateNodeArg(graph_input.name(), &graph_input.type());
+      } else {
+        // subgraph inputs can have type inferred later. need to create a NodeArg in case this input is only used in
+        // a nested subgraph (a NodeArg won't be added by AddNode for the nodes in this subgraph)
+        if (IsSubgraph()) {
+          GetOrCreateNodeArg(graph_input.name(), nullptr);
+        }
+      }
     }
   }
 
@@ -813,7 +821,7 @@ Graph::Graph(const Model& owning_model,
       } else {
         LOGS(logger_, WARNING) << "Initializer " << tensor.name()
                                << " appears in graph inputs and will not be treated as constant value/weight. "
-                               << "This may fail some of the graph optimizations, like const folding. "
+                               << "This may prevent some of the graph optimizations, like const folding. "
                                << "Move it out of graph inputs if there is no need to override it, "
                                << "by either re-generating the model with latest exporter/converter "
                                << "or with the tool onnxruntime/tools/python/remove_initializer_from_input.py.";
@@ -880,7 +888,7 @@ void Graph::InitializeStateFromModelFileGraphProto() {
   for (auto& graph_input : graph_proto_->input()) {
     auto& name = graph_input.name();
     const auto* node_arg = GetNodeArg(name);
-    ORT_ENFORCE(node_arg, "Graph ctor should have created NodeArg for initializer.");
+    ORT_ENFORCE(node_arg, "Graph ctor should have created NodeArg for initializer. Missing:", name);
     graph_inputs.insert({name, node_arg});
     graph_inputs_including_initializers_.push_back(node_arg);
     if (graph_initializers.end() == graph_initializers.find(name)) {

onnxruntime/core/providers/cpu/controlflow/if.cc

@@ -248,24 +248,27 @@ Status IfImpl::AllocateOutputTensors() {
 
   for (auto& graph_output : info_.subgraph.GetOutputs()) {
     auto* graph_output_shape = graph_output->Shape();
-    if (!graph_output_shape) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Subgraph must have the shape set for all outputs but ",
-                             graph_output->Name(), " did not.");
+    bool symbolic_dim_in_shape = false;
+
+    if (graph_output_shape) {
+      TensorShape output_shape = onnxruntime::utils::GetTensorShapeFromTensorShapeProto(*graph_output_shape);
+
+      // if size < 0 we have a symbolic dimension and need to use a temporary OrtValue in the subgraph execution
+      if (output_shape.Size() < 0) {
+        symbolic_dim_in_shape = true;
+      } else {
+        auto* tensor = context_.Output(index, output_shape);
+
+        if (!tensor)
+          return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to create output tensor for ", graph_output->Name());
+
+        outputs_.push_back({AllocationType::IfOutput, *context_.GetOutputMLValue(index)});
+      }
     }
 
-    TensorShape output_shape = onnxruntime::utils::GetTensorShapeFromTensorShapeProto(*graph_output_shape);
-
-    // if size < 0 we have a symbolic dimension and need to use a temporary OrtValue in the subgraph execution
-    if (output_shape.Size() < 0) {
+    if (!graph_output_shape || symbolic_dim_in_shape) {
       // we still need a value to put in the feeds we give to the execution frame, so just use an empty MLValue
       outputs_.push_back({AllocationType::Delayed, {}});
-    } else {
-      auto* tensor = context_.Output(index, output_shape);
-
-      if (!tensor)
-        return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to create output tensor for ", graph_output->Name());
-
-      outputs_.emplace_back(AllocationType::IfOutput, *context_.GetOutputMLValue(index));
     }
 
     ++index;

onnxruntime/test/providers/cpu/controlflow/loop_test.cc

@@ -575,7 +575,7 @@ TEST(Loop, InfiniteLoopTermination) {
 }
 
 // Add basic test to trigger types override logic in Graph::InferAndVerifySubgraphTypes as well as
-// type/shape inferencing for subgraph to flow the type/shape info through 
+// type/shape inferencing for subgraph to flow the type/shape info through
 // subgraph.PerformTypeAndShapeInferencing(options).
 // In this test, main graph has original input/expected output defined as "double" where the subgraph as "float".
 // Expectation is types should get propagated properly in subgraph and yield correct output
@@ -824,6 +824,111 @@ TEST(Loop, Opset11WithNoVariadicInputsAndOutputs) {
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
 }
 
+// Test a combination of things:
+// Subgraph input for loop state var has no type and is not used in the Loop subgraph (used in nested If subgraph)
+// Loop subgraph calls an If where the loop state var is an implicit input so it has no shape due to a loop state
+// var being able to change shape on each iteration.
+TEST(Loop, PassThroughSubgraphInputNoTypeOrShape) {
+  // both subgraphs of the If just pass through an outer scope value
+  auto create_if_subgraph = [](bool is_then) {
+    Model model("if_branch_subgraph", true, DefaultLoggingManager().DefaultLogger());
+    auto& graph = model.MainGraph();
+
+    auto& outer_scope_0 = graph.GetOrCreateNodeArg("loop_state_var", nullptr);
+    graph.AddOuterScopeNodeArg("loop_state_var");
+
+    TypeProto float_tensor;
+    float_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
+
+    auto& if_out = graph.GetOrCreateNodeArg(is_then ? "if_then_out" : "if_else_out", &float_tensor);
+    graph.AddNode("if_out", "Identity", "pass through", {&outer_scope_0}, {&if_out});
+
+    auto status = graph.Resolve();
+    // Resolve will have actually errored out as we don't have type info for the input. That's valid for a subgraph
+    // but not for a main graph but the Resolve doesn't know that it's handling a subgraph. We could add a way to
+    // tell it but generally it's only our unit tests creating a graph this way.
+    // The GraphProto will still be correct.
+    EXPECT_NE(status, Status::OK());
+
+    return graph.ToGraphProto();
+  };
+
+  auto create_subgraph = [&create_if_subgraph]() {
+    Model model("loop_subgraph", true, DefaultLoggingManager().DefaultLogger());
+    auto& graph = model.MainGraph();
+
+    std::vector<NodeArg*> inputs;
+    std::vector<NodeArg*> outputs;
+
+    /*  Inputs: iter_num, cond_in, loop carried state variables.
+
+         iter_num_in    cond_in     [loop_state_var]
+           (unused)        |               |
+                       [Identity]         [If]   (both branches in If return loop_state_var via Identity node)
+                           |               |
+                        cond_out     loop_var_0_out
+    */
+
+    // graph inputs types.
+    TypeProto int64_scalar;
+    int64_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_INT64);
+    int64_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1);
+
+    TypeProto bool_scalar;
+    bool_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL);
+    bool_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1);
+
+    // graph inputs
+    auto& iter_num_in = graph.GetOrCreateNodeArg("iter_num_in", &int64_scalar);
+    auto& cond_in = graph.GetOrCreateNodeArg("cond_in", &bool_scalar);
+
+    auto& loop_state_var = graph.GetOrCreateNodeArg("loop_state_var", nullptr);
+
+    // graph outputs
+    auto& cond_out = graph.GetOrCreateNodeArg("cond_out", &bool_scalar);
+    auto& loop_var_0_out = graph.GetOrCreateNodeArg("loop_var_0_out", nullptr);
+
+    // cond_in -> cond_out
+    {
+      inputs = {&cond_in};
+      outputs = {&cond_out};
+
+      graph.AddNode("cond_in_identity", "Identity", "Forward cond_in to cond_out", inputs, outputs);
+    }
+
+    // loop_state_var -> If(cond_in) -> loop_var_0_out
+    {
+      inputs = {&cond_in};
+      outputs = {&loop_var_0_out};
+
+      auto& node = graph.AddNode("loop_var_out", "If", "If with loop_state_var as implicit_input", inputs, outputs);
+      node.AddAttribute("then_branch", create_if_subgraph(true));
+      node.AddAttribute("else_branch", create_if_subgraph(false));
+    }
+
+    graph.SetInputs({&iter_num_in, &cond_in, &loop_state_var});
+    graph.SetOutputs({&cond_out, &loop_var_0_out});
+
+    auto status = graph.Resolve();
+    // same reason this will fail as in create_if_subgraph - still creates a valid subgraph GraphProto
+    EXPECT_NE(status, Status::OK());
+
+    return graph.ToGraphProto();
+  };
+
+  OpTester test("Loop", 11);
+  auto body = create_subgraph();
+  test.AddAttribute<GraphProto>("body", body);
+  test.AddInput<int64_t>("M", {1}, {1});
+  test.AddInput<bool>("cond", {1}, {true});
+  test.AddInput<float>("initial_value", {1}, {123.f});
+
+  test.AddOutput<float>("loop_var_0_final", {1}, {123.f});
+
+  // Disable TensorRT on unsupported data type BOOL
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
+}
+
 #ifdef USE_CUDA
 // test that when part of the subgraph run on CUDA it executes successfully
 TEST(Loop, MixedExecutionProviders) {