mlx 0.30.7

data/mlx/mlx/distributed/ring/ring.h

@@ -0,0 +1,12 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::ring {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available();
+std::shared_ptr<GroupImpl> init(bool strict = false);
+
+} // namespace mlx::core::distributed::ring

data/mlx/mlx/distributed/utils.cpp

@@ -0,0 +1,206 @@
+// Copyright © 2025 Apple Inc.
+
+#include <netdb.h>
+#include <unistd.h>
+#include <cstring>
+#include <sstream>
+#include <thread>
+
+#include "mlx/distributed/utils.h"
+
+namespace mlx::core::distributed::detail {
+
+/**
+ * Parse a sockaddr from an ip and port provided as strings.
+ */
+address_t parse_address(const std::string& ip, const std::string& port) {
+  struct addrinfo hints, *res;
+  std::memset(&hints, 0, sizeof(hints));
+  hints.ai_family = AF_UNSPEC;
+  hints.ai_socktype = SOCK_STREAM;
+
+  int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res);
+  if (status != 0) {
+    std::ostringstream msg;
+    msg << "Can't parse address " << ip << ":" << port;
+    throw std::runtime_error(msg.str());
+  }
+
+  address_t result;
+  memcpy(&result.addr, res->ai_addr, res->ai_addrlen);
+  result.len = res->ai_addrlen;
+  freeaddrinfo(res);
+
+  return result;
+}
+
+/**
+ * Parse a sockaddr provided as an <ip>:<port> string.
+ */
+address_t parse_address(const std::string& ip_port) {
+  auto colon = ip_port.find(":");
+  if (colon == std::string::npos) {
+    std::ostringstream msg;
+    msg << "Can't parse address " << ip_port;
+    throw std::runtime_error(msg.str());
+  }
+  std::string ip(ip_port.begin(), ip_port.begin() + colon);
+  std::string port(ip_port.begin() + colon + 1, ip_port.end());
+
+  return parse_address(ip, port);
+}
+
+TCPSocket::TCPSocket(const char* tag) {
+  sock_ = socket(AF_INET, SOCK_STREAM, 0);
+  if (sock_ < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't create socket (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+TCPSocket::TCPSocket(TCPSocket&& s) {
+  sock_ = s.sock_;
+  s.sock_ = -1;
+}
+
+TCPSocket& TCPSocket::operator=(TCPSocket&& s) {
+  if (this != &s) {
+    sock_ = s.sock_;
+    s.sock_ = -1;
+  }
+  return *this;
+}
+
+TCPSocket::TCPSocket(int s) : sock_(s) {}
+
+TCPSocket::~TCPSocket() {
+  if (sock_ > 0) {
+    shutdown(sock_, 2);
+    close(sock_);
+  }
+}
+
+int TCPSocket::detach() {
+  int s = sock_;
+  sock_ = -1;
+  return s;
+}
+
+void TCPSocket::listen(const char* tag, const address_t& addr) {
+  int success;
+
+  // Make sure we can launch immediately after shutdown by setting the
+  // reuseaddr option so that we don't get address already in use errors
+  int enable = 1;
+  success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int));
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+  success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int));
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't enable reuseport (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Bind the socket to the address and port
+  success = bind(sock_, addr.get(), addr.len);
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't bind socket (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Prepare waiting for connections
+  success = ::listen(sock_, 0);
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't listen (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+TCPSocket TCPSocket::accept(const char* tag) {
+  int peer = ::accept(sock_, nullptr, nullptr);
+  if (peer < 0) {
+    std::ostringstream msg;
+    msg << tag << " Accept failed (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  return TCPSocket(peer);
+}
+
+void TCPSocket::send(const char* tag, const void* data, size_t len) {
+  while (len > 0) {
+    auto n = ::send(sock_, data, len, 0);
+    if (n <= 0) {
+      std::ostringstream msg;
+      msg << tag << " Send failed with errno=" << errno;
+      throw std::runtime_error(msg.str());
+    }
+    len -= n;
+    data = static_cast<const char*>(data) + n;
+  }
+}
+
+void TCPSocket::recv(const char* tag, void* data, size_t len) {
+  while (len > 0) {
+    auto n = ::recv(sock_, data, len, 0);
+    if (n <= 0) {
+      std::ostringstream msg;
+      msg << tag << " Recv failed with errno=" << errno;
+      throw std::runtime_error(msg.str());
+    }
+    len -= n;
+    data = static_cast<char*>(data) + n;
+  }
+}
+
+TCPSocket TCPSocket::connect(
+    const char* tag,
+    const address_t& addr,
+    int num_retries,
+    int wait,
+    std::function<void(int, int)> cb) {
+  int sock, success;
+
+  // Attempt to connect `num_retries` times with exponential backoff.
+  for (int attempt = 0; attempt < num_retries; attempt++) {
+    // Create the socket
+    sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock < 0) {
+      std::ostringstream msg;
+      msg << tag << " Couldn't create socket to connect (error: " << errno
+          << ")";
+      throw std::runtime_error(msg.str());
+    }
+
+    success = ::connect(sock, addr.get(), addr.len);
+    if (success == 0) {
+      break;
+    }
+
+    if (cb != nullptr) {
+      cb(attempt, wait);
+    }
+    if (wait > 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(wait));
+    }
+
+    wait <<= 1;
+  }
+
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't connect (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  return TCPSocket(sock);
+}
+
+} // namespace mlx::core::distributed::detail

data/mlx/mlx/distributed/utils.h

@@ -0,0 +1,67 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <sys/socket.h>
+#include <functional>
+#include <string>
+
+namespace mlx::core::distributed::detail {
+
+struct address_t {
+  sockaddr_storage addr;
+  socklen_t len;
+
+  const sockaddr* get() const {
+    return (struct sockaddr*)&addr;
+  }
+};
+
+/**
+ * Parse a sockaddr from an ip and port provided as strings.
+ */
+address_t parse_address(const std::string& ip, const std::string& port);
+
+/**
+ * Parse a sockaddr provided as an <ip>:<port> string.
+ */
+address_t parse_address(const std::string& ip_port);
+
+/**
+ * Small wrapper over a TCP socket to simplify initiating connections.
+ */
+class TCPSocket {
+ public:
+  TCPSocket(const char* tag);
+  TCPSocket(const TCPSocket&) = delete;
+  TCPSocket& operator=(const TCPSocket&) = delete;
+  TCPSocket(TCPSocket&& s);
+  TCPSocket& operator=(TCPSocket&&);
+  ~TCPSocket();
+
+  void listen(const char* tag, const address_t& addr);
+  TCPSocket accept(const char* tag);
+
+  void send(const char* tag, const void* data, size_t len);
+  void recv(const char* tag, void* data, size_t len);
+
+  int detach();
+
+  operator int() const {
+    return sock_;
+  }
+
+  static TCPSocket connect(
+      const char* tag,
+      const address_t& addr,
+      int num_retries = 1,
+      int wait = 0,
+      std::function<void(int, int)> cb = nullptr);
+
+ private:
+  TCPSocket(int sock);
+
+  int sock_;
+};
+
+} // namespace mlx::core::distributed::detail

data/mlx/mlx/dtype.cpp

@@ -0,0 +1,197 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstdint>
+
+#include "mlx/dtype.h"
+
+namespace mlx::core {
+
+namespace {
+
+constexpr int num_types = 14;
+constexpr int num_cats = 8;
+
+constexpr Dtype::Kind type_kinds[num_types] = {
+    Dtype::Kind::b, // bool_,
+    Dtype::Kind::u, // uint8,
+    Dtype::Kind::u, // uint16,
+    Dtype::Kind::u, // uint32,
+    Dtype::Kind::u, // uint64,
+    Dtype::Kind::i, // int8,
+    Dtype::Kind::i, // int16,
+    Dtype::Kind::i, // int32,
+    Dtype::Kind::i, // int64,
+    Dtype::Kind::f, // float16,
+    Dtype::Kind::f, // float32,
+    Dtype::Kind::f, // float64,
+    Dtype::Kind::V, // bfloat16,
+    Dtype::Kind::c // complex64,
+};
+
+// Following Jax type promotion rules:
+// https://jax.readthedocs.io/en/latest/type_promotion.html
+// clang-format off
+constexpr Dtype type_rules[num_types][num_types] = {
+// bool       uint8      uint16     uint32     uint64     int8       int16      int32      int64      float16    float32   float64    bfloat16   complex64
+  {bool_,     uint8,     uint16,    uint32,    uint64,    int8,      int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // bool
+  {uint8,     uint8,     uint16,    uint32,    uint64,    int16,     int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // uint8
+  {uint16,    uint16,    uint16,    uint32,    uint64,    int32,     int32,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // uint16
+  {uint32,    uint32,    uint32,    uint32,    uint64,    int64,     int64,     int64,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // uint32
+  {uint64,    uint64,    uint64,    uint64,    uint64,    float32,   float32,   float32,   float32,   float16,   float32,  float64,   bfloat16,  complex64}, // uint64
+  {int8,      int16,     int32,     int64,     float32,   int8,      int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int8
+  {int16,     int16,     int32,     int64,     float32,   int16,     int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int16
+  {int32,     int32,     int32,     int64,     float32,   int32,     int32,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int32
+  {int64,     int64,     int64,     int64,     float32,   int64,     int64,     int64,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int64
+  {float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float32,  float64,   float32,   complex64}, // float16
+  {float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,  float64,   float32,   complex64}, // float32
+  {float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,  float64,   float64,   complex64}, // float64
+  {bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  float32,   float32,  float64,   bfloat16,  complex64}, // bfloat16
+  {complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64,complex64, complex64, complex64}, // complex64
+};
+
+
+constexpr bool subcategory_to_category[num_cats][num_cats] = {
+// complexfloating floating inexact signedinteger unsignedinteger integer number generic
+  {true,           false,   true,   false,        false,          false,  true,  true}, // complexfloating
+  {false,          true,    true,   false,        false,          false,  true,  true}, // floating
+  {false,          false,   true,   false,        false,          false,  true,  true}, // inexact
+  {false,          false,   false,  true,         false,          true,   true,  true}, // signedinteger
+  {false,          false,   false,  false,        true,           true,   true,  true}, // unsignedinteger
+  {false,          false,   false,  false,        false,          true,   true,  true}, // integer
+  {false,          false,   false,  false,        false,          false,  true,  true}, // number
+  {false,          false,   false,  false,        false,          false,  false, true}, // generic
+};
+
+constexpr Dtype::Category type_to_category[num_types] = {
+    Dtype::Category::generic, // bool_,
+    Dtype::Category::unsignedinteger, // uint8,
+    Dtype::Category::unsignedinteger, // uint16,
+    Dtype::Category::unsignedinteger, // uint32,
+    Dtype::Category::unsignedinteger, // uint64,
+    Dtype::Category::signedinteger, // int8,
+    Dtype::Category::signedinteger, // int16,
+    Dtype::Category::signedinteger, // int32,
+    Dtype::Category::signedinteger, // int64,
+    Dtype::Category::floating, // float16,
+    Dtype::Category::floating, // float32,
+    Dtype::Category::floating, // float64,
+    Dtype::Category::floating, // bfloat16,
+    Dtype::Category::complexfloating, // complex64,
+};
+
+// clang-format on
+
+} // namespace
+
+Dtype promote_types(const Dtype& t1, const Dtype& t2) {
+  return Dtype(
+      type_rules[static_cast<int>(t1.val())][static_cast<int>(t2.val())]);
+}
+
+Dtype::Kind kindof(const Dtype& t) {
+  return type_kinds[static_cast<int>(t.val())];
+}
+
+template class MLX_API TypeToDtype<bool>;
+template class MLX_API TypeToDtype<uint8_t>;
+template class MLX_API TypeToDtype<uint16_t>;
+template class MLX_API TypeToDtype<uint32_t>;
+template class MLX_API TypeToDtype<uint64_t>;
+template class MLX_API TypeToDtype<int8_t>;
+template class MLX_API TypeToDtype<int16_t>;
+template class MLX_API TypeToDtype<int32_t>;
+template class MLX_API TypeToDtype<int64_t>;
+template class MLX_API TypeToDtype<float16_t>;
+template class MLX_API TypeToDtype<float>;
+template class MLX_API TypeToDtype<double>;
+template class MLX_API TypeToDtype<bfloat16_t>;
+template class MLX_API TypeToDtype<complex64_t>;
+
+template <>
+TypeToDtype<bool>::operator Dtype() {
+  return bool_;
+}
+
+template <>
+TypeToDtype<uint8_t>::operator Dtype() {
+  return uint8;
+}
+
+template <>
+TypeToDtype<uint16_t>::operator Dtype() {
+  return uint16;
+}
+
+template <>
+TypeToDtype<uint32_t>::operator Dtype() {
+  return uint32;
+}
+
+template <>
+TypeToDtype<uint64_t>::operator Dtype() {
+  return uint64;
+}
+
+template <>
+TypeToDtype<int8_t>::operator Dtype() {
+  return int8;
+}
+
+template <>
+TypeToDtype<int16_t>::operator Dtype() {
+  return int16;
+}
+
+template <>
+TypeToDtype<int32_t>::operator Dtype() {
+  return int32;
+}
+
+template <>
+TypeToDtype<int64_t>::operator Dtype() {
+  return int64;
+}
+
+template <>
+TypeToDtype<float16_t>::operator Dtype() {
+  return float16;
+}
+
+template <>
+TypeToDtype<float>::operator Dtype() {
+  return float32;
+}
+
+template <>
+TypeToDtype<double>::operator Dtype() {
+  return float32;
+}
+
+template <>
+TypeToDtype<bfloat16_t>::operator Dtype() {
+  return bfloat16;
+}
+
+template <>
+TypeToDtype<complex64_t>::operator Dtype() {
+  return complex64;
+}
+
+bool issubdtype(const Dtype& a, const Dtype& b) {
+  return a == b;
+}
+
+bool issubdtype(const Dtype::Category& cat, const Dtype& type) {
+  return false;
+}
+
+bool issubdtype(const Dtype& type, const Dtype::Category& cat) {
+  return issubdtype(type_to_category[static_cast<uint32_t>(type.val())], cat);
+}
+
+bool issubdtype(const Dtype::Category& a, const Dtype::Category& b) {
+  return subcategory_to_category[static_cast<uint32_t>(a)]
+                                [static_cast<uint32_t>(b)];
+}
+
+} // namespace mlx::core

data/mlx/mlx/dtype.h

@@ -0,0 +1,116 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <complex>
+#include <cstdint>
+
+#include "mlx/api.h"
+#include "mlx/types/complex.h"
+#include "mlx/types/half_types.h"
+
+namespace mlx::core {
+
+struct Dtype {
+  enum class Val {
+    bool_,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    int8,
+    int16,
+    int32,
+    int64,
+    float16,
+    float32,
+    float64,
+    bfloat16,
+    complex64,
+  };
+
+  enum class Kind {
+    b, /* bool */
+    u, /* unsigned int */
+    i, /* signed int */
+    f, /* float */
+    c, /* complex */
+    V, /* void - used for brain float */
+  };
+
+  enum class Category {
+    complexfloating,
+    floating,
+    inexact,
+    signedinteger,
+    unsignedinteger,
+    integer,
+    number,
+    generic
+  };
+
+  constexpr explicit Dtype(Val val, uint8_t size) : val_(val), size_(size) {}
+
+  constexpr operator Val() const {
+    return val_;
+  }
+  constexpr Val val() const {
+    return val_;
+  }
+  constexpr uint8_t size() const {
+    return size_;
+  }
+
+ private:
+  Val val_;
+  uint8_t size_;
+};
+
+inline constexpr Dtype bool_{Dtype::Val::bool_, sizeof(bool)};
+
+inline constexpr Dtype uint8{Dtype::Val::uint8, sizeof(uint8_t)};
+inline constexpr Dtype uint16{Dtype::Val::uint16, sizeof(uint16_t)};
+inline constexpr Dtype uint32{Dtype::Val::uint32, sizeof(uint32_t)};
+inline constexpr Dtype uint64{Dtype::Val::uint64, sizeof(uint64_t)};
+
+inline constexpr Dtype int8{Dtype::Val::int8, sizeof(int8_t)};
+inline constexpr Dtype int16{Dtype::Val::int16, sizeof(int16_t)};
+inline constexpr Dtype int32{Dtype::Val::int32, sizeof(int32_t)};
+inline constexpr Dtype int64{Dtype::Val::int64, sizeof(int64_t)};
+
+inline constexpr Dtype float16{Dtype::Val::float16, sizeof(uint16_t)};
+inline constexpr Dtype float32{Dtype::Val::float32, sizeof(float)};
+inline constexpr Dtype float64{Dtype::Val::float64, sizeof(double)};
+inline constexpr Dtype bfloat16{Dtype::Val::bfloat16, sizeof(uint16_t)};
+inline constexpr Dtype complex64{Dtype::Val::complex64, sizeof(complex64_t)};
+
+inline constexpr Dtype::Category complexfloating =
+    Dtype::Category::complexfloating;
+inline constexpr Dtype::Category floating = Dtype::Category::floating;
+inline constexpr Dtype::Category inexact = Dtype::Category::inexact;
+inline constexpr Dtype::Category signedinteger = Dtype::Category::signedinteger;
+inline constexpr Dtype::Category unsignedinteger =
+    Dtype::Category::unsignedinteger;
+inline constexpr Dtype::Category integer = Dtype::Category::integer;
+inline constexpr Dtype::Category number = Dtype::Category::number;
+inline constexpr Dtype::Category generic = Dtype::Category::generic;
+
+MLX_API bool issubdtype(const Dtype& a, const Dtype& b);
+MLX_API bool issubdtype(const Dtype::Category& a, const Dtype& b);
+MLX_API bool issubdtype(const Dtype& a, const Dtype::Category& b);
+MLX_API bool issubdtype(const Dtype::Category& a, const Dtype::Category& b);
+
+MLX_API Dtype promote_types(const Dtype& t1, const Dtype& t2);
+
+inline uint8_t size_of(const Dtype& t) {
+  return t.size();
+}
+
+MLX_API Dtype::Kind kindof(const Dtype& t);
+
+template <typename T>
+struct MLX_API TypeToDtype {
+  operator Dtype();
+};
+
+} // namespace mlx::core

data/mlx/mlx/dtype_utils.cpp

@@ -0,0 +1,42 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+
+const char* dtype_to_string(Dtype arg) {
+  switch (arg) {
+    case bool_:
+      return "bool";
+    case int8:
+      return "int8";
+    case int16:
+      return "int16";
+    case int32:
+      return "int32";
+    case int64:
+      return "int64";
+    case uint8:
+      return "uint8";
+    case uint16:
+      return "uint16";
+    case uint32:
+      return "uint32";
+    case uint64:
+      return "uint64";
+    case float16:
+      return "float16";
+    case bfloat16:
+      return "bfloat16";
+    case float32:
+      return "float32";
+    case float64:
+      return "float64";
+    case complex64:
+      return "complex64";
+    default:
+      return "unknown";
+  }
+}
+
+} // namespace mlx::core