avx512 based int8 -> bf16 dequantization (#4912)

Summary:
X-link: facebookresearch/FBGEMM#1949


Use AVX512-bf16 intrinsics for int8 -> bf16 dequantization

Differential Revision: D82507938

Author: seanx92

fbgemm_gpu/cmake/Fbgemm.cmake

@@ -26,7 +26,8 @@ set(fbgemm_sources_avx2
   "${FBGEMM}/src/QuantUtilsAvx2.cc")
 
 set(fbgemm_sources_avx512
-  "${FBGEMM}/src/EmbeddingSpMDMAvx512.cc")
+  "${FBGEMM}/src/EmbeddingSpMDMAvx512.cc"
+  "${FBGEMM}/src/QuantUtilsAvx512.cc")
 
 if(CXX_AVX2_FOUND)
   set_source_files_properties(${fbgemm_sources_avx2}
@@ -46,7 +47,7 @@ if(CXX_AVX2_FOUND)
     ${fbgemm_sources}
     ${fbgemm_sources_avx2})
 endif()
-if((NOT FBGEMM_BUILD_VARIANT STREQUAL BUILD_VARIANT_ROCM) AND CXX_AVX512_FOUND)
+if(CXX_AVX512_FOUND)
   set(fbgemm_sources
     ${fbgemm_sources}
     ${fbgemm_sources_avx2}

fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h

@@ -411,6 +411,7 @@ at::Tensor FP8rowwise_to_float_cpu(
     const bool forward = true,
     const int64_t output_dtype = 0);
 at::Tensor fused8bitrowwise_to_half_cpu(const at::Tensor& input);
+at::Tensor fused8bitrowwise_to_bfloat16_cpu(const at::Tensor& input);
 at::Tensor fused8bitrowwise_to_float_or_half_cpu(
     const at::Tensor& input,
     const int64_t output_dtype,
@@ -469,6 +470,9 @@ at::Tensor _fusednbitrowwise_to_float_or_half_gpu(
 at::Tensor& _fused8bitrowwise_to_float_cpu_out(
     at::Tensor& output,
     const at::Tensor& input);
+at::Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
+    at::Tensor& output,
+    const at::Tensor& input);
 at::Tensor& _float_to_fused8bitrowwise_cpu_out(
     at::Tensor& output,
     const at::Tensor& input);

fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp

@@ -55,7 +55,7 @@ Tensor& _float_to_fused8bitrowwise_cpu_out_t(
   return output;
 }
 
-template <typename output_t>
+template <typename output_t, bool is_uint16_t_of_type_bf16 = false>
 Tensor& _fused8bitrowwise_to_float_cpu_out_t(
     Tensor& output,
     const Tensor& input) {
@@ -78,7 +78,9 @@ Tensor& _fused8bitrowwise_to_float_cpu_out_t(
   auto output_data = static_cast<output_t*>(
       output.data_ptr()); // output.data_ptr<output_t>(); -> Yields
                           // unresolved data_ptr symbol.
-  fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<output_t>(
+  fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<
+      output_t,
+      is_uint16_t_of_type_bf16>(
       input.data_ptr<uint8_t>(), nrows, ncols, output_data);
 
   return output;
@@ -217,11 +219,19 @@ Tensor _fusednbitrowwise_sbfront_to_float_or_half_cpu(
 Tensor& _fused8bitrowwise_to_float_cpu_out(
     Tensor& output,
     const Tensor& input) {
-  return _fused8bitrowwise_to_float_cpu_out_t<float>(output, input);
+  return _fused8bitrowwise_to_float_cpu_out_t<float, false>(output, input);
 }
 
 Tensor& fused8bitrowwise_to_half_cpu_out(Tensor& output, const Tensor& input) {
-  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16>(output, input);
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16, false>(
+      output, input);
+}
+
+Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
+    Tensor& output,
+    const Tensor& input) {
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::bfloat16, true>(
+      output, input);
 }
 
 /// @ingroup quantize-data-cpu
@@ -285,6 +295,13 @@ Tensor fused8bitrowwise_to_half_cpu(const Tensor& input) {
   return fused8bitrowwise_to_half_cpu_out(output, input);
 }
 
+/// @ingroup quantize-data-cpu
+///
+Tensor fused8bitrowwise_to_bfloat16_cpu(const Tensor& input) {
+  auto output = at::empty({0}, input.options().dtype(at::kBFloat16));
+  return _fused8bitrowwise_to_bfloat16_cpu_out(output, input);
+}
+
 /// @ingroup quantize-data-cpu
 ///
 Tensor fused8bitrowwise_to_float_or_half_cpu(
@@ -305,6 +322,10 @@ Tensor fused8bitrowwise_to_float_or_half_cpu(
       output = at::empty({0}, input.options().dtype(at::kHalf));
       output = fused8bitrowwise_to_half_cpu_out(output, input);
       break;
+    case SparseType::BF16:
+      output = at::empty({0}, input.options().dtype(at::kBFloat16));
+      output = _fused8bitrowwise_to_bfloat16_cpu_out(output, input);
+      break;
     default:
       TORCH_CHECK(false);
   }
@@ -582,6 +603,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
       "FP8RowwiseQuantizedToFloat(Tensor input, bool forward, int output_dtype=0) -> Tensor",
       {PT2_COMPLIANT_TAG});
   m.def("Fused8BitRowwiseQuantizedToHalf(Tensor input) -> Tensor");
+  m.def("Fused8BitRowwiseQuantizedToBfloat16(Tensor input) -> Tensor");
   m.def(
       "Fused8BitRowwiseQuantizedToFloatOrHalf(Tensor input, int output_dtype=0, bool scale_bias_last=True, bool quant_padding_float_type=True) -> Tensor");
   m.def(
@@ -648,6 +670,9 @@ TORCH_LIBRARY_IMPL(fbgemm, CPU, m) {
   DISPATCH_TO_CPU(
       "Fused8BitRowwiseQuantizedToHalf",
       fbgemm_gpu::fused8bitrowwise_to_half_cpu);
+  DISPATCH_TO_CPU(
+      "Fused8BitRowwiseQuantizedToBfloat16",
+      fbgemm_gpu::fused8bitrowwise_to_bfloat16_cpu);
   DISPATCH_TO_CPU(
       "Fused8BitRowwiseQuantizedToFloatOrHalf",
       fbgemm_gpu::fused8bitrowwise_to_float_or_half_cpu);

fbgemm_gpu/test/quantize/fused_8bit_rowwise_test.py

@@ -141,9 +141,9 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
 
         assume(ncols % (2 * num_elem_per_byte) == 0)
         if not test_cuda:
-            # cpu path does not support bf16
+            # cpu path only supports bf16 dequantization
             if output_dtype == SparseType.BF16:
-                return
+                input_data = input_data.float()
             if test_generic_op:
                 quantized_data = (
                     torch.ops.fbgemm.FloatOrHalfToFused8BitRowwiseQuantized(input_data)
@@ -171,6 +171,15 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
                     dequantized_data = torch.ops.fbgemm.Fused8BitRowwiseQuantizedToHalf(
                         quantized_data
                     )
+                elif output_dtype == SparseType.BF16:
+                    quantized_data = torch.ops.fbgemm.FloatToFused8BitRowwiseQuantized(
+                        input_data,
+                    )
+                    dequantized_data = (
+                        torch.ops.fbgemm.Fused8BitRowwiseQuantizedToBfloat16(
+                            quantized_data,
+                        )
+                    )
                 else:
                     raise NotImplementedError("Unsupported dtype")
 
@@ -185,6 +194,10 @@ def quantize_and_dequantize_op_test_helper(  # noqa: C901
                 torch.testing.assert_close(dequantized_data.float(), reference.float())
             elif output_dtype == SparseType.FP16:
                 torch.testing.assert_close(dequantized_data.half(), reference.half())
+            elif output_dtype == SparseType.BF16:
+                torch.testing.assert_close(
+                    dequantized_data.bfloat16(), reference.bfloat16()
+                )
         if test_cuda and gpu_available:
             if nrows == 0 or ncols == 0:
                 return

include/fbgemm/QuantUtils.h

@@ -10,6 +10,7 @@
 
 #include "./FbgemmBuild.h" // @manual
 #include "./QuantUtilsAvx2.h" // @manual
+#include "./QuantUtilsAvx512.h" // @manual
 #include "./QuantUtilsNeon.h" // @manual
 #include "./Types.h" // @manual
 #include "./Utils.h" // @manual
@@ -330,7 +331,7 @@ FBGEMM_API void FloatOrHalfToFused8BitRowwiseQuantizedSBFloat(
  * This version intentionally supports only 8-bit because
  * the corresponding quantize version only supports 8-bit.
  */
-template <typename OutputType>
+template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
 FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
     const uint8_t* input,
     size_t input_rows,
@@ -377,7 +378,7 @@ FBGEMM_API void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef(
  * Same as Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf but unoptimized.
  * This should not be called directly except in testing.
  */
-template <typename OutputType>
+template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
 FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
     const uint8_t* input,
     size_t input_rows,

include/fbgemm/QuantUtilsAvx512.h

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Types.h"
 #if defined(FBGEMM_FBCODE) || !defined(__aarch64__)
 
 #include <cstdint>
@@ -37,6 +38,12 @@ FBGEMM_API void requantizeOutputProcessingGConvAvx512(
     int ld_out,
     int ld_in,
     const requantizationParams_t<BIAS_TYPE>& r);
+
+void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
+    const std::uint8_t* input,
+    size_t input_rows,
+    int input_columns,
+    bfloat16* output);
 } // namespace fbgemm
 
 #endif

include/fbgemm/Utils.h

@@ -177,6 +177,11 @@ FBGEMM_API bool fbgemmHasAvx2Support();
  */
 FBGEMM_API bool fbgemmHasAvx512VnniSupport();
 
+/**
+ * @brief Are we running on a AVX512_BF16 supported cpu?
+ */
+FBGEMM_API bool fbgemmHasAvx512Bf16Support();
+
 /**
  * @brief Are we running on a ARM Neon supported cpu?
  */

src/QuantUtils.cc

@@ -806,7 +806,7 @@ void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalf(
   }
 }
 
-template <typename OutputType>
+template <typename OutputType, bool is_uint16_t_of_type_bf16>
 void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
     const std::uint8_t* input,
     size_t input_rows,
@@ -826,13 +826,17 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
       if constexpr (std::is_same<OutputType, float>()) {
         output_row[col] = output_value;
       } else {
-        output_row[col] = cpu_float2half_rn(output_value);
+        if constexpr (is_uint16_t_of_type_bf16) {
+          output_row[col] = cpu_float2bfloat16(output_value);
+        } else {
+          output_row[col] = cpu_float2half_rn(output_value);
+        }
       }
     }
   }
 }
 
-template <typename OutputType>
+template <typename OutputType, bool is_uint16_t_of_type_bf16>
 void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
     const std::uint8_t* input,
     size_t input_rows,
@@ -844,13 +848,23 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
 #else
   if (cpuinfo_initialize() && fbgemmHasAvx2Support()) {
 #if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
-    Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2<OutputType>(
-        input, input_rows, input_columns, output);
+    if (is_uint16_t_of_type_bf16 && fbgemmHasAvx512Bf16Support()) {
+      Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
+          input,
+          input_rows,
+          input_columns,
+          reinterpret_cast<bfloat16*>(output));
+      return;
+    } else if (!is_uint16_t_of_type_bf16) {
+      Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2<OutputType>(
+          input, input_rows, input_columns, output);
+      return;
+    }
 #endif
-  } else {
-    Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<OutputType>(
-        input, input_rows, input_columns, output);
   }
+  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<
+      OutputType,
+      is_uint16_t_of_type_bf16>(input, input_rows, input_columns, output);
 #endif
 }
 
@@ -906,13 +920,25 @@ void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
       std::uint8_t* output,                                                    \
       const type* rowwise_min_max);                                            \
   template FBGEMM_API void                                                     \
-  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<type>(                      \
+  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<type, false>(               \
+      const uint8_t* input,                                                    \
+      size_t input_rows,                                                       \
+      int input_columns,                                                       \
+      type* output);                                                           \
+  template FBGEMM_API void                                                     \
+  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef<type, true>(                \
+      const uint8_t* input,                                                    \
+      size_t input_rows,                                                       \
+      int input_columns,                                                       \
+      type* output);                                                           \
+  template FBGEMM_API void                                                     \
+  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<type, false>(                  \
       const uint8_t* input,                                                    \
       size_t input_rows,                                                       \
       int input_columns,                                                       \
       type* output);                                                           \
   template FBGEMM_API void                                                     \
-  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<type>(                         \
+  Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<type, true>(                   \
       const uint8_t* input,                                                    \
       size_t input_rows,                                                       \
       int input_columns,                                                       \

src/QuantUtilsAvx512.cc

@@ -6,12 +6,14 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <stdexcept>
 #define FBGEMM_EXPORTS
 #include "fbgemm/QuantUtilsAvx512.h"
 #if defined(__x86_64__) || defined(__i386__) || \
     (defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)))
 #include <immintrin.h>
 #endif
+#include <fbgemm/FloatConversion.h>
 #include <cassert>
 
 namespace fbgemm {
@@ -381,6 +383,50 @@ void requantizeOutputProcessingGConvAvx512(
   } // i loop
 }
 
+void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
+    const std::uint8_t* input,
+    size_t input_rows,
+    int input_columns,
+    bfloat16* output) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+  constexpr int VLEN = 8;
+  int output_columns = input_columns - 2 * sizeof(float);
+
+  for (size_t row = 0; row < input_rows; ++row) {
+    const std::uint8_t* input_row = input + row * input_columns;
+    const float* input_row_scale_bias =
+        reinterpret_cast<const float*>(input_row + output_columns);
+    bfloat16* output_row = output + row * output_columns;
+
+    __m256 scale_v = _mm256_set1_ps(input_row_scale_bias[0]);
+    __m256 bias_v = _mm256_set1_ps(input_row_scale_bias[1]);
+
+    int col = 0;
+    for (col = 0; col < output_columns / VLEN * VLEN; col += VLEN) {
+      __m256 in_v = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(
+          _mm_loadl_epi64(reinterpret_cast<const __m128i*>(input_row + col))));
+#ifdef __FMA__
+      __m256 dequantzed_v = _mm256_fmadd_ps(in_v, scale_v, bias_v);
+#else
+      __m256 dequantzed_v = _mm256_add_ps(_mm256_mul_ps(in_v, scale_v), bias_v);
+#endif
+      _mm_storeu_si128(
+          reinterpret_cast<__m128i*>(output_row + col),
+          (__m128i)(_mm256_cvtneps_pbh(dequantzed_v)));
+    }
+
+    for (; col < output_columns; ++col) {
+      float output_value =
+          input_row[col] * input_row_scale_bias[0] + input_row_scale_bias[1];
+      output_row[col] = cpu_float2bfloat16(output_value);
+    }
+  } // for each row
+#else
+  throw std::runtime_error(
+      "Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512 not implemented for non x86");
+#endif
+}
+
 #define INSTANTIATE_REQUANTIZE_BIAS_TYPE(              \
     A_SYM, B_SYM, Q_GRAN, BIAS, RELU, BIAS_TYPE)       \
   template void requantizeOutputProcessingGConvAvx512< \
@@ -468,4 +514,5 @@ INSTANTIATE_BIAS(false)
 #undef INSTANTIATE_B_SYM
 #undef INSTANTIATE_Q_GRANS
 #undef INSTANTIATE_BIAS
+
 } // namespace fbgemm

src/Utils.cc

@@ -319,6 +319,10 @@ bool fbgemmHasAvx512VnniSupport() {
   return cpuinfo_has_x86_avx512vnni();
 }
 
+bool fbgemmHasAvx512Bf16Support() {
+  return cpuinfo_has_x86_avx512bf16();
+}
+
 bool fbgemmHasArmNeonSupport() {
   return cpuinfo_has_arm_neon();
 }