ai-edge-quantizer-nightly 0.1.0.dev20250415__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/transformation_instruction_generator.py

@@ -23,10 +23,13 @@ from collections.abc import Iterator
 import dataclasses
 from typing import Optional
 from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import constrained_ops_utils
 from ai_edge_quantizer.utils import tfl_flatbuffer_utils
 from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
 
 
+_OpQuantConstraint = common_utils.OpQuantConstraint
 _QuantTransformation = qtyping.QuantTransformation
 
 
@@ -51,6 +54,15 @@ def check_horizontal_optimization(
   Returns:
     True if the two transformations can be merged, False otherwise
   """
+  if (
+      isinstance(param1.parameters, qtyping.UniformQuantParams)
+      and param1.parameters.hadamard is not None
+  ):
+    if (
+        isinstance(param2.parameters, qtyping.UniformQuantParams)
+        and param2.parameters.hadamard is not None
+    ):
+      return True
   return (
       param1.parameters == param2.parameters
       and len(param1.transformations) > index
@@ -165,6 +177,16 @@ class TransformationInstructionsGenerator:
     else:
       self.flatbuffer_model = tfl_flatbuffer_utils.read_model(float_tflite)
       self._create_tensor_name_to_graph_info_map()
+    self._same_as_input_scale_ops = (
+        constrained_ops_utils.get_constrained_op_list(
+            _OpQuantConstraint.SAME_AS_INPUT_SCALE
+        )
+    )
+    self._same_as_output_scale_ops = (
+        constrained_ops_utils.get_constrained_op_list(
+            _OpQuantConstraint.SAME_AS_OUTPUT_SCALE
+        )
+    )
 
   @dataclasses.dataclass(frozen=True)
   class TensorGraphInfo:
@@ -186,11 +208,13 @@ class TransformationInstructionsGenerator:
       A tuple of tensor_name and TensorGraphInfo.
     """
     for tensor_id, tensor in enumerate(subgraph.tensors):
-      consumers = [
-          op_id
-          for (op_id, op) in enumerate(subgraph.operators)
-          if tensor_id in op.inputs
-      ]
+      consumers = []
+      for op_id, op in enumerate(subgraph.operators):
+        # Some ops may use the same input tensor multiple times,
+        # and we should handle each time independently.
+        for op_input in op.inputs:
+          if op_input == tensor_id:
+            consumers.append(op_id)
       producer = -1
       for op_id, op in enumerate(subgraph.operators):
         if tensor_id in op.outputs:
@@ -504,6 +528,89 @@ class TransformationInstructionsGenerator:
       ):
         instructions.pop(i)
 
+  def _is_valid_quantize_requantize_pair(
+      self,
+      instr_0: qtyping.TransformationInst,
+      instr_1: qtyping.TransformationInst,
+  ) -> bool:
+    """Checks if the two instructions form a valid quantize and requantize pair."""
+    return (
+        instr_0.transformation == _QuantTransformation.QUANTIZE_TENSOR
+        and instr_1.transformation == _QuantTransformation.ADD_QUANTIZE
+        and instr_0.consumers == instr_1.consumers
+    )
+
+  def _is_op_constrained(
+      self, subgraph_id: int, op_index: int
+  ) -> bool:
+    """Checks if the op has same as input or output scale constraints."""
+    op_name = tfl_flatbuffer_utils.get_op_name_by_index(
+        self.flatbuffer_model, subgraph_id, op_index
+    )
+    return (
+        op_name in self._same_as_input_scale_ops
+        or op_name in self._same_as_output_scale_ops
+    )
+
+  def _are_quant_params_compatible(
+      self,
+      params_0: qtyping.UniformQuantParams,
+      params_1: qtyping.UniformQuantParams,
+  ) -> bool:
+    """Checks if quant params are the same except for the scale and zero point."""
+    ignore_set = {"scale", "zero_point"}
+    for field_info in dataclasses.fields(qtyping.UniformQuantParams):
+      field_name = field_info.name
+      if field_name in ignore_set:
+        continue
+      if getattr(params_0, field_name) != getattr(params_1, field_name):
+        return False
+    return True
+
+  def _eliminate_requantization_for_nonconstrained_provider(
+      self, tensor_trans_insts: qtyping.TensorTransformationInsts
+  ) -> None:
+    """Removes requantization for tensors with a non-constrained provider.
+
+    Fuses [QUANTIZE_TENSOR, ADD_QUANTIZE] instructions when a tensor has a
+    provider op without same as input/ouput scale constrains. Quant params from
+    the second instruction are copied to the first one and ADD_QUANTIZE is
+    removed.
+
+    Args:
+      tensor_trans_insts: Transformation instructions for a tensor.
+    """
+    instructions = tensor_trans_insts.instructions
+    if instructions is None or len(instructions) != 2:
+      return
+
+    instr_0, instr_1 = instructions
+    params_0 = instr_0.parameters
+    params_1 = instr_1.parameters
+    producer_op_index = instr_0.producer
+    if (
+        not isinstance(params_0, qtyping.UniformQuantParams)
+        or not isinstance(params_1, qtyping.UniformQuantParams)
+        or not self._is_valid_quantize_requantize_pair(instr_0, instr_1)
+        or not self._are_quant_params_compatible(params_0, params_1)
+        # To avoid fusion when subgraph inputs connected to the main subgraph
+        # (e.g. while_body), we skip all tensors with no producer.
+        or producer_op_index == -1
+        # Can't apply fusion to tensors with a constrained producer since that
+        # will break the constraint.
+        or self._is_op_constrained(
+            tensor_trans_insts.subgraph_id, producer_op_index
+        )
+    ):
+      return
+
+    # Fuse the quantize and requantize.
+    instr_0.parameters = dataclasses.replace(
+        params_0, scale=params_1.scale, zero_point=params_1.zero_point
+    )
+    # Remove the requantize instruction.
+    instructions.pop(1)
+
   def _quant_params_to_transformation_insts(
       self,
       param: qtyping.TensorTransformationParams,
@@ -576,6 +683,12 @@ class TransformationInstructionsGenerator:
     # will raise an error if the instructions are not valid.
     self._check_tensor_transformation_instructions_valid(tensor_trans_insts)
 
+    # Remove unnecessary [QUANTIZE_TENSOR, ADD_QUANTIZE] pairs for tensors with
+    # providers without same as input/output scale constraints.
+    self._eliminate_requantization_for_nonconstrained_provider(
+        tensor_trans_insts
+    )
+
     return tensor_trans_insts
 
   def _split_instructions_by_tensor_duplication(
@@ -671,7 +784,6 @@ class TransformationInstructionsGenerator:
     """
     is_tensor_unquantized = False
     is_tensor_quantized = False
-    is_operator_emulated = False
     for instruction in instructions:
       transform_type = instruction.transformation
       if transform_type == qtyping.QuantTransformation.NO_QUANTIZE:
@@ -681,17 +793,10 @@ class TransformationInstructionsGenerator:
           or transform_type == qtyping.QuantTransformation.ADD_DEQUANTIZE
       ):
         is_tensor_quantized = True
-      elif transform_type == qtyping.QuantTransformation.EMULATED_SUBCHANNEL:
-        is_operator_emulated = True
     if is_tensor_unquantized and is_tensor_quantized:
       raise ValueError(
           "Tensor %s can not be both quantized and unquantized" % tensor_name
       )
-    if is_operator_emulated and len(instructions) > 1:
-      raise ValueError(
-          "Tensor %s : op replacement transformation can not be combined with"
-          " other transformations." % tensor_name
-      )
 
   def _check_tensor_transformation_instructions_valid(
       self,

ai_edge_quantizer/transformation_instruction_generator_test.py

@@ -15,7 +15,9 @@
 
 """Tests for instruction_generator."""
 
+from collections.abc import Sequence
 import os
+from typing import Optional
 
 import numpy as np
 
@@ -953,33 +955,6 @@ class InstructionGeneratorTest(parameterized.TestCase):
         instructions["StatefulPartitionedCall:0"], output_transformation
     )
 
-  def test_raise_error_on_op_replacement_transformation_is_not_unique(self):
-    test_model_path = os.path.join(
-        TEST_DATA_PREFIX_PATH, "tests/models/insert_dequant_test.tflite"
-    )
-    quant_parameters = {}
-    quant_parameters["tfl.quantize"] = qtyping.TensorTransformationParams(
-        "tfl.quantize",
-        qtyping.OpToTensorParams(
-            subgraph_op_id=0,
-            transformations=[
-                qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                qtyping.QuantTransformation.EMULATED_SUBCHANNEL,
-            ],
-            parameters=qtyping.UniformQuantParams(
-                8, None, np.array([1]), np.array([0])
-            ),
-        ),
-        [],
-    )
-    ins_gen = instruction_generator.TransformationInstructionsGenerator(
-        test_model_path
-    )
-    with self.assertRaisesRegex(
-        ValueError, "op replacement transformation can not be combined"
-    ):
-      ins_gen.quant_params_to_transformation_insts(quant_parameters)
-
   def test_raise_error_on_no_quant_conflict(self):
     test_model_path = os.path.join(
         TEST_DATA_PREFIX_PATH, "tests/models/insert_dequant_test.tflite"
@@ -1364,5 +1339,166 @@ class InstructionGeneratorTest(parameterized.TestCase):
       )
 
 
+class EliminateUnnecessaryRequantizationTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.ins_gen = instruction_generator.TransformationInstructionsGenerator(
+        os.path.join(TEST_DATA_PREFIX_PATH, "tests/models/conv_fc_mnist.tflite")
+    )
+
+  def _get_test_instruction(
+      self,
+      transformation: qtyping.QuantTransformation,
+      producer: int = -1,
+      consumers: Optional[Sequence[int]] = None,
+      qparams: Optional[qtyping.UniformQuantParams] = None,
+  ) -> qtyping.TransformationInst:
+    if consumers is None:
+      consumers = []
+    if qparams is None:
+      qparams = qtyping.UniformQuantParams(
+          num_bits=8,
+          quantized_dimension=None,
+          scale=np.array([1]),
+          zero_point=np.array([0]),
+      )
+    return qtyping.TransformationInst(
+        transformation=transformation,
+        producer=producer,
+        consumers=consumers,
+        parameters=qparams,
+        # Dummy values below.
+        tensor_id=0,
+    )
+
+  def _create_test_insts(
+      self, instructions: list[qtyping.TransformationInst]
+  ) -> qtyping.TensorTransformationInsts:
+    return qtyping.TensorTransformationInsts(
+        tensor_name="test_tensor", subgraph_id=0, instructions=instructions
+    )
+
+  def test_no_fusion_when_too_few_instructions(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 1)
+
+  def test_no_fusion_when_too_many_instructions(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE),
+        self._get_test_instruction(_QTransf.ADD_DEQUANTIZE),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 3)
+
+  def test_no_fusion_when_invalid_transformation_pair(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.ADD_DEQUANTIZE),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_consumers_mismatch(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR, consumers=[0]),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE, consumers=[1]),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_no_producer(self):
+    producer = -1
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR, producer),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE, producer),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_quant_params_are_incompatible(self):
+    params_8_bits = qtyping.UniformQuantParams(
+        8, None, np.array([1]), np.array([0])
+    )
+    params_16_bits = qtyping.UniformQuantParams(
+        16, None, np.array([1]), np.array([0])
+    )
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(
+            _QTransf.QUANTIZE_TENSOR, qparams=params_8_bits
+        ),
+        self._get_test_instruction(
+            _QTransf.ADD_QUANTIZE, qparams=params_16_bits
+        ),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_producer_constrained(self):
+    # Reshape op (op index 2) has same as input scale constraint.
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR, producer=2),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE, producer=2),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_fusion_succeeds(self):
+    producer = 0
+    consumers = [1]
+    params_0 = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.array([1]),
+        zero_point=np.array([0]),
+    )
+    params_1 = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.array([2]),
+        zero_point=np.array([1]),
+    )
+    inst_0 = self._get_test_instruction(
+        _QTransf.QUANTIZE_TENSOR, producer, consumers, params_0
+    )
+    inst_1 = self._get_test_instruction(
+        _QTransf.ADD_QUANTIZE, producer, consumers, params_1
+    )
+    tensor_insts = self._create_test_insts([inst_0, inst_1])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+
+    self.assertLen(tensor_insts.instructions, 1)
+    result_inst = tensor_insts.instructions[0]
+    self.assertEqual(result_inst.transformation, _QTransf.QUANTIZE_TENSOR)
+
+    result_params = result_inst.parameters
+    # Explicitly narrow the type for pytype.
+    if not isinstance(result_params, qtyping.UniformQuantParams):
+      self.fail("Fused instruction parameters are not UniformQuantParams")
+
+    self.assertEqual(result_params.scale, params_1.scale)
+    self.assertEqual(result_params.zero_point, params_1.zero_point)
+
+
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/transformation_performer.py

@@ -24,7 +24,8 @@ from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.transformations import dequant_insert
 from ai_edge_quantizer.transformations import duplicate_buffer
 from ai_edge_quantizer.transformations import duplicate_tensor
-from ai_edge_quantizer.transformations import emulated_subchannel
+from ai_edge_quantizer.transformations import insert_decomposed_hadamard_rotation
+from ai_edge_quantizer.transformations import insert_hadamard_rotation
 from ai_edge_quantizer.transformations import quant_insert
 from ai_edge_quantizer.transformations import quantize_tensor
 from ai_edge_quantizer.transformations import transformation_utils
@@ -71,7 +72,7 @@ class TransformationPerformer:
             quantize_tensor.quantize_tensor
         ),
         qtyping.QuantTransformation.EMULATED_SUBCHANNEL: (
-            emulated_subchannel.emulated_subchannel
+            transformation_utils.raise_deprecated_error
         ),
         qtyping.QuantTransformation.ADD_QUANTIZE: quant_insert.insert_quant,
         qtyping.QuantTransformation.DUPLICATE_BUFFER: (
@@ -80,6 +81,12 @@ class TransformationPerformer:
         qtyping.QuantTransformation.DUPLICATE_TENSOR: (
             duplicate_tensor.duplicate_tensor
         ),
+        qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION: (
+            insert_hadamard_rotation.insert_hadamard_rotation
+        ),
+        qtyping.QuantTransformation.INSERT_DECOMPOSED_HADAMARD_ROTATION: (
+            insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation
+        ),
     }
     # transformations are seprated in two categories:
     # op_insertion_transformations are transformations that only insert ops
@@ -91,6 +98,8 @@ class TransformationPerformer:
         qtyping.QuantTransformation.ADD_QUANTIZE,
         qtyping.QuantTransformation.DUPLICATE_BUFFER,
         qtyping.QuantTransformation.DUPLICATE_TENSOR,
+        qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
+        qtyping.QuantTransformation.INSERT_DECOMPOSED_HADAMARD_ROTATION,
     ])
     self._op_replacement_transformations = set(
         [qtyping.QuantTransformation.EMULATED_SUBCHANNEL]
@@ -180,6 +189,38 @@ class TransformationPerformer:
             )
           transformation.tensor_id = trans_info.output_tensor_id
 
+  def _get_updated_producer_id(
+      self, original_producer_id: int, subgraph_id: int
+  ) -> int:
+    """Update the producer of a transformation instruction."""
+    if original_producer_id is None or original_producer_id < 0:
+      producer = -1
+    elif original_producer_id < len(self._original_op_id_map[subgraph_id]):
+      producer = self._original_op_id_map[subgraph_id][original_producer_id]
+    else:
+      # If the producer id is not in the original op map, it's an added op,
+      # go the added op map to find the producer.
+      producer = self._added_op_id_map[subgraph_id][
+          original_producer_id - len(self._original_op_id_map[subgraph_id])
+      ]
+    return producer
+
+  def _get_updated_consumer_ids(
+      self,
+      original_consumer_ids: list[int],
+      subgraph_id: int,
+  ) -> list[int]:
+    """Update the consumers of a transformation instruction."""
+    consumers = []
+    for original_op_id in original_consumer_ids:
+      new_consumer_id = (
+          -1
+          if original_op_id == -1
+          else self._original_op_id_map[subgraph_id][original_op_id]
+      )
+      consumers.append(new_consumer_id)
+    return consumers
+
   def _apply_single_transformation(
       self,
       transformation_inst: qtyping.TensorTransformationInsts,
@@ -198,28 +239,12 @@ class TransformationPerformer:
       None, update the transformation_inst & tflite_model in place
     """
     instruction = transformation_inst.instructions[transformation_index]
-    if not instruction.producer or instruction.producer < 0:
-      producer = -1
-    elif instruction.producer < len(
-        self._original_op_id_map[transformation_inst.subgraph_id]
-    ):
-      producer = self._original_op_id_map[transformation_inst.subgraph_id][
-          instruction.producer
-      ]
-    else:
-      # if the producer id is not in the original op map, it's an added op,
-      # go the corresponding new maps
-      producer = self._added_op_id_map[transformation_inst.subgraph_id][
-          instruction.producer
-          - len(self._original_op_id_map[transformation_inst.subgraph_id])
-      ]
-    consumers = []
-    for original_op_id in instruction.consumers:
-      consumers.append(
-          self._original_op_id_map[transformation_inst.subgraph_id][
-              original_op_id
-          ]
-      )
+    producer = self._get_updated_producer_id(
+        instruction.producer, transformation_inst.subgraph_id
+    )
+    consumers = self._get_updated_consumer_ids(
+        instruction.consumers, transformation_inst.subgraph_id
+    )
     trans_info = self._transformation_registration[instruction.transformation](
         transformation_utils.TransformationInput(
             instruction.tensor_id,
@@ -239,7 +264,12 @@ class TransformationPerformer:
     )
     self._update_op_id_map(
         transformation_inst.subgraph_id,
-        min(instruction.consumers),
+        # The added op must be right before the most immediate consumer, unless
+        # the consumer is the graph output (id=-1), then use the producer's
+        # index instead.
+        min(instruction.consumers)
+        if min(instruction.consumers) >= 0
+        else instruction.producer + 1,
         trans_info.num_ops_added,
     )