[PT2] Detect missed inputs for noed based on TensorMeta in arguments (#3360)

### Changes

For new tracing detect missed inputs based on count of TensorMeta
instance in op inputs.

### Reason for changes

Failed on detection missed input for concat nodes.

Author: AlexanderDokuchaev

nncf/torch/graph/graph.py

@@ -9,13 +9,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from itertools import chain
 from typing import Dict, List, Tuple
 
 import nncf
 from nncf.common.graph import NNCFGraph
 from nncf.common.graph import NNCFNode
 from nncf.common.graph import NNCFNodeName
 from nncf.common.graph.layer_attributes import MultipleInputLayerAttributes
+from nncf.experimental.common.check_feature import is_experimental_torch_tracing_enabled
+from nncf.experimental.torch2.function_hook.graph.graph_utils import TensorMeta
+from nncf.experimental.torch2.function_hook.nncf_graph.layer_attributes import PT2OpLayerAttributes
 from nncf.torch.dynamic_graph.scope import Scope
 from nncf.torch.graph.transformations.commands import PTTargetPoint
 
@@ -102,16 +106,33 @@ def get_nodes_with_missed_input_edges(self) -> List[NNCFNode]:
         :return: List of NNCFNodes that are identified as disconnected.
         """
         input_nodes = set()
-        for node in self.get_all_nodes():
-            num_expected_input_edges = None
-            if hasattr(node.metatype, "num_expected_input_edges"):
-                num_expected_input_edges = node.metatype.num_expected_input_edges
-            if node.layer_attributes is not None and isinstance(node.layer_attributes, MultipleInputLayerAttributes):
-                num_expected_input_edges = node.layer_attributes.num_inputs
-            if num_expected_input_edges:
-                input_edges = self.get_input_edges(node)
-                if len(input_edges) < num_expected_input_edges:
-                    # If node has missed input edges we assume this node is an input node
-                    # that was disconnected from an activation input.
+        if is_experimental_torch_tracing_enabled():
+            # Check expected number of input edges by counting TensorMeta in op_args and op_kwargs.
+            for node in self.get_all_nodes():
+                input_edges = len(self.get_input_edges(node))
+                if not isinstance(node.layer_attributes, PT2OpLayerAttributes):
+                    continue
+                num_expected_input_edges = 0
+                for val in chain(node.layer_attributes.op_args, node.layer_attributes.op_kwargs.values()):
+                    if isinstance(val, TensorMeta):
+                        num_expected_input_edges += 1
+                    if isinstance(val, (list, tuple)):
+                        num_expected_input_edges += sum(isinstance(v, TensorMeta) for v in val)
+                if input_edges < num_expected_input_edges:
                     input_nodes.add(node)
+        else:
+            for node in self.get_all_nodes():
+                num_expected_input_edges = None
+                if hasattr(node.metatype, "num_expected_input_edges"):
+                    num_expected_input_edges = node.metatype.num_expected_input_edges
+                if node.layer_attributes is not None and isinstance(
+                    node.layer_attributes, MultipleInputLayerAttributes
+                ):
+                    num_expected_input_edges = node.layer_attributes.num_inputs
+                if num_expected_input_edges:
+                    input_edges = self.get_input_edges(node)
+                    if len(input_edges) < num_expected_input_edges:
+                        # If node has missed input edges we assume this node is an input node
+                        # that was disconnected from an activation input.
+                        input_nodes.add(node)
         return list(input_nodes)

tests/torch2/function_hook/nncf_graph/test_nncf_graph.py

@@ -11,11 +11,12 @@
 
 from dataclasses import dataclass
 from functools import partial
-from typing import List, Tuple, Union
+from typing import Callable, List, Tuple, Union
 
 import networkx as nx
 import pytest
 import torch
+import torch.nn.functional as F
 import torchvision.models as models
 
 from nncf.common.graph.layer_attributes import Dtype
@@ -24,13 +25,18 @@
 from nncf.experimental.torch2.function_hook.graph.graph_utils import FunctionMeta
 from nncf.experimental.torch2.function_hook.graph.graph_utils import InOutMeta
 from nncf.experimental.torch2.function_hook.graph.graph_utils import NodeType
+from nncf.experimental.torch2.function_hook.graph.graph_utils import TensorMeta
+from nncf.experimental.torch2.function_hook.nncf_graph.layer_attributes import PT2OpLayerAttributes
 from nncf.experimental.torch2.function_hook.nncf_graph.nncf_graph_builder import build_nncf_graph
 from nncf.experimental.torch2.function_hook.nncf_graph.nncf_graph_builder import convert_to_nncf_graph
 from nncf.experimental.torch2.function_hook.nncf_graph.nncf_graph_builder import get_dtype
 from nncf.experimental.torch2.function_hook.nncf_graph.nncf_graph_builder import get_name_of_node
 from nncf.experimental.torch2.function_hook.nncf_graph.nncf_graph_builder import get_node_type
 from nncf.experimental.torch2.function_hook.wrapper import register_post_function_hook
 from nncf.experimental.torch2.function_hook.wrapper import wrap_model
+from nncf.torch.graph.graph import PTNNCFGraph
+from nncf.torch.graph.operator_metatypes import PTCatMetatype
+from nncf.torch.graph.operator_metatypes import PTConv2dMetatype
 from tests.cross_fw.shared.paths import TEST_ROOT
 from tests.torch2.function_hook import helpers
 from tests.torch2.utils import compare_with_reference_file
@@ -155,3 +161,81 @@ def test_model_graph_with_shared_parameters(regen_ref_data):
     nx_nncf_graph = nx.nx_pydot.to_pydot(graph)
     ref_file = REF_DIR / "model_graph_with_shared_parameters.dot"
     compare_with_reference_file(str(nx_nncf_graph), ref_file, regen_ref_data)
+
+
+def _missed_input_edge_for_conv() -> PTNNCFGraph:
+    graph = PTNNCFGraph()
+    graph.add_nncf_node(
+        node_name="conv",
+        node_type="conv",
+        node_metatype=PTConv2dMetatype,
+        layer_attributes=PT2OpLayerAttributes(
+            func=F.conv2d,
+            op_args=(
+                TensorMeta(shape=(1,), dtype=torch.float),
+                TensorMeta(shape=(1,), dtype=torch.float),
+            ),
+            op_kwargs={},
+            constant_port_ids=set([1]),
+        ),
+    )
+    return graph
+
+
+def _missed_input_edge_for_concat() -> PTNNCFGraph:
+    graph = PTNNCFGraph()
+    graph.add_nncf_node(
+        node_name="concat",
+        node_type="concat",
+        node_metatype=PTCatMetatype,
+        layer_attributes=PT2OpLayerAttributes(
+            func=torch.concat,
+            op_args=(
+                [
+                    TensorMeta(shape=(1,), dtype=torch.float),
+                    TensorMeta(shape=(1,), dtype=torch.float),
+                ]
+            ),
+            op_kwargs={},
+            constant_port_ids=set(),
+        ),
+    )
+    return graph
+
+
+def _no_missed_input_edge_for_conv() -> PTNNCFGraph:
+    graph = PTNNCFGraph()
+    node_input = graph.add_nncf_node("input", "input", None, None)
+    node_weight = graph.add_nncf_node("weight", "weight", None, None)
+    node_conv = graph.add_nncf_node(
+        node_name="conv",
+        node_type="conv",
+        node_metatype=PTConv2dMetatype,
+        layer_attributes=PT2OpLayerAttributes(
+            func=F.conv2d,
+            op_args=(
+                TensorMeta(shape=(1, 1, 1, 1), dtype=torch.float),
+                TensorMeta(shape=(1, 1, 1, 1), dtype=torch.float),
+            ),
+            op_kwargs={},
+            constant_port_ids=set([1]),
+        ),
+    )
+    graph.add_edge_between_nncf_nodes(node_input.node_id, node_conv.node_id, (1,), 0, 0, Dtype.FLOAT)
+    graph.add_edge_between_nncf_nodes(node_weight.node_id, node_conv.node_id, (1,), 1, 0, Dtype.FLOAT)
+    return graph
+
+
+@pytest.mark.parametrize(
+    "graph_builder, ref",
+    (
+        (_missed_input_edge_for_conv, ["conv"]),
+        (_missed_input_edge_for_concat, ["concat"]),
+        (_no_missed_input_edge_for_conv, []),
+    ),
+)
+def test_get_nodes_with_missed_input_edges(graph_builder: Callable[[], PTNNCFGraph], ref: List[str]):
+    graph = graph_builder()
+    ret = graph.get_nodes_with_missed_input_edges()
+    ret_names = [node.node_name for node in ret]
+    assert ret_names == ref