Make float16/bfloat16 distinct types

Author: cyyever

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)                            \

src/EmbeddingSpMDMAutovec.cc

@@ -69,13 +69,17 @@ static inline void fill_output(
   } else if constexpr (std::is_same_v<OutType, uint16_t>) {
     if (is_bf16_out) {
       for (int j = 0; j < block_size; ++j) {
-        out[j] = cpu_float2bfloat16(src[j]);
+        out[j] = cpu_float2bfloat16(src[j]).val;
       }
     } else {
       for (int j = 0; j < block_size; ++j) {
-        out[j] = cpu_float2half(src[j]);
+        out[j] = cpu_float2half(src[j]).val;
       }
     }
+  } else if constexpr (std::is_same_v<OutType, float16>) {
+    for (int j = 0; j < block_size; ++j) {
+      out[j] = cpu_float2half(src[j]);
+    }
   }
 }
 
@@ -1053,18 +1057,24 @@ static bool ALWAYS_INLINE EmbeddingSpMDMRowWiseSparse_autovec(
 #ifdef FBGEMM_VECTOR_WIDTH
         for (; j < block_size - (block_size % FBGEMM_VECTOR_WIDTH); ++j) {
           const InType* inptr = input_row++;
-          out[j] = std::fma(
-              weight,
-              std::is_same_v<InType, float16> ? cpu_half2float(*inptr) : *inptr,
-              out[j]);
+          float in_val = 0.f;
+          if constexpr (std::is_same_v<InType, float16>) {
+            in_val = cpu_half2float(*inptr);
+          } else {
+            in_val = *inptr;
+          }
+          out[j] = std::fma(weight, in_val, out[j]);
         }
 #endif
         for (; j < block_size; ++j) {
           const InType* inptr = input_row++;
-          out[j] = std::fma(
-              weight,
-              std::is_same_v<InType, float16> ? cpu_half2float(*inptr) : *inptr,
-              out[j]);
+          float in_val = 0.f;
+          if constexpr (std::is_same_v<InType, float16>) {
+            in_val = cpu_half2float(*inptr);
+          } else {
+            in_val = *inptr;
+          }
+          out[j] = std::fma(weight, in_val, out[j]);
         }
       }
       if (normalize_by_lengths && len) {
@@ -2303,9 +2313,10 @@ GenerateEmbeddingSpMDMRowWiseSparse_autovec(
   INSTANTIATE_SPMDM_NBIT_WITH_STRIDES(INDEX_TYPE, OFFSET_TYPE, OUT_TYPE) \
   INSTANTIATE_SPMDM_FP8(INDEX_TYPE, OFFSET_TYPE, OUT_TYPE)
 
-#define INSTANTIATE_SPMDM_OUT_T(INDEX_TYPE, OFFSET_TYPE)   \
-  INSTANTIATE_SPMDM_BASE(INDEX_TYPE, OFFSET_TYPE, float)   \
-  INSTANTIATE_SPMDM_BASE(INDEX_TYPE, OFFSET_TYPE, float16) \
+#define INSTANTIATE_SPMDM_OUT_T(INDEX_TYPE, OFFSET_TYPE)    \
+  INSTANTIATE_SPMDM_BASE(INDEX_TYPE, OFFSET_TYPE, float)    \
+  INSTANTIATE_SPMDM_BASE(INDEX_TYPE, OFFSET_TYPE, float16)  \
+  INSTANTIATE_SPMDM_BASE(INDEX_TYPE, OFFSET_TYPE, uint16_t) \
   INSTANTIATE_SPMDM_BASE(INDEX_TYPE, OFFSET_TYPE, uint8_t)
 
 #define INSTANTIATE_SPMDM_OFFSET_T(INDEX_TYPE) \
@@ -2356,10 +2367,11 @@ INSTANTIATE_SPMDM_OFFSET_T(int64_t)
       bool is_bf16_out,                                                    \
       bool is_bf16_in);
 
-#define INSTANTIATE_SPMDM_OUT_T(IN_TYPE, INDEX_TYPE, OFFSET_TYPE)        \
-  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, float)        \
-  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, float16)      \
-  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, std::uint8_t) \
+#define INSTANTIATE_SPMDM_OUT_T(IN_TYPE, INDEX_TYPE, OFFSET_TYPE)         \
+  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, float)         \
+  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, float16)       \
+  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, std::uint16_t) \
+  INSTANTIATE_SPMDM_BASE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE, std::uint8_t)  \
   INSTANTIATE_SPMDM_ROWWISE(IN_TYPE, INDEX_TYPE, OFFSET_TYPE)
 
 #define INSTANTIATE_SPMDM_OFFSET_T(IN_TYPE, INDEX_TYPE)      \
@@ -2372,6 +2384,7 @@ INSTANTIATE_SPMDM_OFFSET_T(int64_t)
 
 INSTANTIATE_SPMDM_INDEX_T(float)
 INSTANTIATE_SPMDM_INDEX_T(float16)
+INSTANTIATE_SPMDM_INDEX_T(std::uint16_t)
 INSTANTIATE_SPMDM_INDEX_T(std::uint8_t)
 
 #undef INSTANTIATE_SPMDM_ROWWISE

src/EmbeddingSpMDMAvx2.cc

@@ -151,6 +151,7 @@ bool EmbeddingSpMDMBlockSize1_(
 
 INSTANTIATE_SPMDM_INDEX_T(float)
 INSTANTIATE_SPMDM_INDEX_T(float16)
+INSTANTIATE_SPMDM_INDEX_T(std::uint16_t)
 INSTANTIATE_SPMDM_INDEX_T(std::uint8_t)
 
 #undef INSTANTIATE_SPMDM_INDEX_T

src/EmbeddingSpMDMNBit.cc

@@ -1392,6 +1392,7 @@ GenerateEmbeddingSpMDMNBitRowWiseSparse(
 #define INSTANTIATE_SPMDM_OUT_T(INDEX_TYPE, OFFSET_TYPE)                   \
   INSTANTIATE_SPMDM_THREAD_LOCAL(INDEX_TYPE, OFFSET_TYPE, float)           \
   INSTANTIATE_SPMDM_THREAD_LOCAL(INDEX_TYPE, OFFSET_TYPE, uint16_t)        \
+  INSTANTIATE_SPMDM_THREAD_LOCAL(INDEX_TYPE, OFFSET_TYPE, float16)         \
   INSTANTIATE_SPMDM_THREAD_LOCAL(INDEX_TYPE, OFFSET_TYPE, uint8_t)         \
   template FBGEMM_API typename EmbeddingSpMDMRowWiseSparseKernelSignature< \
       uint8_t,                                                             \