blockwise: enable blockwise quantization in policy

PiperOrigin-RevId: 746515536

Author: zichuan-wei

ai_edge_quantizer/algorithm_manager_api_test.py

@@ -18,6 +18,7 @@
 from absl.testing import parameterized
 from tensorflow.python.platform import googletest
 from ai_edge_quantizer import algorithm_manager_api
+from ai_edge_quantizer import default_policy
 from ai_edge_quantizer import qtyping
 
 _TFLOpName = qtyping.TFLOperationName
@@ -205,6 +206,12 @@ def test_register_config_check_policy_succeeds(self):
         self._alg_manager._config_check_policy_registry[test_algorithm_name]
     )
 
+  def test_default_policy_not_empty(self):
+    """Tests that the default policy is not empty & no empty policy is generated."""
+    self.assertNotEmpty(default_policy.DEFAULT_CONFIG_CHECK_POLICY)
+    for policy in default_policy.DEFAULT_CONFIG_CHECK_POLICY.values():
+      self.assertNotEmpty(policy)
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/default_policy.py

@@ -57,6 +57,17 @@
       "explicit_dequantize": false,
       "compute_precision": "INTEGER"
     },
+    "dynamic_wi4_afp32_blockwise": {
+      "weight_tensor_config": {
+        "num_bits": 4,
+        "symmetric": [true],
+        "granularity": ["BLOCKWISE"],
+        "dtype": "INT",
+        "block_size": [32, 64, 96, 128, 256]
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
     "static_wi8_ai16": {
       "activation_tensor_config": {
         "num_bits": 16,
@@ -216,6 +227,7 @@
       "FULLY_CONNECTED"
     ],
     "dynamic_wi4_afp32": ["FULLY_CONNECTED", "EMBEDDING_LOOKUP", "CONV_2D"],
+    "dynamic_wi4_afp32_blockwise": ["EMBEDDING_LOOKUP", "FULLY_CONNECTED"],
     "weightonly_wi8_afp32": [
       "BATCH_MATMUL",
       "CONV_2D",
@@ -259,6 +271,7 @@ def _unroll_json_config(
 
   # Then unroll weight configs and turn them into quantization configs.
   quant_configs = []
+  weight_configs = []
   for symmetric in json_config["weight_tensor_config"]["symmetric"]:
     for granularity in json_config["weight_tensor_config"]["granularity"]:
       tensor_config = {
@@ -267,6 +280,16 @@ def _unroll_json_config(
           "granularity": granularity,
           "dtype": json_config["weight_tensor_config"]["dtype"],
       }
+      if "block_size" in json_config["weight_tensor_config"]:
+        for block_size in json_config["weight_tensor_config"]["block_size"]:
+          tensor_config["block_size"] = block_size
+          weight_configs.append(
+              qtyping.TensorQuantizationConfig.from_dict(tensor_config)
+          )
+      else:
+        weight_configs.append(
+            qtyping.TensorQuantizationConfig.from_dict(tensor_config)
+        )
 
       if activation_configs:
         for activation_config in activation_configs:
@@ -281,15 +304,14 @@ def _unroll_json_config(
               )
           )
       else:
-        quant_configs.append(
-            qtyping.OpQuantizationConfig(
-                weight_tensor_config=qtyping.TensorQuantizationConfig.from_dict(
-                    tensor_config
-                ),
-                compute_precision=json_config["compute_precision"],
-                explicit_dequantize=json_config["explicit_dequantize"],
-            )
-        )
+        for weight_config in weight_configs:
+          quant_configs.append(
+              qtyping.OpQuantizationConfig(
+                  weight_tensor_config=weight_config,
+                  compute_precision=json_config["compute_precision"],
+                  explicit_dequantize=json_config["explicit_dequantize"],
+              )
+          )
 
   return quant_configs