Fp8 conversion (#2686)

* add fp8 e4m3 converters

* add cuda

* default saturate to min/max

* fix for older OS

* fix no gpu/cpu

* fix saturate

* fix compile

mlx/backend/cpu/quantized.cpp

@@ -1,8 +1,11 @@
 // Copyright © 2023 Apple Inc.
 
+#include "mlx/backend/common/unary.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/backend/cpu/unary.h"
+#include "mlx/backend/cpu/unary_ops.h"
 #include "mlx/fast_primitives.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
@@ -1102,4 +1105,44 @@ void fast::Quantize::eval_cpu(
   });
 }
 
+void fast::ConvertFP8::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& in = inputs[0];
+  auto& out = outputs[0];
+  set_unary_output_data(in, out);
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    to_fp8 = to_fp8_]() mutable {
+    if (to_fp8) {
+      switch (in.dtype()) {
+        case float16:
+          unary_op<float16_t, uint8_t>(in, out, detail::ToFP8());
+          break;
+        case bfloat16:
+          unary_op<bfloat16_t, uint8_t>(in, out, detail::ToFP8());
+          break;
+        default:
+          unary_op<float, uint8_t>(in, out, detail::ToFP8());
+          break;
+      }
+    } else {
+      switch (out.dtype()) {
+        case float16:
+          unary_op<uint8_t, float16_t>(in, out, detail::FromFP8());
+          break;
+        case bfloat16:
+          unary_op<uint8_t, bfloat16_t>(in, out, detail::FromFP8());
+          break;
+        default:
+          unary_op<uint8_t, float>(in, out, detail::FromFP8());
+          break;
+      }
+    }
+  });
+}
+
 } // namespace mlx::core

mlx/backend/cpu/simd/accelerate_simd.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <arm_neon.h>
 #include <simd/math.h>
 #include <simd/vector.h>
 
@@ -200,6 +201,15 @@ SIMD_DEFAULT_COMPARISONS(<=)
 SIMD_DEFAULT_COMPARISONS(==)
 SIMD_DEFAULT_COMPARISONS(!=)
 
+template <typename T, int N>
+Simd<T, N> clz(Simd<T, N> x) {
+  auto a = *(uint32x4_t*)(&x);
+  auto b = *((uint32x4_t*)(&x) + 1);
+  a = vclzq_u32(a);
+  b = vclzq_u32(b);
+  return asd::make_uint8(a, b);
+}
+
 template <typename T, int N>
 Simd<T, N> atan2(Simd<T, N> a, Simd<T, N> b) {
   return asd::atan2(a.value, b.value);

mlx/backend/cpu/simd/base_simd.h

@@ -171,6 +171,11 @@ DEFAULT_BINARY(&)
 DEFAULT_BINARY(&&)
 DEFAULT_BINARY(||)
 
+template <typename T>
+Simd<T, 1> clz(Simd<T, 1> x_) {
+  return __builtin_clz(x_.value);
+}
+
 template <typename T>
 Simd<T, 1> remainder(Simd<T, 1> a_, Simd<T, 1> b_) {
   T a = a_.value;

mlx/backend/cpu/unary.h

@@ -24,9 +24,9 @@ void unary_op(const array& a, array& out, Op) {
   auto ndim = a.ndim();
   if (a.flags().contiguous) {
     auto size = a.data_size();
-    constexpr int N = simd::max_size<T>;
+    constexpr int N = std::min(simd::max_size<T>, simd::max_size<U>);
     while (size >= N) {
-      simd::store(dst, Op{}(simd::load<T, N>(src)));
+      simd::store(dst, simd::Simd<U, N>(Op{}(simd::load<T, N>(src))));
       size -= N;
       src += N;
       dst += N;

mlx/backend/cpu/unary_ops.h

@@ -108,4 +108,73 @@ struct Square {
   SINGLE()
 };
 
+template <int N>
+Simd<float, N> fp32_from_bits(Simd<uint32_t, N> x) {
+  return *(Simd<float, N>*)(&x);
+}
+template <int N>
+Simd<uint32_t, N> fp32_to_bits(Simd<float, N> x) {
+  return *(Simd<uint32_t, N>*)(&x);
+}
+
+struct ToFP8 {
+  template <typename T, int N>
+  Simd<uint8_t, N> operator()(Simd<T, N> f) {
+    uint32_t fp8_max = 1087 << 20;
+    auto denorm_mask = Simd<uint32_t, N>(141 << 23);
+    Simd<uint32_t, N> f_bits;
+    Simd<float, N> f32 = f;
+    f_bits = fp32_to_bits(f32);
+    Simd<uint8_t, N> result = 0u;
+    auto sign = f_bits & 0x80000000;
+    f_bits = f_bits ^ sign;
+
+    auto f_bits_low =
+        fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
+    auto result_low = Simd<uint8_t, N>(f_bits_low - denorm_mask);
+
+    auto mant_odd = Simd<uint8_t, N>((f_bits >> 20) & 1);
+    auto f_bits_high = f_bits + (((uint32_t)(7 - 127) << 23) + 0x7FFFF);
+    f_bits_high = f_bits_high + Simd<uint32_t, N>(mant_odd);
+
+    auto result_high = Simd<uint8_t, N>(f_bits_high >> 20);
+    result = select(f_bits < (121 << 23), result_low, result_high);
+
+    auto result_sat = Simd<uint8_t, N>(0x7E);
+    result = select(f_bits >= fp8_max, result_sat, result);
+    return result | Simd<uint8_t, N>(sign >> 24);
+  }
+
+  template <typename T>
+  uint8_t operator()(T x) {
+    return (*this)(Simd<T, 1>(x)).value;
+  }
+};
+
+struct FromFP8 {
+  template <int N>
+  Simd<float, N> operator()(Simd<uint8_t, N> x) {
+    auto w = Simd<uint32_t, N>(x) << 24;
+    auto sign = w & 0x80000000;
+    auto nonsign = w & 0x7FFFFFFF;
+
+    auto renorm_shift = clz(nonsign);
+    renorm_shift = simd::select(
+        renorm_shift > Simd<uint32_t, N>{4},
+        renorm_shift - Simd<uint32_t, N>{4},
+        Simd<uint32_t, N>{0});
+
+    Simd<int32_t, N> inf_nan_mask =
+        (Simd<int32_t, N>(nonsign + 0x01000000) >> 8) & 0x7F800000;
+    auto zero_mask = Simd<int32_t, N>(nonsign - 1) >> 31;
+    auto result = sign |
+        ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
+          inf_nan_mask) &
+         ~zero_mask);
+    return fp32_from_bits(result);
+  }
+  float operator()(uint8_t x) {
+    return (*this)(Simd<uint8_t, 1>(x)).value;
+  }
+};
 } // namespace mlx::core::detail