Make float16/bfloat16 distinct types

Author: cyyever

bench/EmbeddingQuantizeFloatToFloatOrHalfBenchmark.cc

@@ -93,7 +93,7 @@ static void performance_test_bf16() {
       randFill<uint8_t>(inpVec, 0, 20);
 
       int output_columns = colSize - 2 * sizeof(float);
-      aligned_vector<float16> outVec(rowSize * output_columns);
+      aligned_vector<bfloat16> outVec(rowSize * output_columns);
 
       double duration = 0.0f;
 
@@ -102,7 +102,7 @@ static void performance_test_bf16() {
       duration = measureWithWarmup(
           [&]() {
             for (int i = 0; i < kNumRepeats; ++i) {
-              Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<float16, true>(
+              Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<bfloat16>(
                   inpVec.data(), rowSize, colSize, outVec.data());
             }
           },

bench/EmbeddingQuantizeNBitToFloatOrHalfBenchmark.cc

@@ -107,7 +107,7 @@ static void performance_test_bf16() {
         aligned_vector<uint8_t> inpVec(rowSize * bytes_per_row);
         randFill<uint8_t>(inpVec, 0, 20);
 
-        aligned_vector<float16> outVec(rowSize * colSize);
+        aligned_vector<bfloat16> outVec(rowSize * colSize);
 
         double duration = 0.0f;
 
@@ -116,7 +116,7 @@ static void performance_test_bf16() {
         duration = measureWithWarmup(
             [&]() {
               for (int i = 0; i < kNumRepeats; ++i) {
-                FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef<float16, true>(
+                FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef<bfloat16>(
                     bit_rate,
                     inpVec.data(),
                     rowSize,

bench/EmbeddingSpMDMBenchmark.cc

@@ -146,7 +146,7 @@ static void run_benchmark(
             batch_size,
             lengths_sum,
             num_rows,
-            embedding_table_fp16.data(),
+            reinterpret_cast<const uint16_t*>(embedding_table_fp16.data()),
             indices_32.data(),
             offsets.data(),
             has_weight ? weights.data() : nullptr,
@@ -158,7 +158,7 @@ static void run_benchmark(
             batch_size,
             lengths_sum,
             num_rows,
-            embedding_table_fp16.data(),
+            reinterpret_cast<const uint16_t*>(embedding_table_fp16.data()),
             indices.data(),
             offsets.data(),
             has_weight ? weights.data() : nullptr,
@@ -172,7 +172,7 @@ static void run_benchmark(
             batch_size,
             lengths_sum,
             num_rows,
-            embedding_table_bf16.data(),
+            reinterpret_cast<const uint16_t*>(embedding_table_bf16.data()),
             indices_32.data(),
             offsets.data(),
             has_weight ? weights.data() : nullptr,
@@ -184,7 +184,7 @@ static void run_benchmark(
             batch_size,
             lengths_sum,
             num_rows,
-            embedding_table_bf16.data(),
+            reinterpret_cast<const uint16_t*>(embedding_table_bf16.data()),
             indices.data(),
             offsets.data(),
             has_weight ? weights.data() : nullptr,
@@ -223,19 +223,19 @@ static void run_benchmark(
         embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
     auto kernel_fp32_i64 = GenerateEmbeddingSpMDM<float, int64_t>(
         embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
-    auto kernel_fp16_i32 = GenerateEmbeddingSpMDM<float16, int32_t>(
+    auto kernel_fp16_i32 = GenerateEmbeddingSpMDM<uint16_t, int32_t>(
         embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
-    auto kernel_fp16_i64 = GenerateEmbeddingSpMDM<float16, int64_t>(
+    auto kernel_fp16_i64 = GenerateEmbeddingSpMDM<uint16_t, int64_t>(
         embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
-    auto kernel_bf16_i32 = GenerateEmbeddingSpMDM<bfloat16, int32_t>(
+    auto kernel_bf16_i32 = GenerateEmbeddingSpMDM<uint16_t, int32_t>(
         embedding_dim,
         has_weight,
         normalize_by_lengths,
         prefetch ? 16 : 0,
         /*is_weight_positional=*/false,
         /*use_offsets=*/true,
         /*is_bf16_out=*/true);
-    auto kernel_bf16_i64 = GenerateEmbeddingSpMDM<bfloat16, int64_t>(
+    auto kernel_bf16_i64 = GenerateEmbeddingSpMDM<uint16_t, int64_t>(
         embedding_dim,
         has_weight,
         normalize_by_lengths,
@@ -254,7 +254,7 @@ static void run_benchmark(
                     batch_size,
                     lengths_sum,
                     num_rows,
-                    embedding_table_fp16.data(),
+                    reinterpret_cast<const uint16_t*>(embedding_table_fp16.data()),
                     indices_32.data(),
                     offsets.data(),
                     has_weight ? weights.data() : nullptr,
@@ -264,7 +264,7 @@ static void run_benchmark(
                     batch_size,
                     lengths_sum,
                     num_rows,
-                    embedding_table_fp16.data(),
+                    reinterpret_cast<const uint16_t*>(embedding_table_fp16.data()),
                     indices.data(),
                     offsets.data(),
                     has_weight ? weights.data() : nullptr,
@@ -276,7 +276,7 @@ static void run_benchmark(
                     batch_size,
                     lengths_sum,
                     num_rows,
-                    embedding_table_bf16.data(),
+                    reinterpret_cast<const uint16_t*>(embedding_table_bf16.data()),
                     indices_32.data(),
                     offsets.data(),
                     has_weight ? weights.data() : nullptr,
@@ -286,7 +286,7 @@ static void run_benchmark(
                     batch_size,
                     lengths_sum,
                     num_rows,
-                    embedding_table_bf16.data(),
+                    reinterpret_cast<const uint16_t*>(embedding_table_bf16.data()),
                     indices.data(),
                     offsets.data(),
                     has_weight ? weights.data() : nullptr,

fbgemm_gpu/src/quantize_ops/quantize_ops_cpu.cpp

@@ -25,6 +25,23 @@ using Tensor = at::Tensor;
 
 namespace fbgemm_gpu {
 
+// Map at::Half/at::BFloat16 to the corresponding fbgemm strict types;
+// other types (e.g. float) pass through unchanged.
+template <typename T>
+struct to_fbgemm_type {
+  using type = T;
+};
+template <>
+struct to_fbgemm_type<at::Half> {
+  using type = fbgemm::float16;
+};
+template <>
+struct to_fbgemm_type<at::BFloat16> {
+  using type = fbgemm::bfloat16;
+};
+template <typename T>
+using to_fbgemm_type_t = typename to_fbgemm_type<T>::type;
+
 template <typename input_t>
 Tensor& _float_to_fused8bitrowwise_cpu_out_t(
     Tensor& output,
@@ -55,7 +72,7 @@ Tensor& _float_to_fused8bitrowwise_cpu_out_t(
   return output;
 }
 
-template <typename output_t, bool is_uint16_t_of_type_bf16 = false>
+template <typename output_t>
 Tensor& _fused8bitrowwise_to_float_cpu_out_t(
     Tensor& output,
     const Tensor& input,
@@ -86,9 +103,7 @@ Tensor& _fused8bitrowwise_to_float_cpu_out_t(
   auto output_data = static_cast<output_t*>(
       output.mutable_data_ptr()); // output.mutable_data_ptr<output_t>(); ->
                                   // Yields unresolved data_ptr symbol.
-  fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<
-      output_t,
-      is_uint16_t_of_type_bf16>(
+  fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf<output_t>(
       input.const_data_ptr<uint8_t>(),
       nrows,
       ncols,
@@ -206,17 +221,12 @@ Tensor _fusednbitrowwise_sbfront_to_float_or_half_cpu(
         "Unsupported output dtype for _fusednbitrowwise_sbfront_to_float_or_half_cpu");
   }
 
-  using output_ty = std::
-      conditional_t<std::is_same_v<output_t, float>, float, fbgemm::float16>;
+  using output_ty = to_fbgemm_type_t<output_t>;
   output_ty* output_data = static_cast<output_ty*>(
       output.mutable_data_ptr()); // output.mutable_data_ptr<output_t>(); ->
                                   // Yields unresolved data_ptr symbol.
 
-  constexpr bool is_uint16_t_of_type_bf16 =
-      std::is_same_v<output_t, at::BFloat16>;
-  fbgemm::FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef<
-      output_ty,
-      is_uint16_t_of_type_bf16>(
+  fbgemm::FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef<output_ty>(
       bit_rate,
       input.const_data_ptr<uint8_t>(),
       nrows,
@@ -234,7 +244,7 @@ Tensor& _fused8bitrowwise_to_float_cpu_out(
     const Tensor& input,
     const bool scale_bias_last,
     const bool quant_padding_float_type) {
-  return _fused8bitrowwise_to_float_cpu_out_t<float, false>(
+  return _fused8bitrowwise_to_float_cpu_out_t<float>(
       output, input, scale_bias_last, quant_padding_float_type);
 }
 
@@ -243,7 +253,7 @@ Tensor& fused8bitrowwise_to_half_cpu_out(
     const Tensor& input,
     const bool scale_bias_last,
     const bool quant_padding_float_type) {
-  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16, false>(
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::float16>(
       output, input, scale_bias_last, quant_padding_float_type);
 }
 
@@ -252,7 +262,7 @@ Tensor& _fused8bitrowwise_to_bfloat16_cpu_out(
     const Tensor& input,
     const bool scale_bias_last,
     const bool quant_padding_float_type) {
-  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::bfloat16, true>(
+  return _fused8bitrowwise_to_float_cpu_out_t<fbgemm::bfloat16>(
       output, input, scale_bias_last, quant_padding_float_type);
 }
 

include/fbgemm/FloatConversion.h

@@ -210,18 +210,18 @@ inline typename Tgt::value_type ieee754_trunc(typename Src::value_type value) {
 
 inline float16 cpu_float2half_rn(float f) {
   uint32_t f_u32 = std::bit_cast<uint32_t>(f);
-  return detail::ieee754_trunc<
+  return {detail::ieee754_trunc<
       /*Src=*/detail::IEEE754Single,
       /*Tgt=*/detail::IEEE754Half,
-      detail::RoundingMode::ToNearestTiesToEven>(f_u32);
+      detail::RoundingMode::ToNearestTiesToEven>(f_u32)};
 }
 
 inline float16 cpu_float2half_rz(float f) {
   uint32_t f_u32 = std::bit_cast<uint32_t>(f);
-  return detail::ieee754_trunc<
+  return {detail::ieee754_trunc<
       /*Src=*/detail::IEEE754Single,
       /*Tgt=*/detail::IEEE754Half,
-      detail::RoundingMode::ToZero>(f_u32);
+      detail::RoundingMode::ToZero>(f_u32)};
 }
 
 // Converts a 16-bit unsigned integer representation of a IEEE754 half-precision
@@ -245,10 +245,10 @@ inline float cpu_half2float_ref(const float16 h) {
   constexpr uint32_t f32_most_significant_bit = 1u << 22;
 
   // Get sign and exponent alone by themselves
-  uint32_t sign_bit = (h >> f16_num_non_sign_bits) & 1;
-  uint32_t exponent = (h >> f16_num_mantissa_bits) & f16_exponent_mask;
+  uint32_t sign_bit = (h.val >> f16_num_non_sign_bits) & 1;
+  uint32_t exponent = (h.val >> f16_num_mantissa_bits) & f16_exponent_mask;
   // Shift mantissa so that it fills the most significant bits of a float32
-  uint32_t mantissa = (h & f16_mantissa_mask)
+  uint32_t mantissa = (h.val & f16_mantissa_mask)
       << (f32_num_mantissa_bits - f16_num_mantissa_bits);
 
   if (exponent == f16_exponent_mask) { // NaN or Inf
@@ -280,10 +280,10 @@ inline float cpu_half2float_ref(const float16 h) {
 
 inline float cpu_half2float(const float16 h) {
 #ifdef HAS_NATIVE_FP16_TYPE
-  return std::bit_cast<__fp16>(h);
+  return std::bit_cast<__fp16>(h.val);
 #elif defined(HAS_F16C)
   // Use F16C VCVTPH2PS instruction
-  __m128i v = _mm_cvtsi32_si128(static_cast<int>(h));
+  __m128i v = _mm_cvtsi32_si128(static_cast<int>(h.val));
   return _mm_cvtss_f32(_mm_cvtph_ps(v));
 #else
   return cpu_half2float_ref(h);
@@ -293,25 +293,25 @@ inline float cpu_half2float(const float16 h) {
 inline float16 cpu_float2half(const float f) {
 #ifdef HAS_NATIVE_FP16_TYPE
   __fp16 h = f;
-  return std::bit_cast<float16>(h);
+  return {std::bit_cast<uint16_t>(h)};
 #elif defined(HAS_F16C)
   // Use F16C VCVTPS2PH instruction
   __m128 v = _mm_set_ss(f);
-  return static_cast<float16>(
-      _mm_extract_epi16(_mm_cvtps_ph(v, _MM_FROUND_TO_NEAREST_INT), 0));
+  return {static_cast<uint16_t>(
+      _mm_extract_epi16(_mm_cvtps_ph(v, _MM_FROUND_TO_NEAREST_INT), 0))};
 #else
   return cpu_float2half_rn(f);
 #endif
 }
 
 inline float cpu_bf162float(bfloat16 src) {
-  uint32_t val_fp32 = static_cast<uint32_t>(src) << 16;
+  uint32_t val_fp32 = static_cast<uint32_t>(src.val) << 16;
   return std::bit_cast<float>(val_fp32);
 }
 
 inline bfloat16 cpu_float2bfloat16(float src) {
   uint32_t temp = std::bit_cast<uint32_t>(src);
-  return (temp + (1u << 15)) >> 16;
+  return {static_cast<uint16_t>((temp + (1u << 15)) >> 16)};
 }
 
 } // namespace fbgemm

include/fbgemm/QuantUtils.h

@@ -324,7 +324,7 @@ FBGEMM_API void FloatOrHalfToFused8BitRowwiseQuantizedSBFloat(
  * This version intentionally supports only 8-bit because
  * the corresponding quantize version only supports 8-bit.
  */
-template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
+template <typename OutputType>
 FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
     const uint8_t* input,
     size_t input_rows,
@@ -360,7 +360,7 @@ FBGEMM_API void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatRef(
  * Same as FusedNBitRowwiseQuantizedSBHalfToFloat but unoptimized.
  * This should not be called directly except in testing.
  */
-template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
+template <typename OutputType>
 FBGEMM_API void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef(
     int bit_rate,
     const uint8_t* input,
@@ -373,7 +373,7 @@ FBGEMM_API void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef(
  * Same as Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf but unoptimized.
  * This should not be called directly except in testing.
  */
-template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
+template <typename OutputType>
 FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
     const uint8_t* input,
     size_t input_rows,

include/fbgemm/Types.h

@@ -12,8 +12,18 @@
 
 namespace fbgemm {
 
-using float16 = std::uint16_t;
-using bfloat16 = std::uint16_t;
+struct float16 {
+  uint16_t val;
+  bool operator==(const float16&) const = default;
+};
+
+struct bfloat16 {
+  uint16_t val;
+  bool operator==(const bfloat16&) const = default;
+};
+
+static_assert(sizeof(float16) == 2);
+static_assert(sizeof(bfloat16) == 2);
 
 constexpr int64_t round_up(int64_t val, int64_t unit) {
   return (val + unit - 1) / unit * unit;

src/EmbeddingSpMDM.cc

@@ -1660,6 +1660,7 @@ GenerateEmbeddingSpMDMRowWiseSparse(
   INSTANTIATE_SPMDMFP8_BASE(INDEX_TYPE, OFFSET_TYPE, OUT_TYPE)
 #define INSTANTIATE_SPMDMFP8_BASE_float(INDEX_TYPE, OFFSET_TYPE, OUT_TYPE)
 #define INSTANTIATE_SPMDMFP8_BASE_uint16_t(INDEX_TYPE, OFFSET_TYPE, OUT_TYPE)
+#define INSTANTIATE_SPMDMFP8_BASE_float16(INDEX_TYPE, OFFSET_TYPE, OUT_TYPE)
 
 #define INSTANTIATE_SPMDM_BASE_THREAD_LOCAL(                               \
     IN_TYPE, INDEX_TYPE, OFFSET_TYPE, OUT_TYPE)                            \
@@ -1695,6 +1696,7 @@ GenerateEmbeddingSpMDMRowWiseSparse(
   INSTANTIATE_SPMDM_OFFSET_T(IN_TYPE, int64_t)
 
 INSTANTIATE_SPMDM_INDEX_T(float)
+INSTANTIATE_SPMDM_INDEX_T(float16)
 INSTANTIATE_SPMDM_INDEX_T(uint16_t)
 INSTANTIATE_SPMDM_INDEX_T(uint8_t)