Refactor fusing

.github/workflows/run_pytorch_tests.yml

@@ -33,10 +33,11 @@ jobs:
           pip install torch==${{ inputs.torch-version }} torchvision onnx onnxruntime "onnxruntime-extensions<0.14"
           pip install pytest pytest-mock
           pip check
-      - name: Run unittests
-        run: |
-          python -m unittest discover tests/pytorch_tests -v
       - name: Run pytest
         run: |
           pytest tests_pytest/pytorch_tests
+      - name: Run unittests
+        run: |
+          python -m unittest discover tests/pytorch_tests -v
+
 

model_compression_toolkit/core/common/fusion/fusing_metadata_wrapper.py

@@ -0,0 +1,159 @@
+#  Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+import types
+
+from functools import wraps
+
+from typing import Any, Iterator
+
+from model_compression_toolkit.core.common import BaseNode, Graph
+from model_compression_toolkit.core.common.fusion.fusing_info import FusingInfo
+
+
+
+class FusingMetadataWrapper:
+    def __init__(self, graph: Graph, fusing_info: FusingInfo):
+        """
+        Initialize with a graph and its fusing information.
+
+        Args:
+            graph: The neural network graph (e.g., a networkx.DiGraph or similar).
+            fusing_info: Dict mapping fused operation IDs to sets of node objects.
+        """
+        assert isinstance(graph, Graph)
+        self._internal_graph = graph
+        self._fusing_info = fusing_info
+        self._fusing_info.validate(graph)  # Ensure initial consistency
+        # TODO: temp disable activation quantization to keep similar functionality. This will be removed in the future
+        self._disable_nodes_activation_quantization()
+
+    # We added __getstate__ and __setstate__ to FusingMetadataWrapper to fix a recursion error during copy.deepcopy. Without
+    # these, deepcopy endlessly traverses attributes via __getattr__, causing a loop. Now, __getstate__ defines what
+    # to copy (self._graph and self._fusing_info), and __setstate__ rebuilds the object, ensuring a clean copy
+    # without recursion, assuming Graph and FusingInfo are copyable.
+    def __getstate__(self):
+        """
+        Define how the object is serialized for copying.
+        Returns a dictionary of the essential attributes.
+        """
+        self._fusing_info.validate(self._internal_graph)
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        """
+        Reconstruct the object from the serialized state.
+
+        Args:
+            state: Dictionary containing the serialized attributes.
+        """
+        self.__dict__.update(state)
+        self._fusing_info.validate(self._internal_graph)
+
+    def __getattr__(self, name: str) -> Any:
+        """
+        Delegate attribute access to the underlying graph if not found in FusingMetadataWrapper.
+
+        Ensures that if the accessed attribute is a callable (e.g., a method like remove_node),
+        it is wrapped so that the fusing information is validated after execution.
+        Non-callable attributes are returned directly without validation.
+
+        Args:
+            name: The name of the attribute being accessed.
+
+        Returns:
+            The attribute or a wrapped method from self._graph.
+
+        Raises:
+            AttributeError: If the attribute doesn't exist in self._graph.
+        """
+
+        # TODO: Optimize validation by restricting it to known modifying methods to improve efficiency. For now,
+        #  validating after every method call ensures correctness. In the
+        #  future, define explicit modification methods (e.g., remove_node)
+        #  in FusingMetadataWrapper for better efficiency.
+
+        graph_attr = getattr(self._internal_graph, name)
+        # Only wrap methods or functions, excluding properties and descriptors
+        if isinstance(graph_attr, (types.MethodType, types.FunctionType)):
+            @wraps(graph_attr)
+            def wrapper(*args, **kwargs):
+                result = graph_attr(*args, **kwargs)
+                self._fusing_info.validate(self._internal_graph)
+                return result
+            return wrapper
+
+        return graph_attr
+
+    def __iter__(self) -> Iterator[BaseNode]:
+        """
+        Make FusingMetadataWrapper iterable by delegating to the wrapped graph's iterator.
+
+        This allows FusingMetadataWrapper to be used in contexts expecting an iterable of nodes,
+        such as topological_sort, without requiring changes to external code.
+
+        Returns:
+            An iterator over the nodes in the underlying graph.
+        """
+        return iter(self._internal_graph)
+
+    def __getitem__(self, key: Any) -> Any:
+        """
+        Delegate subscripting to the wrapped graph.
+
+        This enables FusingMetadataWrapper to support dictionary-like access (e.g., graph[node][child])
+        as required by operations like topological_generations in NetworkX, maintaining
+        compatibility with code expecting a subscriptable Graph object.
+
+        Args:
+            key: The key (e.g., node) to look up in the graph.
+
+        Returns:
+            The value associated with the key in the wrapped graph.
+
+        Raises:
+            KeyError: If the key doesn't exist in self._graph.
+        """
+        return self._internal_graph[key]
+
+    def update_fusing_info(self, new_fusing_info: FusingInfo):
+        self._fusing_info = new_fusing_info
+
+    def get_internal_graph(self):
+        """Return the original graph."""
+        return self._internal_graph
+
+    def get_fusing_info(self):
+        """Return the fusing information."""
+        return self._fusing_info
+
+    def is_part_of_fused_op(self, node):
+        """Check if a node is part of any fused operation."""
+        return self._fusing_info.is_node_in_fused_op(node)
+
+    def _disable_nodes_activation_quantization(self):
+        """
+        Disable activation quantization for nodes inside the fused operators.
+        """
+        nodes_to_disable = [node for nodes in self._fusing_info.get_all_fused_operations().values() for node in nodes[:-1]]
+        for node in nodes_to_disable:
+            for qc in node.candidates_quantization_cfg:
+                qc.activation_quantization_cfg.enable_activation_quantization = False
+
+    def validate(self):
+        """
+        Check if the internal graph and fusing data are consistent.
+        """
+        return self._fusing_info.validate(self._internal_graph)

model_compression_toolkit/core/common/fusion/graph_fuser.py

@@ -1,4 +1,4 @@
-#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#  Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -13,10 +13,14 @@
 #  limitations under the License.
 #  ==============================================================================
 
-from typing import Dict, List
+import copy
+from typing import List
 
-from model_compression_toolkit.core.common import Graph, BaseNode
-from model_compression_toolkit.core.common.graph.base_graph import OutTensor
+from model_compression_toolkit.core.common.fusion.fusing_info import FusingInfoGenerator
+from model_compression_toolkit.core.common.fusion.fusing_metadata_wrapper import FusingMetadataWrapper
+from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNode, OutTensor
+from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import CandidateNodeQuantizationConfig
+from itertools import product
 
 
 class FusedLayerType:
@@ -27,35 +31,42 @@ class FusedLayerType:
     def __init__(self):
         self.__name__ = 'FusedLayer'
 
-
 class GraphFuser:
-
-    def create_fused_graph(self, graph: Graph) -> Dict[str, str]:
+    def apply_node_fusion(self, wrapper: FusingMetadataWrapper) -> Graph:
         """
-        GraphFuser is responsible for fusing nodes in a networkx graph.
-        The fusion process involves:
-            1. Creating new fused nodes to represent these groups.
-            2. Updating the graph structure to replace the original nodes with fused nodes.
-            3. Maintaining mapping of original node names to their fused node names.
+        Applies node fusion to the graph contained in the given FusingMetadataWrapper.
+
+        The fusion process includes:
+            1. Generating new fused nodes to replace groups of original nodes.
+            2. Updating the graph structure to replace those nodes with the fused representations.
 
         Args:
-            graph: Graph to fuse its nodes.
+            wrapper: A wrapper object that contains the graph and its fusing metadata.
 
         Returns:
-            Mapping of original node names to their fused node names
+            The updated graph with fused nodes replacing the original node groups.
         """
-        fused_nodes_mapping = {}
-        # Iterate through each group of nodes to be fused
-        for fused_nodes_list in graph.fused_nodes:
-            new_fused_node = self._create_fused_node(fused_nodes_list)
-            self._replace_nodes_with_fused_node(graph, fused_nodes_list, new_fused_node)
-            # Update the mapping to keep track of which original nodes are now part of which fused nodes
-            for node in fused_nodes_list:
-                fused_nodes_mapping[node.name] = new_fused_node.name
-        return fused_nodes_mapping
+        wrapper_copy = copy.deepcopy(wrapper)
+        expected_fusing_info = FusingInfoGenerator(wrapper_copy.get_fusing_info().fusing_patterns).generate_fusing_info(wrapper_copy.get_internal_graph())
+
+        if expected_fusing_info != wrapper_copy.get_fusing_info():
+            raise ValueError(
+                f"Mismatch between expected and existing fusing information.\n"
+                f"Expected:\n{expected_fusing_info}\nExisting:\n{wrapper_copy.get_fusing_info()}"
+            )
+
+        for fused_node_id, original_nodes in wrapper_copy.get_fusing_info().get_all_fused_operations().items():
+            fused_node = self._create_fused_node(fused_node_id,
+                                                 original_nodes)
+            self._replace_nodes_with_fused_node(wrapper_copy.get_internal_graph(),
+                                                original_nodes,
+                                                fused_node)
+
+        return wrapper_copy.get_internal_graph()
+
 
     @staticmethod
-    def _create_fused_node(nodes: List[BaseNode]) -> BaseNode:
+    def _create_fused_node(fused_node_id: str, nodes: List[BaseNode]) -> BaseNode:
         """
         Create a new node that represents the fusion of the given nodes.
 
@@ -67,15 +78,34 @@ def _create_fused_node(nodes: List[BaseNode]) -> BaseNode:
         """
         # Create a new node with a name that reflects its components
         # Use the input shape of the first node and output shape of the last node
-        fused_node = BaseNode(name='FusedNode_' + '_'.join([node.name for node in nodes]),
+        # TODO: consider replacing the fused node with a sub-model to allow inference on it, etc.
+        fused_node = BaseNode(name=fused_node_id,
                               framework_attr={},
                               input_shape=nodes[0].input_shape,
                               output_shape=nodes[-1].output_shape,
-                              weights={},
+                              weights={}, # TODO: update with weights of all nodes
                               layer_class=FusedLayerType)
 
-        # Preserve the final activation quantization configuration
-        # This is important for maintaining the correct behavior of the fused node
+        # Ensure weights quantization config is None for all nodes except the first
+        for n in nodes[1:]:
+            for c in n.candidates_quantization_cfg:
+                assert len(c.weights_quantization_cfg.all_weight_attrs) == 0, (
+                    f"Expected weights_quantization_cfg to have no attributes to quantize, but found attributes {c.weights_quantization_cfg.all_weight_attrs} for node {n}"
+                )
+
+        # Create candidates for this node (we assume that the weights configuration should be taken from the first node, and the activaion configuration
+        # is the output quantization configuration of the last node. We ignore all configurations of middle nodes.
+        weight_cfgs = [c.weights_quantization_cfg for c in nodes[0].candidates_quantization_cfg]
+        activation_cfgs = [c.activation_quantization_cfg for c in nodes[-1].candidates_quantization_cfg]
+        if weight_cfgs and activation_cfgs:
+            combinations = list(product(weight_cfgs, activation_cfgs))
+            fused_node.candidates_quantization_cfg = [
+                CandidateNodeQuantizationConfig(weights_quantization_cfg=w, activation_quantization_cfg=a)
+                for w, a in combinations
+            ]
+
+        # Keep the final configurations if they were set already.
+        fused_node.final_weights_quantization_cfg = nodes[0].final_weights_quantization_cfg
         fused_node.final_activation_quantization_cfg = nodes[-1].final_activation_quantization_cfg
 
         return fused_node
@@ -136,3 +166,5 @@ def _replace_nodes_with_fused_node(graph: Graph,
 
         # Finally, add the new fused node to the graph
         graph.add_node(fused_node)
+
+