Fix PR comments

Author: elad-c

tests_pytest/_test_util/graph_builder_utils.py

@@ -70,7 +70,7 @@ def full_attr_name(canonical_name: Union[str, dict, Iterable]):
 
 
 def build_nbits_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, ()),
-                   convert_canonical_attr=True) -> CandidateNodeQuantizationConfig:
+                   convert_canonical_attr=True, q_preserving=False) -> CandidateNodeQuantizationConfig:
     """
     Build quantization config with configurable nbits and enabling/disabling quantization only.
 
@@ -87,6 +87,8 @@ def build_nbits_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, (
     Returns:
 
     """
+    assert not(a_enable and q_preserving)
+
     w_attr = w_attr or {}
     attr_weights_configs_mapping = {
         k: AttributeQuantizationConfig(weights_n_bits=v[0], enable_weights_quantization=v[1])
@@ -102,7 +104,7 @@ def build_nbits_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, (
         default_weight_attr_config=AttributeQuantizationConfig(weights_n_bits=pos_attr[0],
                                                                enable_weights_quantization=pos_attr[1]),
         activation_quantization_method=QuantizationMethod.POWER_OF_TWO,
-        quantization_preserving=False,
+        quantization_preserving=q_preserving,
         supported_input_activation_n_bits=[2, 4, 8],
         fixed_scale=None,
         fixed_zero_point=None,

tests_pytest/common_tests/unit_tests/core/graph/test_quantization_preserving_node.py

@@ -15,70 +15,18 @@
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.graph.edge import Edge
 
-from mct_quantizers import QuantizationMethod
-from model_compression_toolkit.core import QuantizationConfig
-
-from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
-    CandidateNodeQuantizationConfig
-from model_compression_toolkit.core.common.quantization.node_quantization_config import \
-    NodeActivationQuantizationConfig, NodeWeightsQuantizationConfig
-from model_compression_toolkit.target_platform_capabilities import AttributeQuantizationConfig, OpQuantizationConfig, \
-    Signedness
-
-from tests_pytest._test_util.graph_builder_utils import build_node
-
-
-def build_quant_preserving_qc(a_enable=True, qp_enable=False) -> CandidateNodeQuantizationConfig:
-    """
-    Build quantization config with enabling/disabling quantization and quantization preserving only.
-
-    Args:
-        a_enable: whether to enable activation quantization.
-        qp_enable: whether to enable activation quantization preserving flag.
-
-    Returns:
-        CandidateNodeQuantizationConfig object.
-
-    """
-    qc = QuantizationConfig()
-    # positional attrs are set via default weight config (so all pos attrs have the same q config)
-    op_cfg = OpQuantizationConfig(
-        # canonical names (as 'kernel')
-        attr_weights_configs_mapping={},
-        activation_n_bits=8,
-        enable_activation_quantization=a_enable,
-        default_weight_attr_config=AttributeQuantizationConfig(weights_n_bits=8,
-                                                               enable_weights_quantization=False),
-        activation_quantization_method=QuantizationMethod.POWER_OF_TWO,
-        quantization_preserving=qp_enable,
-        supported_input_activation_n_bits=[2, 4, 8],
-        fixed_scale=None,
-        fixed_zero_point=None,
-        simd_size=None,
-        signedness=Signedness.AUTO
-    )
-    a_qcfg = NodeActivationQuantizationConfig(qc=qc, op_cfg=op_cfg,
-                                              activation_quantization_fn=None,
-                                              activation_quantization_params_fn=None)
-    # full names from the layers
-    w_qcfg = NodeWeightsQuantizationConfig(qc=qc, op_cfg=op_cfg,
-                                           weights_channels_axis=None,
-                                           node_attrs_list=[])
-    qc = CandidateNodeQuantizationConfig(activation_quantization_cfg=a_qcfg,
-                                         weights_quantization_cfg=w_qcfg)
-
-    return qc
+from tests_pytest._test_util.graph_builder_utils import build_node, build_nbits_qc
 
 
 class TestQuantizationPreservingNode:
 
     def test_activation_preserving_candidate(self):
         """ Tests that the correct activation quantization candidate is selected. """
-        n1 = build_node('qact_node', qcs=[build_quant_preserving_qc()])
-        n2 = build_node('qp1a_node', qcs=[build_quant_preserving_qc(a_enable=False, qp_enable=True)])
-        n3 = build_node('qp1b_node', qcs=[build_quant_preserving_qc(a_enable=False, qp_enable=True)])
-        n4 = build_node('qp2a_node', qcs=[build_quant_preserving_qc()])
-        n5 = build_node('qp2b_node', qcs=[build_quant_preserving_qc(a_enable=False, qp_enable=True)])
+        n1 = build_node('qact_node', qcs=[build_nbits_qc()])
+        n2 = build_node('qp1a_node', qcs=[build_nbits_qc(a_enable=False, q_preserving=True)])
+        n3 = build_node('qp1b_node', qcs=[build_nbits_qc(a_enable=False, q_preserving=True)])
+        n4 = build_node('qp2a_node', qcs=[build_nbits_qc()])
+        n5 = build_node('qp2b_node', qcs=[build_nbits_qc(a_enable=False, q_preserving=True)])
         graph = Graph('g', input_nodes=[n1], nodes=[n2, n4], output_nodes=[n3, n5],
                       edge_list=[Edge(n1, n2, 0, 0), Edge(n2, n3, 0, 0),
                                  Edge(n1, n4, 0, 0), Edge(n4, n5, 0, 0)])

tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py

@@ -298,9 +298,10 @@ class TestComputeActivationTensorsUtilization:
     """ Tests for activation tensors utilization public apis. """
     def test_compute_node_activation_tensor_utilization(self, graph_mock, fw_impl_mock, fw_info_mock):
         mp_reuse = build_node('mp_reuse', output_shape=(None, 3, 14), qcs=[build_qc(4), build_qc(16)], reuse=True)
+        qp = build_node('qp', output_shape=(None, 15, 9), qcs=[build_qc(a_enable=False, q_preserving=True)])
         noq = build_node('noq', output_shape=(None, 15, 9), qcs=[build_qc(a_enable=False)])
-        graph_mock.nodes = [mp_reuse, noq]
-        graph_mock.retrieve_preserved_quantization_node = _identity_func
+        graph_mock.nodes = [mp_reuse, qp, noq]
+        graph_mock.retrieve_preserved_quantization_node = lambda n: mp_reuse if n is qp else n
 
         ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock)
         # _get_activation_nbits is already fully checked, just make sure we use it, and use correctly
@@ -313,6 +314,9 @@ def test_compute_node_activation_tensor_utilization(self, graph_mock, fw_impl_mo
         # reused is not ignored
         res = ru_calc.compute_node_activation_tensor_utilization(mp_reuse, TIC.QConfigurable, BM.QMinBit)
         assert res == Utilization(42, 21.)
+        # quantization preserving uses custom_qc.
+        res = ru_calc.compute_node_activation_tensor_utilization(qp, TIC.AnyQuantized, BM.QCustom, custom_qc)
+        assert res == Utilization(135, 270.)
         # not a target node
         res = ru_calc.compute_node_activation_tensor_utilization(noq, TIC.AnyQuantized, BM.QCustom, custom_qc)
         assert res == Utilization(0, 0)
@@ -394,11 +398,14 @@ def test_compute_cuts_integration(self, graph_mock, fw_impl_mock, fw_info_mock,
         """ Test integration with max cut computation. """
         # Test a simple linear dummy graph with the real max cut computation.
         n1 = build_node('n1', qcs=[build_qc()], input_shape=(None, 10, 20, 3), output_shape=(None, 10, 20, 3))
+        n1_qp = build_node('n1_qp', qcs=[build_qc(a_enable=False, q_preserving=True)],
+                           input_shape=(None, 10, 20, 3), output_shape=(None, 10, 20, 3))
         n2 = build_node('n2', qcs=[build_qc()], input_shape=(None, 10, 20, 3), output_shape=(None, 5, 10))
         n3 = build_node('n3', qcs=[build_qc()], input_shape=(None, 5, 10), output_shape=(None, 5, 10))
         n4 = build_node('n4', qcs=[build_qc()], input_shape=(None, 5, 10, 32), output_shape=(None, 5, 10, 32))
-        edges = [Edge(n1, n2, 0, 0), Edge(n2, n3, 0, 0), Edge(n3, n4, 0, 0)]
-        graph = Graph('g', input_nodes=[n1], nodes=[n2, n3], output_nodes=[n4], edge_list=edges)
+        edges = [Edge(n1, n1_qp, 0, 0), Edge(n1_qp, n2, 0, 0),
+                 Edge(n2, n3, 0, 0), Edge(n3, n4, 0, 0)]
+        graph = Graph('g', input_nodes=[n1], nodes=[n1_qp, n2, n3], output_nodes=[n4], edge_list=edges)
         ru_calc = ResourceUtilizationCalculator(graph, fw_impl_mock, fw_info_mock)
         # wrap the real implementation
         maxcut_spy = mocker.patch('model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.'
@@ -408,11 +415,11 @@ def test_compute_cuts_integration(self, graph_mock, fw_impl_mock, fw_info_mock,
         cuts_cache = ru_calc.cuts
 
         # verify the cache
-        assert len(cuts_cache) == 5
+        assert len(cuts_cache) == 6
         assert all(isinstance(k, Cut) for k in cuts_cache.keys())
         # for each cut we save a list of its nodes
         cuts_nodes = {tuple(sorted(n.name for n in nodes)) for nodes in cuts_cache.values()}
-        assert cuts_nodes == {('n1',), ('n4',), ('n1', 'n2'), ('n2', 'n3'), ('n3', 'n4')}
+        assert cuts_nodes == {('n1',), ('n4',), ('n1', 'n1_qp'), ('n1_qp', 'n2'), ('n2', 'n3'), ('n3', 'n4')}
 
         # verify cuts computation only happens the first time
         cuts_cache2 = ru_calc.cuts
@@ -423,7 +430,8 @@ def test_compute_cuts_integration(self, graph_mock, fw_impl_mock, fw_info_mock,
         nodes_to_cuts = {tuple(sorted(elem.node_name for elem in cut.mem_elements.elements)): cut
                          for cut in cuts_cache.keys()}
         cut1 = nodes_to_cuts[('n1',)]
-        cut12 = nodes_to_cuts[('n1', 'n2')]
+        cut11 = nodes_to_cuts[('n1', 'n1_qp')]
+        cut12 = nodes_to_cuts[('n1_qp', 'n2')]
         cut23 = nodes_to_cuts[('n2', 'n3')]
         cut34 = nodes_to_cuts[('n3', 'n4')]
         cut4 = nodes_to_cuts[('n4',)]
@@ -433,7 +441,8 @@ def test_compute_cuts_integration(self, graph_mock, fw_impl_mock, fw_info_mock,
                                                                                          bitwidth_mode=BM.QDefaultSP)
 
         assert per_cut_per_node == {cut1: {'n1': Utilization(10 * 20 * 3, 600)},
-                                    cut12: {'n1': Utilization(10 * 20 * 3, 600),
+                                    cut11: {'n1': Utilization(10 * 20 * 3, 600), 'n1_qp': Utilization(10 * 20 * 3, 600)},
+                                    cut12: {'n1_qp': Utilization(10 * 20 * 3, 600),
                                             'n2': Utilization(5 * 10, 50)},
                                     cut23: {'n2': Utilization(5*10, 50),
                                             'n3': Utilization(5*10, 50)},
@@ -442,7 +451,8 @@ def test_compute_cuts_integration(self, graph_mock, fw_impl_mock, fw_info_mock,
                                     cut4: {'n4': Utilization(5 * 10 * 32, 1600)}}
         assert per_cut == {
             nodes_to_cuts[('n1',)]: Utilization(600, 600),
-            nodes_to_cuts[('n1', 'n2')]: Utilization(650, 650),
+            nodes_to_cuts[('n1', 'n1_qp')]: Utilization(1200, 1200),
+            nodes_to_cuts[('n1_qp', 'n2')]: Utilization(650, 650),
             nodes_to_cuts[('n2', 'n3')]: Utilization(100, 100),
             nodes_to_cuts[('n3', 'n4')]: Utilization(1650, 1650),
             nodes_to_cuts[('n4',)]: Utilization(1600, 1600)