Quantization with min & max bounds support - 4-bit & 2-bit on X86-64 (#4833)

Summary:
Pull Request resolved: #4833

X-link: facebookresearch/FBGEMM#1860

In D78181177 we have added support for row-wise min/max for 8-bit quantization. In this diff similar thing is done for 4-bit & 2-bit quantization as well.

Reviewed By: spcyppt, excelle08

Differential Revision: D81858256

fbshipit-source-id: b6fbb5914b3a832143ebacf34dfaf2573d74fd85

Author: sampathvic

include/fbgemm/QuantUtils.h

@@ -285,7 +285,8 @@ FBGEMM_API void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(
     const InputType* input,
     size_t input_rows,
     int input_columns,
-    std::uint8_t* output);
+    std::uint8_t* output,
+    const InputType* rowwise_min_max = nullptr);
 
 /**
  * Convert fused rowwise quantized inputs to float (fp32 or fp16).

include/fbgemm/QuantUtilsAvx2.h

@@ -17,6 +17,10 @@
 
 namespace fbgemm {
 
+/// Number of columns in the rowwise min/max buffer passed to the quantization
+/// function(s)
+constexpr int kRowwiseMinMaxNumCols = 2;
+
 /// Struct from <a href="https://github.com/google/gemmlowp">`gemmlowp`</a>
 ///
 /// A structure to hold quantization parameters `scale` and `zero_point`.
@@ -144,7 +148,8 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
     const InputType* input,
     size_t input_rows,
     int input_columns,
-    std::uint8_t* output);
+    std::uint8_t* output,
+    const InputType* rowwise_min_max = nullptr);
 
 template <typename InputType>
 void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatAvx2(

src/QuantUtils.cc

@@ -626,7 +626,8 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(
     const InputType* input,
     size_t input_rows,
     int input_columns,
-    std::uint8_t* output) {
+    std::uint8_t* output,
+    const InputType* rowwise_min_max) {
   // Currenlty we can only dequantize if the number of input columns
   // is a multiple of number of elements_per_byte
 
@@ -640,15 +641,15 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(
     switch (bit_rate) {
       case 2:
         FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2<InputType, 2>(
-            input, input_rows, input_columns, output);
+            input, input_rows, input_columns, output, rowwise_min_max);
         break;
       case 4:
         FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2<InputType, 4>(
-            input, input_rows, input_columns, output);
+            input, input_rows, input_columns, output, rowwise_min_max);
         break;
       case 8:
         FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2<InputType, 8>(
-            input, input_rows, input_columns, output);
+            input, input_rows, input_columns, output, rowwise_min_max);
         break;
       default:
         FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfRef<InputType>(
@@ -866,7 +867,8 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
       const type* input,                                                       \
       size_t input_rows,                                                       \
       int input_columns,                                                       \
-      std::uint8_t* output);                                                   \
+      std::uint8_t* output,                                                    \
+      const type* rowwise_min_max);                                            \
   template FBGEMM_API void                                                     \
   FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef<type, false>(                \
       int bit_rate,                                                            \

src/QuantUtilsAvx2.cc

@@ -1563,7 +1563,8 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
     const InputType* input,
     size_t input_rows,
     int input_columns,
-    std::uint8_t* output) {
+    std::uint8_t* output,
+    const InputType* rowwise_min_max) {
   static_assert(
       std::is_same<InputType, float>() || std::is_same<InputType, float16>(),
       "Only float and float16 types are allowed.");
@@ -1574,6 +1575,8 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
       2 * sizeof(std::uint16_t);
 
   float* input_row_float_for_fp16 = nullptr;
+  float min_max_row_float_for_fp16[kRowwiseMinMaxNumCols];
+  const auto is_valid_rowwise_min_max = (rowwise_min_max != nullptr);
   if constexpr (std::is_same<InputType, float16>()) {
     input_row_float_for_fp16 = static_cast<float*>(
         fbgemmAlignedAlloc(64, input_columns * sizeof(float)));
@@ -1591,48 +1594,72 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
       input_row_float = input_row_float_for_fp16;
     }
 
+    const float* min_max_row_float = nullptr;
+    if (is_valid_rowwise_min_max) {
+      const InputType* min_max_row =
+          rowwise_min_max + row * kRowwiseMinMaxNumCols;
+
+      if constexpr (std::is_same_v<InputType, float>) {
+        min_max_row_float = reinterpret_cast<const float*>(min_max_row);
+      } else {
+        min_max_row_float_for_fp16[0] = halfToFloat(min_max_row[0]);
+        min_max_row_float_for_fp16[1] = halfToFloat(min_max_row[1]);
+        min_max_row_float = min_max_row_float_for_fp16;
+      }
+    }
+
     std::uint8_t* output_row = output + row * output_columns;
     std::uint16_t* output_row_scale_bias = reinterpret_cast<std::uint16_t*>(
         output_row +
         (input_columns + NUM_ELEM_PER_BYTE - 1) / NUM_ELEM_PER_BYTE);
 
     float minimum_element = FLT_MAX;
     float maximum_element = -FLT_MAX;
-    __m256 min_v = _mm256_set1_ps(minimum_element);
-    __m256 max_v = _mm256_set1_ps(maximum_element);
+    if (is_valid_rowwise_min_max) {
+      minimum_element = min_max_row_float[0];
+      maximum_element = min_max_row_float[1];
 
-    int col = 0;
-    for (col = 0; col < input_columns / VLEN * VLEN; col += VLEN) {
-      __m256 in_v;
-      if constexpr (std::is_same<InputType, float>()) {
-        in_v = _mm256_loadu_ps(input_row_float + col);
-      } else {
-        __m128i in_half_v =
-            _mm_loadu_si128(reinterpret_cast<const __m128i*>(input_row + col));
-        in_v = _mm256_cvtph_ps(in_half_v);
-        _mm256_store_ps(input_row_float_for_fp16 + col, in_v);
+      for (int col = 0; col < input_columns; ++col) {
+        if constexpr (std::is_same<InputType, float16>()) {
+          input_row_float_for_fp16[col] = halfToFloat(input_row[col]);
+        }
       }
+    } else {
+      __m256 min_v = _mm256_set1_ps(minimum_element);
+      __m256 max_v = _mm256_set1_ps(maximum_element);
+      int col = 0;
+      for (col = 0; col < input_columns / VLEN * VLEN; col += VLEN) {
+        __m256 in_v;
+        if constexpr (std::is_same<InputType, float>()) {
+          in_v = _mm256_loadu_ps(input_row_float + col);
+        } else {
+          __m128i in_half_v = _mm_loadu_si128(
+              reinterpret_cast<const __m128i*>(input_row + col));
+          in_v = _mm256_cvtph_ps(in_half_v);
+          _mm256_store_ps(input_row_float_for_fp16 + col, in_v);
+        }
 
-      min_v = _mm256_min_ps(min_v, in_v);
-      max_v = _mm256_max_ps(max_v, in_v);
-    }
-    alignas(64) float min_buf[VLEN], max_buf[VLEN];
-    _mm256_store_ps(min_buf, min_v);
-    _mm256_store_ps(max_buf, max_v);
-    for (int i = 0; i < VLEN; ++i) {
-      minimum_element = std::min(minimum_element, min_buf[i]);
-      maximum_element = std::max(maximum_element, max_buf[i]);
-    }
+        min_v = _mm256_min_ps(min_v, in_v);
+        max_v = _mm256_max_ps(max_v, in_v);
+      }
+      alignas(64) float min_buf[VLEN], max_buf[VLEN];
+      _mm256_store_ps(min_buf, min_v);
+      _mm256_store_ps(max_buf, max_v);
+      for (int i = 0; i < VLEN; ++i) {
+        minimum_element = std::min(minimum_element, min_buf[i]);
+        maximum_element = std::max(maximum_element, max_buf[i]);
+      }
 
-    for (; col < input_columns; ++col) {
-      if constexpr (std::is_same<InputType, float>()) {
-        minimum_element = std::min(minimum_element, input_row_float[col]);
-        maximum_element = std::max(maximum_element, input_row_float[col]);
-      } else {
-        float element = halfToFloat(input_row[col]);
-        input_row_float_for_fp16[col] = element;
-        minimum_element = std::min(minimum_element, element);
-        maximum_element = std::max(maximum_element, element);
+      for (; col < input_columns; ++col) {
+        if constexpr (std::is_same<InputType, float>()) {
+          minimum_element = std::min(minimum_element, input_row_float[col]);
+          maximum_element = std::max(maximum_element, input_row_float[col]);
+        } else {
+          float element = halfToFloat(input_row[col]);
+          input_row_float_for_fp16[col] = element;
+          minimum_element = std::min(minimum_element, element);
+          maximum_element = std::max(maximum_element, element);
+        }
       }
     }
 
@@ -1657,12 +1684,12 @@ void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
 
     output_row_scale_bias[0] = floatToHalf(scale);
 
-    col = 0;
+    int col = 0;
     if constexpr (BIT_RATE == 2 || BIT_RATE == 4) {
       __m256i permute_mask1_v =
           _mm256_set_epi32(0x07, 0x03, 0x06, 0x02, 0x05, 0x01, 0x04, 0x00);
       __m256 inverse_scale_v = _mm256_set1_ps(inverse_scale);
-      min_v = _mm256_set1_ps(minimum_element);
+      __m256 min_v = _mm256_set1_ps(minimum_element);
 
       for (; col + 4 * VLEN <= input_columns; col += 4 * VLEN) {
         __m256i x_rounded_v = _mm256_cvtps_epi32(_mm256_mul_ps(
@@ -1778,7 +1805,7 @@ void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatAvx2(
 
   const int64_t output_columns = input_columns + 2 * sizeof(float);
   float* input_row_float_for_fp16 = nullptr;
-  float min_max_row_float_for_fp16[2];
+  float min_max_row_float_for_fp16[kRowwiseMinMaxNumCols];
   const auto is_valid_rowwise_min_max = (rowwise_min_max != nullptr);
   if constexpr (std::is_same_v<InputType, float16>) {
     input_row_float_for_fp16 = static_cast<float*>(
@@ -1798,7 +1825,8 @@ void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatAvx2(
 
     const float* min_max_row_float = nullptr;
     if (is_valid_rowwise_min_max) {
-      const InputType* min_max_row = rowwise_min_max + row * 2;
+      const InputType* min_max_row =
+          rowwise_min_max + row * kRowwiseMinMaxNumCols;
 
       if constexpr (std::is_same_v<InputType, float>) {
         min_max_row_float = reinterpret_cast<const float*>(min_max_row);
@@ -2205,7 +2233,8 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2(
       const type* input,                                            \
       size_t input_rows,                                            \
       int input_columns,                                            \
-      std::uint8_t* output);                                        \
+      std::uint8_t* output,                                         \
+      const type* rowwise_min_max);                                 \
   template void                                                     \
   FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfAvx2<type, bit_rate>( \
       const std::uint8_t* input,                                    \

test/QuantUtilsTest.cc

@@ -594,6 +594,14 @@ class EmbeddingQuantizeFixedNumberTest : public testing::TestWithParam<int> {
       1, 1,
       -64, 191,
     };
+
+    float16_test_input_rowwise_min_max.resize(
+        float_test_input_rowwise_min_max.size());
+    std::transform(
+        float_test_input_rowwise_min_max.begin(),
+        float_test_input_rowwise_min_max.end(),
+        float16_test_input_rowwise_min_max.begin(),
+        [](float input) { return cpu_float2half_rn(input); });
   }
   // clang-format on
 
@@ -607,6 +615,7 @@ class EmbeddingQuantizeFixedNumberTest : public testing::TestWithParam<int> {
       expected_output_half;
   std::vector<uint8_t> expected_output_float;
   std::vector<float> float_test_input_rowwise_min_max;
+  std::vector<float16> float16_test_input_rowwise_min_max;
 };
 
 INSTANTIATE_TEST_SUITE_P(
@@ -861,3 +870,119 @@ TEST_P(
       col));
 #endif
 }
+
+TEST_P(
+    EmbeddingQuantizeFixedNumberTest,
+    embeddingFloatOrHalfToFusedNBitRowwiseQuantizedSBHalfTest) {
+  const int bit_rate = GetParam();
+
+  // Confirm that quantization with rowwise_min_max produces expected results.
+  vector<uint8_t> outVectFloatTest(row * out_cols_half);
+  FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf<float>(
+      bit_rate,
+      float_test_input.data(),
+      row,
+      col,
+      outVectFloatTest.data(),
+      float_test_input_rowwise_min_max.data());
+  EXPECT_TRUE(isQEmbeddingClose<float16>(
+      expected_output_half[bit_rate], outVectFloatTest, row, col));
+
+  vector<uint8_t> outVectHalfTest(row * out_cols_half);
+  FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf<float16>(
+      bit_rate,
+      float16_test_input.data(),
+      row,
+      col,
+      outVectHalfTest.data(),
+      float16_test_input_rowwise_min_max.data());
+  EXPECT_TRUE(isQEmbeddingClose<float16>(
+      expected_output_half[bit_rate], outVectHalfTest, row, col));
+
+  // Confirm that quantization with and without rowwise_min_max produces
+  // similar results.
+  vector<uint8_t> outVecFloatTestNoRowwiseMinMax(row * out_cols_half);
+  FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf<float>(
+      bit_rate,
+      float_test_input.data(),
+      row,
+      col,
+      outVecFloatTestNoRowwiseMinMax.data());
+  EXPECT_TRUE(isQEmbeddingClose<float16>(
+      expected_output_half[bit_rate],
+      outVecFloatTestNoRowwiseMinMax,
+      row,
+      col));
+  EXPECT_TRUE(isQEmbeddingClose<float16>(
+      outVectFloatTest, outVecFloatTestNoRowwiseMinMax, row, col));
+
+  vector<uint8_t> outVecHalfTestNoRowwiseMinMax(row * out_cols_half);
+  FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf<float16>(
+      bit_rate,
+      float16_test_input.data(),
+      row,
+      col,
+      outVecHalfTestNoRowwiseMinMax.data());
+  EXPECT_TRUE(isQEmbeddingClose<float16>(
+      expected_output_half[bit_rate], outVecHalfTestNoRowwiseMinMax, row, col));
+  EXPECT_TRUE(isQEmbeddingClose<float16>(
+      outVectHalfTest, outVecHalfTestNoRowwiseMinMax, row, col));
+
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+  // Confirm that incorrect min and max values for each row in the input
+  // of rowwise_min_max produces different results.
+  // Since Windows & ARM are not yet supported, only run this test on x86_64.
+  std::vector<float> float_test_input_incorrect_rowwise_min_max = {
+      10,
+      15,
+      -14,
+      1,
+  };
+  std::vector<float16> float16_test_input_incorrect_rowwise_min_max;
+  float16_test_input_incorrect_rowwise_min_max.resize(
+      float_test_input_incorrect_rowwise_min_max.size());
+  std::transform(
+      float_test_input_incorrect_rowwise_min_max.begin(),
+      float_test_input_incorrect_rowwise_min_max.end(),
+      float16_test_input_incorrect_rowwise_min_max.begin(),
+      [](float input) { return cpu_float2half_rn(input); });
+
+  vector<uint8_t> outVecFloatTestIncorrectRowwiseMinMax(row * out_cols_half);
+  FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf<float>(
+      bit_rate,
+      float_test_input.data(),
+      row,
+      col,
+      outVecFloatTestIncorrectRowwiseMinMax.data(),
+      float_test_input_incorrect_rowwise_min_max.data());
+  EXPECT_FALSE(isQEmbeddingClose<float16>(
+      expected_output_half[bit_rate],
+      outVecFloatTestIncorrectRowwiseMinMax,
+      row,
+      col));
+  EXPECT_FALSE(isQEmbeddingClose<float16>(
+      outVecFloatTestIncorrectRowwiseMinMax,
+      outVecFloatTestNoRowwiseMinMax,
+      row,
+      col));
+
+  vector<uint8_t> outVecHalfTestIncorrectRowwiseMinMax(row * out_cols_half);
+  FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf<float16>(
+      bit_rate,
+      float16_test_input.data(),
+      row,
+      col,
+      outVecHalfTestIncorrectRowwiseMinMax.data(),
+      float16_test_input_incorrect_rowwise_min_max.data());
+  EXPECT_FALSE(isQEmbeddingClose<float16>(
+      expected_output_half[bit_rate],
+      outVecHalfTestIncorrectRowwiseMinMax,
+      row,
+      col));
+  EXPECT_FALSE(isQEmbeddingClose<float16>(
+      outVecHalfTestIncorrectRowwiseMinMax,
+      outVecHalfTestNoRowwiseMinMax,
+      row,
+      col));
+#endif
+}