mlx 1.0.0

data/mlx/mlx/backend/cpu/unary.cpp

@@ -0,0 +1,238 @@
+// Copyright © 2024 Apple Inc.
+
+// Required for using M_LN2 in MSVC.
+#define _USE_MATH_DEFINES
+
+#include <cassert>
+
+#include "mlx/backend/cpu/unary.h"
+#include "mlx/backend/cpu/unary_ops.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), unsignedinteger) || in.dtype() == bool_) {
+    // No-op for unsigned types
+    out.copy_shared_buffer(in);
+  } else {
+    unary_signed(in, out, detail::Abs(), stream());
+  }
+}
+
+void ArcCos::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcCos(), stream());
+}
+
+void ArcCosh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcCosh(), stream());
+}
+
+void ArcSin::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcSin(), stream());
+}
+
+void ArcSinh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcSinh(), stream());
+}
+
+void ArcTan::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcTan(), stream());
+}
+
+void ArcTanh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcTanh(), stream());
+}
+
+void BitwiseInvert::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_int(in, out, detail::BitwiseInvert(), stream());
+}
+
+void Ceil::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_fp(in, out, detail::Ceil(), stream());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+void Conjugate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  unary_complex(inputs[0], out, detail::Conjugate(), stream());
+}
+
+void Cos::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Cos(), stream());
+}
+
+void Cosh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Cosh(), stream());
+}
+
+void Erf::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_real_fp(in, out, detail::Erf(), stream());
+}
+
+void ErfInv::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_real_fp(in, out, detail::ErfInv(), stream());
+}
+
+void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Exp(), stream());
+}
+
+void Expm1::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Expm1(), stream());
+}
+
+void Floor::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_fp(in, out, detail::Floor(), stream());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+void Imag::eval_cpu(const std::vector<array>& inputs, array& out) {
+  unary_complex_to_float(inputs[0], out, detail::Imag(), stream());
+}
+
+void Log::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  switch (base_) {
+    case Base::e:
+      unary_fp(in, out, detail::Log(), stream());
+      break;
+    case Base::two:
+      unary_fp(in, out, detail::Log2(), stream());
+      break;
+    case Base::ten:
+      unary_fp(in, out, detail::Log10(), stream());
+      break;
+  }
+}
+
+void Log1p::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Log1p(), stream());
+}
+
+void LogicalNot::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  unary(in, out, detail::LogicalNot(), stream());
+}
+
+void Negative::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  unary(in, out, detail::Negative(), stream());
+}
+
+void Real::eval_cpu(const std::vector<array>& inputs, array& out) {
+  unary_complex_to_float(inputs[0], out, detail::Real(), stream());
+}
+
+void Round::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_fp(in, out, detail::Round(), stream());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+void Sigmoid::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Sigmoid(), stream());
+}
+
+void Sign::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (in.dtype() == bool_) {
+    out.copy_shared_buffer(in);
+  } else {
+    unary(in, out, detail::Sign(), stream());
+  }
+}
+
+void Sin::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Sin(), stream());
+}
+
+void Sinh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Sinh(), stream());
+}
+
+void Square::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  unary(in, out, detail::Square(), stream());
+}
+
+void Sqrt::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (recip_) {
+    unary_fp(in, out, detail::Rsqrt(), stream());
+  } else {
+    unary_fp(in, out, detail::Sqrt(), stream());
+  }
+}
+
+void Tan::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Tan(), stream());
+}
+
+void Tanh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Tanh(), stream());
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/unary.h

@@ -0,0 +1,281 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/unary.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+template <typename T, typename U = T, typename Op>
+void unary_op(const T* a, U* out, size_t shape, size_t stride) {
+  for (size_t i = 0; i < shape; i += 1) {
+    out[i] = Op{}(*a);
+    a += stride;
+  }
+}
+
+template <typename T, typename U = T, typename Op>
+void unary_op(const array& a, array& out, Op) {
+  const T* src = a.data<T>();
+  U* dst = out.data<U>();
+  auto ndim = a.ndim();
+  if (a.flags().contiguous) {
+    auto size = a.data_size();
+    constexpr int N = std::min(simd::max_size<T>, simd::max_size<U>);
+    while (size >= N) {
+      simd::store(dst, simd::Simd<U, N>(Op{}(simd::load<T, N>(src))));
+      size -= N;
+      src += N;
+      dst += N;
+    }
+    while (size > 0) {
+      *dst = Op{}(*src);
+      size--;
+      dst++;
+      src++;
+    }
+  } else {
+    size_t shape = ndim > 0 ? a.shape().back() : 1;
+    size_t stride = ndim > 0 ? a.strides().back() : 1;
+    if (ndim <= 1) {
+      unary_op<T, U, Op>(src, dst, shape, stride);
+      return;
+    }
+    auto it = ContiguousIterator(a.shape(), a.strides(), ndim - 1);
+    for (size_t elem = 0; elem < a.size(); elem += shape) {
+      unary_op<T, U, Op>(src + it.loc, dst + elem, shape, stride);
+      it.step();
+    }
+  }
+}
+
+template <typename Op>
+void unary(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        unary_op<bool>(a, out, op);
+        break;
+      case uint8:
+        unary_op<uint8_t>(a, out, op);
+        break;
+      case uint16:
+        unary_op<uint16_t>(a, out, op);
+        break;
+      case uint32:
+        unary_op<uint32_t>(a, out, op);
+        break;
+      case uint64:
+        unary_op<uint64_t>(a, out, op);
+        break;
+      case int8:
+        unary_op<int8_t>(a, out, op);
+        break;
+      case int16:
+        unary_op<int16_t>(a, out, op);
+        break;
+      case int32:
+        unary_op<int32_t>(a, out, op);
+        break;
+      case int64:
+        unary_op<int64_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case complex64:
+        unary_op<complex64_t>(a, out, op);
+        break;
+    }
+  });
+}
+
+template <typename Op>
+void unary_real_fp(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      default:
+        std::ostringstream err;
+        err << "[unary_real] Does not support " << out.dtype();
+        throw std::runtime_error(err.str());
+    }
+  });
+}
+template <typename Op>
+void unary_fp(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      case complex64:
+        unary_op<complex64_t>(a, out, op);
+        break;
+      default:
+        std::ostringstream err;
+        err << "[unary_fp] Does not support " << out.dtype();
+        throw std::runtime_error(err.str());
+    }
+  });
+}
+
+template <typename Op>
+void unary_signed(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case int8:
+        unary_op<int8_t>(a, out, op);
+        break;
+      case int16:
+        unary_op<int16_t>(a, out, op);
+        break;
+      case int32:
+        unary_op<int32_t>(a, out, op);
+        break;
+      case int64:
+        unary_op<int64_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case complex64:
+        unary_op<complex64_t>(a, out, op);
+        break;
+      default:
+        throw std::runtime_error("[Abs] Called on unsigned type");
+    }
+  });
+}
+
+template <typename Op>
+void unary_complex(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable { unary_op<complex64_t>(a, out, op); });
+}
+
+template <typename Op>
+void unary_complex_to_float(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch(
+      [a = array::unsafe_weak_copy(a),
+       out = array::unsafe_weak_copy(out),
+       op = op]() mutable { unary_op<complex64_t, float>(a, out, op); });
+}
+
+template <typename Op>
+void unary_int(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case uint8:
+        unary_op<uint8_t>(a, out, op);
+        break;
+      case uint16:
+        unary_op<uint16_t>(a, out, op);
+        break;
+      case uint32:
+        unary_op<uint32_t>(a, out, op);
+        break;
+      case uint64:
+        unary_op<uint64_t>(a, out, op);
+        break;
+      case int8:
+        unary_op<int8_t>(a, out, op);
+        break;
+      case int16:
+        unary_op<int16_t>(a, out, op);
+        break;
+      case int32:
+        unary_op<int32_t>(a, out, op);
+        break;
+      case int64:
+        unary_op<int64_t>(a, out, op);
+        break;
+      default:
+        std::ostringstream err;
+        err << "[unary_int] Does not support " << out.dtype();
+        throw std::runtime_error(err.str());
+    }
+  });
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/unary_ops.h

@@ -0,0 +1,175 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <stdint.h>
+#include <cmath>
+#include <complex>
+
+#include "mlx/backend/cpu/simd/simd.h"
+
+namespace mlx::core::detail {
+
+using namespace mlx::core::simd;
+
+#define SINGLE()                         \
+  template <typename T>                  \
+  T operator()(T x) {                    \
+    return (*this)(Simd<T, 1>(x)).value; \
+  }
+
+#define DEFAULT_OP(Op, op)                \
+  struct Op {                             \
+    template <int N, typename T>          \
+    Simd<T, N> operator()(Simd<T, N> x) { \
+      return simd::op(x);                 \
+    }                                     \
+    SINGLE()                              \
+  };
+
+DEFAULT_OP(Abs, abs)
+DEFAULT_OP(ArcCos, acos)
+DEFAULT_OP(ArcCosh, acosh)
+DEFAULT_OP(ArcSin, asin)
+DEFAULT_OP(ArcSinh, asinh)
+DEFAULT_OP(ArcTan, atan)
+DEFAULT_OP(ArcTanh, atanh)
+DEFAULT_OP(BitwiseInvert, operator~)
+DEFAULT_OP(Ceil, ceil)
+DEFAULT_OP(Conjugate, conj)
+DEFAULT_OP(Cos, cos)
+DEFAULT_OP(Cosh, cosh)
+DEFAULT_OP(Erf, erf)
+DEFAULT_OP(ErfInv, erfinv)
+DEFAULT_OP(Exp, exp)
+DEFAULT_OP(Expm1, expm1)
+DEFAULT_OP(Floor, floor);
+DEFAULT_OP(Log, log);
+DEFAULT_OP(Log2, log2);
+DEFAULT_OP(Log10, log10);
+DEFAULT_OP(Log1p, log1p);
+DEFAULT_OP(LogicalNot, operator!)
+DEFAULT_OP(Negative, operator-)
+DEFAULT_OP(Round, rint);
+DEFAULT_OP(Sin, sin)
+DEFAULT_OP(Sinh, sinh)
+DEFAULT_OP(Sqrt, sqrt)
+DEFAULT_OP(Rsqrt, rsqrt)
+DEFAULT_OP(Tan, tan)
+DEFAULT_OP(Tanh, tanh)
+
+struct Imag {
+  template <int N>
+  Simd<float, N> operator()(Simd<complex64_t, N> x) {
+    return simd::imag(x);
+  }
+  SINGLE()
+};
+
+struct Real {
+  template <int N>
+  Simd<float, N> operator()(Simd<complex64_t, N> x) {
+    return simd::real(x);
+  }
+  SINGLE()
+};
+
+struct Sigmoid {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x) {
+    auto y = 1.0f / (1.0f + simd::exp(simd::abs(x)));
+    return simd::select(x < Simd<T, N>{0}, y, Simd<T, N>{1} - y);
+  }
+  SINGLE()
+};
+
+struct Sign {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x) {
+    auto z = Simd<T, N>{0};
+    auto o = Simd<T, N>{1};
+    auto m = Simd<T, N>{-1};
+    if constexpr (std::is_unsigned_v<T>) {
+      return simd::select(x == z, z, o);
+    } else if constexpr (std::is_same_v<T, complex64_t>) {
+      return simd::select(x == z, x, Simd<T, N>(x / simd::abs(x)));
+    } else {
+      return simd::select(x < z, m, simd::select(x > z, o, z));
+    }
+  }
+  SINGLE()
+};
+
+struct Square {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x) {
+    return x * x;
+  }
+  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 = 543 << 21;
+    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 v = Simd<uint16_t, N>(x & 127) << 7;
+    Simd<float, N> out;
+    if constexpr (simd::max_size<float16_t> >= N) {
+      auto converted = *(Simd<float16_t, N>*)(&v);
+      out = converted * 256.0;
+    } else {
+      for (int i = 0; i < N; ++i) {
+        auto converted = *(float16_t*)(&v[i]);
+        out[i] = converted * 256.0;
+      }
+    }
+    auto sign = Simd<bool, N>(x & 128);
+    return select(sign, -out, out);
+  }
+  float operator()(uint8_t x) {
+    return (*this)(Simd<uint8_t, 1>(x)).value;
+  }
+};
+} // namespace mlx::core::detail