mct-nightly 2.3.0.20250505.616__py3-none-any.whl → 2.3.0.20250507.555__py3-none-any.whl

{mct_nightly-2.3.0.20250505.616.dist-info → mct_nightly-2.3.0.20250507.555.dist-info}/METADATA

@@ -1,7 +1,8 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250505.616
+Version: 2.3.0.20250507.555
 Summary: A Model Compression Toolkit for neural networks
+Author-email: ssi-dnn-dev@sony.com
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
@@ -23,6 +24,7 @@ Requires-Dist: protobuf
 Requires-Dist: mct-quantizers-nightly
 Requires-Dist: pydantic>=2.0
 Requires-Dist: sony-custom-layers-dev==0.4.0.dev6
+Dynamic: author-email
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
@@ -51,7 +53,7 @@ ______________________________________________________________________
 </p>
 <p align="center">
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/pytorch-2.2%20%7C%202.3%20%7C%202.4%20%7C%202.5-blue" /></a>
-  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-02.14%20%7C%202.15-blue" /></a>
+  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-2.14%20%7C%202.15-blue" /></a>
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%20%7C%203.12-blue" /></a>
   <a href="https://github.com/sony/model_optimization/releases"><img src="https://img.shields.io/github/v/release/sony/model_optimization" /></a>
   <a href="https://github.com/sony/model_optimization/blob/main/LICENSE.md"><img src="https://img.shields.io/badge/license-Apache%202.0-blue" /></a>

{mct_nightly-2.3.0.20250505.616.dist-info → mct_nightly-2.3.0.20250507.555.dist-info}/RECORD

@@ -1,5 +1,5 @@
-mct_nightly-2.3.0.20250505.616.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-model_compression_toolkit/__init__.py,sha256=blX207LzIjS7kQeI15kYyjJnPY-XhwGQPlrTjt2S0CY,1557
+mct_nightly-2.3.0.20250507.555.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+model_compression_toolkit/__init__.py,sha256=QzNbJcOvpHUdWIDaA4UVEDr7PLGs8z-feuEZ7nopltg,1557
 model_compression_toolkit/constants.py,sha256=iJ6vfTjC2oFIZWt8wvHoxEw5YJi3yl0Hd4q30_8q0Zc,3958
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=L3q7tn3Uht0i_7phnlOWMR2Te2zvzrt2HOz9vYEInts,4529
@@ -40,7 +40,7 @@ model_compression_toolkit/core/common/graph/edge.py,sha256=buoSEUZwilWBK3WeBKpJ-
 model_compression_toolkit/core/common/graph/functional_node.py,sha256=GH5wStmw8SoAj5IdT_-ItN1Meo_P5NUTt_5bgJC4fak,3935
 model_compression_toolkit/core/common/graph/graph_matchers.py,sha256=CrDoHYq4iPaflgJWmoJ1K4ziLrRogJvFTVWg8P0UcDU,4744
 model_compression_toolkit/core/common/graph/graph_searches.py,sha256=2oKuW6L8hP-oL0lFO9PhQFt9fEFgVJwpc1u4fHExAtE,5128
-model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py,sha256=JH33qnHTaqFXcYSzTVMpDc9N93503y2pY3hiVJELuZI,10704
+model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py,sha256=ahuvX2H7__hwTrtR02QbadlDJjagvKovFg6KKNU9svo,10443
 model_compression_toolkit/core/common/graph/memory_graph/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/core/common/graph/memory_graph/bipartite_graph.py,sha256=X6FK3C3y8ixFRPjC_wm3ClloCX8_06SOdA1TRi7o_LA,3800
 model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py,sha256=oyz260JXDbvL8aI-DVtUvLHtLRWC2Yu4SBYlGL68c2Y,3498
@@ -69,13 +69,13 @@ model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=onHgDwfw8CUbZFNU-RYit9eqA6FrzAtFA3akVZ2d7IM,4533
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_ru_helper.py,sha256=-hOMBucYn12ePyLd0b1KxniPOIRu4b53SwEzv0bWToI,4943
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=d5-3j2e_rdcQOT7c4s0p7640i3nSetjJ6MgMhhMM7dc,6152
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=J8io_axti6gRoch9QR0FmKOP8JSHGeKqX95rf-nG6fI,37719
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=Lk5cftihGpgFQoyqnRGiwJFFqkI8dkx0l1q0sVJi2CE,27505
 model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=R3UIO9lKf-lpEGfJOqgpQAXdP1IWMatWxXKYDkhWj_E,28096
 model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py,sha256=P8QtKgFXtt5b2RoubzI5OGlCfbEfZsAirjyrkFzK26A,2846
 model_compression_toolkit/core/common/mixed_precision/solution_refinement_procedure.py,sha256=S1ChgxtUjzXJufNWyRbKoNdyNC6fGUjPeComDMx8ZCo,9479
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py,sha256=PKkhc5q8pEPnNLXwo3U56EOCfYnPXIvPs0LlCGZOoKU,4426
-model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py,sha256=PXBuUUuYDmukjhgyrwEe71egpT_iu-LQt5SqddgkRHo,40793
+model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py,sha256=-kNcmQQFVHRPizInaRrCEIuh_q_57CWxC6CIV6azF4g,39640
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py,sha256=QQwtl08DiDxUOQGpYPnek_RlZjWm1Ky7tL2ESHXMK78,4050
 model_compression_toolkit/core/common/mixed_precision/search_methods/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
 model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py,sha256=6Z6nQL9UH7B8dbcUR0cuCTEYFOKZAlvOb-SCk_cAZFA,6670
@@ -528,7 +528,7 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=UVN_S9ULHBEldBpShCOt8-soT8YTQ5oE362y96qF_FA,3950
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=mkoEktLFFHtEKzzFRn_jCnxjhJolK12TZ5AQeDHzUO8,9767
-mct_nightly-2.3.0.20250505.616.dist-info/METADATA,sha256=StJnfqa-V7mHpnvKCpoRrFkWKX2TVY2tXMbSAK5vkA0,25101
-mct_nightly-2.3.0.20250505.616.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-mct_nightly-2.3.0.20250505.616.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.3.0.20250505.616.dist-info/RECORD,,
+mct_nightly-2.3.0.20250507.555.dist-info/METADATA,sha256=hIfm1mpPLcDseqKnO60EFxdc5f3T66WfAIb0gwT7TEk,25157
+mct_nightly-2.3.0.20250507.555.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+mct_nightly-2.3.0.20250507.555.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.3.0.20250507.555.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.3.0.20250505.000616"
+__version__ = "2.3.0.20250507.000555"

model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py

@@ -12,22 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import abc
 import uuid
 
-from typing import Dict, Any, Tuple
-
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.constants import VIRTUAL_ACTIVATION_WEIGHTS_NODE_PREFIX, \
     VIRTUAL_WEIGHTS_SUFFIX, VIRTUAL_ACTIVATION_SUFFIX, FLOAT_BITWIDTH
-
+from model_compression_toolkit.core.common.framework_info import DEFAULT_KERNEL_ATTRIBUTES
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
-import numpy as np
-
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
+from model_compression_toolkit.core.common.quantization.node_quantization_config import ActivationQuantizationMode
+
 
+class VirtualNode(BaseNode, abc.ABC):
+    """ Base class for all virtual nodes. """
+    pass
 
-class VirtualSplitNode(BaseNode):
+
+class VirtualSplitNode(VirtualNode, abc.ABC):
     """
     A class that represents a node that was split from a kernel node (node with weights).
     """
@@ -73,14 +76,11 @@ class VirtualSplitWeightsNode(VirtualSplitNode):
         super().__init__(origin_node)
 
         self.name = origin_node.name + VIRTUAL_WEIGHTS_SUFFIX
-        # Virtual weights node is created only to be absorbed into virtual composed node right away.
-        # However, in some cases composition is impossible and virtual weights node can remain in the graph.
-        # In such case it messes up resource utilization computation, specifically activation cuts. In order to minimize
-        # the impact, we preserve the behavior of the original node wrt activation (shape and quantization),
-        # so that prev - virtualW cut is identical to prev-origin_node. Only the cut virtualW-virtualA will be different
-        # from the original graph, so in the worst case the utilization will be higher in virtual graph.
-        # This should guarantee that the utilization of the original graph does not exceed the requested target.
-        self.candidates_quantization_cfg = origin_node.candidates_quantization_cfg
+
+        self.candidates_quantization_cfg = origin_node.get_unique_weights_candidates(kernel_attr)
+        for c in self.candidates_quantization_cfg:
+            c.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.NO_QUANT
+            c.activation_quantization_cfg.activation_n_bits = FLOAT_BITWIDTH
 
 
 class VirtualSplitActivationNode(VirtualSplitNode):
@@ -113,7 +113,7 @@ class VirtualSplitActivationNode(VirtualSplitNode):
             c.weights_quantization_cfg.weights_n_bits = FLOAT_BITWIDTH
 
 
-class VirtualActivationWeightsNode(BaseNode):
+class VirtualActivationWeightsNode(VirtualNode):
     """
     A node that represents a composition of pair of sequential activation node and weights (kernel) node.
     This structure is used for mixed-precision search with bit-operation constraint.
@@ -149,7 +149,7 @@ class VirtualActivationWeightsNode(BaseNode):
         weights = weights_node.weights.copy()
         act_node_w_rename = {}
         if act_node.weights:
-            if not fw_info.get_kernel_op_attributes(act_node)[0] is None:
+            if fw_info.get_kernel_op_attributes(act_node) != DEFAULT_KERNEL_ATTRIBUTES:
                 raise NotImplementedError(f'Node {act_node} with kernel cannot be used as activation for '
                                           f'VirtualActivationWeightsNode.')
             if act_node.has_any_configurable_weight():

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -19,7 +19,7 @@ from collections import defaultdict
 
 from tqdm import tqdm
 
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 
 import numpy as np
 
@@ -28,7 +28,7 @@ from model_compression_toolkit.core.common.framework_implementation import Frame
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualActivationWeightsNode, \
-    VirtualSplitWeightsNode, VirtualSplitActivationNode
+    VirtualSplitWeightsNode, VirtualSplitActivationNode, VirtualNode
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
     RUTarget, ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_calculator import \
@@ -83,10 +83,9 @@ class MixedPrecisionSearchManager:
         self.min_ru_config: Dict[BaseNode, int] = self.mp_graph.get_min_candidates_config(fw_info)
         self.max_ru_config: Dict[BaseNode, int] = self.mp_graph.get_max_candidates_config(fw_info)
 
-        self.config_reconstruction_helper = ConfigReconstructionHelper(virtual_graph=self.mp_graph,
-                                                                       original_graph=self.original_graph)
+        self.config_reconstruction_helper = ConfigReconstructionHelper(self.original_graph)
         if self.using_virtual_graph:
-            real_min_ru_config: Dict[BaseNode, int] = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(self.min_ru_config)
+            real_min_ru_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.min_ru_config)
             self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, real_min_ru_config)
         else:
             self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, self.min_ru_config)
@@ -101,7 +100,7 @@ class MixedPrecisionSearchManager:
         mp_config = self._prepare_and_run_solver()
 
         if self.using_virtual_graph:
-            mp_config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(mp_config)
+            mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_config)
 
         return mp_config
 
@@ -112,9 +111,9 @@ class MixedPrecisionSearchManager:
         Returns:
             Mapping from nodes to indices of the selected bit-widths candidate.
         """
-        layers_candidates_sensitivity: Dict[BaseNode, List[float]] = self._build_sensitivity_mapping()
         candidates_ru = self._compute_relative_ru_matrices()
         rel_target_ru = self._get_relative_ru_constraint_per_mem_element()
+        layers_candidates_sensitivity: Dict[BaseNode, List[float]] = self._build_sensitivity_mapping()
         solver = MixedPrecisionIntegerLPSolver(layers_candidates_sensitivity, candidates_ru, rel_target_ru)
         mp_config = solver.run()
         return mp_config
@@ -171,8 +170,7 @@ class MixedPrecisionSearchManager:
                                                              topo_cfg(baseline_cfg) if baseline_cfg else None)
 
         if self.using_virtual_graph:
-            origin_max_config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(
-                self.max_ru_config)
+            origin_max_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.max_ru_config)
             max_config_value = compute_metric(origin_max_config)
         else:
             max_config_value = compute_metric(self.max_ru_config)
@@ -192,22 +190,12 @@ class MixedPrecisionSearchManager:
                 # Build a distance matrix using the function we got from the framework implementation.
                 if self.using_virtual_graph:
                     # Reconstructing original graph's configuration from virtual graph's configuration
-                    origin_mp_model_configuration = \
-                        self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(
-                            mp_model_configuration,
-                            changed_virtual_nodes_idx=[node_idx],
-                            original_base_config=origin_max_config)
-                    origin_changed_nodes_indices = [i for i, (n, c) in enumerate(origin_max_config.items()) if
-                                                    c != origin_mp_model_configuration[n]]
-                    metric_value = compute_metric(
-                        origin_mp_model_configuration,
-                        origin_changed_nodes_indices,
-                        origin_max_config)
+                    orig_mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_model_configuration)
+                    changed_nodes = [orig_sorted_nodes.index(n) for n, ind in orig_mp_config.items()
+                                     if origin_max_config[n] != ind]
+                    metric_value = compute_metric(orig_mp_config, changed_nodes, origin_max_config)
                 else:
-                    metric_value = compute_metric(
-                        mp_model_configuration,
-                        [node_idx],
-                        self.max_ru_config)
+                    metric_value = compute_metric(mp_model_configuration, [node_idx], self.max_ru_config)
                 metric_value = max(metric_value, max_config_value + eps)
                 layer_to_metrics_mapping[node].append(metric_value)
 
@@ -256,7 +244,7 @@ class MixedPrecisionSearchManager:
                 else:
                     cfg = self.min_ru_config.copy()
                     cfg[node] = candidate_idx
-                    real_cfg = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(cfg)
+                    real_cfg = self.config_reconstruction_helper.reconstruct_full_configuration(cfg)
                     candidate_rus = self.ru_helper.compute_utilization(self.ru_targets, real_cfg)
 
                 for target, ru in candidate_rus.items():
@@ -326,353 +314,192 @@ class MixedPrecisionSearchManager:
 
 
 class ConfigReconstructionHelper:
-    """
-    A class to help reconstruct an original mixed-precision configuration from a virtual one,
-    when running mixed-precision search with BOPS utilization.
-    It provides a reconstruct_config_from_virtual_graph which allows to translate a bit-width config of a virtual graph
-    to a config of the original configurable nodes.
-    """
-
-    def __init__(self, virtual_graph: Graph, original_graph: Graph):
-        """
-        Init a ConfigReconstructionHelper object.
-        It holds a dictionary variable named origin_node_idx_to_cfg which holds the mapping from an original graph's
-        configurable node to its actual bit-width index (this data structure is being cleared
-        after every reconstruction call).
-
-        Args:
-            virtual_graph: The virtual graph.
-            original_graph: The original graph.
-        """
-
-        self.virtual_graph = virtual_graph
-        self.original_graph = original_graph
-        self.fw_info = original_graph.fw_info
-
-        self.virtual_sorted_nodes_names = self.virtual_graph.get_configurable_sorted_nodes_names(self.fw_info)
-        self.origin_sorted_conf_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
-        self.origin_sorted_conf_nodes_names = [n.name for n in self.origin_sorted_conf_nodes]
-
-        self.origin_node_idx_to_cfg = {}
-
-    def _clear_reconstruction_dict(self):
-        """
-        Clears the origin_node_idx_to_cfg data structure.
-        """
-
-        self.origin_node_idx_to_cfg = {}
-
-    def reconstruct_config_from_virtual_graph(self,
-                                              virtual_mp_cfg: Dict[BaseNode, int],
-                                              changed_virtual_nodes_idx: List[int] = None,
-                                              original_base_config: Dict[BaseNode, int] = None) -> Dict[BaseNode, int]:
-        """
-        Reconstructs the original config for a given virtual graph mixed-precision config.
-        It iterates over all virtual configurable node (that has some chosen bit-width virtual candidate)
-        and translates its chosen candidate to a candidate index of configurable nodes in the original graph.
-        The translation is based of the virtual node's type. Note that if the node is a split activation node
-        for instance, then we need to find its matching weights node in order to construct the original linear node's
-        chosen config.
-
-        Args:
-            virtual_mp_cfg: A mixed-precision configuration (list of candidates indices) of the virtual graph.
-            changed_virtual_nodes_idx: Provide an optional list of virtual nodes indices for which the
-                config reconstruction will be computed.
-            original_base_config: If changed_virtual_nodes_idx is provided, need to provide a base config from which the
-                bit-width for all un-changed original nodes will be taken.
-
-        Returns: A mixed-precision configuration (list of candidates indices) of the original graph.
-
-        """
-
-        if changed_virtual_nodes_idx is not None:
-            if original_base_config is None:
-                Logger.critical("To run config reconstruction for a partial set of nodes, a base original config must be provided.")  # pragma: no cover
-
-            updated_virtual_nodes = \
-                [(idx, self.virtual_graph.get_configurable_sorted_nodes(self.fw_info)[idx]) for idx in changed_virtual_nodes_idx]
-            # Iterating only over the virtual nodes that have updated config
-            for virtual_node_idx, n in updated_virtual_nodes:
-                self.reconstruct_node_config(n, list(virtual_mp_cfg.values()), virtual_node_idx)
-            # Updating reconstructed config for all other nodes based on provided base_config
-            original_sorted_conf_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
-            for i, (n, qc_ind) in enumerate(original_base_config.items()):
-                if i not in list(self.origin_node_idx_to_cfg.keys()):
-                    self.update_config_at_original_idx(n=n, origin_cfg_idx=qc_ind)
-        else:
-            # Reconstruct entire config
-            for virtual_node_idx, n in enumerate(self.virtual_graph.get_configurable_sorted_nodes(self.fw_info)):
-                self.reconstruct_node_config(n, list(virtual_mp_cfg.values()), virtual_node_idx)
-
-        res_config = [self.origin_node_idx_to_cfg[key] for key in sorted(self.origin_node_idx_to_cfg.keys())]
-        self._clear_reconstruction_dict()
-        assert len(res_config) == len(self.origin_sorted_conf_nodes)
-        return {n: candidate_idx for n, candidate_idx in zip(self.origin_sorted_conf_nodes, res_config)}
-
-    def reconstruct_node_config(self,
-                                n: BaseNode,
-                                virtual_mp_cfg: List[int],
-                                virtual_node_idx: int):
-        """
-        Reconstructs the original configuration for a single node. Updates the mapping inplace.
-
-        Args:
-            n: The node to reconstruct the configuration for.
-            virtual_mp_cfg: A mixed-precision configuration (list of candidates indices) of the virtual graph.
-            virtual_node_idx: The index of the virtual node in the virtual mixed-precision configuration.
-        """
-
-        virtual_cfg_idx = virtual_mp_cfg[virtual_node_idx]
-
-        if isinstance(n, VirtualActivationWeightsNode):
-            weights_node = n.original_weights_node
-            if isinstance(weights_node, VirtualSplitWeightsNode):
-                self.get_activation_for_split_weights(weights_node, n, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                Logger.critical(f"Virtual graph construction error: Expected all weights nodes to be split into weights and activation nodes. Found node '{n.name}' not split as expected. Every weights node should correspond to a VirtualSplitWeightsNode type.")  # pragma: no cover
-
-            activation_node = n.original_activation_node
-            if isinstance(activation_node, VirtualSplitActivationNode):
-                self.get_weights_for_split_activation(activation_node, n, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                if activation_node.name in self.origin_sorted_conf_nodes_names:
-                    # It is possible that the original activation node is not configurable,
-                    # in this case we don't need to retrieve its bit-width config
-                    self.retrieve_activation_only_config(activation_node, n, virtual_cfg_idx)
-        elif isinstance(n, VirtualSplitWeightsNode):
-            # If the node's predecessor have multiple outgoing edges then it is possible that this weights
-            # node is not composed with an activation, but otherwise there is something wrong, and we need
-            # to raise an exception
-            predecessor = self.virtual_graph.get_prev_nodes(n)
-            assert len(predecessor) == 1  # Sanity check
-            predecessor = predecessor[0]
-            if len(self.virtual_graph.out_edges(predecessor)) > 1:
-                # It's ok, need to find the node's configuration
-                self.get_activation_for_split_weights(n, n, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                Logger.critical(f"Virtual graph configuration error: Expected the predecessor of node '{n.name}' to have multiple outputs when not composed with an activation node.")  # pragma: no cover
-        elif isinstance(n, VirtualSplitActivationNode):
-            self.get_weights_for_split_activation(n, n, virtual_cfg_idx, virtual_mp_cfg)
-        else:
-            # Node didn't change in virtual graph - candidates list is similar to original
-            if n.name not in self.origin_sorted_conf_nodes_names:
-                Logger.critical(f"Configuration mismatch: Node '{n.name}' is configurable in the virtual graph but not in the original graph. Verify node configurations.")  # pragma: no cover
-            origin_idx = self.origin_sorted_conf_nodes_names.index(n.name)
-            self.origin_node_idx_to_cfg[origin_idx] = virtual_cfg_idx
-
-    def retrieve_weights_only_config(self, weights_node: BaseNode, virtual_node: BaseNode, virtual_cfg_idx: int):
-        """
-        Retrieves the configuration of an original weights configurable node based on a
-        virtual weights configurable node's chosen config idx, and updates (inplace) the origin_cfg_idx mapping dict.
-        If the original node is not configurable, nothing will be updated.
+    def __init__(self, original_graph):
+        # mapping in order to return the actual node objects from the original graph
+        self.orig_nodes = {n.name: n for n in original_graph.nodes}
 
-        Args:
-            weights_node: The original weights (possibly configurable) node.
-            virtual_node: The virtual weights configurable node.
-            virtual_cfg_idx: The virtual node's chosen config index.
-        """
-
-        if weights_node.name in self.origin_sorted_conf_nodes_names:
-            # It is possible that the original weights node is not configurable,
-            # in this case we don't need to retrieve its bit-width config
-            kernel_attr = self.fw_info.get_kernel_op_attributes(weights_node.type)[0]
-            weights_bitwidth = (virtual_node.candidates_quantization_cfg[virtual_cfg_idx].weights_quantization_cfg
-                                .get_attr_config(kernel_attr).weights_n_bits)
-            origin_cfg_idx = [i for i, c in
-                              enumerate(weights_node.candidates_quantization_cfg) if
-                              c.weights_quantization_cfg.get_attr_config(kernel_attr).weights_n_bits == weights_bitwidth]
-
-            self.update_config_at_original_idx(weights_node, origin_cfg_idx[0])
-
-    def retrieve_activation_only_config(self, activation_node: BaseNode, virtual_node: BaseNode, virtual_cfg_idx: int):
+    def reconstruct_full_configuration(self,
+                                       virtual_cfg: Dict[BaseNode, int],
+                                       include_non_configurable: bool = False) -> Dict[BaseNode, int]:
         """
-        Retrieves the configuration of an original activation configurable node based on a
-        virtual activation configurable node's chosen config idx, and updates (inplace) the origin_cfg_idx mapping dict.
-        If the original node is not configurable, nothing will be updated.
+        Convert a configuration of a virtual graph into the corresponding configuration of the original graph.
+        Note that a configurable VirtualActivationWeightsNode might comprise one configurable and one non-configurable
+        original nodes.
 
         Args:
-            activation_node: The original activation (possibly configurable) node.
-            virtual_node: The virtual activation configurable node.
-            virtual_cfg_idx: The virtual node's chosen config index.
-        """
+            virtual_cfg: a mapping from nodes in the virtual graph to selected candidate index. Should contain all
+                configurable nodes of the virtual graph, and only configurable nodes.
+            include_non_configurable: whether to return configs for non-configurable original nodes.
 
-        if activation_node.name in self.origin_sorted_conf_nodes_names:
-            # It is possible that the original activation node is not configurable,
-            # in this case we don't need to retrieve its bit-width config
-            activation_bitwidth = virtual_node.candidates_quantization_cfg[
-                virtual_cfg_idx].activation_quantization_cfg.activation_n_bits
-            origin_cfg_idx = [i for i, c in
-                              enumerate(activation_node.candidates_quantization_cfg) if
-                              c.activation_quantization_cfg.activation_n_bits == activation_bitwidth]
-
-            self.update_config_at_original_idx(activation_node, origin_cfg_idx[0])
-
-    def retrieve_activation_weights_config(self,
-                                           activation_node: BaseNode,
-                                           weights_node: BaseNode,
-                                           virtual_node: BaseNode,
-                                           virtual_cfg_idx: int,
-                                           virtual_mp_cfg: List[int]):
-        """
-        Retrieves the configuration of an original weights and activation (possibly) configurable node based on a given
-        virtual split weights node and a virtual split activation node which represents its matching in the original graph.
-        it updates (inplace) the origin_cfg_idx mapping dict.
+        Returns:
+            A mapping from configurable nodes in the original graph to their candidate indices.
+        """
+        # Original candidate of a node that has been split might be determined by two different virtual nodes, one
+        # determines activation and one - weights. First, for each virtual node we collect the original
+        # activation / weights nodes, with all original candidates that match the virtual candidate
+        # activation / weights config. If both activation and weights of the original node are determined by virtual
+        # candidates, we look for a common candidate.
+        orig_nodes_a_candidates = {}
+        orig_nodes_w_candidates = {}
+        for virtual_node, virtual_qc_ind in virtual_cfg.items():
+            assert virtual_node.has_configurable_activation() or virtual_node.has_any_configurable_weight()
+            orig_a_node, orig_a_candidates = self._retrieve_matching_orig_a_candidates(virtual_node, virtual_qc_ind)
+            if orig_a_node and (include_non_configurable or orig_a_node.has_configurable_activation()):
+                assert orig_a_node not in orig_nodes_a_candidates
+                orig_nodes_a_candidates[orig_a_node] = orig_a_candidates
+            orig_w_node, orig_w_candidates = self._retrieve_matching_orig_w_candidates(virtual_node, virtual_qc_ind)
+            if orig_w_node and (include_non_configurable or orig_w_node.has_any_configurable_weight()):
+                assert orig_w_node not in orig_nodes_w_candidates
+                orig_nodes_w_candidates[orig_w_node] = orig_w_candidates
+
+        orig_cfg = {}
+        common_orig_nodes = set(orig_nodes_a_candidates.keys()).intersection(set(orig_nodes_w_candidates))
+        for orig_node in common_orig_nodes:
+            a_candidates = orig_nodes_a_candidates[orig_node]
+            w_candidates = orig_nodes_w_candidates[orig_node]
+            # find the common candidate
+            common_candidates = set(a_candidates).intersection(set(w_candidates))
+            if len(common_candidates) != 1:
+                raise ValueError(f'Expected to find exactly one candidate with the required activation and weights '
+                                 f'quantization configuration for node {orig_node}. Found {len(common_candidates)}')
+            # in theory it's possible that original non-configurable node gets split and each part is combined
+            # with a configurable part of another node and we end up here
+            if orig_node.has_configurable_activation() or orig_node.has_any_configurable_weight():
+                orig_cfg[orig_node] = common_candidates.pop()
+            del orig_nodes_a_candidates[orig_node]
+            del orig_nodes_w_candidates[orig_node]
+
+        # remaining a nodes
+        for orig_node, a_candidates in orig_nodes_a_candidates.items():
+            assert not orig_node.has_any_configurable_weight()  # if it had we should have caught it above
+            assert len(a_candidates) == 1
+            assert orig_node not in orig_cfg
+            if include_non_configurable or orig_node.has_configurable_activation():
+                orig_cfg[orig_node] = a_candidates[0]
+
+        # remaining w nodes
+        for orig_node, w_candidates in orig_nodes_w_candidates.items():
+            assert not orig_node.has_configurable_activation()  # if it had we should have caught it above
+            assert len(w_candidates) == 1
+            assert orig_node not in orig_cfg
+            if include_non_configurable or orig_node.has_any_configurable_weight():
+                orig_cfg[orig_node] = w_candidates[0]
+
+        return orig_cfg
+
+    def reconstruct_separate_aw_configs(self, virtual_cfg: Dict[BaseNode, int], include_non_configurable: bool) \
+            -> Tuple[Dict[BaseNode, int], Dict[BaseNode, int]]:
+        """
+        Retrieves original activation and weights nodes and corresponding candidates for a given configuration of the
+        virtual graph. Only returns configuration specified by the virtual config, per configurable target (activation
+        or weights). For example, if 'virtual_cfg' contains a single VirtualActivationWeightsNode, the returned
+        configuration will contain only activation config for the original activation node, and only weights config
+        for the original weights node).
+        In practice, we return candidate index in both cases, instead of actual activation or weights config, since
+        sensitivity evaluator heavily depends on it, so we must ignore activation config in weights candidate and vice
+        versa. This is bad!!! TODO
 
         Args:
-            activation_node: The virtual node that contains the activation that matches the weights node in the original graph.
-            weights_node: The virtual node that contains the weights representation of an original node.
-            virtual_node: The virtual node that contains the virtual weights node (either a composed node or a split weights node).
-            virtual_cfg_idx: The virtual node's chosen config index.
-            virtual_mp_cfg: The virtual graph's chosen mp config.
-        """
+            virtual_cfg: a mapping from nodes in the virtual graph to selected candidate index.
+            include_non_configurable: whether to return configs for non-configurable target (i.e. activation config
+              for non-configurable activation, and weights config for non-configurable weight).
 
-        activation_bitwidth = activation_node.candidates_quantization_cfg[virtual_mp_cfg[
-            self.virtual_sorted_nodes_names.index(activation_node.name)]].activation_quantization_cfg.activation_n_bits
-
-        kernel_attr = self.fw_info.get_kernel_op_attributes(weights_node.type)[0]
-
-        weights_bitwidth = (virtual_node.candidates_quantization_cfg[virtual_cfg_idx].weights_quantization_cfg
-                            .get_attr_config(kernel_attr).weights_n_bits)
-
-        origin_cfg_idx = [i for i, c in
-                          enumerate(weights_node.origin_node.candidates_quantization_cfg) if
-                          c.weights_quantization_cfg.get_attr_config(kernel_attr).weights_n_bits == weights_bitwidth and
-                          c.activation_quantization_cfg.activation_n_bits == activation_bitwidth]
-
-        self.update_config_at_original_idx(weights_node.origin_node, origin_cfg_idx[0])
-
-    def retrieve_weights_activation_config(self,
-                                           activation_node: BaseNode,
-                                           weights_node: BaseNode,
-                                           virtual_node: BaseNode,
-                                           virtual_cfg_idx: int,
-                                           virtual_mp_cfg: List[int]):
-        """
-        Retrieves the configuration of an original weights and activation (possibly) configurable node based on a given
-        virtual split activation node and a virtual split weights node which represents its matching in the original graph.
-        it updates (inplace) the origin_cfg_idx mapping dict.
-
-        Args:
-            activation_node: The virtual node that contains the activation representation of an original node.
-            weights_node: The virtual node that contains the weights that matches the activation node in the original graph.
-            virtual_node: The virtual node that contains the virtual activation node (either a composed node or a split activation node).
-            virtual_cfg_idx: The virtual node's chosen config index.
-            virtual_mp_cfg: The virtual graph's chosen mp config.
+        Returns:
+            Configuration for original activation nodes and a separate configuration for original weights nodes.
         """
+        a_cfg = {}
+        w_cfg = {}
+        for virtual_node, virtual_qc_ind in virtual_cfg.items():
+            orig_a_node, orig_a_candidates = self._retrieve_matching_orig_a_candidates(virtual_node, virtual_qc_ind)
+            if orig_a_node and (include_non_configurable or orig_a_node.has_configurable_activation()):
+                # we may have retrieved multiple candidates with different weights configs and identical activation
+                # configs, so we just take the first
+                a_cfg[orig_a_node] = orig_a_candidates[0]
 
-        kernel_attr = self.fw_info.get_kernel_op_attributes(weights_node.type)[0]
-
-        weights_bitwidth = (weights_node.candidates_quantization_cfg[virtual_mp_cfg[
-            self.virtual_sorted_nodes_names.index(weights_node.name)]]
-                            .weights_quantization_cfg.get_attr_config(kernel_attr).weights_n_bits)
-
-        activation_bitwidth = virtual_node.candidates_quantization_cfg[
-            virtual_cfg_idx].activation_quantization_cfg.activation_n_bits
-
-        origin_cfg_idx = [i for i, c in enumerate(activation_node.origin_node.candidates_quantization_cfg) if
-                          c.weights_quantization_cfg.get_attr_config(kernel_attr).weights_n_bits == weights_bitwidth and
-                          c.activation_quantization_cfg.activation_n_bits == activation_bitwidth]
+            orig_w_node, orig_w_candidates = self._retrieve_matching_orig_w_candidates(virtual_node, virtual_qc_ind)
+            if orig_w_node and (include_non_configurable or orig_w_node.has_any_configurable_weight()):
+                # we may have retrieved multiple candidates with different activation configs and identical weights
+                # configs, so we just take the first
+                w_cfg[orig_w_node] = orig_w_candidates[0]
 
-        self.update_config_at_original_idx(activation_node.origin_node, origin_cfg_idx[0])
+        return a_cfg, w_cfg
 
-    def get_activation_for_split_weights(self,
-                                         weights_node: BaseNode,
-                                         virtual_node: BaseNode,
-                                         virtual_cfg_idx: int,
-                                         virtual_mp_cfg: List[int]):
+    def _retrieve_matching_orig_a_candidates(self,
+                                             virtual_node: BaseNode,
+                                             virtual_qc_ind: int) -> Tuple[Optional[BaseNode], Optional[List[int]]]:
         """
-        Finds the matching activation node in the virtual graph for a given split weights node,
-        and calls the relevant method for updating the configuration mapping.
+        Retrieve the original activation node and all its candidates matching activation quantization config of the
+        given virtual candidate (candidate of a node in the virtual graph).
+        Note that we do simple matching, without any filtering, so disabled activation quantization will be also matched.
 
         Args:
-            weights_node: A virtual weights node.
-            virtual_node: A virtual node that contains the virtual weights node (either a composed node or a split weights node).
-            virtual_cfg_idx: The virtual node's chosen config index.
-            virtual_mp_cfg: The virtual graph's chosen mp config.
+            virtual_node: node in the virtual graph (can be virtual or regular).
+            virtual_qc_ind: candidate index of the virtual node.
 
-        """
-
-        # This is a weights node that was split, means it has an activation node that should follow it,
-        # and we need its configuration in order to reconstruct the original node's configuration.
-        matching_activation_node = self.virtual_graph.get_next_nodes(virtual_node)
-        assert len(matching_activation_node) == 1
-        activation_node = matching_activation_node[0]
-
-        if isinstance(activation_node, VirtualActivationWeightsNode):
-            if activation_node.original_activation_node.is_activation_quantization_enabled() and not \
-                    activation_node.original_activation_node.is_all_activation_candidates_equal():
-                assert activation_node.name in self.virtual_sorted_nodes_names  # Sanity check
-                # The original node is both weights and activation configurable
-                self.retrieve_activation_weights_config(activation_node, weights_node, virtual_node, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                # weights_node here is a split weights node therefore must have 'origin_node'
-                self.retrieve_weights_only_config(weights_node.origin_node, virtual_node, virtual_cfg_idx)
+        Returns:
+            The original activation node (actual object from the original graph) and a list of its matching candidates.
+        """
+        if not isinstance(virtual_node, VirtualNode):
+            return self.orig_nodes[virtual_node.name], [virtual_qc_ind]
+        if isinstance(virtual_node, VirtualSplitWeightsNode):
+            return None, None
+        if isinstance(virtual_node, VirtualActivationWeightsNode):
+            orig_a_node = virtual_node.original_activation_node
+            if isinstance(orig_a_node, VirtualSplitActivationNode):
+                orig_a_node = orig_a_node.origin_node
         else:
-            assert isinstance(activation_node, VirtualSplitActivationNode)  # Sanity check
-            if activation_node.name in self.virtual_sorted_nodes_names:
-                self.retrieve_activation_weights_config(activation_node, weights_node, virtual_node, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                # The original node is only weights configurable
-                # weights_node here is a split weights node therefore must have 'origin_node'
-                self.retrieve_weights_only_config(weights_node.origin_node, virtual_node, virtual_cfg_idx)
-
-    def get_weights_for_split_activation(self,
-                                         activation_node: BaseNode,
-                                         virtual_node: BaseNode,
-                                         virtual_cfg_idx: int,
-                                         virtual_mp_cfg: List[int]):
-        """
-        Finds the matching weights node in the virtual graph for a given split activation node,
-        and calls the relevant method for updating the configuration mapping.
+            assert isinstance(virtual_node, VirtualSplitActivationNode)
+            orig_a_node = virtual_node.origin_node
 
-        Args:
-            activation_node: A virtual activation node.
-            virtual_node: A virtual node that contains the virtual activation node (either a composed node or a split activation node).
-            virtual_cfg_idx: The virtual node's chosen config index.
-            virtual_mp_cfg: The virtual graph's chosen mp config.
-
-        """
+        virtual_qc = virtual_node.candidates_quantization_cfg[virtual_qc_ind]
+        matching_orig_a_cfgs = [i for i, orig_qc in enumerate(orig_a_node.candidates_quantization_cfg)
+                                if orig_qc.activation_quantization_cfg == virtual_qc.activation_quantization_cfg]
+        if not matching_orig_a_cfgs:    # pragma: no cover
+            raise ValueError(f'Could not find matching activation quantization config in the original node '
+                             f'{orig_a_node} for candidate {virtual_qc_ind} of the virtual node {virtual_node}')
+        return self.orig_nodes[orig_a_node.name], matching_orig_a_cfgs
 
-        # This is an activation node that was split, means it has a weights node that should come before it,
-        # and we need its configuration in order to reconstruct the original node's configuration.
-        matching_weights_node = self.virtual_graph.get_prev_nodes(virtual_node)
-        assert len(matching_weights_node) == 1
-        weights_node = matching_weights_node[0]
-
-        if isinstance(weights_node, VirtualActivationWeightsNode):
-            kernel_attr = self.fw_info.get_kernel_op_attributes(weights_node.type)[0]
-            if weights_node.original_weights_node.is_weights_quantization_enabled(kernel_attr) and not \
-                    weights_node.original_weights_node.is_all_weights_candidates_equal(kernel_attr):
-                assert weights_node.name in self.virtual_sorted_nodes_names  # Sanity check
-                # The original node is both weights and activation configurable
-                self.retrieve_weights_activation_config(activation_node, weights_node, virtual_node, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                # The original node is only activation configurable
-                # activation_node here is a split activation node therefore must have 'origin_node'
-                self.retrieve_activation_only_config(activation_node.origin_node, virtual_node, virtual_cfg_idx)
-        else:
-            # If the node's predecessor e multiple outgoing edges than it is possible that this weights
-            # node is not composed with an activation, but otherwise this is something wrong and we need
-            # to raise an exception
-            predecessor = self.virtual_graph.get_prev_nodes(weights_node)
-            assert len(predecessor) == 1  # Sanity check
-            predecessor = predecessor[0]
-            if len(self.virtual_graph.out_edges(predecessor)) > 1:
-                # It's ok, need to find the node's configuration
-                self.retrieve_weights_activation_config(activation_node, weights_node, virtual_node, virtual_cfg_idx, virtual_mp_cfg)
-            else:
-                Logger.critical(f"Virtual graph configuration error: Expected the predecessor of node '{weights_node.name}' to have multiple outputs when not composed with an activation node.")  # pragma: no cover
-
-    def update_config_at_original_idx(self, n: BaseNode, origin_cfg_idx: int):
+    def _retrieve_matching_orig_w_candidates(self,
+                                             virtual_node: BaseNode,
+                                             virtual_qc_ind: int) -> Tuple[Optional[BaseNode], Optional[List[int]]]:
         """
-        Updates (inplace) the origin_node_idx_to_cfg mapping wit hthe given index for a given original node index
-        (in the original graph's sorted configurable nodes list).
+        Retrieve the original weights node and all its candidates matching weights quantization config of the
+        given virtual candidate (candidate of a node in the virtual graph).
 
         Args:
-            n: An original graph's node
-            origin_cfg_idx: A candidate index.
+            virtual_node: node in the virtual graph (can be virtual or regular).
+            virtual_qc_ind: candidate index of the virtual node.
 
-        """
-
-        origin_idx = self.origin_sorted_conf_nodes_names.index(n.name)
-        self.origin_node_idx_to_cfg[origin_idx] = origin_cfg_idx
+        Returns:
+            The original weights node (actual object from the original graph) and a list of all its matching candidates.
+        """
+        if not isinstance(virtual_node, VirtualNode):
+            if virtual_node.weights:
+                return self.orig_nodes[virtual_node.name], [virtual_qc_ind]
+            return None, None
+        if isinstance(virtual_node, VirtualSplitActivationNode):
+            return None, None
+
+        if isinstance(virtual_node, VirtualActivationWeightsNode):
+            assert isinstance(virtual_node.original_weights_node, VirtualSplitWeightsNode)
+            orig_w_node = virtual_node.original_weights_node.origin_node
+        else:
+            assert isinstance(virtual_node, VirtualSplitWeightsNode)
+            orig_w_node = virtual_node.origin_node
+
+        virtual_qc = virtual_node.candidates_quantization_cfg[virtual_qc_ind]
+
+        # Matching candidate is a candidate with matching configs for configurable weights. We cannot compare the entire
+        # weights config since the virtual node may contain additional non-configurable weights from the activation node
+        orig_configurable_attrs = [attr for attr in orig_w_node.weights if virtual_node.is_configurable_weight(attr)]
+        assert all(virtual_node.is_configurable_weight(attr) for attr in orig_configurable_attrs)
+
+        def get_configurable_attrs_cfgs(qc):
+            return {attr: qc.weights_quantization_cfg.get_attr_config(attr) for attr in orig_configurable_attrs}
+        virtual_cfg = get_configurable_attrs_cfgs(virtual_qc)
+        matching_orig_w_cfgs = [i for i, orig_qc in enumerate(orig_w_node.candidates_quantization_cfg)
+                                if get_configurable_attrs_cfgs(orig_qc) == virtual_cfg]
+        if not matching_orig_w_cfgs:    # pragma: no cover
+            raise ValueError(f'Could not find matching weights quantization config in the original node '
+                             f'{orig_w_node} for candidate {virtual_qc_ind} of the virtual node {virtual_node}')
+        return self.orig_nodes[orig_w_node.name], matching_orig_w_cfgs

model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py

@@ -29,7 +29,7 @@ from model_compression_toolkit.core.common.graph.memory_graph.compute_graph_max_
 from model_compression_toolkit.core.common.graph.memory_graph.cut import Cut
 from model_compression_toolkit.core.common.graph.memory_graph.memory_graph import MemoryGraph
 from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualActivationWeightsNode, \
-    VirtualSplitWeightsNode
+    VirtualSplitWeightsNode, VirtualNode
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
     RUTarget, ResourceUtilization
 from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig, \
@@ -531,6 +531,7 @@ class ResourceUtilizationCalculator:
         Returns:
             Node's BOPS count.
         """
+        assert not isinstance(n, VirtualNode), 'Use original graph to compute BOPS.'
         if target_criterion is None:
             target_criterion = TargetInclusionCriterion.Any
         if target_criterion not in [TargetInclusionCriterion.AnyQuantized, TargetInclusionCriterion.Any]:
@@ -539,20 +540,6 @@ class ResourceUtilizationCalculator:
         self._validate_custom_qcs(act_qcs, bitwidth_mode)
         self._validate_custom_qcs(w_qc, bitwidth_mode)
 
-        if isinstance(n, VirtualSplitWeightsNode):
-            # Virtual weights node can only be present if it couldn't be merged into VirtualActivationWeightsNode.
-            # This means that during MP search we cannot compute bops for all A/W nbits combinations. To prevent
-            # inconsistencies we ignore such nodes for bops computation.
-            return 0
-
-        # Fetch the original weights node for mac computation (VirtualActivationWeightsNode input/output shapes are
-        # based on the activation original node, not weights original node)
-        orig_w_node = n
-        if isinstance(n, VirtualActivationWeightsNode):
-            orig_w_node = n.original_weights_node
-            if isinstance(orig_w_node, VirtualSplitWeightsNode):
-                orig_w_node = orig_w_node.origin_node
-
         # check if the node has kernel
         kernel_attrs = self.fw_info.get_kernel_op_attributes(n.type)
         if len(kernel_attrs) > 1:  # pragma: no cover
@@ -561,21 +548,13 @@ class ResourceUtilizationCalculator:
             return 0
 
         kernel_attr = kernel_attrs[0]
-        node_mac = self.fw_impl.get_node_mac_operations(orig_w_node, self.fw_info)
+        node_mac = self.fw_impl.get_node_mac_operations(n, self.fw_info)
         if node_mac == 0:
             return node_mac
 
-        # find the activation node from which to get quantization info and for which to look in custom configuration
-        if isinstance(n, VirtualActivationWeightsNode):
-            # we don't need the original node (and cannot use it for custom configuration anyway)
-            a_node = n
-        else:
-            # if we are running on the original (non-virtual) graph, we only compute bops if it would be computed in an
-            # equivalent virtual graph for consistency.
-            a_node = get_input_activation_if_composable(self.graph, n, warn=False)
-            if a_node is None:
-                return 0
-
+        prev_nodes = self.graph.get_prev_nodes(n)
+        assert len(prev_nodes) == 1, f'Weights node is expected to have exactly one input, {n} has {len(prev_nodes)}'
+        a_node = prev_nodes[0]
         if (target_criterion == TargetInclusionCriterion.AnyQuantized and
                 not (a_node.is_activation_quantization_enabled() or n.is_weights_quantization_enabled(kernel_attr))):
             return 0