mct-nightly 2.3.0.20250322.517__tar.gz → 2.3.0.20250324.606__tar.gz

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250322.517
+Version: 2.3.0.20250324.606
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250322.517
+Version: 2.3.0.20250324.606
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/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.20250322.000517"
+__version__ = "2.3.0.20250324.000606"

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/model_compression_toolkit/core/common/graph/base_graph.py

@@ -706,14 +706,24 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         """
         self.fused_nodes.append(fusion)
 
-    def is_single_activation_cfg(self):
+    def has_any_configurable_activation(self) -> bool:
         """
-        Checks whether all nodes in the graph that have activation quantization are quantized with the same bit-width.
+        Checks whether any node in the graph has a configurable activation quantization.
 
-        Returns: True if all quantization config candidates of all nodes have the same activation quantization bit-width.
+        Returns:
+            Whether any node in the graph has a configurable activation quantization.
+        """
+        return any([n.has_configurable_activation() for n in self.nodes])
+
+    def has_any_configurable_weights(self):
+        """
+        Checks whether any node in the graph has any configurable weights quantization.
 
+        Returns:
+            Whether any node in the graph has any configurable weights quantization.
         """
-        return all([n.is_all_activation_candidates_equal() for n in self.nodes])
+
+        return any([n.has_any_configurable_weight() for n in self.nodes])
 
     def replace_node(self, node_to_replace: BaseNode, new_node: BaseNode):
         """

mct_nightly-2.3.0.20250324.606/model_compression_toolkit/core/common/mixed_precision/mixed_precision_ru_helper.py

@@ -0,0 +1,93 @@
+# Copyright 2022 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.
+# ==============================================================================
+from typing import List, Set, Dict, Tuple
+
+import numpy as np
+
+from model_compression_toolkit.core import FrameworkInfo
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
+    RUTarget
+from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_calculator import \
+    ResourceUtilizationCalculator, BitwidthMode, TargetInclusionCriterion
+from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig, \
+    NodeActivationQuantizationConfig
+
+
+class MixedPrecisionRUHelper:
+    """ Helper class for resource utilization computations for mixed precision optimization. """
+
+    def __init__(self, graph: Graph, fw_info: FrameworkInfo, fw_impl: FrameworkImplementation):
+        self.graph = graph
+        self.fw_info = fw_info
+        self.fw_impl = fw_impl
+        self.ru_calculator = ResourceUtilizationCalculator(graph, fw_impl, fw_info)
+
+    def compute_utilization(self, ru_targets: Set[RUTarget], mp_cfg: List[int]) -> Dict[RUTarget, np.ndarray]:
+        """
+        Compute utilization of requested targets for a specific configuration:
+          for weights and bops - total utilization,
+          for activations and total - utilization per cut.
+
+        Args:
+            ru_targets: resource utilization targets to compute.
+            mp_cfg: a list of candidates indices for configurable layers.
+
+        Returns:
+            Dict of the computed utilization per target, as 1d vector.
+        """
+        act_qcs, w_qcs = self.get_quantization_candidates(mp_cfg)
+
+        ru, detailed_ru = self.ru_calculator.compute_resource_utilization(TargetInclusionCriterion.AnyQuantized,
+                                                                          BitwidthMode.QCustom,
+                                                                          act_qcs=act_qcs,
+                                                                          w_qcs=w_qcs,
+                                                                          ru_targets=ru_targets,
+                                                                          allow_unused_qcs=True,
+                                                                          return_detailed=True)
+
+        ru_dict = {k: np.array([v]) for k, v in ru.get_resource_utilization_dict(restricted_only=True).items()}
+        # For activation and total we need utilization per cut, as different mp configurations might result in
+        # different cuts to be maximal.
+        for target in [RUTarget.ACTIVATION, RUTarget.TOTAL]:
+            if target in ru_dict:
+                ru_dict[target] = np.array(list(detailed_ru[target].values()))
+
+        assert all(v.ndim == 1 for v in ru_dict.values())
+        if RUTarget.ACTIVATION in ru_targets and RUTarget.TOTAL in ru_targets:
+            assert ru_dict[RUTarget.ACTIVATION].shape == ru_dict[RUTarget.TOTAL].shape
+
+        assert len(ru_dict) == len(ru_targets), (f'Mismatch between the number of computed and requested metrics.'
+                                                 f'Requested {ru_targets}')
+        return ru_dict
+
+    def get_quantization_candidates(self, mp_cfg) \
+            -> Tuple[Dict[str, NodeActivationQuantizationConfig], Dict[str, NodeWeightsQuantizationConfig]]:
+        """
+        Retrieve quantization candidates objects for weights and activations from the configuration list.
+
+        Args:
+            mp_cfg: a list of candidates indices for configurable layers.
+
+        Returns:
+            A mapping between nodes to weights quantization config, and a mapping between nodes and activation
+            quantization config.
+        """
+        mp_nodes = self.graph.get_configurable_sorted_nodes(self.fw_info)
+        node_qcs = {n: n.candidates_quantization_cfg[mp_cfg[i]] for i, n in enumerate(mp_nodes)}
+        act_qcs = {n.name: cfg.activation_quantization_cfg for n, cfg in node_qcs.items()}
+        w_qcs = {n.name: cfg.weights_quantization_cfg for n, cfg in node_qcs.items()}
+        return act_qcs, w_qcs

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -13,37 +13,27 @@
 # limitations under the License.
 # ==============================================================================
 
-import copy
 from enum import Enum
-import numpy as np
-from typing import List, Callable, Dict
+from typing import List, Callable
 
 from model_compression_toolkit.core import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.core.common.hessian import HessianInfoService
-from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization, RUTarget
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_search_manager import MixedPrecisionSearchManager
-from model_compression_toolkit.core.common.mixed_precision.search_methods.linear_programming import \
-    mp_integer_programming_search
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.common.hessian import HessianInfoService
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_search_manager import \
+    MixedPrecisionSearchManager
+from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
+    ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.solution_refinement_procedure import \
     greedy_solution_refinement_procedure
-from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
-from model_compression_toolkit.logger import Logger
 
 
 class BitWidthSearchMethod(Enum):
-    # When adding a new search_methods MP configuration method, these enum and factory dictionary
-    # should be updated with it's kind and a search_method implementation.
     INTEGER_PROGRAMMING = 0
 
 
-search_methods = {
-    BitWidthSearchMethod.INTEGER_PROGRAMMING: mp_integer_programming_search}
-
-
-def search_bit_width(graph_to_search_cfg: Graph,
+def search_bit_width(graph: Graph,
                      fw_info: FrameworkInfo,
                      fw_impl: FrameworkImplementation,
                      target_resource_utilization: ResourceUtilization,
@@ -60,7 +50,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
     target_resource_utilization have to be passed. If it was not passed, the facade is not supposed to get here by now.
 
     Args:
-        graph_to_search_cfg: Graph to search a MP configuration for.
+        graph: Graph to search a MP configuration for.
         fw_info: FrameworkInfo object about the specific framework (e.g., attributes of different layers' weights to quantize).
         fw_impl: FrameworkImplementation object with specific framework methods implementation.
         target_resource_utilization: Target Resource Utilization to bound our feasible solution space s.t the configuration does not violate it.
@@ -75,17 +65,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
         bit-width index on the node).
 
     """
-
-    # target_resource_utilization have to be passed. If it was not passed, the facade is not supposed to get here by now.
-    if target_resource_utilization is None:
-        Logger.critical("Target ResourceUtilization is required for the bit-width search method's configuration.")  # pragma: no cover
-
-    # Set graph for MP search
-    graph = copy.deepcopy(graph_to_search_cfg)  # Copy graph before searching
-    if target_resource_utilization.bops_restricted():
-        # TODO: we only need the virtual graph is both activations and weights are configurable
-        # Since Bit-operations count target resource utilization is set, we need to reconstruct the graph for the MP search
-        graph = substitute(graph, fw_impl.get_substitutions_virtual_weights_activation_coupling())
+    assert target_resource_utilization.is_any_restricted()
 
     # If we only run weights compression with MP than no need to consider activation quantization when computing the
     # MP metric (it adds noise to the computation)
@@ -93,33 +73,28 @@ def search_bit_width(graph_to_search_cfg: Graph,
     weight_only_restricted = tru.weight_restricted() and not (tru.activation_restricted() or
                                                               tru.total_mem_restricted() or
                                                               tru.bops_restricted())
-    disable_activation_for_metric = weight_only_restricted or graph_to_search_cfg.is_single_activation_cfg()
+    disable_activation_for_metric = weight_only_restricted or not graph.has_any_configurable_activation()
 
     # Set Sensitivity Evaluator for MP search. It should always work with the original MP graph,
     # even if a virtual graph was created (and is used only for BOPS utilization computation purposes)
     se = fw_impl.get_sensitivity_evaluator(
-        graph_to_search_cfg,
+        graph,
         mp_config,
         representative_data_gen=representative_data_gen,
         fw_info=fw_info,
         disable_activation_for_metric=disable_activation_for_metric,
         hessian_info_service=hessian_info_service)
 
-    # Instantiate a manager object
+    if search_method != BitWidthSearchMethod.INTEGER_PROGRAMMING:
+        raise NotImplementedError()
+
+    # Search manager and LP are highly coupled, so LP search method was moved inside search manager.
     search_manager = MixedPrecisionSearchManager(graph,
                                                  fw_info,
                                                  fw_impl,
                                                  se,
-                                                 target_resource_utilization,
-                                                 original_graph=graph_to_search_cfg)
-
-    if search_method not in search_methods:
-        raise NotImplementedError()  # pragma: no cover
-
-    search_method_fn = search_methods[search_method]
-    # Search for the desired mixed-precision configuration
-    result_bit_cfg = search_method_fn(search_manager,
-                                      target_resource_utilization)
+                                                 target_resource_utilization)
+    result_bit_cfg = search_manager.search()
 
     if mp_config.refine_mp_solution:
         result_bit_cfg = greedy_solution_refinement_procedure(result_bit_cfg, search_manager, target_resource_utilization)

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -12,11 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import copy
+from collections import defaultdict
 
-from typing import Callable, Dict, List
+from tqdm import tqdm
+
+from typing import Dict, List, Tuple
 
 import numpy as np
 
+from model_compression_toolkit.constants import EPS
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
@@ -29,7 +34,10 @@ from model_compression_toolkit.core.common.mixed_precision.resource_utilization_
     TargetInclusionCriterion, BitwidthMode
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_ru_helper import \
     MixedPrecisionRUHelper
+from model_compression_toolkit.core.common.mixed_precision.search_methods.linear_programming import \
+    MixedPrecisionIntegerLPSolver
 from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
+from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 from model_compression_toolkit.logger import Logger
 
 
@@ -43,8 +51,7 @@ class MixedPrecisionSearchManager:
                  fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation,
                  sensitivity_evaluator: SensitivityEvaluation,
-                 target_resource_utilization: ResourceUtilization,
-                 original_graph: Graph = None):
+                 target_resource_utilization: ResourceUtilization):
         """
 
         Args:
@@ -54,96 +61,208 @@ class MixedPrecisionSearchManager:
             sensitivity_evaluator: A SensitivityEvaluation which provides a function that evaluates the sensitivity of
                 a bit-width configuration for the MP model.
             target_resource_utilization: Target Resource Utilization to bound our feasible solution space s.t the configuration does not violate it.
-            original_graph: In case we have a search over a virtual graph (if we have BOPS utilization target), then this argument
-                will contain the original graph (for config reconstruction purposes).
         """
 
-        self.graph = graph
-        self.original_graph = graph if original_graph is None else original_graph
         self.fw_info = fw_info
         self.fw_impl = fw_impl
+
+        self.original_graph = graph
+        # graph for mp search
+        self.mp_graph, self.using_virtual_graph = self._get_mp_graph(graph, target_resource_utilization)
+        del graph  # so that it's not used by mistake
+
         self.sensitivity_evaluator = sensitivity_evaluator
+        self.target_resource_utilization = target_resource_utilization
+
+        self.mp_topo_configurable_nodes = self.mp_graph.get_configurable_sorted_nodes(fw_info)
         self.layer_to_bitwidth_mapping = self.get_search_space()
-        self.compute_metric_fn = self.get_sensitivity_metric()
-        self._cuts = None
 
-        # To define RU Total constraints we need to compute weights and activations even if they have no constraints
-        # TODO currently this logic is duplicated in linear_programming.py
-        targets = target_resource_utilization.get_restricted_targets()
-        if RUTarget.TOTAL in targets:
-            targets = targets.union({RUTarget.ACTIVATION, RUTarget.WEIGHTS}) - {RUTarget.TOTAL}
-        self.ru_targets_to_compute = targets
+        self.ru_targets = target_resource_utilization.get_restricted_targets()
+        self.ru_helper = MixedPrecisionRUHelper(self.mp_graph, fw_info, fw_impl)
 
-        self.ru_helper = MixedPrecisionRUHelper(graph, fw_info, fw_impl)
-        self.target_resource_utilization = target_resource_utilization
-        self.min_ru_config = self.graph.get_min_candidates_config(fw_info)
-        self.max_ru_config = self.graph.get_max_candidates_config(fw_info)
-        self.min_ru = self.ru_helper.compute_utilization(self.ru_targets_to_compute, self.min_ru_config)
-        self.non_conf_ru_dict = self.ru_helper.compute_utilization(self.ru_targets_to_compute, None)
+        self.min_ru_config = self.mp_graph.get_min_candidates_config(fw_info)
+        self.max_ru_config = self.mp_graph.get_max_candidates_config(fw_info)
+        self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, self.min_ru_config)
 
-        self.config_reconstruction_helper = ConfigReconstructionHelper(virtual_graph=self.graph,
+        self.config_reconstruction_helper = ConfigReconstructionHelper(virtual_graph=self.mp_graph,
                                                                        original_graph=self.original_graph)
 
-    def get_search_space(self) -> Dict[int, List[int]]:
+    def search(self) -> List[int]:
         """
-        The search space is a mapping from a node's index to a list of integers (possible bitwidths candidates indeces
-        for the node).
+        Run mixed precision search.
 
         Returns:
-            The entire search space of the graph.
+            Indices of the selected bit-widths candidates.
         """
+        candidates_sensitivity = self._build_sensitivity_mapping()
+        candidates_ru = self._compute_relative_ru_matrices()
+        rel_target_ru = self._get_relative_ru_constraint_per_mem_element()
+        solver = MixedPrecisionIntegerLPSolver(candidates_sensitivity, candidates_ru, rel_target_ru)
+        config = solver.run()
 
-        indices_mapping = {}
-        nodes_to_configure = self.graph.get_configurable_sorted_nodes(self.fw_info)
-        for idx, n in enumerate(nodes_to_configure):
-            # For each node, get all possible bitwidth indices for it
-            # (which is a list from 0 to the length of the candidates mp_config list of the node).
-            indices_mapping[idx] = list(range(len(n.candidates_quantization_cfg)))  # all search_methods space
-        return indices_mapping
+        if self.using_virtual_graph:
+            config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(config)
+        return config
 
-    def get_sensitivity_metric(self) -> Callable:
+    def _get_relative_ru_constraint_per_mem_element(self) -> Dict[RUTarget, np.ndarray]:
         """
+        Computes resource utilization constraint with respect to the minimal bit configuration, i.e. corresponding
+        constraint for each memory element is the relative utilization between the target utilization and
+        element's utilization for min-bit configuration.
+
+        Returns:
+            A dictionary of relative resource utilization constraints per ru target.
+
+        Raises:
+            ValueError: if target resource utilization cannot be satisfied (utilization for the minimal bit
+              configuration exceeds the requested target utilization for any target).
+        """
+        target_ru = self.target_resource_utilization.get_resource_utilization_dict(restricted_only=True)
+        rel_target_ru = {
+            ru_target: ru - self.min_ru[ru_target] for ru_target, ru in target_ru.items()
+        }
+        unsatisfiable_targets = {
+            ru_target.value: target_ru[ru_target] for ru_target, ru in rel_target_ru.items() if any(ru < 0)
+        }
+        if unsatisfiable_targets:
+            raise ValueError(f"The model cannot be quantized to meet the specified resource utilization for the "
+                             f"following targets: {unsatisfiable_targets}")
+        return rel_target_ru
+
+    def _build_sensitivity_mapping(self, eps: float = EPS) -> Dict[int, Dict[int, float]]:
+        """
+        This function measures the sensitivity of a change in a bitwidth of a layer on the entire model.
+        It builds a mapping from a node's index, to its bitwidht's effect on the model sensitivity.
+        For each node and some possible node's bitwidth (according to the given search space), we use
+        the framework function compute_metric_fn in order to infer
+        a batch of images, and compute (using the inference results) the sensitivity metric of
+        the configured mixed-precision model.
+
+        Args:
+            eps: Epsilon value to manually increase metric value (if necessary) for numerical stability
 
-        Returns: Return a function (from the framework implementation) to compute a metric that
-        indicates the similarity of the mixed-precision model (to the float model) for a given
-        mixed-precision configuration.
+        Returns:
+            Mapping from each node's index in a graph, to a dictionary from the bitwidth index (of this node) to
+            the sensitivity of the model.
 
         """
-        # Get from the framework an evaluation function on how a MP configuration,
-        # affects the expected loss.
 
-        return self.sensitivity_evaluator.compute_metric
+        Logger.info('Starting to evaluate metrics')
+        layer_to_metrics_mapping = {}
+
+        compute_metric = self.sensitivity_evaluator.compute_metric
+        if self.using_virtual_graph:
+            origin_max_config = self.config_reconstruction_helper.reconstruct_config_from_virtual_graph(
+                self.max_ru_config)
+            max_config_value = compute_metric(origin_max_config)
+        else:
+            max_config_value = compute_metric(self.max_ru_config)
+
+        for node_idx, layer_possible_bitwidths_indices in tqdm(self.layer_to_bitwidth_mapping.items(),
+                                                               total=len(self.layer_to_bitwidth_mapping)):
+            layer_to_metrics_mapping[node_idx] = {}
+
+            for bitwidth_idx in layer_possible_bitwidths_indices:
+                if self.max_ru_config[node_idx] == bitwidth_idx:
+                    # This is a computation of the metric for the max configuration, assign pre-calculated value
+                    layer_to_metrics_mapping[node_idx][bitwidth_idx] = max_config_value
+                    continue
+
+                # Create a configuration that differs at one layer only from the baseline model
+                mp_model_configuration = self.max_ru_config.copy()
+                mp_model_configuration[node_idx] = bitwidth_idx
+
+                # 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, c in enumerate(origin_max_config) if
+                                                    c != origin_mp_model_configuration[i]]
+                    metric_value = compute_metric(
+                        origin_mp_model_configuration,
+                        origin_changed_nodes_indices,
+                        origin_max_config)
+                else:
+                    metric_value = compute_metric(
+                        mp_model_configuration,
+                        [node_idx],
+                        self.max_ru_config)
+
+                layer_to_metrics_mapping[node_idx][bitwidth_idx] = max(metric_value, max_config_value + eps)
 
-    def compute_resource_utilization_matrix(self, target: RUTarget) -> np.ndarray:
+        # Finalize distance metric mapping
+        self.finalize_distance_metric(layer_to_metrics_mapping)
+
+        return layer_to_metrics_mapping
+
+    def _get_mp_graph(self, graph: Graph, target_resource_utilization: ResourceUtilization) -> Tuple[Graph, bool]:
         """
-        Computes and builds a resource utilization matrix, to be used for the mixed-precision search problem formalization.
-        Utilization is computed relative to the minimal configuration, i.e. utilization for it will be 0.
+        Get graph for mixed precision search. Virtual graph is built if bops is restricted and both activation and
+        weights are configurable.
 
         Args:
-            target: The resource target for which the resource utilization is calculated (a RUTarget value).
+            graph: input graph.
+            target_resource_utilization: target resource utilization.
+
+        Returns:
+            Graph for mixed precision search (virtual or original), and a boolean flag whether a virtual graph has been
+            constructed.
+        """
+        if (target_resource_utilization.bops_restricted() and
+                graph.has_any_configurable_activation() and
+                graph.has_any_configurable_weights()):
+            mp_graph = substitute(copy.deepcopy(graph),
+                                  self.fw_impl.get_substitutions_virtual_weights_activation_coupling())
+            return mp_graph, True
+
+        return graph, False
+
+    def get_search_space(self) -> Dict[int, List[int]]:
+        """
+        The search space is a mapping from a node's index to a list of integers (possible bitwidths candidates indeces
+        for the node).
 
         Returns:
-            A resource utilization matrix of shape (num configurations, num memory elements). Num memory elements
-            depends on the target, e.g. num nodes or num cuts, for which utilization is computed.
+            The entire search space of the graph.
         """
-        assert isinstance(target, RUTarget), f"{target} is not a valid resource target"
 
-        configurable_sorted_nodes = self.graph.get_configurable_sorted_nodes(self.fw_info)
+        indices_mapping = {}
+        for idx, n in enumerate(self.mp_topo_configurable_nodes):
+            # For each node, get all possible bitwidth indices for it
+            # (which is a list from 0 to the length of the candidates mp_config list of the node).
+            indices_mapping[idx] = list(range(len(n.candidates_quantization_cfg)))  # all search_methods space
+        return indices_mapping
+
+    def _compute_relative_ru_matrices(self) -> Dict[RUTarget, np.ndarray]:
+        """
+        Computes and builds a resource utilization matrix for all restricted targets, to be used for the
+        mixed-precision search problem formalization.
+        Utilization is computed relative to the minimal configuration, i.e. utilization for it will be 0.
 
-        ru_matrix = []
-        for c, c_n in enumerate(configurable_sorted_nodes):
+        Returns:
+            A dictionary containing resource utilization matrix of shape (num configurations, num memory elements)
+            per ru target. Num memory elements depends on the target, e.g. num cuts or 1 for cumulative metrics.
+        """
+        rus_per_candidate = defaultdict(list)
+        for c, c_n in enumerate(self.mp_topo_configurable_nodes):
             for candidate_idx in range(len(c_n.candidates_quantization_cfg)):
                 if candidate_idx == self.min_ru_config[c]:
-                    candidate_rus = self.min_ru[target]
+                    candidate_rus = self.min_ru
                 else:
-                    candidate_rus = self.compute_node_ru_for_candidate(c, candidate_idx, target)
+                    candidate_rus = self.compute_ru_for_candidate(c, candidate_idx)
 
-                ru_matrix.append(np.asarray(candidate_rus))
+                for target, ru in candidate_rus.items():
+                    rus_per_candidate[target].append(ru)
 
-        np_ru_matrix = np.array(ru_matrix) - self.min_ru[target]    # num configurations X num elements
-        return np_ru_matrix
+        # Each target contains a matrix of num configurations X num elements
+        relative_rus = {target: np.array(ru) - self.min_ru[target] for target, ru in rus_per_candidate.items()}
+        return relative_rus
 
-    def compute_node_ru_for_candidate(self, conf_node_idx: int, candidate_idx: int, target: RUTarget) -> np.ndarray:
+    def compute_ru_for_candidate(self, conf_node_idx: int, candidate_idx: int) -> Dict[RUTarget, np.ndarray]:
         """
         Computes a resource utilization vector after replacing the given node's configuration candidate in the minimal
         target configuration with the given candidate index.
@@ -151,13 +270,13 @@ class MixedPrecisionSearchManager:
         Args:
             conf_node_idx: The index of a node in a sorted configurable nodes list.
             candidate_idx: Quantization config candidate to be used for the node's resource utilization computation.
-            target: The target for which the resource utilization is calculated (a RUTarget value).
 
-        Returns: Node's resource utilization vector.
+        Returns:
+            Node's resource utilization vector.
 
         """
         cfg = self.replace_config_in_index(self.min_ru_config, conf_node_idx, candidate_idx)
-        return self.ru_helper.compute_utilization({target}, cfg)[target]
+        return self.ru_helper.compute_utilization(self.ru_targets, cfg)
 
     @staticmethod
     def replace_config_in_index(mp_cfg: List[int], idx: int, value: int) -> List[int]:
@@ -191,7 +310,7 @@ class MixedPrecisionSearchManager:
         act_qcs, w_qcs = self.ru_helper.get_quantization_candidates(config)
         ru = self.ru_helper.ru_calculator.compute_resource_utilization(
             target_criterion=TargetInclusionCriterion.AnyQuantized, bitwidth_mode=BitwidthMode.QCustom, act_qcs=act_qcs,
-            w_qcs=w_qcs, ru_targets=self.ru_targets_to_compute, allow_unused_qcs=True)
+            w_qcs=w_qcs, ru_targets=self.ru_targets, allow_unused_qcs=True)
         return ru
 
     def finalize_distance_metric(self, layer_to_metrics_mapping: Dict[int, Dict[int, float]]):

{mct_nightly-2.3.0.20250322.517 → mct_nightly-2.3.0.20250324.606}/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py

@@ -51,25 +51,34 @@ class ResourceUtilization:
     bops: float = np.inf
 
     def weight_restricted(self):
-        return self.weights_memory < np.inf
+        return self._is_restricted(self.weights_memory)
 
     def activation_restricted(self):
-        return self.activation_memory < np.inf
+        return self._is_restricted(self.activation_memory)
 
     def total_mem_restricted(self):
-        return self.total_memory < np.inf
+        return self._is_restricted(self.total_memory)
 
     def bops_restricted(self):
-        return self.bops < np.inf
+        return self._is_restricted(self.bops)
 
-    def get_resource_utilization_dict(self) -> Dict[RUTarget, float]:
+    def get_resource_utilization_dict(self, restricted_only: bool = False) -> Dict[RUTarget, float]:
         """
-        Returns: a dictionary with the ResourceUtilization object's values for each resource utilization target.
+        Get resource utilization as a dictionary.
+
+        Args:
+            restricted_only: whether to include only targets with restricted utilization.
+
+        Returns:
+            A dictionary containing the resource utilization with targets as keys.
         """
-        return {RUTarget.WEIGHTS: self.weights_memory,
-                RUTarget.ACTIVATION: self.activation_memory,
-                RUTarget.TOTAL: self.total_memory,
-                RUTarget.BOPS: self.bops}
+        ru_dict = {RUTarget.WEIGHTS: self.weights_memory,
+                   RUTarget.ACTIVATION: self.activation_memory,
+                   RUTarget.TOTAL: self.total_memory,
+                   RUTarget.BOPS: self.bops}
+        if restricted_only:
+            ru_dict = {k: v for k, v in ru_dict.items() if self._is_restricted(v)}
+        return ru_dict
 
     def is_satisfied_by(self, ru: 'ResourceUtilization') -> bool:
         """
@@ -114,3 +123,6 @@ class ResourceUtilization:
         if RUTarget.BOPS in targets:
             summary.append(f"BOPS: {self.bops}")
         return ', '.join(summary)
+
+    def _is_restricted(self, v):
+        return v < np.inf