fix: correct spirv node incoming edge mapping

Author: jmahbs

src/vgf_adapter_model_explorer/builder/builder.py

@@ -71,27 +71,25 @@ def _build_segment_spirv_nodes(
         spirv_nodes = self._build_spirv_nodes(self.vgf_data, module.index)
         if not spirv_nodes:
             return []
-        self._clean_spirv_edges(spirv_nodes)
-        self._connect_spirv_node(spirv_nodes[0], input_nodes)
+        self._connect_spirv_nodes(spirv_nodes, input_nodes)
         return spirv_nodes
 
-    def _clean_spirv_edges(self, spirv_nodes: list[gb.GraphNode]) -> None:
-        """Clean incoming edges of SPIR-V nodes."""
-        for node in spirv_nodes:
-            new_incoming_edges = []
-            for e in node.incomingEdges:
-                if e.sourceNodeId != GRAPH_INPUT_ANNOTATION:
-                    new_incoming_edges.append(e)
-            node.incomingEdges[:] = new_incoming_edges
-
-    def _connect_spirv_node(
-        self, spirv_node: gb.GraphNode, input_nodes: list[gb.GraphNode]
+    def _connect_spirv_nodes(
+        self, spirv_nodes: list[gb.GraphNode], input_nodes: list[gb.GraphNode]
     ) -> None:
-        """Connect SPIR-V node to the graph."""
-        for node in input_nodes:
-            spirv_node.incomingEdges.append(
-                gb.IncomingEdge(sourceNodeId=node.id, targetNodeInputId="0")
-            )
+        """Connect SPIR-V nodes to the graph by mapping block arguments to input nodes."""
+        for spirv_node in spirv_nodes:
+            new_incoming_edges = []
+            for edge in spirv_node.incomingEdges:
+                if edge.sourceNodeId == GRAPH_INPUT_ANNOTATION:
+                    arg_number = int(edge.sourceNodeOutputId or "0")
+                    if arg_number < len(input_nodes):
+                        edge.sourceNodeId = input_nodes[arg_number].id
+                        edge.sourceNodeOutputId = "0"
+                        new_incoming_edges.append(edge)
+                else:
+                    new_incoming_edges.append(edge)
+            spirv_node.incomingEdges[:] = new_incoming_edges
 
     def _build_segment_input_nodes(
         self, inputs: list[IOBase], module: Module

src/vgf_adapter_model_explorer/tests/fixtures/mobilenet_v2/expected.json

@@ -73,6 +73,11 @@
               "sourceNodeOutputId": "0",
               "targetNodeInputId": "4"
             },
+            {
+              "sourceNodeId": "mrt_115",
+              "sourceNodeOutputId": "0",
+              "targetNodeInputId": "5"
+            },
             {
               "sourceNodeId": "5",
               "sourceNodeOutputId": "0",
@@ -92,11 +97,6 @@
               "sourceNodeId": "8",
               "sourceNodeOutputId": "0",
               "targetNodeInputId": "9"
-            },
-            {
-              "sourceNodeId": "mrt_115",
-              "sourceNodeOutputId": "0",
-              "targetNodeInputId": "0"
             }
           ],
           "outputsMetadata": [

src/vgf_adapter_model_explorer/tests/fixtures/vww4/expected.json

@@ -73,6 +73,11 @@
               "sourceNodeOutputId": "0",
               "targetNodeInputId": "4"
             },
+            {
+              "sourceNodeId": "mrt_95",
+              "sourceNodeOutputId": "0",
+              "targetNodeInputId": "5"
+            },
             {
               "sourceNodeId": "5",
               "sourceNodeOutputId": "0",
@@ -92,11 +97,6 @@
               "sourceNodeId": "8",
               "sourceNodeOutputId": "0",
               "targetNodeInputId": "9"
-            },
-            {
-              "sourceNodeId": "mrt_95",
-              "sourceNodeOutputId": "0",
-              "targetNodeInputId": "0"
             }
           ],
           "outputsMetadata": [

src/vgf_adapter_model_explorer/tests/test_vgf_builder.py

@@ -91,6 +91,8 @@
 
 
 def test_builder():
+    from ..constants import GRAPH_INPUT_ANNOTATION
+
     mock_get_spirv_nodes = Mock(
         return_value=[
             gb.GraphNode(
@@ -99,7 +101,13 @@ def test_builder():
                 namespace="",
                 subgraphIds=[],
                 attrs=[],
-                incomingEdges=[],
+                incomingEdges=[
+                    gb.IncomingEdge(
+                        sourceNodeId=GRAPH_INPUT_ANNOTATION,
+                        sourceNodeOutputId="0",
+                        targetNodeInputId="0",
+                    )
+                ],
                 outputsMetadata=[
                     gb.MetadataItem(
                         id="0",