mlx 0.30.7

data/mlx/mlx/compile.h

@@ -0,0 +1,45 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+enum class CompileMode { disabled, no_simplify, no_fuse, enabled };
+
+/** Compile takes a function and returns a compiled function. */
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    bool shapeless = false);
+
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::vector<array> (*fun)(const std::vector<array>&),
+    bool shapeless = false);
+
+// Convert capture-less lambdas to function pointers.
+template <
+    typename F,
+    typename = std::enable_if_t<
+        std::is_convertible_v<F, decltype(+std::declval<F>())>>>
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    F&& f,
+    bool shapeless = false) {
+  return compile(+f, shapeless);
+}
+
+/** Globally disable compilation.
+ * Setting the environment variable ``MLX_DISABLE_COMPILE`` can also
+ * be used to disable compilation.
+ */
+MLX_API void disable_compile();
+
+/** Globally enable compilation.
+ * This will override the environment variable ``MLX_DISABLE_COMPILE``.
+ */
+MLX_API void enable_compile();
+
+/** Set the compiler mode to the given value. */
+MLX_API void set_compile_mode(CompileMode mode);
+} // namespace mlx::core

data/mlx/mlx/compile_impl.h

@@ -0,0 +1,70 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core::detail {
+
+using ArraysAndExtra = std::pair<std::vector<array>, std::shared_ptr<void>>;
+using ArrayFnWithExtra =
+    std::function<ArraysAndExtra(const std::vector<array>&)>;
+
+// This is not part of the general C++ API as calling with a bad id is a bad
+// idea.
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::uintptr_t fun_id,
+    bool shapeless = false,
+    std::vector<uint64_t> constants = {});
+
+MLX_API ArrayFnWithExtra compile(
+    ArrayFnWithExtra fun,
+    std::uintptr_t fun_id,
+    bool shapeless,
+    std::vector<uint64_t> constants);
+
+// Erase cached compile functions
+MLX_API void compile_erase(std::uintptr_t fun_id);
+
+// Clear the compiler cache causing a recompilation of all compiled functions
+// when called again.
+MLX_API void compile_clear_cache();
+
+bool compile_available_for_device(const Device& device);
+
+std::tuple<std::vector<array>, std::vector<array>, std::shared_ptr<void>>
+compile_trace(
+    const ArrayFnWithExtra& fun,
+    const std::vector<array>& inputs,
+    bool shapeless);
+
+using ParentsMap =
+    std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>;
+
+// Traverses the graph to build a tape and a map of array ids to their parents
+std::pair<std::vector<array>, ParentsMap> compile_dfs(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const std::vector<array>& original_inputs);
+
+// Simplify the tape.
+void compile_simplify(
+    std::vector<array>& tape,
+    ParentsMap& parents_map,
+    std::vector<array>& outputs,
+    int passes);
+
+std::vector<array> compile_replace(
+    const std::vector<array>& tape,
+    const std::vector<array>& trace_inputs,
+    const std::vector<array>& trace_outputs,
+    const std::vector<array>& inputs,
+    bool shapeless);
+
+void compile_validate_shapeless(const std::vector<array>& tape);
+
+} // namespace mlx::core::detail

data/mlx/mlx/device.cpp

@@ -0,0 +1,72 @@
+// Copyright © 2023-2026 Apple Inc.
+
+#include <stdexcept>
+
+#include "mlx/backend/cpu/device_info.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/device.h"
+
+namespace mlx::core {
+
+Device& mutable_default_device() {
+  static Device default_device{gpu::is_available() ? Device::gpu : Device::cpu};
+  return default_device;
+}
+
+const Device& default_device() {
+  return mutable_default_device();
+}
+
+void set_default_device(const Device& d) {
+  if (!gpu::is_available() && d == Device::gpu) {
+    throw std::invalid_argument(
+        "[set_default_device] Cannot set gpu device without gpu backend.");
+  }
+  mutable_default_device() = d;
+}
+
+bool operator==(const Device& lhs, const Device& rhs) {
+  return lhs.type == rhs.type && lhs.index == rhs.index;
+}
+
+bool operator!=(const Device& lhs, const Device& rhs) {
+  return !(lhs == rhs);
+}
+
+bool is_available(const Device& d) {
+  switch (d.type) {
+    case Device::cpu:
+      return cpu::is_available() && (d.index < cpu::device_count());
+    case Device::gpu:
+      return gpu::is_available() && (d.index < gpu::device_count());
+  }
+  // appease compiler
+  return false;
+}
+
+int device_count(Device::DeviceType type) {
+  switch (type) {
+    case Device::cpu:
+      return cpu::device_count();
+    case Device::gpu:
+      return gpu::device_count();
+  }
+  // appease compiler
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(const Device& d) {
+  switch (d.type) {
+    case Device::cpu:
+      return cpu::device_info(d.index);
+    case Device::gpu:
+      return gpu::device_info(d.index);
+  }
+  // appease compiler
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
+} // namespace mlx::core

data/mlx/mlx/device.h

@@ -0,0 +1,56 @@
+// Copyright © 2023-2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+namespace mlx::core {
+
+struct MLX_API Device {
+  enum class DeviceType {
+    cpu,
+    gpu,
+  };
+
+  static constexpr DeviceType cpu = DeviceType::cpu;
+  static constexpr DeviceType gpu = DeviceType::gpu;
+
+  Device(DeviceType type, int index = 0) : type(type), index(index) {}
+
+  DeviceType type;
+  int index;
+};
+
+MLX_API const Device& default_device();
+
+MLX_API void set_default_device(const Device& d);
+
+MLX_API bool operator==(const Device& lhs, const Device& rhs);
+MLX_API bool operator!=(const Device& lhs, const Device& rhs);
+
+MLX_API bool is_available(const Device& d);
+
+/** Get the number of available devices for the given device type. */
+MLX_API int device_count(Device::DeviceType type);
+
+/**
+ * Get information about a device.
+ *
+ * Returns a map of device properties. Keys vary by backend:
+ *   - device_name (string): Device name
+ *   - architecture (string): Architecture identifier
+ *   - total_memory/memory_size (size_t): Total device memory
+ *   - free_memory (size_t): Available memory (CUDA only)
+ *   - uuid (string): Device UUID (CUDA only)
+ *   - pci_bus_id (string): PCI bus ID (CUDA only)
+ *   - compute_capability_major/minor (size_t): Compute capability (CUDA only)
+ */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(const Device& d = default_device());
+
+} // namespace mlx::core

data/mlx/mlx/distributed/CMakeLists.txt

@@ -0,0 +1,14 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp)
+
+if(MLX_BUILD_CPU AND NOT WIN32)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mpi)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ring)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/nccl)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/jaccl)

data/mlx/mlx/distributed/distributed.cpp

@@ -0,0 +1,197 @@
+// Copyright © 2024 Apple Inc.
+
+#include <unordered_map>
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/jaccl/jaccl.h"
+#include "mlx/distributed/mpi/mpi.h"
+#include "mlx/distributed/nccl/nccl.h"
+#include "mlx/distributed/ring/ring.h"
+
+namespace mlx::core::distributed {
+
+namespace detail {
+
+Stream communication_stream(Group group, StreamOrDevice s /* = {} */) {
+  return group.raw_group()->communication_stream(s);
+}
+
+void all_sum(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_sum(input, output, stream);
+}
+
+void all_max(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_max(input, output, stream);
+}
+
+void all_min(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_min(input, output, stream);
+}
+
+void all_gather(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_gather(input, output, stream);
+}
+
+void send(Group group, const array& input, int dst, Stream stream) {
+  group.raw_group()->send(input, dst, stream);
+}
+
+void recv(Group group, array& out, int src, Stream stream) {
+  group.raw_group()->recv(out, src, stream);
+}
+
+void sum_scatter(
+    Group group,
+    const array& input,
+    array& output,
+    Stream stream) {
+  group.raw_group()->sum_scatter(input, output, stream);
+}
+
+class EmptyGroup : public GroupImpl {
+ public:
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s);
+  }
+
+  int rank() override {
+    return 0;
+  }
+
+  int size() override {
+    return 1;
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("Cannot split the distributed group further.");
+  }
+
+  void all_sum(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void all_gather(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void send(const array&, int, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void recv(array&, int, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+
+  void all_max(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+
+  void all_min(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void sum_scatter(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+};
+
+} // namespace detail
+
+bool is_available() {
+  return mpi::is_available() || ring::is_available() || nccl::is_available() ||
+      jaccl::is_available();
+}
+
+bool is_available(const std::string& bk) {
+  if (bk == "any") {
+    return is_available();
+  }
+  if (bk == "mpi") {
+    return mpi::is_available();
+  }
+  if (bk == "ring") {
+    return ring::is_available();
+  }
+  if (bk == "nccl") {
+    return nccl::is_available();
+  }
+  if (bk == "jaccl") {
+    return jaccl::is_available();
+  }
+  return false;
+}
+
+int Group::rank() const {
+  return group_->rank();
+}
+
+int Group::size() const {
+  return group_->size();
+}
+
+Group Group::split(int color, int key /* = -1 */) const {
+  return Group(group_->split(color, key));
+}
+
+Group init(bool strict /* = false */, const std::string& bk /* = "any" */) {
+  static std::unordered_map<std::string, std::shared_ptr<detail::GroupImpl>>
+      backends;
+
+  // Already initialized so return the group.
+  if (auto g = backends.find(bk); g != backends.end()) {
+    return Group(g->second);
+  }
+
+  // Create the requested communication group
+  std::shared_ptr<detail::GroupImpl> group{nullptr};
+  std::string bk_ = bk;
+  if (bk == "mpi") {
+    group = mpi::init(strict);
+  } else if (bk == "ring") {
+    group = ring::init(strict);
+  } else if (bk == "nccl") {
+    group = nccl::init(strict);
+  } else if (bk == "jaccl") {
+    group = jaccl::init(strict);
+  } else if (bk == "any") {
+    if (mlx::core::cu::is_available()) {
+      group = nccl::init(false);
+      bk_ = "nccl";
+    }
+    if (group == nullptr) {
+      group = ring::init(false);
+      bk_ = "ring";
+    }
+    if (group == nullptr) {
+      group = mpi::init(false);
+      bk_ = "mpi";
+    }
+    if (group == nullptr) {
+      group = jaccl::init(false);
+      bk_ = "jaccl";
+    }
+    if (group == nullptr && strict) {
+      throw std::runtime_error("[distributed] Couldn't initialize any backend");
+    }
+  } else {
+    std::ostringstream msg;
+    msg << "[distributed] The only valid values for backend are 'any', 'mpi', 'nccl', "
+        << "'jaccl' and 'ring' but '" << bk << "' was provided.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (group == nullptr) {
+    group = std::make_shared<detail::EmptyGroup>();
+  } else {
+    backends.insert({"any", group});
+  }
+  backends.insert({std::move(bk_), group});
+  return Group(group);
+}
+
+} // namespace mlx::core::distributed

data/mlx/mlx/distributed/distributed.h

@@ -0,0 +1,61 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <memory>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::distributed {
+
+// Forward declaration of the base group implementation.
+namespace detail {
+class GroupImpl;
+};
+
+/* Check if a communication backend is available */
+MLX_API bool is_available();
+MLX_API bool is_available(const std::string& bk);
+
+/**
+ * A distributed::Group represents a group of independent mlx processes that
+ * can communicate. We must also be able to create sub-groups from a group in
+ * order to define more granular communication.
+ */
+struct MLX_API Group {
+  Group(std::shared_ptr<detail::GroupImpl> group) : group_(std::move(group)) {}
+
+  int rank() const;
+  int size() const;
+
+  /**
+   * Split the group according to the provided color. Namely processes that use
+   * the same color will go to the same group.
+   *
+   * The key defines the rank of the processes in the new group. The smaller
+   * the key the smaller the rank. If the provided key is negative, then the
+   * rank in the current group is used.
+   */
+  Group split(int color, int key = -1) const;
+
+  const std::shared_ptr<detail::GroupImpl>& raw_group() const {
+    return group_;
+  }
+
+ private:
+  std::shared_ptr<detail::GroupImpl> group_{nullptr};
+};
+
+/**
+ * Initialize the distributed backend and return the group containing all
+ * discoverable processes.
+ *
+ * If strict is true then throw an error if we couldn't initialize the
+ * distributed subsystem. Otherwise simply return a singleton group which will
+ * render communication operations as no-op.
+ */
+MLX_API Group init(bool strict = false, const std::string& bk = "any");
+
+} // namespace mlx::core::distributed

data/mlx/mlx/distributed/distributed_impl.h

@@ -0,0 +1,59 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::detail {
+
+/**
+ * Abstract base class of a distributed group implementation.
+ */
+class GroupImpl {
+ public:
+  virtual ~GroupImpl() {}
+
+  // Choose the stream this communication group can operate on
+  virtual Stream communication_stream(StreamOrDevice s = {}) = 0;
+
+  // Group operations
+  virtual int rank() = 0;
+  virtual int size() = 0;
+  virtual std::shared_ptr<GroupImpl> split(int color, int key = -1) = 0;
+
+  // Actual communication operations
+  virtual void all_sum(const array& input, array& output, Stream stream) = 0;
+  virtual void all_gather(const array& input, array& output, Stream stream) = 0;
+  virtual void send(const array& input, int dst, Stream stream) = 0;
+  virtual void recv(array& out, int src, Stream stream) = 0;
+  virtual void all_max(const array& input, array& output, Stream stream) = 0;
+  virtual void all_min(const array& input, array& output, Stream stream) = 0;
+  virtual void
+  sum_scatter(const array& input, array& output, Stream stream) = 0;
+};
+
+/* Define the MLX stream that the communication should happen in. */
+Stream communication_stream(Group group, StreamOrDevice s = {});
+
+/* Perform an all reduce sum operation */
+void all_sum(Group group, const array& input, array& output, Stream stream);
+
+/* Perform an all gather operation */
+void all_gather(Group group, const array& input, array& output, Stream stream);
+
+/** Send an array to the dst rank */
+void send(Group group, const array& input, int dst, Stream stream);
+
+/** Recv an array from the src rank */
+void recv(Group group, array& out, int src, Stream stream);
+
+/** Max reduction */
+void all_max(Group group, const array& input, array& output, Stream stream);
+
+/** Min reduction */
+void all_min(Group group, const array& input, array& output, Stream stream);
+
+/** Reduce scatter with average operation */
+void sum_scatter(Group group, const array& input, array& output, Stream stream);
+
+} // namespace mlx::core::distributed::detail

data/mlx/mlx/distributed/jaccl/CMakeLists.txt

@@ -0,0 +1,12 @@
+if(MLX_BUILD_CPU
+   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
+   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
+  target_sources(
+    mlx
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/mesh.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp)
+endif()