[NNCFGraph] Fix get_input_edges for parallel edges (#2107)

### Changes

* `parallel_input_port_ids` parameter plugged in to `OVNNCFGraph`
* `NNCFGraph.get_input_edges` unwraps `parallel_input_port_ids` and
returns separate entities of `NNCFGraphEdge` for each parallel edge

### Reason for changes

* To make possible to insert statistics collectors on parallel edges in
OV backend

### Related tickets

119666

### Tests

* tests/common/graph/test_nncf_graph.py is updated
* tests/openvino/native/test_nncf_graph_builder.py is updated

nncf/common/graph/graph.py

@@ -148,17 +148,23 @@ def __init__(
         self.parallel_input_port_ids = parallel_input_port_ids
 
     def __str__(self):
-        return str(self.from_node) + " -> " + str(self.tensor_shape) + " -> " + str(self.to_node)
+        return f"{self.from_node}:{self.output_port_id} -> {self.tensor_shape} -> {self.to_node}:{self.input_port_id}"
 
     def __hash__(self):
-        return hash(str(self))
+        return hash(
+            (
+                self.from_node,
+                self.to_node,
+                self.input_port_id,
+                self.output_port_id,
+                tuple(self.tensor_shape),
+                self.dtype,
+                tuple(self.parallel_input_port_ids),
+            )
+        )
 
     def __eq__(self, other):
-        return (
-            self.from_node == other.from_node
-            and self.to_node == other.to_node
-            and self.tensor_shape == other.tensor_shape
-        )
+        return isinstance(other, NNCFGraphEdge) and self.__dict__ == other.__dict__
 
 
 class NNCFGraphPatternIO:
@@ -331,7 +337,9 @@ def get_input_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]:
         :return: List of input edges for the node sorted by input port ID.
         """
         input_nodes = self.get_previous_nodes(node)
-        edges = [self.get_edge(from_node, node) for from_node in input_nodes]
+        edges = []
+        for from_node in input_nodes:
+            edges.extend(self._get_edges(from_node, node))
         return sorted(edges, key=lambda x: x.input_port_id)
 
     def get_output_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]:
@@ -343,9 +351,31 @@ def get_output_edges(self, node: NNCFNode) -> List[NNCFGraphEdge]:
         """
 
         output_nodes = self.get_next_nodes(node)
-        edges = [self.get_edge(node, to_node) for to_node in output_nodes]
+        edges = []
+        for to_node in output_nodes:
+            edges.extend(self._get_edges(node, to_node))
         return sorted(edges, key=lambda x: x.output_port_id)
 
+    def _get_edges(self, from_node: NNCFNode, to_node: NNCFNode) -> List[NNCFGraphEdge]:
+        edges = []
+        edge = self.get_edge(from_node, to_node)
+        parallel_input_port_ids = edge.parallel_input_port_ids
+        edge.parallel_input_port_ids = []
+        edges.append(edge)
+        for input_port_id in parallel_input_port_ids:
+            edges.append(
+                NNCFGraphEdge(
+                    from_node=edge.from_node,
+                    to_node=edge.to_node,
+                    input_port_id=input_port_id,
+                    output_port_id=edge.output_port_id,
+                    tensor_shape=edge.tensor_shape,
+                    dtype=edge.dtype,
+                    parallel_input_port_ids=[],
+                )
+            )
+        return edges
+
     def traverse_graph(
         self,
         curr_node: NNCFNode,

nncf/experimental/common/tensor_statistics/collectors.py

@@ -366,7 +366,7 @@ def get_tensor_collector_inputs(
 
         :param outputs: Target model outputs.
         :param output_info: Output info collected by a `TensorCollector.get_output_info` method.
-        :returns: Model outputs in a format required by `TensorCollector.register_input` method.
+        :returns: Model outputs in a format required by `TensorCollector.register_inputs` method.
         """
         target_inputs = {}
         for reducer, names in output_info:

nncf/openvino/graph/nncf_graph_builder.py

@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import defaultdict
 from typing import List, Type
 
 import openvino.runtime as ov
@@ -88,19 +89,28 @@ def _add_edges_to_nncf_graph(model: ov.Model, graph: NNCFGraph) -> None:
         for op in model.get_ops():
             in_node_id = graph.get_node_by_name(op.get_friendly_name()).node_id
             for output_port_id, out in enumerate(op.outputs()):
+                node_vs_target_inputs = defaultdict(list)
                 for inp in out.get_target_inputs():
-                    out_node = inp.get_node()
+                    node_vs_target_inputs[inp.get_node()].append(inp)
+
+                for out_node, inputs in node_vs_target_inputs.items():
                     tensor_shape = list(out.partial_shape.get_max_shape())
                     output_node_id = graph.get_node_by_name(out_node.get_friendly_name()).node_id
                     ov_dtype = out.get_element_type().get_type_name()
                     nncf_dtype = GraphConverter.convert_to_nncf_dtype(ov_dtype)
+
+                    parallel_inputs = None
+                    if len(inputs) > 1:
+                        parallel_inputs = [inp.get_index() for inp in inputs[1:]]
+
                     graph.add_edge_between_nncf_nodes(
                         from_node_id=in_node_id,
                         to_node_id=output_node_id,
                         tensor_shape=tensor_shape,
-                        input_port_id=inp.get_index(),
+                        input_port_id=inputs[0].get_index(),
                         output_port_id=output_port_id,
                         dtype=Dtype(nncf_dtype),
+                        parallel_input_port_ids=parallel_inputs,
                     )
 
     @staticmethod

tests/common/graph/test_nncf_graph.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 
 from nncf.common.graph.graph import NNCFGraph
+from nncf.common.graph.graph import NNCFGraphEdge
 from nncf.common.graph.layer_attributes import Dtype
 from nncf.common.graph.patterns import GraphPattern
 
@@ -48,3 +49,61 @@ def test_find_matching_subgraphs():
             continue
         assert len(match) == 2
         assert match == nodes[:2]
+
+
+def test_parallel_edges():
+    def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_id):
+        return NNCFGraphEdge(
+            from_node,
+            to_node,
+            input_port_id=input_port_id,
+            output_port_id=output_port_id,
+            parallel_input_port_ids=[],
+            tensor_shape=(1, 2, 3),
+            dtype="dummy",
+        )
+
+    nncf_graph = NNCFGraph()
+    nodes = []
+    for node in "abc":
+        nodes.append(nncf_graph.add_nncf_node(node, f"type_{node}", f"metatype_{node}"))
+
+    nncf_graph.add_edge_between_nncf_nodes(
+        nodes[0].node_id,
+        nodes[1].node_id,
+        input_port_id=0,
+        output_port_id=0,
+        parallel_input_port_ids=list(range(1, 5)),
+        tensor_shape=(1, 2, 3),
+        dtype="dummy",
+    )
+    nncf_graph.add_edge_between_nncf_nodes(
+        nodes[0].node_id,
+        nodes[2].node_id,
+        input_port_id=10,
+        output_port_id=15,
+        parallel_input_port_ids=[],
+        tensor_shape=(1, 2, 3),
+        dtype="dummy",
+    )
+    output_edges = nncf_graph.get_output_edges(nodes[0])
+    input_edges = nncf_graph.get_input_edges(nodes[1])
+    assert len(input_edges) == 5
+    assert len(output_edges) == 6
+    assert input_edges == output_edges[:-1]
+    for input_port_id, edge in enumerate(input_edges):
+        ref_edge = _get_default_nncf_graph_edge(
+            nodes[0],
+            nodes[1],
+            input_port_id=input_port_id,
+            output_port_id=0,
+        )
+        assert ref_edge == edge
+
+    ordinary_edge = _get_default_nncf_graph_edge(
+        nodes[0],
+        nodes[2],
+        input_port_id=10,
+        output_port_id=15,
+    )
+    assert ordinary_edge == output_edges[-1]

tests/openvino/native/models.py

@@ -650,6 +650,19 @@ def _create_ov_model(self):
         return model
 
 
+class ParallelEdgesModel(OVReferenceModel):
+    def _create_ov_model(self) -> ov.Model:
+        input_shape = [1, 3, 3]
+
+        input_1 = opset.parameter(input_shape, name="Input")
+        mm = opset.matmul(input_1, input_1, False, False, name="Mm")
+        add = opset.add(input_1, np.array(1.0, dtype=np.float32), name="Add")
+        result_0 = opset.result(mm, name="Result_mm")
+        result_1 = opset.result(add, name="Result_add")
+        model = ov.Model([result_0, result_1], [input_1])
+        return model
+
+
 @SYNTHETIC_MODELS.register()
 class UnifiedEmbeddingModel(OVReferenceModel):
     def _create_ov_model(self):

tests/openvino/native/test_nncf_graph_builder.py

@@ -12,9 +12,13 @@
 import openvino.runtime as ov
 import pytest
 
+from nncf.common.graph.graph import NNCFGraphEdge
+from nncf.common.graph.layer_attributes import Dtype
+from nncf.openvino.graph.nncf_graph_builder import GraphConverter
 from tests.openvino.conftest import OPENVINO_NATIVE_TEST_ROOT
 from tests.openvino.native.common import compare_nncf_graphs
 from tests.openvino.native.models import SYNTHETIC_MODELS
+from tests.openvino.native.models import ParallelEdgesModel
 from tests.openvino.omz_helpers import convert_model
 from tests.openvino.omz_helpers import download_model
 
@@ -47,3 +51,47 @@ def test_compare_nncf_graph_omz_models(tmp_path, model_name):
 
     path_to_dot = REFERENCE_GRAPHS_DIR / f"{model_name}.dot"
     compare_nncf_graphs(model, path_to_dot)
+
+
+def test_parallel_edges():
+    def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_id):
+        return NNCFGraphEdge(
+            from_node,
+            to_node,
+            input_port_id=input_port_id,
+            output_port_id=output_port_id,
+            tensor_shape=[1, 3, 3],
+            dtype=Dtype.FLOAT,
+            parallel_input_port_ids=[],
+        )
+
+    model = ParallelEdgesModel().ov_model
+    nncf_graph = GraphConverter.create_nncf_graph(model)
+    input_node = nncf_graph.get_node_by_name("Input")
+    mm_node = nncf_graph.get_node_by_name("Mm")
+    add_node = nncf_graph.get_node_by_name("Add")
+    ref_input_edges = {
+        _get_default_nncf_graph_edge(
+            input_node,
+            mm_node,
+            input_port_id=0,
+            output_port_id=0,
+        ),
+        _get_default_nncf_graph_edge(
+            input_node,
+            mm_node,
+            input_port_id=1,
+            output_port_id=0,
+        ),
+    }
+    ref_output_edges = ref_input_edges.copy()
+    ref_output_edges.add(
+        _get_default_nncf_graph_edge(
+            input_node,
+            add_node,
+            input_port_id=0,
+            output_port_id=0,
+        )
+    )
+    assert set(nncf_graph.get_input_edges(mm_node)) == ref_input_edges
+    assert set(nncf_graph.get_output_edges(input_node)) == ref_output_edges

tests/torch/test_graph_analysis.py

@@ -78,15 +78,15 @@ def get_node(name: NNCFNodeName):
             ["/C_0"],
             NNCFGraphPatternIO(
                 input_edges=[make_mock_edge("/B_0", "/C_0", input_port_id=0, output_port_id=0)],
-                output_edges=[make_mock_edge("/C_0", "/D_0", input_port_id=0, output_port_id=0)],
+                output_edges=[make_mock_edge("/C_0", "/D_0", input_port_id=1, output_port_id=0)],
             ),
         ),
         (
             ["/A_0", "/B_0", "/C_0"],
             NNCFGraphPatternIO(
                 input_edges=[],
                 output_edges=[
-                    make_mock_edge("/C_0", "/D_0", input_port_id=0, output_port_id=0),
+                    make_mock_edge("/C_0", "/D_0", input_port_id=1, output_port_id=0),
                     make_mock_edge("/A_0", "/D_0", input_port_id=0, output_port_id=1),
                 ],
             ),
@@ -95,7 +95,7 @@ def get_node(name: NNCFNodeName):
             ["/D_0"],
             NNCFGraphPatternIO(
                 input_edges=[
-                    make_mock_edge("/C_0", "/D_0", input_port_id=0, output_port_id=0),
+                    make_mock_edge("/C_0", "/D_0", input_port_id=1, output_port_id=0),
                     make_mock_edge("/A_0", "/D_0", input_port_id=0, output_port_id=1),
                 ],
                 output_edges=[

tests/torch/test_graph_building.py

@@ -21,6 +21,7 @@
 from nncf.common.graph.definitions import MODEL_INPUT_OP_NAME
 from nncf.common.graph.definitions import MODEL_OUTPUT_OP_NAME
 from nncf.common.graph.definitions import NNCFGraphNodeType
+from nncf.common.graph.graph import NNCFGraphEdge
 from nncf.common.graph.layer_attributes import GetItemLayerAttributes
 from nncf.common.graph.layer_attributes import MultipleInputLayerAttributes
 from nncf.common.graph.layer_attributes import MultipleOutputLayerAttributes
@@ -407,10 +408,7 @@ def test_split_attributes(input_shape):
     graph = graph_builder.build_graph(model)
     chunk_nodes_with_attributes = {
         "ModelForTestWithSplit/chunk_0": {"chunks": 2, "axis": 1},
-        "ModelForTestWithSplit/unbind_0": {"chunks": 2, "axis": 1}
-        # TODO: fix NNCFGraph tracing so valid reference below
-        # will be generated by NNCF
-        #'ModelForTestWithSplit/unbind_0': {'chunks': 20, 'axis': 1}
+        "ModelForTestWithSplit/unbind_0": {"chunks": 20, "axis": 1},
     }
 
     for node in graph.get_all_nodes():
@@ -460,6 +458,55 @@ def test_getitem_attributes(input_shape):
                 assert getitem_nodes_with_attributes[node.node_name] is None
 
 
+class ParallelEdgesModel(nn.Module):
+    def forward(self, x):
+        mm_res = torch.mm(x, x)
+        return mm_res, x + mm_res
+
+
+def test_parallel_edges_in_nncf_graph():
+    def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_id):
+        return NNCFGraphEdge(
+            from_node,
+            to_node,
+            input_port_id=input_port_id,
+            output_port_id=output_port_id,
+            tensor_shape=(3, 3),
+            dtype=Dtype.FLOAT,
+            parallel_input_port_ids=[],
+        )
+
+    input_shape = (3, 3)
+    model = ParallelEdgesModel()
+    input_info = ModelInputInfo(input_shape)
+    graph_builder = GraphBuilder(
+        create_dummy_forward_fn(
+            [
+                input_info,
+            ],
+            with_input_tracing=True,
+            with_output_tracing=True,
+        )
+    )
+
+    nncf_graph = graph_builder.build_graph(model)
+
+    input_node = nncf_graph.get_node_by_name("/nncf_model_input_0")
+    mm_node = nncf_graph.get_node_by_name("ParallelEdgesModel/mm_0")
+    ref_input_edges = {
+        _get_default_nncf_graph_edge(input_node, mm_node, input_port_id=0, output_port_id=0),
+        _get_default_nncf_graph_edge(input_node, mm_node, input_port_id=1, output_port_id=0),
+    }
+    mm_node_input_edges = nncf_graph.get_input_edges(mm_node)
+    assert set(mm_node_input_edges) == ref_input_edges
+    ref_output_edges = ref_input_edges.copy()
+
+    add_node = nncf_graph.get_node_by_name("ParallelEdgesModel/__add___0")
+    ref_output_edges.add(_get_default_nncf_graph_edge(input_node, add_node, input_port_id=0, output_port_id=0))
+    input_node_output_edges = nncf_graph.get_output_edges(input_node)
+    assert set(input_node_output_edges) == ref_output_edges
+
+
 TEST_KEYWORD_1 = "keyword1"
 TEST_KEYWORD_2 = "keyword2"
 INPUT_INFO_CONFIG_VS_FORWARD_ARGS = [