From 63f029fe4d399cab81da91c166432757209195c5 Mon Sep 17 00:00:00 2001 From: gangmul12 Date: Mon, 21 Aug 2023 11:28:58 -0400 Subject: [PATCH 1/3] [ONNX][BugFix] Support If body with free variable from graph input When graph inputs are used in an inner body of If node, the original TVM ONNX frontend did not set the span properly. Because of the wrong or partial span, new relay.Var is introduced and failed to match identical Var. Firstly, there was an issue where the free variable of the inner body was updated in _node but not applied to _input. Secondly, although the free variable of the then body successfully updated to relay.Var in _node, but this was obscured by the update of _node in the else body. This commit fixes the ONNX importer and adds an ONNX import testcase for the revised code. --- python/tvm/relay/frontend/onnx.py | 10 ++- tests/python/frontend/onnx/test_forward.py | 88 ++++++++++++++++++++++ 2 files changed, 95 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 13609704ccb7..1e6c9fbd8bef 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -4510,14 +4510,19 @@ def _impl_v1(cls, inputs, attr, params): # Add constants from both branches to parent graph. graph_scope._params.update(then_graph._params) graph_scope._nodes.update(then_graph._nodes) + graph_scope._params.update(else_graph._params) + graph_scope._nodes.update(else_graph._nodes) + then_free_vars = analysis.free_vars(then_expr) for var in then_free_vars: graph_scope._nodes.update({var.name_hint: var}) - graph_scope._params.update(else_graph._params) - graph_scope._nodes.update(else_graph._nodes) + if var.name_hint in graph_scope._inputs: + graph_scope._inputs.update({var.name_hint: var}) else_free_vars = analysis.free_vars(else_expr) for var in else_free_vars: graph_scope._nodes.update({var.name_hint: var}) + if var.name_hint in graph_scope._inputs: + graph_scope._inputs.update({var.name_hint: var}) # Sometimes pytorch to onnx will insert silly if statements that produce dynamic ranks. # Often these dont contribute anything. If we see a dynamic rank output, try to unify @@ -6954,7 +6959,6 @@ def _construct_nodes(self, graph): op = fold_constant(op) else: op = _expr.TupleWrapper(fold_constant(op.astuple()), len(op)) - op = set_span(op, node_source_name) if outputs_num > 1: diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 216732343028..ac17b0c81f87 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -5147,6 +5147,94 @@ def append_constant_nodes(nodes, outputs, expected, name): verify_if(cond_array=True, num_outputs=2) +@tvm.testing.parametrize_targets +def test_graph_input_use_in_if(target, dev): + """test_graph_input_use_in_if""" + + def verify_if(num_nested, cond): + # return "graph input" if cond is True, else return constant(-1). + + input_tensor = helper.make_tensor_value_info("graph_input", TensorProto.FLOAT, [1]) + output_tensor = helper.make_tensor_value_info("graph_output", TensorProto.FLOAT, [1]) + constant_node = make_constant_node("const_val", TensorProto.FLOAT, [1], [-1]) + cond_tensor = helper.make_tensor_value_info("cond", TensorProto.BOOL, []) + inner_if_node = None + for i in range(num_nested): + identity_node = helper.make_node( + "Identity", + inputs=["const_val"], + outputs=[f"const{i}"], + name=f"depth{i}'th else identity", + ) + else_branch = helper.make_graph( + [identity_node], + f"else{i}_body", + inputs=[], + outputs=[helper.make_tensor_value_info(f"const{i}", TensorProto.FLOAT, [1])], + ) + out_name = f"if_output{i}" if i != (num_nested - 1) else "graph_output" + + if i == 0: + identity_node = helper.make_node( + "Identity", + inputs=["graph_input"], + outputs=[f"input_identity{i}"], + name=f"depth{i}'th then identity", + ) + then_branch = helper.make_graph( + [identity_node], + f"then{i}_body", + inputs=[], + outputs=[ + helper.make_tensor_value_info(f"input_identity{i}", TensorProto.FLOAT, [1]) + ], + ) + if_node = helper.make_node( + "If", + inputs=["cond"], + outputs=[out_name], + then_branch=then_branch, + else_branch=else_branch, + name=f"depth{i}'s If node", + ) + inner_if_node = if_node + else: + then_branch = helper.make_graph( + [inner_if_node], f"then{i}_body", [], [f"if_output{i-1}"] + ) + if_node = helper.make_node( + "If", + inputs=["cond"], + outputs=[out_name], + then_branch=then_branch, + else_branch=else_branch, + name=f"depth{i}'s If node", + ) + inner_if_node = if_node + graph_nodes = [constant_node, inner_if_node] + graph = helper.make_graph( + graph_nodes, + "input_use_in_if_test", + inputs=[input_tensor, cond_tensor], + outputs=[output_tensor], + ) + model = helper.make_model(graph, producer_name="input_use_in_if_test") + + verify_with_ort_with_inputs( + model, + [np.array([3.0], dtype="float32"), [cond]], + dtype="float32", + use_vm=True, + opset=14, + ) + + # Confirm that if works with cond as an array or scalar. + verify_if(num_nested=1, cond=True) + verify_if(num_nested=1, cond=False) + verify_if(num_nested=2, cond=True) + verify_if(num_nested=2, cond=False) + + @tvm.testing.parametrize_targets def test_size(target, dev): """test_size""" From c5b028fa6d275f4e0841c0fd5e962edb006508c3 Mon Sep 17 00:00:00 2001 From: gangmul12 Date: Mon, 21 Aug 2023 11:59:37 -0400 Subject: [PATCH 2/3] remove meaningless line change --- python/tvm/relay/frontend/onnx.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 1e6c9fbd8bef..9934d4f13269 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -6959,6 +6959,7 @@ def _construct_nodes(self, graph): op = fold_constant(op) else: op = _expr.TupleWrapper(fold_constant(op.astuple()), len(op)) + op = set_span(op, node_source_name) if outputs_num > 1: From 21c6a079fc22d15f4f6233df192ef599cac4a064 Mon Sep 17 00:00:00 2001 From: gangmul12 Date: Tue, 22 Aug 2023 11:26:21 -0400 Subject: [PATCH 3/3] fix test_graph_input_use_in_if work on llvm test --- tests/python/frontend/onnx/test_forward.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index ac17b0c81f87..b9f2d14b7888 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -5157,7 +5157,7 @@ def verify_if(num_nested, cond): input_tensor = helper.make_tensor_value_info("graph_input", TensorProto.FLOAT, [1]) output_tensor = helper.make_tensor_value_info("graph_output", TensorProto.FLOAT, [1]) constant_node = make_constant_node("const_val", TensorProto.FLOAT, [1], [-1]) - cond_tensor = helper.make_tensor_value_info("cond", TensorProto.BOOL, []) + cond_tensor = helper.make_tensor_value_info("cond", TensorProto.BOOL, [1]) inner_if_node = None for i in range(num_nested): identity_node = helper.make_node( @@ -5200,7 +5200,12 @@ def verify_if(num_nested, cond): inner_if_node = if_node else: then_branch = helper.make_graph( - [inner_if_node], f"then{i}_body", [], [f"if_output{i-1}"] + [inner_if_node], + f"then{i}_body", + inputs=[], + outputs=[ + helper.make_tensor_value_info(f"if_output{i-1}", TensorProto.FLOAT, [1]) + ], ) if_node = helper.make_node( "If", @@ -5222,10 +5227,12 @@ def verify_if(num_nested, cond): verify_with_ort_with_inputs( model, - [np.array([3.0], dtype="float32"), [cond]], + [np.array([3.0], dtype="float32"), np.array([cond])], dtype="float32", use_vm=True, opset=14, + target=target, + dev=dev, ) # Confirm that if works with cond as an array or scalar.