mlx 1.0.0

data/mlx/mlx/types/fp16.h

@@ -0,0 +1,234 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <vector>
+
+#define __MLX_HALF_NAN__ 0x7D00
+
+namespace mlx::core {
+
+namespace {
+union float_bits_fp16 {
+  float f;
+  uint32_t u;
+};
+} // namespace
+
+struct _MLX_Float16 {
+  uint16_t bits_;
+
+  // Default constructor
+  _MLX_Float16() = default;
+
+  // Default copy constructor
+  _MLX_Float16(_MLX_Float16 const&) = default;
+
+  // Appease std::vector<bool> for being special
+  _MLX_Float16& operator=(std::vector<bool>::reference x) {
+    bits_ = x;
+    return *this;
+  }
+
+  _MLX_Float16& operator=(const float& x) {
+    return (*this = _MLX_Float16(x));
+  }
+
+  // From float32
+  _MLX_Float16(const float& x) : bits_(0) {
+    // Conversion following
+    // https://github.com/Maratyszcza/FP16/blob/master/include/fp16/fp16.h
+
+    // Union
+    float_bits_fp16 in;
+
+    // Take fp32 bits
+    in.f = x;
+
+    // Find and take sign bit
+    uint32_t x_sign_32 = in.u & uint32_t(0x80000000);
+    uint16_t x_sign_16 = (x_sign_32 >> 16);
+
+    if (std::isnan(x)) {
+      bits_ = x_sign_16 | uint16_t(__MLX_HALF_NAN__);
+    } else {
+      // Union
+      float_bits_fp16 inf_scale, zero_scale, magic_bits;
+
+      // Find exponent bits and take the max supported by half
+      uint32_t x_expo_32 = in.u & uint32_t(0x7f800000);
+      uint32_t max_expo_32 = uint32_t(0x38800000);
+      x_expo_32 = x_expo_32 < max_expo_32 ? max_expo_32 : x_expo_32;
+      x_expo_32 += uint32_t(15) << 23;
+
+      // Handle scaling to inf as needed
+      inf_scale.u = uint32_t(0x77800000);
+      zero_scale.u = uint32_t(0x08800000);
+
+      // Combine with magic and let addition do rounding
+      magic_bits.u = x_expo_32;
+      magic_bits.f += (std::abs(x) * inf_scale.f) * zero_scale.f;
+
+      // Take the lower 5 bits of the exponent
+      uint32_t x_expo_16 = ((magic_bits.u >> 13) & uint32_t(0x7c00));
+
+      // Collect the lower 12 bits which have the mantissa
+      uint32_t x_mant_16 = magic_bits.u & uint32_t(0x0fff);
+
+      // Combine sign, exp and mantissa
+      bits_ = (x_sign_16 | uint16_t(x_expo_16 + x_mant_16));
+    }
+  }
+
+  // To float32
+  operator float() const {
+    // Conversion following
+    // https://github.com/Maratyszcza/FP16/blob/master/include/fp16/fp16.h
+
+    // Union
+    float_bits_fp16 out;
+
+    uint32_t x_sign_32 = (bits_ << 16) & uint32_t(0x80000000);
+    uint32_t base = (bits_ << 16);
+    uint32_t two_base = base + base;
+
+    uint32_t denorm_max = 1u << 27;
+    if (two_base < denorm_max) {
+      out.u = uint32_t(126) << 23; // magic mask
+      out.u |= (two_base >> 17); // Bits from fp16
+      out.f -= 0.5f; // magic bias
+    } else {
+      out.u = uint32_t(0xE0) << 23; // exponent offset
+      out.u += (two_base >> 4); // Bits from fp16
+      float out_unscaled = out.f; // Store value
+      out.u = uint32_t(0x7800000); // exponent scale
+      out.f *= out_unscaled;
+    }
+
+    // Add sign
+    out.u |= x_sign_32;
+
+    return out.f;
+  }
+};
+
+#define half_binop_base(__op__, __operator__, otype, atype, btype, ctype) \
+  inline otype __operator__(atype lhs, btype rhs) {                       \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);        \
+  }
+
+#define half_binop_helper(__op__, __operator__, otype, itype, ctype) \
+  inline otype __operator__(_MLX_Float16 lhs, itype rhs) {           \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);   \
+  }                                                                  \
+  inline otype __operator__(itype lhs, _MLX_Float16 rhs) {           \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);   \
+  }
+
+// Operators
+#define half_binop(__op__, __operator__)                                      \
+  half_binop_base(                                                            \
+      __op__, __operator__, _MLX_Float16, _MLX_Float16, _MLX_Float16, float); \
+  half_binop_helper(__op__, __operator__, float, float, float);               \
+  half_binop_helper(__op__, __operator__, double, double, double);            \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, bool, float);         \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, int32_t, float);      \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, uint32_t, float);     \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, int64_t, float);      \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, uint64_t, float);
+
+half_binop(+, operator+);
+half_binop(-, operator-);
+half_binop(*, operator*);
+half_binop(/, operator/);
+
+#undef half_binop
+
+// Comparison ops
+#define half_compop(__op__, __operator__)                             \
+  half_binop_base(                                                    \
+      __op__, __operator__, bool, _MLX_Float16, _MLX_Float16, float); \
+  half_binop_helper(__op__, __operator__, bool, float, float);        \
+  half_binop_helper(__op__, __operator__, bool, double, double);      \
+  half_binop_helper(__op__, __operator__, bool, int32_t, float);      \
+  half_binop_helper(__op__, __operator__, bool, uint32_t, float);     \
+  half_binop_helper(__op__, __operator__, bool, int64_t, float);      \
+  half_binop_helper(__op__, __operator__, bool, uint64_t, float);
+
+half_compop(>, operator>);
+half_compop(<, operator<);
+half_compop(>=, operator>=);
+half_compop(<=, operator<=);
+half_compop(==, operator==);
+half_compop(!=, operator!=);
+
+#undef half_compop
+
+// Negative
+inline _MLX_Float16 operator-(_MLX_Float16 lhs) {
+  return -static_cast<float>(lhs);
+}
+
+// Inplace ops
+#define half_inplace_op(__op__, __operator__)                              \
+  inline _MLX_Float16& __operator__(_MLX_Float16& lhs, const float& rhs) { \
+    lhs = lhs __op__ rhs;                                                  \
+    return lhs;                                                            \
+  }                                                                        \
+  inline float& __operator__(float& lhs, _MLX_Float16 rhs) {               \
+    lhs = lhs __op__ rhs;                                                  \
+    return lhs;                                                            \
+  }
+
+half_inplace_op(+, operator+=);
+half_inplace_op(-, operator-=);
+half_inplace_op(*, operator*=);
+half_inplace_op(/, operator/=);
+
+#undef half_inplace_op
+
+// Bitwise ops
+
+#define half_bitop(__op__, __operator__)                                 \
+  inline _MLX_Float16 __operator__(_MLX_Float16 lhs, _MLX_Float16 rhs) { \
+    _MLX_Float16 out;                                                    \
+    out.bits_ = lhs.bits_ __op__ rhs.bits_;                              \
+    return out;                                                          \
+  }                                                                      \
+  inline _MLX_Float16 __operator__(_MLX_Float16 lhs, uint16_t rhs) {     \
+    _MLX_Float16 out;                                                    \
+    out.bits_ = lhs.bits_ __op__ rhs;                                    \
+    return out;                                                          \
+  }                                                                      \
+  inline _MLX_Float16 __operator__(uint16_t lhs, _MLX_Float16 rhs) {     \
+    _MLX_Float16 out;                                                    \
+    out.bits_ = lhs __op__ rhs.bits_;                                    \
+    return out;                                                          \
+  }
+
+half_bitop(|, operator|);
+half_bitop(&, operator&);
+half_bitop(^, operator^);
+
+#undef half_bitop
+
+#define half_inplace_bitop(__op__, __operator__)                           \
+  inline _MLX_Float16& __operator__(_MLX_Float16& lhs, _MLX_Float16 rhs) { \
+    lhs.bits_ = lhs.bits_ __op__ rhs.bits_;                                \
+    return lhs;                                                            \
+  }                                                                        \
+  inline _MLX_Float16& __operator__(_MLX_Float16& lhs, uint16_t rhs) {     \
+    lhs.bits_ = lhs.bits_ __op__ rhs;                                      \
+    return lhs;                                                            \
+  }
+
+half_inplace_bitop(|, operator|=);
+half_inplace_bitop(&, operator&=);
+half_inplace_bitop(^, operator^=);
+
+#undef half_inplace_bitop
+
+} // namespace mlx::core

data/mlx/mlx/types/half_types.h

@@ -0,0 +1,58 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+
+#include <arm_fp16.h>
+namespace mlx::core {
+using ::float16_t;
+} // namespace mlx::core
+
+#else
+
+#define ADD_HALF_BINOPS
+#include "mlx/types/fp16.h"
+namespace mlx::core {
+typedef struct _MLX_Float16 float16_t;
+} // namespace mlx::core
+
+#endif // __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+
+#ifdef __ARM_FEATURE_BF16
+
+#include <arm_bf16.h>
+namespace mlx::core {
+using ::bfloat16_t;
+} // namespace mlx::core
+
+#else
+
+#define ADD_HALF_BINOPS
+#include "mlx/types/bf16.h"
+namespace mlx::core {
+typedef struct _MLX_BFloat16 bfloat16_t;
+} // namespace mlx::core
+
+#endif // __ARM_FEATURE_BF16
+
+#ifdef ADD_HALF_BINOPS
+namespace mlx::core {
+
+// clang-format off
+#define fp16_bf16_binop_helper(__op__, __operator__)               \
+  inline float __operator__(float16_t lhs, bfloat16_t rhs) {       \
+    return static_cast<float>(lhs) __op__ static_cast<float>(rhs); \
+  }                                                                \
+  inline float __operator__(bfloat16_t lhs, float16_t rhs) {       \
+    return static_cast<float>(lhs) __op__ static_cast<float>(rhs); \
+  }
+
+fp16_bf16_binop_helper(+, operator+)
+fp16_bf16_binop_helper(-, operator-)
+fp16_bf16_binop_helper(*, operator*)
+fp16_bf16_binop_helper(/, operator/)
+// clang-format on
+
+} // namespace mlx::core
+#endif

data/mlx/mlx/types/limits.h

@@ -0,0 +1,70 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <limits>
+#include "mlx/types/half_types.h"
+
+namespace mlx::core {
+
+template <typename T>
+struct numeric_limits;
+
+template <>
+struct numeric_limits<float> : public std::numeric_limits<float> {};
+
+template <>
+struct numeric_limits<double> : public std::numeric_limits<double> {};
+
+template <>
+struct numeric_limits<float16_t> {
+ private:
+  union half_or_bits {
+    uint16_t bits;
+    float16_t value;
+  };
+  constexpr static float16_t bits_to_half(uint16_t v) {
+    return half_or_bits{v}.value;
+  }
+
+ public:
+  constexpr static float16_t lowest() {
+    return bits_to_half(0xFBFF);
+  }
+  static constexpr float16_t max() {
+    return bits_to_half(0x7BFF);
+  }
+  static constexpr float16_t epsilon() {
+    return bits_to_half(0x1400);
+  }
+  static constexpr float16_t infinity() {
+    return bits_to_half(0x7C00);
+  }
+};
+
+template <>
+struct numeric_limits<bfloat16_t> {
+ private:
+  union bfloat_or_bits {
+    uint16_t bits;
+    bfloat16_t value;
+  };
+  constexpr static bfloat16_t bits_to_bfloat(uint16_t v) {
+    return bfloat_or_bits{v}.value;
+  }
+
+ public:
+  constexpr static bfloat16_t lowest() {
+    return bits_to_bfloat(0xFF7F);
+  }
+  static constexpr bfloat16_t max() {
+    return bits_to_bfloat(0x7F7F);
+  }
+  static constexpr bfloat16_t epsilon() {
+    return bits_to_bfloat(0x3C00);
+  }
+  static constexpr bfloat16_t infinity() {
+    return bits_to_bfloat(0x7F80);
+  }
+};
+
+} // namespace mlx::core

data/mlx/mlx/utils.cpp

@@ -0,0 +1,302 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cstdlib>
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+#include "mlx/dtype_utils.h"
+#include "mlx/types/limits.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+Stream to_stream(StreamOrDevice s) {
+  if (std::holds_alternative<std::monostate>(s)) {
+    return default_stream(default_device());
+  } else if (std::holds_alternative<Device>(s)) {
+    return default_stream(std::get<Device>(s));
+  } else {
+    return std::get<Stream>(s);
+  }
+}
+
+Stream to_stream(StreamOrDevice s, Device default_) {
+  if (std::holds_alternative<std::monostate>(s)) {
+    return default_stream(default_);
+  } else if (std::holds_alternative<Device>(s)) {
+    return default_stream(std::get<Device>(s));
+  } else {
+    return std::get<Stream>(s);
+  }
+}
+
+void PrintFormatter::print(std::ostream& os, bool val) {
+  if (capitalize_bool) {
+    os << (val ? "True" : "False");
+  } else {
+    os << val;
+  }
+}
+inline void PrintFormatter::print(std::ostream& os, int16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, uint16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, int32_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, uint32_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, int64_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, uint64_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, float16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, bfloat16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, float val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, double val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, complex64_t val) {
+  os << val.real();
+  if (val.imag() >= 0 || std::isnan(val.imag())) {
+    os << "+" << val.imag() << "j";
+  } else {
+    os << "-" << -val.imag() << "j";
+  }
+}
+
+PrintFormatter& get_global_formatter() {
+  static PrintFormatter formatter;
+  return formatter;
+}
+
+void abort_with_exception(const std::exception& error) {
+  std::ostringstream msg;
+  msg << "Terminating due to uncaught exception: " << error.what();
+  std::cerr << msg.str() << std::endl;
+  std::abort();
+}
+
+Dtype result_type(const std::vector<array>& arrays) {
+  Dtype t = bool_;
+  for (auto& arr : arrays) {
+    t = promote_types(t, arr.dtype());
+  }
+  return t;
+}
+
+Shape broadcast_shapes(const Shape& s1, const Shape& s2) {
+  // Use the same broadcasting rules as numpy
+  // https://numpy.org/doc/1.20/user/theory.broadcasting.html
+  // "The size of the trailing axes for both arrays in an operation must
+  // either be the same size or one of them must be one."
+  int ndim1 = s1.size();
+  int ndim2 = s2.size();
+  int ndim = std::max(ndim1, ndim2);
+  int diff = std::abs(ndim1 - ndim2);
+  const auto& big = ndim1 > ndim2 ? s1 : s2;
+  const auto& small = ndim1 > ndim2 ? s2 : s1;
+  Shape out_shape(ndim);
+  for (int i = ndim - 1; i >= diff; --i) {
+    auto a = big[i];
+    auto b = small[i - diff];
+    if (b == a) {
+      out_shape[i] = a;
+    } else if (a == 1 || b == 1) {
+      // 0 if a or b is 0 otherwise max(a, b)
+      out_shape[i] = a * b;
+    } else {
+      std::ostringstream msg;
+      msg << "[broadcast_shapes] Shapes " << s1 << " and " << s2
+          << " cannot be broadcast.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  for (int i = diff - 1; i >= 0; --i) {
+    out_shape[i] = big[i];
+  }
+  return out_shape;
+}
+
+int normalize_axis_index(
+    int axis,
+    int ndim,
+    const std::string& msg_prefix /* = "" */) {
+  if (axis < -ndim || axis >= ndim) {
+    std::ostringstream msg;
+    msg << msg_prefix << "Axis " << axis << " is out of bounds for array with "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  return axis < 0 ? axis + ndim : axis;
+}
+
+std::ostream& operator<<(std::ostream& os, const Device& d) {
+  os << "Device(";
+  switch (d.type) {
+    case Device::cpu:
+      os << "cpu";
+      break;
+    case Device::gpu:
+      os << "gpu";
+      break;
+  }
+  os << ", " << d.index << ")";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, const Stream& s) {
+  os << "Stream(";
+  os << s.device;
+  os << ", " << s.index << ")";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, int8_t x) {
+  os << static_cast<int>(x);
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, uint8_t x) {
+  os << static_cast<unsigned int>(x);
+  return os;
+}
+
+namespace {
+
+template <typename T>
+void print_subarray(std::ostream& os, const array& a, size_t index, int dim) {
+  int num_print = 3;
+  int n = a.shape(dim);
+  size_t s = a.strides()[dim];
+  bool is_last = dim == a.ndim() - 1;
+  auto prefix = is_last ? "" : std::string(7 + dim, ' ');
+  auto postfix = is_last ? ", " : ",\n";
+  os << "[";
+  for (int i = 0; i < n; ++i) {
+    os << (i == 0 ? "" : prefix);
+    if (i == num_print && n > 2 * num_print) {
+      os << "...";
+      i = n - num_print - 1;
+      index += s * (n - 2 * num_print - 1);
+    } else if (is_last) {
+      get_global_formatter().print(os, a.data<T>()[index]);
+    } else {
+      print_subarray<T>(os, a, index, dim + 1);
+    }
+    os << (i == n - 1 ? "" : postfix);
+    index += s;
+  }
+  os << "]";
+}
+
+template <typename T>
+void print_array(std::ostream& os, const array& a) {
+  os << std::boolalpha;
+  os << "array(";
+  if (a.ndim() == 0) {
+    auto data = a.data<T>();
+    get_global_formatter().print(os, data[0]);
+  } else {
+    print_subarray<T>(os, a, 0, 0);
+  }
+  os << ", dtype=" << a.dtype() << ")";
+  os << std::noboolalpha;
+}
+
+} // namespace
+
+std::ostream& operator<<(std::ostream& os, const Dtype& dtype) {
+  return os << dtype_to_string(dtype);
+}
+
+std::ostream& operator<<(std::ostream& os, const Dtype::Kind& k) {
+  switch (k) {
+    case Dtype::Kind::b:
+      return os << "b";
+    case Dtype::Kind::i:
+      return os << "i";
+    case Dtype::Kind::u:
+      return os << "u";
+    case Dtype::Kind::f:
+      return os << "f";
+    case Dtype::Kind::c:
+      return os << "c";
+    case Dtype::Kind::V:
+      return os << "V";
+  }
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, array a) {
+  a.eval();
+  dispatch_all_types(a.dtype(), [&](auto type_tag) {
+    print_array<MLX_GET_TYPE(type_tag)>(os, a);
+  });
+  return os;
+}
+
+namespace env {
+
+int get_var(const char* name, int default_value) {
+  if (const char* buff_str = std::getenv(name)) {
+    return atoi(buff_str);
+  } else {
+    return default_value;
+  }
+}
+
+} // namespace env
+
+template <typename T>
+void set_finfo_limits(double& min, double& max, double& eps) {
+  min = numeric_limits<T>::lowest();
+  max = numeric_limits<T>::max();
+  eps = numeric_limits<T>::epsilon();
+}
+
+finfo::finfo(Dtype dtype) : dtype(dtype) {
+  if (!issubdtype(dtype, inexact)) {
+    std::ostringstream msg;
+    msg << "[finfo] dtype " << dtype << " is not inexact.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (dtype == float32) {
+    set_finfo_limits<float>(min, max, eps);
+  } else if (dtype == float16) {
+    set_finfo_limits<float16_t>(min, max, eps);
+  } else if (dtype == bfloat16) {
+    set_finfo_limits<bfloat16_t>(min, max, eps);
+  } else if (dtype == float64) {
+    set_finfo_limits<double>(min, max, eps);
+  } else if (dtype == complex64) {
+    this->dtype = float32;
+    set_finfo_limits<float>(min, max, eps);
+  }
+}
+
+template <typename T>
+void set_iinfo_limits(int64_t& min, uint64_t& max) {
+  min = std::numeric_limits<T>::min();
+  max = std::numeric_limits<T>::max();
+}
+
+iinfo::iinfo(Dtype dtype) : dtype(dtype) {
+  dispatch_int_types(dtype, "[iinfo]", [&](auto type_tag) {
+    set_iinfo_limits<MLX_GET_TYPE(type_tag)>(min, max);
+  });
+}
+
+} // namespace mlx::core