Adds efficient special cases to the default bitpacking utils in tensor_encoding.

PiperOrigin-RevId: 352917915

Author: PraChetit

tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/tf_utils.py

@@ -311,6 +311,10 @@ def pack_into_int(value, input_bitrange, target_bitrange):
   """
   # TODO(b/161433177): Provide a general solution without memory overhead.
   # Special cases implemented without extra memory overhead.
+  if input_bitrange == 6 and target_bitrange == 28:
+    return _pack_into_int_6_28(value)
+  if input_bitrange == 7 and target_bitrange == 28:
+    return _pack_into_int_7_28(value)
   if input_bitrange == 8 and target_bitrange == 28:
     return _pack_into_int_8_28(value)
   if input_bitrange == 12 and target_bitrange == 28:
@@ -349,6 +353,10 @@ def unpack_from_int(value, original_bitrange, target_bitrange, shape):
   """
   # TODO(b/161433177): Provide a general solution without memory overhead.
   # Special cases implemented without extra memory overhead.
+  if original_bitrange == 6 and target_bitrange == 28:
+    return _unpack_from_int_6_28(value, shape)
+  if original_bitrange == 7 and target_bitrange == 28:
+    return _unpack_from_int_7_28(value, shape)
   if original_bitrange == 8 and target_bitrange == 28:
     return _unpack_from_int_8_28(value, shape)
   if original_bitrange == 12 and target_bitrange == 28:
@@ -386,6 +394,85 @@ def _expand_to_binary_form(value, input_bits):
   return tf.reshape(bits, [-1])
 
 
+def _pack_into_int_6_28(value):
+  """Implementation of `pack_into_int` for specific bitranges.
+
+  This method corresponts to `(input_bitrange, target_bitrange)` form the
+  `pack_into_int` method equal to `(6, 28)`. This method relies on the fact that
+  14 values in 6-bit bitrange can be packed into 3 values in 28-bitrange
+  (14 = least_common_multiple(6, 28) / 6).
+
+  It reshapes the input into matrix of 14 columns and performs operations on the
+  columns of the matrix, thus vectorizing the operations and avoiding memory
+  overhead of an earlier general implementation.
+
+  Args:
+    value: An integer Tensor to be packed with values in [0, 2**6 - 1].
+
+  Returns:
+    An integer Tensor representing `value` of the same dtype as `value`.
+  """
+  value = tf.reshape(value, [-1])
+  extra_zeros = tf.zeros(tf.math.mod(-tf.shape(value), 14), value.dtype)
+  val = tf.reshape(tf.concat([value, extra_zeros], 0), [-1, 14])
+
+  a = (val[:, 0] +
+       val[:, 1] * 2**6 +
+       val[:, 2] * 2**12 +
+       val[:, 3] * 2**18 +
+       tf.math.mod(val[:, 4], 2**4) * 2**24)
+  b = (tf.math.floordiv(val[:, 4], 2**4) +
+       val[:, 5] * 2**2 +
+       val[:, 6] * 2**8 +
+       val[:, 7] * 2**14 +
+       val[:, 8] * 2**20 +
+       tf.math.mod(val[:, 9], 2**2) * 2**26)
+  c = (tf.math.floordiv(val[:, 9], 2**2) +
+       val[:, 10] * 2**4 +
+       val[:, 11] * 2**10 +
+       val[:, 12] * 2**16 +
+       val[:, 13] * 2**22)
+
+  packed_val = tf.reshape(tf.stack([a, b, c], 1), [-1, 1])
+  if extra_zeros.shape[0] in [5, 6, 7, 8, 9]:
+    # We added unnecessary product of zeros to the representation.
+    packed_val = tf.slice(packed_val, [0, 0], [packed_val.shape[0] - 1, 1])
+  if extra_zeros.shape[0] in [10, 11, 12, 13]:
+    # We added unnecessary two products of zeros to the representation.
+    packed_val = tf.slice(packed_val, [0, 0], [packed_val.shape[0] - 2, 1])
+  return packed_val
+
+
+def _pack_into_int_7_28(value):
+  """Implementation of `pack_into_int` for specific bitranges.
+
+  This method corresponts to `(input_bitrange, target_bitrange)` form the
+  `pack_into_int` method equal to `(7, 28)`. This method relies on the fact that
+  4 values in 7-bit bitrange can be packed into 1 value in 28-bitrange
+  (4 = least_common_multiple(7, 28) / 7).
+
+  It reshapes the input into matrix of 4 columns and performs operations on the
+  columns of the matrix, thus vectorizing the operations and avoiding memory
+  overhead of an earlier general implementation.
+
+  Args:
+    value: An integer Tensor to be packed with values in [0, 2**7 - 1].
+
+  Returns:
+    An integer Tensor representing `value` of the same dtype as `value`.
+  """
+  value = tf.reshape(value, [-1])
+  extra_zeros = tf.zeros(tf.math.mod(-tf.shape(value), 4), value.dtype)
+  val = tf.reshape(tf.concat([value, extra_zeros], 0), [-1, 4])
+
+  packed_val = (val[:, 0] +
+                val[:, 1] * 2**7 +
+                val[:, 2] * 2**14 +
+                val[:, 3] * 2**21)
+
+  return tf.reshape(packed_val, [-1, 1])
+
+
 def _pack_into_int_8_28(value):
   """Implementation of `pack_into_int` for specific bitranges.
 
@@ -466,6 +553,47 @@ def _pack_into_int_12_28(value):
   return packed_val
 
 
+def _unpack_from_int_6_28(value, shape):
+  """Inverse operation of `_pack_into_int_6_28`."""
+  value = tf.reshape(value, [-1])
+  extra_zeros = tf.zeros(tf.math.mod(-tf.shape(value), 3), value.dtype)
+  val = tf.reshape(tf.concat([value, extra_zeros], 0), [-1, 3])
+
+  a = tf.math.mod(val[:, 0], 2**6)
+  b = tf.math.mod(tf.math.floordiv(val[:, 0], 2**6), 2**6)
+  c = tf.math.mod(tf.math.floordiv(val[:, 0], 2**12), 2**6)
+  d = tf.math.mod(tf.math.floordiv(val[:, 0], 2**18), 2**6)
+  e = tf.math.floordiv(val[:, 0], 2**24) + tf.math.mod(val[:, 1], 2**2) * 2**4
+  f = tf.math.mod(tf.math.floordiv(val[:, 1], 2**2), 2**6)
+  g = tf.math.mod(tf.math.floordiv(val[:, 1], 2**8), 2**6)
+  h = tf.math.mod(tf.math.floordiv(val[:, 1], 2**14), 2**6)
+  i = tf.math.mod(tf.math.floordiv(val[:, 1], 2**20), 2**6)
+  j = tf.math.floordiv(val[:, 1], 2**26) + tf.math.mod(val[:, 2], 2**4) * 2**2
+  k = tf.math.mod(tf.math.floordiv(val[:, 2], 2**4), 2**6)
+  l = tf.math.mod(tf.math.floordiv(val[:, 2], 2**10), 2**6)
+  m = tf.math.mod(tf.math.floordiv(val[:, 2], 2**16), 2**6)
+  n = tf.math.mod(tf.math.floordiv(val[:, 2], 2**22), 2**6)
+
+  unpacked_val = tf.reshape(
+      tf.stack([a, b, c, d, e, f, g, h, i, j, k, l, m, n], 1), [-1,])
+  unpacked_val = tf.slice(unpacked_val, [0], [tf.reduce_prod(shape)])
+  return tf.reshape(unpacked_val, shape)
+
+
+def _unpack_from_int_7_28(value, shape):
+  """Inverse operation of `_pack_into_int_7_28`."""
+  val = tf.reshape(value, [-1, 1])
+
+  a = tf.math.mod(val[:, 0], 2**7)
+  b = tf.math.mod(tf.math.floordiv(val[:, 0], 2**7), 2**7)
+  c = tf.math.mod(tf.math.floordiv(val[:, 0], 2**14), 2**7)
+  d = tf.math.mod(tf.math.floordiv(val[:, 0], 2**21), 2**7)
+
+  unpacked_val = tf.reshape(tf.stack([a, b, c, d], 1), [-1,])
+  unpacked_val = tf.slice(unpacked_val, [0], [tf.reduce_prod(shape)])
+  return tf.reshape(unpacked_val, shape)
+
+
 def _unpack_from_int_8_28(value, shape):
   """Inverse operation of `_pack_into_int_8_28`."""
   value = tf.reshape(value, [-1])

tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/tf_utils_test.py

@@ -375,7 +375,8 @@ def test_unpack_from_int_different_outputs(self):
         packed_value, original_bitrange=4, target_bitrange=28, shape=(2,))
     self.assertAllEqual([9, 0], self.evaluate(unpacked_value))
 
-  @parameterized.parameters([(1, 28), (2, 28), (8, 28), (12, 28)])
+  @parameterized.parameters([(1, 28), (2, 28), (6, 28), (7, 28), (8, 28),
+                             (12, 28)])
   def test_boundary_conditions(self, input_bitrange, target_bitrange):
     max_v = 2**input_bitrange - 1
     input_value = tf.constant(
@@ -390,7 +391,8 @@ def test_boundary_conditions(self, input_bitrange, target_bitrange):
 
     self.assertAllEqual(self.evaluate(value), self.evaluate(unpacked_value))
 
-  @parameterized.parameters([(1, 28), (2, 28), (8, 28), (12, 28)])
+  @parameterized.parameters([(1, 28), (2, 28), (6, 28), (7, 28), (8, 28),
+                             (12, 28)])
   def test_random_input(self, input_bitrange, target_bitrange):
     # Tests that packing/unpacking amounts to identity, regardless of the input.
     num_elements = np.random.randint(low=1, high=50)
@@ -407,6 +409,42 @@ def test_random_input(self, input_bitrange, target_bitrange):
     except:  # pylint: disable=bare-except
       self.fail(f'Random input test failed with input value: {value}')
 
+  def test_pack_into_int_special_case_6_28(self):
+    value = tf.constant(
+        [50, 19, 51, 59, 10, 53, 36, 44, 31, 44, 31, 10, 31, 56, 49, 48, 35])
+    packed_value = tf_utils.pack_into_int(
+        value, input_bitrange=6, target_bitrange=28)
+    expected_packed_value = tf.constant([[183448818], [33236180], [236923387],
+                                         [146481]])
+    self.assertAllEqual(self.evaluate(expected_packed_value),
+                        self.evaluate(packed_value))
+
+  def test_unpack_from_int_special_case_6_28(self):
+    packed_value = tf.constant([[183448818], [33236180], [236923387], [146481]])
+    unpacked_value = tf_utils.unpack_from_int(
+        packed_value, original_bitrange=6, target_bitrange=28, shape=(17,))
+    expected_unpacked_value = tf.constant(
+        [50, 19, 51, 59, 10, 53, 36, 44, 31, 44, 31, 10, 31, 56, 49, 48, 35])
+    self.assertAllEqual(self.evaluate(expected_unpacked_value),
+                        self.evaluate(unpacked_value))
+
+  def test_pack_into_int_special_case_7_28(self):
+    value = tf.constant([117, 86, 42, 69, 9, 70, 66, 8, 112, 116])
+    packed_value = tf_utils.pack_into_int(
+        value, input_bitrange=7, target_bitrange=28)
+    expected_packed_value = tf.constant([[145402741], [17867529], [14960]])
+    self.assertAllEqual(self.evaluate(expected_packed_value),
+                        self.evaluate(packed_value))
+
+  def test_unpack_from_int_special_case_7_28(self):
+    packed_value = tf.constant([[145402741], [17867529], [14960]])
+    unpacked_value = tf_utils.unpack_from_int(
+        packed_value, original_bitrange=7, target_bitrange=28, shape=(10,))
+    expected_unpacked_value = tf.constant(
+        [117, 86, 42, 69, 9, 70, 66, 8, 112, 116])
+    self.assertAllEqual(self.evaluate(expected_unpacked_value),
+                        self.evaluate(unpacked_value))
+
   def test_pack_into_int_special_case_8_28(self):
     value = tf.constant([38, 147, 1, 201, 205, 36, 155, 78, 163, 98])
     packed_value = tf_utils.pack_into_int(