Enable lowering from FQ Composite for 2-bit

This also adds an additional test for this lowering.

LiteRT-Converter-PiperOrigin-RevId: 820534395

Author: majiddadashi

tflite/converter/quantization/common/quantization_lib/quantization.td

@@ -56,6 +56,7 @@ class Int8UniformQuantizedType<int zero_pt, int smantissa, int sexp>
 
 // General uniform quantized types. The definitions can be used to specify
 // operand's tensor types.
+def QI2 : QuantizedType<"Uniform", [2], 1>;
 def QI4 : QuantizedType<"Uniform", [4], 1>;
 def QUI8 : QuantizedType<"Uniform", [8], 0>;
 def QI8 : QuantizedType<"Uniform", [8], 1>;

tflite/converter/tests/lower_quant_annotations.mlir

@@ -3,6 +3,8 @@
 func.func private @XlaCallModule_quant.fake_quant.impl_0(tensor<1x28x28x3xf32>) -> tensor<1x28x28x3xf32>
 func.func private @XlaCallModule_quant.fake_quant.impl_5_0(tensor<2x1x1x1xf32>) -> tensor<2x1x1x1xf32>
 func.func private @XlaCallModule_quant.fake_quant.impl_17_0(tensor<1x30x30x2xf32>) -> tensor<1x30x30x2xf32>
+func.func private @XlaCallModule_quant.fake_quant.impl_i2_0(tensor<1x4xf32>) -> tensor<1x4xf32>
+func.func private @XlaCallModule_quant.fake_quant.impl_i2_1(tensor<1x4xf32>) -> tensor<1x4xf32>
 // CHECK-LABEL: func.func @serving_default
 func.func @serving_default(%arg0: tensor<1x28x28x3xf32>) -> (tensor<1x30x30x2xf32>) {
   %cst = arith.constant dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi32>
@@ -22,4 +24,15 @@ func.func @serving_default(%arg0: tensor<1x28x28x3xf32>) -> (tensor<1x30x30x2xf3
   // CHECK-OFF: %[[DEQUANT2:.+]] = "tfl.dequantize"(%[[QUANT2]]) : (tensor<1x30x30x2x!quant.uniform<i8:f32, 0.018049469217658043:8>>) -> tensor<1x30x30x2xf32>
   %5 = stablehlo.composite "quant.fake_quant" %4 {composite_attributes = {dtype = "i8", narrow_range = false, scale = dense<0.0180494692> : tensor<1xf32>, zero_point = dense<8> : tensor<1xi32>}, decomposition = @XlaCallModule_quant.fake_quant.impl_17_0} : (tensor<1x30x30x2xf32>) -> tensor<1x30x30x2xf32>
   return %5 : tensor<1x30x30x2xf32>
+}
+
+// CHECK-LABEL: func.func @i2_test
+func.func @i2_test(%arg0: tensor<1x4xf32>) -> (tensor<1x4xf32>) {
+  // CHECK: %[[QUANT0:.+]] = "tfl.quantize"(%arg0) <{qtype = tensor<1x4x!quant.uniform<i2:f32, 1.000000e+00>>}> : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i2:f32, 1.000000e+00>>
+  // CHECK: %[[DEQUANT0:.+]] = "tfl.dequantize"(%[[QUANT0]]) : (tensor<1x4x!quant.uniform<i2:f32, 1.000000e+00>>) -> tensor<1x4xf32>
+  %0 = stablehlo.composite "quant.fake_quant" %arg0 {composite_attributes = {dtype = "i2", narrow_range = false, scale = dense<1.0> : tensor<1xf32>, zero_point = dense<0> : tensor<1xi32>}, decomposition = @XlaCallModule_quant.fake_quant.impl_i2_0} : (tensor<1x4xf32>) -> tensor<1x4xf32>
+  // CHECK: %[[QUANT1:.+]] = "tfl.quantize"(%[[DEQUANT0]]) <{qtype = tensor<1x4x!quant.uniform<i2<-1:1>:f32:1, {1.000000e+00,2.000000e+00,3.000000e+00,4.000000e+00}>>}> : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i2<-1:1>:f32:1, {1.000000e+00,2.000000e+00,3.000000e+00,4.000000e+00}>>
+  // CHECK: %[[DEQUANT1:.+]] = "tfl.dequantize"(%[[QUANT1]]) : (tensor<1x4x!quant.uniform<i2<-1:1>:f32:1, {1.000000e+00,2.000000e+00,3.000000e+00,4.000000e+00}>>) -> tensor<1x4xf32>
+  %1 = stablehlo.composite "quant.fake_quant" %0 {composite_attributes = {dtype = "i2", narrow_range = true, quantization_dimension = 1 : i32, scale = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>}, decomposition = @XlaCallModule_quant.fake_quant.impl_i2_1} : (tensor<1x4xf32>) -> tensor<1x4xf32>
+  return %1 : tensor<1x4xf32>
 }

tflite/converter/transforms/lower_quant_annotations_helper.cc

@@ -71,12 +71,15 @@ LogicalResult FillCompositeParams(stablehlo::CompositeOp op,
     return failure();
   }
   std::string dtype = dtype_attr.getValue().str();
-  if (dtype == "i8") {
-    num_bits = 8;
+  if (dtype == "i2") {
+    num_bits = 2;
     is_signed = true;
   } else if (dtype == "i4") {
     num_bits = 4;
     is_signed = true;
+  } else if (dtype == "i8") {
+    num_bits = 8;
+    is_signed = true;
   } else {
     return failure();
   }
@@ -110,7 +113,16 @@ LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range,
                                bool is_signed, MLIRContext* ctx,
                                Type& storage_type, int64_t& qmin,
                                int64_t& qmax) {
-  if (num_bits <= 4) {
+  if (num_bits == 2) {
+    storage_type = IntegerType::get(ctx, 2);
+    if (is_signed) {
+      qmin = -2;
+      qmax = 1;
+    } else {
+      qmin = 0;
+      qmax = 3;
+    }
+  } else if (num_bits <= 4) {
     storage_type = IntegerType::get(ctx, 4);
     if (is_signed) {
       qmin = -8;