Add max_hadamard_size parameter for Hadamard rotations.

PiperOrigin-RevId: 888789202

Author: ai-edge-bot

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py

@@ -61,20 +61,24 @@ def _make_hadamard_matrix(size: int) -> np.ndarray:
     h_int8 = np.kron(h_int8, h2)
     current_size *= 2
     if current_size not in _HADAMARD_MATRIX_CACHE:
-      h_norm = h_int8 / np.sqrt(size, dtype=np.float32)
+      h_norm = h_int8 / np.sqrt(current_size, dtype=np.float32)
       _HADAMARD_MATRIX_CACHE[current_size] = h_norm
+    else:
+      h_norm = _HADAMARD_MATRIX_CACHE[current_size]
   return h_norm
 
 
 def _rotate_with_diagonal_hadamard(
     tensor_content: np.ndarray,
     axis: int,
+    max_size: int | None = None,
 ):
   """Quantizes the given float array using the diagonal Hadamard algorithm.
 
   Args:
     tensor_content: The float array to quantize.
     axis: The axis of the tensor to rotate.
+    max_size: The maximum size of the Hadamard matrix.
 
   Returns:
     A tuple containing the quantized array and the recovered array.
@@ -93,6 +97,8 @@ def _rotate_with_diagonal_hadamard(
   # tile this Hadamard matrix along the diagonal. 2**30 is just a large power
   # of 2 to calculate this factor.
   hadamard_size = np.gcd(tensor_content.shape[axis], 2**30)
+  if max_size:
+    hadamard_size = min(hadamard_size, 1 << (max_size.bit_length() - 1))
   random_vector = np.ones(hadamard_size, dtype=np.int8)
 
   # Use a canonical Hadamard matrix.
@@ -150,7 +156,9 @@ def get_tensor_quant_params(
 
   # Rotate the tensor with a Hadamard matrix.
   w_rotated, hadamard_size, random_vector = _rotate_with_diagonal_hadamard(
-      tensor_content, axis=reduce_axis
+      tensor_content,
+      axis=reduce_axis,
+      max_size=tensor_quant_config.algorithm_params.get("max_hadamard_size"),
   )
 
   # Get the quantized values of the rotated tensor.

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py

@@ -15,6 +15,7 @@
 
 """Test Hadamard rotation materialization."""
 
+import dataclasses
 import pathlib
 
 from absl.testing import absltest
@@ -240,6 +241,36 @@ def test_get_tensor_quant_params_basic(self):
     if qparams.hadamard is not None:
       self.assertEqual(qparams.hadamard.hadamard_size, 32)
 
+  def test_get_tensor_quant_params_max_size(self):
+    input_tensor = self._subgraph.tensors[self._fc_op.inputs[1]]
+    buffer = self._graph_info.buffers[self._fc_buffer_id]
+    np_buffer = np.frombuffer(buffer.data, dtype=np.float32).reshape(
+        input_tensor.shape
+    )
+    # The original dimension is 32. The largest power of 2 factor is 32.
+    # If we set algorithm_params to {'max_hadamard_size': 16}, then it should be
+    # 16.
+    new_op_quant_config = dataclasses.replace(
+        self._op_info.op_quant_config,
+        weight_tensor_config=qtyping.TensorQuantizationConfig(
+            num_bits=8,
+            symmetric=True,
+            granularity=qtyping.QuantGranularity.CHANNELWISE,
+            algorithm_params={"max_hadamard_size": 16},
+        ),
+    )
+    self._op_info = dataclasses.replace(
+        self._op_info, op_quant_config=new_op_quant_config
+    )
+    qparams = hadamard_rotation.get_tensor_quant_params(
+        self._op_info,
+        self._op_info.op_quant_config.weight_tensor_config,
+        np_buffer,
+        self._tensor_name_to_qsv,
+    )
+    self.assertIsNotNone(qparams.hadamard)
+    self.assertEqual(qparams.hadamard.hadamard_size, 16)
+
   def test_get_tensor_quant_params_golden_1(self):
     test_data = np.ones((6, 6))
     # expected:

ai_edge_quantizer/algorithms/utils/common_utils.py

@@ -126,21 +126,36 @@ def check_if_valid_op_config(
         f"No policy was specified for op: {op_name} with config:"
         f" {op_quant_config}."
     )
-  # The config_check_policy contains all possible valid configs, except for
-  # variations in the min_weight_elements field (it's set to 0 for all of them).
-  # min_weight_elements has to be ignored during policy check here because it
-  # can be any non-negative integer, which means we can't list all possible
-  # values in the policy.
-  elif (
-      dataclasses.replace(op_quant_config, min_weight_elements=0)
-      not in config_check_policy[op_name]
-  ):
-    error_msg = (
-        f"Quantization config for op: {op_name} with config:"
-        f" {op_quant_config} was not found in the policy."
-    )
   else:
-    check_passed = True
+    # min_weight_elements and algorithm_params have to be ignored during
+    # policy check here because they can be any non-negative integer or dict,
+    # which means we can't list all possible values in the policy.
+    op_quant_config_to_check = dataclasses.replace(
+        op_quant_config, min_weight_elements=0
+    )
+    if op_quant_config_to_check.weight_tensor_config is not None:
+      op_quant_config_to_check = dataclasses.replace(
+          op_quant_config_to_check,
+          weight_tensor_config=dataclasses.replace(
+              op_quant_config_to_check.weight_tensor_config, algorithm_params={}
+          ),
+      )
+    if op_quant_config_to_check.activation_tensor_config is not None:
+      op_quant_config_to_check = dataclasses.replace(
+          op_quant_config_to_check,
+          activation_tensor_config=dataclasses.replace(
+              op_quant_config_to_check.activation_tensor_config,
+              algorithm_params={},
+          ),
+      )
+
+    if op_quant_config_to_check not in config_check_policy[op_name]:
+      error_msg = (
+          f"Quantization config for op: {op_name} with config:"
+          f" {op_quant_config!r} was not found in the policy."
+      )
+    else:
+      check_passed = True
 
   if not check_passed:
     raise ValueError(

ai_edge_quantizer/algorithms/utils/common_utils_test.py

@@ -224,6 +224,21 @@ def test_check_drq_config_with_non_default_min_weight_elements_succeeds(self):
         _TFLOpName.CONV_2D, op_quant_config, _DEFAULT_CONFIG_CHECK_POLICY
     )
 
+  def test_check_config_with_non_default_algorithm_params_succeeds(self):
+    op_quant_config = _OpQuantConfig(
+        weight_tensor_config=_TensorQuantConfig(
+            num_bits=8,
+            granularity=qtyping.QuantGranularity.CHANNELWISE,
+            algorithm_params={"max_hadamard_size": 1024},
+        ),
+        compute_precision=_ComputePrecision.INTEGER,  # DRQ.
+    )
+    common_utils.check_if_valid_op_config(
+        _TFLOpName.FULLY_CONNECTED,
+        op_quant_config,
+        _DEFAULT_CONFIG_CHECK_POLICY,
+    )
+
   @parameterized.product(
       op_name=(_TFLOpName.FULLY_CONNECTED, _TFLOpName.CONV_2D),
       act_num_bits=(8, 16),

ai_edge_quantizer/qtyping.py

@@ -20,10 +20,9 @@
 import copy
 import dataclasses
 import enum
-from typing import Any, Callable, Optional, Union
-
+from typing import Any, Callable, Mapping, Optional, Union, TypeAlias
+from immutabledict import immutabledict
 import numpy as np
-from typing_extensions import TypeAlias
 
 from ai_edge_litert.tools import flatbuffer_utils
 
@@ -355,22 +354,36 @@ class TensorQuantizationConfig:
       quantization.
     dtype: The data type of the tensor.
     algorithm_key: The algorithm key to use for quantization.
+    algorithm_params: Additional parameters for the quantization algorithm.
   """
 
   num_bits: int
   symmetric: bool = True
   granularity: QuantGranularity = QuantGranularity.TENSORWISE
   dtype: TensorDataType = TensorDataType.INT
+  algorithm_params: Mapping[str, Any] = dataclasses.field(
+      default_factory=immutabledict
+  )
+
+  def __post_init__(self):
+    if not isinstance(self.algorithm_params, immutabledict):
+      object.__setattr__(
+          self, 'algorithm_params', immutabledict(self.algorithm_params)
+      )
 
   def to_dict(self) -> dict[str, Any]:
     """Converts ActivationQuantizationConfig to dict."""
     return dataclasses.asdict(
         self,
         dict_factory=lambda x: {  # pylint: disable=g-long-lambda
-            k: v
+            k: (
+                dict(v)
+                if isinstance(v, Mapping) and not isinstance(v, dict)
+                else v
+            )
             for (k, v) in x
             # Skip None and empty dict values.
-            if v is not None and not (isinstance(v, dict) and not v)
+            if v is not None and not (isinstance(v, Mapping) and not v)
         },
     )
 
@@ -380,7 +393,15 @@ def from_dict(cls, params: dict[str, Any]) -> 'TensorQuantizationConfig':
     params_copy = copy.deepcopy(params)
     # Process block_size config from legacy recipe.
     params_copy = _process_block_size(params_copy)
-    return cls(**params_copy)
+
+    # Move any unknown fields to algorithm_params for backward compatibility.
+    known_fields = {f.name for f in dataclasses.fields(cls)}
+    algorithm_params = params_copy.pop('algorithm_params', {})
+    for key in list(params_copy.keys()):
+      if key not in known_fields:
+        algorithm_params[key] = params_copy.pop(key)
+
+    return cls(algorithm_params=algorithm_params, **params_copy)
 
 
 def _process_block_size(params: dict[str, Any]) -> dict[str, Any]:
@@ -462,10 +483,14 @@ def to_dict(self) -> dict[str, Any]:
     return dataclasses.asdict(
         self,
         dict_factory=lambda x: {  # pylint: disable=g-long-lambda
-            k: v
+            k: (
+                dict(v)
+                if isinstance(v, Mapping) and not isinstance(v, dict)
+                else v
+            )
             for (k, v) in x
             # Skip None and empty dict values.
-            if v is not None and not (isinstance(v, dict) and not v)
+            if v is not None and not (isinstance(v, Mapping) and not v)
         },
     )
 

ai_edge_quantizer/recipe_manager.py

@@ -18,7 +18,7 @@
 import collections
 import dataclasses
 import re
-from typing import Any, Optional
+from typing import Any, Mapping, Optional
 from absl import logging
 from ai_edge_quantizer import algorithm_manager
 from ai_edge_quantizer import qtyping
@@ -205,16 +205,28 @@ def get_quantization_recipe(self) -> ModelQuantizationRecipe:
     Returns:
       A list of quantization configs in the recipe.
     """
-    ret = []
-    for _, scope_config in self._scope_configs.items():
-      for quant_config in scope_config:
-        config = dict()
-        config['regex'] = quant_config.regex
-        config['operation'] = quant_config.operation
-        config['algorithm_key'] = quant_config.algorithm_key
-        config['op_config'] = quant_config.op_config.to_dict()
-        ret.append(config)
-    return ret
+    recipe = []
+    for _, op_recipes in self._scope_configs.items():
+      for op_recipe in op_recipes:
+        recipe_dict = dataclasses.asdict(
+            op_recipe,
+            dict_factory=lambda x: {  # pylint: disable=g-long-lambda
+                k: (
+                    dict(v)
+                    if isinstance(v, Mapping)
+                    and not isinstance(v, dict)
+                    else v
+                )
+                for (k, v) in x
+                # Skip None and empty dict values.
+                if v is not None
+                and not (
+                    isinstance(v, (dict, Mapping)) and not v
+                )
+            },
+        )
+        recipe.append(recipe_dict)
+    return recipe
 
   def load_quantization_recipe(
       self, quantization_recipe: ModelQuantizationRecipe