Generalize PackInt8IntoDenseInt4 to support 2-bit and 4-bit packing.

Rename `PackInt8IntoDenseInt4` to `PackInt8IntoDenseInt` and add a `bit_width` parameter. Implement packing logic for both 2-bit and 4-bit integers. Update existing call sites in `quantize.cc` and `transpose.cc` to use the new function signature. Add new unit tests for both 2-bit and 4-bit packing.

PiperOrigin-RevId: 818850218

Author: majiddadashi

tflite/kernels/internal/portable_tensor_utils.cc

@@ -143,23 +143,40 @@ void UnpackPackedIntToInt8(const int8_t* src_buffer, int num_elements,
   }
 }
 
-void PackInt8IntoDenseInt4(const int8_t* src_buffer, int num_elements,
-                           int8_t* dst_buffer) {
-  // num_elements means the number of elements regardless of packed or unpacked.
-  // For example, 3 elements means both
-  //   1) Packed: 3 int4's = 12 bit -> 16 bits (padded) = 2 bytes.
-  //      stored in src_buffer[0] and src_buffer[1] (i = 0..1)
-  //   2) Unpacked: 3 int8's = 3 bytes.
-  //      stored in dst_buffer[0], dst_buffer[1] and dst_buffer[2] (j = 0..2)
-  for (int i = 0; i < num_elements - 1; i += 2) {
-    dst_buffer[i / 2] = src_buffer[i] & 0x0F;
-    dst_buffer[i / 2] |= src_buffer[i + 1] << 4;
-  }
-  auto packed_size = (num_elements + 1) / 2;
-
-  // Copy the final nibble if the buffer is odd-lengthed
-  if (num_elements % 2 != 0) {
-    dst_buffer[packed_size - 1] = src_buffer[num_elements - 1] & 0x0F;
+void PackInt8IntoDenseInt(const int8_t* src_buffer, int num_elements,
+                          int bit_width, int8_t* dst_buffer) {
+  assert(bit_width == 2 || bit_width == 4);
+  if (bit_width == 4) {
+    // num_elements means the number of elements regardless of packed or
+    // unpacked. For example, 3 elements means both
+    //   1) Unpacked: 3 int8's = 3 bytes.
+    //      stored in src_buffer[0], src_buffer[1] and src_buffer[2] (j = 0..2)
+    //   2) Packed: 3 int4's = 12 bit -> 16 bits (padded) = 2 bytes.
+    //      stored in dst_buffer[0] and dst_buffer[1] (i = 0..1)
+    for (int i = 0; i < num_elements / 2; ++i) {
+      dst_buffer[i] = (src_buffer[2 * i] & 0x0F) | (src_buffer[2 * i + 1] << 4);
+    }
+    // If the buffer size is odd, pack the final nibble.
+    if (num_elements % 2 != 0) {
+      dst_buffer[num_elements / 2] = src_buffer[num_elements - 1] & 0x0F;
+    }
+  } else if (bit_width == 2) {
+    for (int i = 0; i < num_elements / 4; ++i) {
+      dst_buffer[i] = (src_buffer[4 * i] & 0x03) |
+                      ((src_buffer[4 * i + 1] & 0x03) << 2) |
+                      ((src_buffer[4 * i + 2] & 0x03) << 4) |
+                      ((src_buffer[4 * i + 3] & 0x03) << 6);
+    }
+    // Handle the remaining elements.
+    int remaining_elements = num_elements % 4;
+    if (remaining_elements > 0) {
+      int8_t packed_val = 0;
+      for (int i = 0; i < remaining_elements; ++i) {
+        packed_val |= (src_buffer[num_elements - remaining_elements + i] & 0x03)
+                      << (i * 2);
+      }
+      dst_buffer[num_elements / 4] = packed_val;
+    }
   }
 }
 

tflite/kernels/internal/portable_tensor_utils.h

@@ -635,20 +635,24 @@ void UnpackDenseInt4IntoInt8(const int8_t* src_buffer, int num_elements,
 void UnpackPackedIntToInt8(const int8_t* src_buffer, int num_elements,
                            int bit_width, int8_t* dst_buffer);
 
-// Pack `src_buffer` into a densely packed buffer of int4 values.
+// Pack `src_buffer` into a densely packed buffer of int2 or int4 values.
 // Parameters:
-//   src_buffer   : Buffer containing int4 values stored in int8 memory.
+//   src_buffer   : Buffer containing int2 or int4 values stored in int8
+//                  memory.
 //   num_elements : Number of elements stored in the buffer. Note that this can
 //                  be smaller than the size of `src_buffer` by 1 if it's odd,
 //                  in which case the last nibble in `src_buffer` is ignored.
 //                  This should be equal to the size of `dst_buffer`.
+//   bit_width    : The bit width of the packed elements (either 2 or 4).
 //   dst_buffer   : Buffer to pack into. Should be allocated by the caller.
 //                  Size should be at least `num_elements`.
 // Notes:
-//   For example, given `src_buffer = {0x02, 0x01, 0x04, 0x03}`, calling this
-//   function will return `dst_buffer = {0x12, 0x34}`.
-void PackInt8IntoDenseInt4(const int8_t* src_buffer, int num_elements,
-                           int8_t* dst_buffer);
+//   For 4-bit packing: e.g., given `src_buffer = {0x02, 0x01, 0x04, 0x03}`,
+//   calling this function will return `dst_buffer = {0x12, 0x34}`.
+//   For 2-bit packing: e.g., given `src_buffer = {0x00, 0x01, 0x00, 0x02}`,
+//   calling this function will return `dst_buffer = {0x84}`.
+void PackInt8IntoDenseInt(const int8_t* src_buffer, int num_elements,
+                          int bit_width, int8_t* dst_buffer);
 }  // namespace tensor_utils
 
 }  // namespace tflite

tflite/kernels/internal/tensor_utils_test.cc

@@ -2149,6 +2149,44 @@ TEST(uKernels, UnpackInt2OddLength) {
               testing::Pointwise(testing::Eq(), expected_output));
 }
 
+TEST(uKernels, PackInt4Basic) {
+  const int8_t input[4] = {-8, 3, -2, -5};
+  const int8_t expected_output[2] = {0x38, static_cast<int8_t>(0xBE)};
+  int8_t actual_output[2];
+  PackInt8IntoDenseInt(input, 4, 4, actual_output);
+  EXPECT_THAT(actual_output,
+              testing::Pointwise(testing::Eq(), expected_output));
+}
+
+TEST(uKernels, PackInt4OddLength) {
+  // `num_elements` is odd, so the last element 0x4 should be ignored
+  const int8_t input[3] = {1, 2, 3};
+  const int8_t expected_output[2] = {0x21, 0x03};
+  int8_t actual_output[2];
+  PackInt8IntoDenseInt(input, 3, 4, actual_output);
+  EXPECT_THAT(actual_output,
+              testing::Pointwise(testing::Eq(), expected_output));
+}
+
+TEST(uKernels, PackInt2Basic) {
+  const int8_t input[4] = {0, -1, -2, 1};
+  const int8_t expected_output[1] = {0x6C};
+  int8_t actual_output[1];
+  PackInt8IntoDenseInt(input, 4, 2, actual_output);
+  EXPECT_THAT(actual_output,
+              testing::Pointwise(testing::Eq(), expected_output));
+}
+
+TEST(uKernels, PackInt2OddLength) {
+  // `num_elements` is odd
+  const int8_t input[3] = {0, -2, 1};
+  const int8_t expected_output[1] = {0x18};
+  int8_t actual_output[1];
+  PackInt8IntoDenseInt(input, 3, 2, actual_output);
+  EXPECT_THAT(actual_output,
+              testing::Pointwise(testing::Eq(), expected_output));
+}
+
 }  // namespace tensor_utils
 }  // namespace tflite
 

tflite/kernels/quantize.cc

@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "tflite/core/c/common.h"
 #include "tflite/kernels/internal/optimized/optimized_ops.h"
+#include "tflite/kernels/internal/portable_tensor_utils.h"
 #include "tflite/kernels/internal/quantization_util.h"
 #include "tflite/kernels/internal/reference/reference_ops.h"
 #include "tflite/kernels/internal/reference/requantize.h"
@@ -109,8 +110,8 @@ void AffineQuantizeToInt4(const tflite::QuantizationParams& op_params,
     int32_t clamped = std::min(std::max(unclamped, min_val), max_val);
     quantized_buffer[i] = clamped;
   }
-  tensor_utils::PackInt8IntoDenseInt4(quantized_buffer.data(), flat_size,
-                                      output_data);
+  tensor_utils::PackInt8IntoDenseInt(quantized_buffer.data(), flat_size,
+                                     /*bit_width=*/4, output_data);
 }
 
 void ReportError(TfLiteContext* context, TfLiteType input_type,

tflite/kernels/transpose.cc

@@ -142,10 +142,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           params, GetTensorShape(op_context.input), unpacked_input_data.get(),
           GetTensorShape(op_context.output), unpacked_output_data.get());
       // Pack the output back to int4.
-      tflite::tensor_utils::PackInt8IntoDenseInt4(
+      tflite::tensor_utils::PackInt8IntoDenseInt(
           unpacked_output_data.get(),
           GetTensorShape(op_context.input).FlatSize(),
-          GetTensorData<int8_t>(op_context.output));
+          /*bit_width=*/4, GetTensorData<int8_t>(op_context.output));
       break;
     }
     case kTfLiteInt16: