[ONNX][BugFix] Support If body with free variable from graph input (#…

…15602)

* [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.

* remove meaningless line change

* fix test_graph_input_use_in_if work on llvm test

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

tests/python/frontend/onnx/test_forward.py

@@ -5147,6 +5147,101 @@ 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, [1])
+        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",
+                    inputs=[],
+                    outputs=[
+                        helper.make_tensor_value_info(f"if_output{i-1}", 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
+        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"), 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.
+    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"""