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

ai_edge_quantizer/transformation_performer_test.py

@@ -112,6 +112,32 @@ class TransformationPerformerTest(parameterized.TestCase):
     for index, op_id in enumerate(op_id_map[0]):
       self.assertEqual(op_id, index)
 
+  def test_update_op_id_map_not_changing_value_single_op_model(self):
+    """test for _update_op_id_map."""
+    model = tfl_flatbuffer_utils.read_model(
+        os.path.join(
+            TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias.tflite"
+        )
+    )
+    self._transformation_performer._create_op_id_map(model)
+    instruction = qtyping.TransformationInst(
+        transformation=qtyping.QuantTransformation.QUANTIZE_TENSOR,
+        tensor_id=0,
+        producer=0,
+        consumers=[-1],
+        parameters=qtyping.UniformQuantParams(
+            8, None, np.array([1]), np.array([0])
+        ),
+    )
+    producer = self._transformation_performer._get_updated_producer_id(
+        instruction.producer, 0
+    )
+    consumers = self._transformation_performer._get_updated_consumer_ids(
+        instruction.consumers, 0
+    )
+    self.assertEqual(producer, 0)
+    self.assertEqual(consumers, [-1])
+
   @parameterized.named_parameters(
       dict(
           testcase_name="test_no_update",
@@ -271,7 +297,7 @@ class TransformationPerformerTest(parameterized.TestCase):
         expected_added_op_id_map,
     )
 
-  def test__update_instructions_updates_tensor_id_after_duplicate_tensor(self):
+  def test_update_instructions_updates_tensor_id_after_duplicate_tensor(self):
     def get_test_instruction(transformation, consumers):
       return qtyping.TransformationInst(
           transformation=transformation,
@@ -325,6 +351,8 @@ class TransformationPerformerTest(parameterized.TestCase):
             tensor_name="sequential/conv2d/Relu;sequential/conv2d/BiasAdd;"
             + "sequential/conv2d/Conv2D;sequential/conv2d/BiasAdd/ReadVariableOp1",
             subgraph_id=0,
+            # Conv2d: op_id=0, output_tensor_id=7.
+            # This should add two sequential dequants after the conv2d.
             instructions=[
                 qtyping.TransformationInst(
                     transformation=qtyping.QuantTransformation.ADD_DEQUANTIZE,
@@ -349,6 +377,8 @@ class TransformationPerformerTest(parameterized.TestCase):
         "sequential/average_pooling2d/AvgPool": qtyping.TensorTransformationInsts(
             tensor_name="sequential/average_pooling2d/AvgPool",
             subgraph_id=0,
+            # Avg_pool: op_id=1, output_tensor_id=8.
+            # This should add two sequential dequants after the avg_pool.
             instructions=[
                 qtyping.TransformationInst(
                     transformation=qtyping.QuantTransformation.ADD_DEQUANTIZE,
@@ -376,19 +406,111 @@ class TransformationPerformerTest(parameterized.TestCase):
     )
     self.assertLen(self._test_model.subgraphs, 1)
     self.assertLen(self._test_model.subgraphs[0].operators, 10)
+    # The original model has 13 tensors, each dequant adds 1 tensor.
     self.assertLen(self._test_model.subgraphs[0].tensors, 17)
+    # Check that the dequant opcode is added to the model.
     self.assertEqual(
         self._test_model.subgraphs[0].operators[1].opcodeIndex,
         len(self._test_model.operatorCodes) - 1,
     )
+    # Conv2d, dequant, dequant, avgpool, dequant, dequant, etc.
+    expected_builtin_op_order = [3, 6, 6, 1, 6, 6, 22, 9, 9, 25]
+    for i, op in enumerate(self._test_model.subgraphs[0].operators):
+      op_code = self._test_model.operatorCodes[op.opcodeIndex].builtinCode
+      self.assertEqual(op_code, expected_builtin_op_order[i])
+    # Check that the first dequant input is connected to the conv2d output.
+    self.assertEqual(self._test_model.subgraphs[0].operators[1].inputs[0], 7)
+    # Output is a new tensor just added.
+    self.assertEqual(self._test_model.subgraphs[0].operators[1].outputs[0], 13)
+    # Second dequant has new tensors.
     self.assertEqual(self._test_model.subgraphs[0].operators[2].inputs[0], 13)
     self.assertEqual(self._test_model.subgraphs[0].operators[2].outputs[0], 14)
-    self.assertEqual(
-        self._test_model.subgraphs[0].operators[2].outputs[0],
-        self._test_model.subgraphs[0].operators[3].inputs[0],
-    )
+    # Avgpool's input is second dequant's output.
+    self.assertEqual(self._test_model.subgraphs[0].operators[3].inputs[0], 14)
+    # Avgpool's output remains the same.
     self.assertEqual(self._test_model.subgraphs[0].operators[3].outputs[0], 8)
+    # Third dequant's output is a new tensor.
     self.assertEqual(self._test_model.subgraphs[0].operators[4].outputs[0], 15)
+    # Fourth dequant.
+    self.assertEqual(self._test_model.subgraphs[0].operators[5].inputs[0], 15)
+    self.assertEqual(self._test_model.subgraphs[0].operators[5].outputs[0], 16)
+
+    # Avgpool (op_id=1) and reshape (op_id=2) are bumped by 2 due to the two
+    # dequants added after it.
+    expected_op_id_map = [0, 3, 6, 7, 8, 9]
+    self.assertEqual(
+        self._transformation_performer._original_op_id_map[0],
+        expected_op_id_map,
+    )
+    # New dequants are added at these indices.
+    expected_added_op_id_map = [1, 2, 4, 5]
+    self.assertEqual(
+        self._transformation_performer._added_op_id_map[0],
+        expected_added_op_id_map,
+    )
+
+  def test_op_insertion_at_input_and_output(self):
+    """test for _update_op_id_map."""
+    model = tfl_flatbuffer_utils.read_model(
+        os.path.join(
+            TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias.tflite"
+        )
+    )
+    self._transformation_performer._create_op_id_map(model)
+    instructions = {
+        # Fully_connected: op_id=0, input_tensor_id=0, output_tensor_id=3.
+        # Add a new quantize op to the input of the fully_connected.
+        "serving_default_input_2:0": qtyping.TensorTransformationInsts(
+            tensor_name="serving_default_input_2:0",
+            subgraph_id=0,
+            instructions=[
+                qtyping.TransformationInst(
+                    transformation=qtyping.QuantTransformation.ADD_QUANTIZE,
+                    tensor_id=0,
+                    producer=-1,
+                    consumers=[0],
+                    parameters=qtyping.UniformQuantParams(
+                        8, None, np.array([1]), np.array([0])
+                    ),
+                ),
+            ],
+        ),
+        # Add a new dequantize op to the output of the fully_connected.
+        "StatefulPartitionedCall:0": qtyping.TensorTransformationInsts(
+            tensor_name="StatefulPartitionedCall:0",
+            subgraph_id=0,
+            instructions=[
+                qtyping.TransformationInst(
+                    transformation=qtyping.QuantTransformation.ADD_DEQUANTIZE,
+                    tensor_id=3,
+                    producer=0,
+                    consumers=[-1],
+                    parameters=qtyping.UniformQuantParams(
+                        8, None, np.array([1]), np.array([0])
+                    ),
+                ),
+            ],
+        ),
+    }
+    self._transformation_performer.transform_graph(instructions, model)
+
+    # Original fc (op_id=0) should be bumped to op_id=1.
+    self.assertEqual(
+        self._transformation_performer._original_op_id_map[0],
+        [1],
+    )
+    # New quantize added at op_id=0, dequantize added at op_id=1.
+    expected_added_op_id_map = [0, 2]
+    self.assertEqual(
+        self._transformation_performer._added_op_id_map[0],
+        expected_added_op_id_map,
+    )
+    # Quantize, fully_connected, dequantize.
+    expected_builtin_op_order = [114, 9, 6]
+    for i, op in enumerate(model.subgraphs[0].operators):
+      op_code = model.operatorCodes[op.opcodeIndex].builtinCode
+      self.assertEqual(op_code, expected_builtin_op_order[i])
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/transformations/duplicate_buffer.py

@@ -34,9 +34,10 @@ def duplicate_buffer(
         f' Tensor {tensor_name} is not constant.'
     )
 
-  duplicated_buffer_id = transformation_utils.add_new_constant_buffer(
+  duplicated_buffer_id = transformation_utils.get_constant_buffer(
       data=buffer_data,
       buffers=transformation_input.buffers,
+      force_duplicate_buffer=True,
   )
   tensor.buffer = duplicated_buffer_id
 

ai_edge_quantizer/transformations/duplicate_tensor.py

@@ -41,6 +41,7 @@ def duplicate_tensor(
       tensor_shape=tensor.shape,
       subgraph=subgraph,
       buffers=transformation_input.buffers,
+      force_duplicate_buffer=True,
   )
   # Update the tensor name to avoid name collision in case when tensor is
   # duplicated mulitple times.

ai_edge_quantizer/transformations/insert_decomposed_hadamard_rotation.py

@@ -0,0 +1,299 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Hadamard rotation decomposed pattern transformation."""
+
+from flatbuffers import flexbuffers
+import numpy as np
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+
+def _to_flexbuffer(
+    hadamard_size: int,
+    random_binary_vector: list[np.int8],
+) -> bytes:
+  """Converts hadamard_size to flexbuffer."""
+  fbb = flexbuffers.Builder()
+  with fbb.Map():
+    fbb.Int('hadamard_size', hadamard_size)
+    fbb.VectorFromElements('random_binary_vector', random_binary_vector)
+  return fbb.Finish()
+
+
+def _update_embedding_lookup_consumers(
+    transformation: transformation_utils.TransformationInput,
+    new_tensor_id: int,
+) -> bool:
+  """Updates the consumers of the embedding lookup op to use the new tensor.
+
+  Args:
+    transformation: The transformation input to update the consumers of.
+    new_tensor_id: The new tensor id to use as the input to the embedding lookup
+      consumers.
+  """
+  for consumer in transformation.consumers:
+    # If the consumer is a graph output and not an op, we can ignore it here
+    # since the graph output will be updated later.
+    if consumer == -1:
+      continue
+    consumer_op = transformation.subgraph.operators[consumer]
+    # Find the input that was attached to the insertion point, and replace it
+    # with the new tensor.
+    for i in range(len(consumer_op.inputs)):
+      if consumer_op.inputs[i] == transformation.tensor_id:
+        consumer_op.inputs[i] = new_tensor_id
+
+
+def _update_fully_connected_consumers(
+    transformation: transformation_utils.TransformationInput,
+    new_tensor_id: int,
+) -> bool:
+  """Updates the fully connected op(s) to use the new tensor.
+
+  Since the new tensor is inserted to the fully_connected's input, we need to
+  scan each consumer (in case of multiple fully_connected ops), and update
+  the input tensor to the new tensor.
+
+  Args:
+    transformation: The transformation input to update the consumers of.
+    new_tensor_id: The new tensor id to use as the input to the fully connected
+      consumers.
+
+  Returns:
+    True if the fully connected op(s) were updated to use the new tensor.
+  """
+  updated = False
+  for consumer in transformation.consumers:
+    if (
+        transformation_utils.get_schema_op_id(transformation, consumer)
+        == schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+    ):
+      transformation.subgraph.operators[consumer].inputs[0] = new_tensor_id
+      updated = True
+  return updated
+
+
+def _make_hadamard_matrix(size: int):
+  """Generates a Hadamard matrix of the given size.
+
+  Args:
+    size: The size of the Hadamard matrix. Must be a power of 2. This represents
+      a single dimension. E.g. if size is 4, then the Hadamard matrix is a 4x4
+      matrix.
+
+  Returns:
+    The Hadamard matrix.
+
+  Raises:
+    ValueError: If the size is not a power of 2.
+  """
+  if size <= 0 or (size & (size - 1)) != 0:
+    raise ValueError('Hadamard matrix size must be a power of 2. ')
+  h = h2 = np.array([[1, 1], [1, -1]])
+  current_size = 2
+  while current_size < size:
+    h = np.kron(h, h2)
+    current_size *= 2
+  return h / np.sqrt(size)
+
+
+def insert_decomposed_hadamard_rotation(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Inserts a decomposed pattern of Hadamard rotation on this tensor.
+
+  This function works for float32 tensors only. Instead of inserting a single
+  custom op (aeq.hadamard_rotation), this inserts the mathematical equivalent
+  expressed in built-in TFLite ops. The mathematical equivalent is:
+    x' = reshape(x, (-1, hadamard_size))
+    x' = x' @ H(hadamard_size)
+    x' = reshape(x, x.shape)
+  where H(n) is a Hadamard matrix of size n.
+
+  Args:
+    transformation_input: The transformation input to insert the ops on.
+
+  Returns:
+    The transformation info of the inserted ops.
+
+  Raises:
+    ValueError: If the transformation input is not a uniform quantization
+    transformation.
+    ValueError: If the Hadamard quantization params are not set.
+    ValueError: If the tensor is not a float32 tensor.
+    ValueError: If no supported ops were found as the tensor's producer or
+    consumers.
+  """
+  if not isinstance(
+      transformation_input.quant_params, qtyping.UniformQuantParams
+  ):
+    raise ValueError('Hadamard rotation supports uniform quantization only')
+
+  if transformation_input.quant_params.hadamard is None:
+    raise ValueError(
+        'Hadamard rotation quantization params are not set but op insertion is'
+        ' requested.'
+    )
+
+  tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
+  if tensor.type != schema_py_generated.TensorType.FLOAT32:
+    raise ValueError(
+        'The Hadamard rotation op supports float32 tensors only. Got'
+        f' {tensor.type} tensor.'
+    )
+
+  # Insert x' = tfl.reshape to reshape x to (-1, hadamard_size)
+  hadamard_size = transformation_input.quant_params.hadamard.hadamard_size
+  tensor_size = np.prod(tensor.shape)
+  num_hadamard_blocks = tensor_size // hadamard_size
+  prerotate_shape = [num_hadamard_blocks, hadamard_size]
+  prerotate_shape_tensor_id = transformation_utils.add_new_constant_tensor(
+      tensor.name + b'_prerotate_shape',
+      np.array(prerotate_shape, dtype=np.int32),
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+  prerotate_reshape_output_tensor_id = (
+      transformation_utils.add_new_activation_tensor(
+          tensor.name + b'_prerotate_reshaped',
+          prerotate_shape,
+          schema_py_generated.TensorType.FLOAT32,
+          transformation_input.subgraph,
+      )
+  )
+
+  prerotate_reshape_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.RESHAPE,
+      transformation_input.op_codes,
+      'RESHAPE',
+  )
+  prerorate_reshape_op = schema_py_generated.OperatorT()
+  prerorate_reshape_op.opcodeIndex = prerotate_reshape_op_code_idx
+  prerorate_reshape_op.inputs = [
+      transformation_input.tensor_id,
+      prerotate_shape_tensor_id,
+  ]
+  prerorate_reshape_op.outputs = [prerotate_reshape_output_tensor_id]
+
+  # Generate hadamard_matrix(hadamard_size).
+  # We could quantize this to INT4 for better memory efficiency, but for large
+  # models the memory overhead is not significant, and floating point
+  # computation does seem to result in better accuracy.
+  hadamard_matrix = _make_hadamard_matrix(hadamard_size)
+  hadamard_matrix_tensor_id = transformation_utils.add_new_constant_tensor(
+      tensor.name + b'_hadamard_matrix',
+      hadamard_matrix.astype(np.float32),
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  # Insert x' = tfl.fully_connected(x', hadamard_matrix)
+  fc_output_tensor_id = transformation_utils.add_new_activation_tensor(
+      tensor.name + b'_rotated',
+      prerotate_shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+
+  fc_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+      transformation_input.op_codes,
+      'FULLY_CONNECTED',
+  )
+  fc_op = schema_py_generated.OperatorT()
+  fc_op.opcodeIndex = fc_op_code_idx
+  fc_op.inputs = [prerotate_reshape_output_tensor_id, hadamard_matrix_tensor_id]
+  fc_op.outputs = [fc_output_tensor_id]
+  fc_options = schema_py_generated.FullyConnectedOptionsT()
+  fc_options.fusedActivationFunction = (
+      schema_py_generated.ActivationFunctionType.NONE
+  )
+  fc_op.builtinOptionsType = (
+      schema_py_generated.BuiltinOptions.FullyConnectedOptions
+  )
+  fc_op.builtinOptions = fc_options
+
+  # Insert x' = tfl.reshape(x', x.shape)
+  post_reshape_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.RESHAPE,
+      transformation_input.op_codes,
+      'RESHAPE',
+  )
+  post_reshape_op = schema_py_generated.OperatorT()
+  post_reshape_op.opcodeIndex = post_reshape_op_code_idx
+  post_reshape_shape_tensor_id = transformation_utils.add_new_constant_tensor(
+      tensor.name + b'_postrotate_shape',
+      np.array(tensor.shape, dtype=np.int32),
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  post_reshape_output_tensor_id = (
+      transformation_utils.add_new_activation_tensor(
+          tensor.name + b'_postrotate_reshaped',
+          tensor.shape,
+          schema_py_generated.TensorType.FLOAT32,
+          transformation_input.subgraph,
+      )
+  )
+  post_reshape_op.inputs = [
+      fc_output_tensor_id,
+      post_reshape_shape_tensor_id,
+  ]
+  post_reshape_op.outputs = [post_reshape_output_tensor_id]
+
+  # Update the users of this tensor to use the new tensor.
+  if (
+      transformation_utils.get_producer_schema_op_id(transformation_input)
+      == schema_py_generated.BuiltinOperator.EMBEDDING_LOOKUP
+  ):
+    _update_embedding_lookup_consumers(
+        transformation_input, post_reshape_output_tensor_id
+    )
+  elif not _update_fully_connected_consumers(
+      transformation_input, post_reshape_output_tensor_id
+  ):
+    raise ValueError(
+        'The Hadamard rotation op supports embedding lookup and fully connected'
+        ' ops only, but no such ops were found.'
+    )
+
+  # If the tensor is a graph output, we need to replace the tensor with the
+  # new tensor.
+  for i, output in enumerate(transformation_input.subgraph.outputs):
+    if output == transformation_input.tensor_id:
+      transformation_input.subgraph.outputs[i] = post_reshape_output_tensor_id
+
+  # Find the actual insertion point. The insertion point should be after the
+  # producer op and before the first consumer op. The max() operation ensures
+  # that we're not using -1 as the insertion point.
+  first_consumer_id = min(transformation_input.consumers)
+  op_id = max(transformation_input.producer + 1, first_consumer_id)
+
+  # Insert the new ops in the correct order.
+  transformation_input.subgraph.operators.insert(op_id, prerorate_reshape_op)
+  transformation_input.subgraph.operators.insert(op_id + 1, fc_op)
+  transformation_input.subgraph.operators.insert(op_id + 2, post_reshape_op)
+
+  return qtyping.TransformationInfo(
+      op_id=op_id,
+      num_ops_added=3,
+      output_tensor_id=post_reshape_output_tensor_id,
+  )