From 2a96be83f629bfbf2aa5c3937e9c6e86e050bdad Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Fri, 29 May 2020 14:48:07 +1000 Subject: [PATCH] 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 | 22 ++-- .../core/providers/cpu/controlflow/if.cc | 31 ++--- .../providers/cpu/controlflow/loop_test.cc | 107 +++++++++++++++++- 3 files changed, 138 insertions(+), 22 deletions(-) diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index eac2e7e9ec..9b144b89fe 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/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)) { diff --git a/onnxruntime/core/providers/cpu/controlflow/if.cc b/onnxruntime/core/providers/cpu/controlflow/if.cc index 01d2675e2d..e3937deb73 100644 --- a/onnxruntime/core/providers/cpu/controlflow/if.cc +++ b/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; diff --git a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc index 5eab21c501..41d29d97aa 100644 --- a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc +++ b/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 inputs; + std::vector 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("body", body); + test.AddInput("M", {1}, {1}); + test.AddInput("cond", {1}, {true}); + test.AddInput("initial_value", {1}, {123.f}); + + test.AddOutput("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) {