mlx 1.0.0

data/mlx/mlx/CMakeLists.txt

@@ -0,0 +1,107 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/dtype_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/export.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/einsum.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/graph_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h)
+
+# Define MLX_VERSION only in the version.cpp file.
+add_library(mlx_version OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp)
+target_compile_definitions(mlx_version PRIVATE MLX_VERSION="${MLX_VERSION}")
+target_include_directories(mlx_version PRIVATE ${PROJECT_SOURCE_DIR})
+target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:mlx_version>)
+
+# Do not export symbols by default.
+set_target_properties(
+  mlx mlx_version
+  PROPERTIES VISIBILITY_INLINES_HIDDEN ON
+             CXX_VISIBILITY_PRESET hidden
+             CUDA_VISIBILITY_PRESET hidden)
+
+# Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries.
+set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
+if(BUILD_SHARED_LIBS)
+  target_compile_definitions(mlx_version PUBLIC MLX_EXPORT)
+else()
+  target_compile_definitions(mlx PUBLIC MLX_STATIC)
+  target_compile_definitions(mlx_version PUBLIC MLX_STATIC)
+endif()
+
+if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+  # Supress warnings: note: parameter passing for argument of type
+  # 'std::pair<float, float>' when C++17 is enabled changed to match C++14 in
+  # GCC 10.1
+  target_compile_options(mlx PRIVATE -Wno-psabi)
+endif()
+
+if(MSVC)
+  # Some of CUDA's headers include windows.h, which defines min/max macros.
+  target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN)
+  # Unicode support in fmt does not compile in .cu files.
+  target_compile_definitions(mlx PRIVATE FMT_UNICODE=0)
+  # Disable some MSVC warnings to speed up compilation.
+  target_compile_options(
+    mlx
+    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4244 /wd4267>
+    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/wd4068
+            /wd4146
+            /wd4700
+            /wd4804
+            /wd4805>
+            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4244
+            -Xcompiler=/wd4267>)
+  # Enable /bigobj for heavily templated code (e.g., binary.cpp) that exceeds
+  # the default 65,535 section limit in COFF object files.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>
+                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/bigobj>)
+  # Use modern preprocessor, otherwise CCCL would complain.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>
+                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
+
+if(MLX_BUILD_CPU)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cpu)
+else()
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_cpu)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/distributed)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
+
+if(MLX_BUILD_METAL)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/metal)
+else()
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/no_metal.cpp)
+endif()
+
+if(MLX_BUILD_CUDA)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda)
+else()
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda/no_cuda.cpp)
+endif()
+
+if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/gpu)
+else()
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_gpu)
+endif()

data/mlx/mlx/allocator.h

@@ -0,0 +1,75 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "mlx/api.h"
+
+namespace mlx::core::allocator {
+
+// Simple wrapper around buffer pointers
+// WARNING: Only Buffer objects constructed from and those that wrap
+//          raw pointers from mlx::allocator are supported.
+class MLX_API Buffer {
+ private:
+  void* ptr_;
+
+ public:
+  explicit Buffer(void* ptr) : ptr_(ptr) {};
+
+  // Get the raw data pointer from the buffer
+  void* raw_ptr();
+
+  // Get the buffer pointer from the buffer
+  const void* ptr() const {
+    return ptr_;
+  };
+  void* ptr() {
+    return ptr_;
+  };
+};
+
+class MLX_API Allocator {
+  /** Abstract base class for a memory allocator. */
+ public:
+  virtual Buffer malloc(size_t size) = 0;
+  virtual void free(Buffer buffer) = 0;
+  virtual size_t size(Buffer buffer) const = 0;
+  virtual Buffer make_buffer(void* ptr, size_t size) {
+    return Buffer{nullptr};
+  };
+  virtual void release(Buffer buffer) {}
+
+  Allocator() = default;
+  Allocator(const Allocator& other) = delete;
+  Allocator(Allocator&& other) = delete;
+  Allocator& operator=(const Allocator& other) = delete;
+  Allocator& operator=(Allocator&& other) = delete;
+  virtual ~Allocator() = default;
+};
+
+MLX_API Allocator& allocator();
+
+inline Buffer malloc(size_t size) {
+  return allocator().malloc(size);
+}
+
+inline void free(Buffer buffer) {
+  allocator().free(buffer);
+}
+
+// Make a Buffer from a raw pointer of the given size without a copy.  If a
+// no-copy conversion is not possible then the returned buffer.ptr() will be
+// nullptr. Any buffer created with this function must be released with
+// release(buffer)
+inline Buffer make_buffer(void* ptr, size_t size) {
+  return allocator().make_buffer(ptr, size);
+};
+
+// Release a buffer from the allocator made with make_buffer
+inline void release(Buffer buffer) {
+  allocator().release(buffer);
+}
+
+} // namespace mlx::core::allocator

data/mlx/mlx/api.h

@@ -0,0 +1,29 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+// MLX_API macro for controlling symbol visibility, must add for public APIs.
+//
+// Usage:
+//   MLX_API void some_function(...);
+//   class MLX_API SomeClass { ... };
+
+#if defined(MLX_STATIC)
+
+// Static library build - no import/export decorations needed
+#define MLX_API
+
+#else
+
+// Shared library build.
+#if defined(_WIN32)
+#if defined(MLX_EXPORT)
+#define MLX_API __declspec(dllexport)
+#else
+#define MLX_API __declspec(dllimport)
+#endif // defined(MLX_EXPORT)
+#else
+#define MLX_API __attribute__((visibility("default")))
+#endif // defined(_WIN32)
+
+#endif // defined(MLX_STATIC)

data/mlx/mlx/array.cpp

@@ -0,0 +1,354 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <functional>
+#include <unordered_map>
+
+#include "mlx/array.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+
+namespace mlx::core {
+
+array::array(const std::complex<float>& val, Dtype dtype /* = complex64 */)
+    : array_desc_(std::make_shared<ArrayDesc>(Shape{}, dtype)) {
+  auto cval = static_cast<complex64_t>(val);
+  init(&cval);
+}
+
+array::array(
+    Shape shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array> inputs)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              std::move(shape),
+              dtype,
+              std::move(primitive),
+              std::move(inputs))) {
+  if (has_primitive() && this->primitive().stream().device == Device::gpu) {
+    for (auto& in : this->inputs()) {
+      if (in.dtype() == float64) {
+        throw std::invalid_argument("float64 is not supported on the GPU");
+      }
+    }
+    if (this->dtype() == float64) {
+      throw std::invalid_argument("float64 is not supported on the GPU");
+    }
+  }
+}
+
+std::vector<array> array::make_arrays(
+    std::vector<Shape> shapes,
+    const std::vector<Dtype>& dtypes,
+    const std::shared_ptr<Primitive>& primitive,
+    const std::vector<array>& inputs) {
+  std::vector<array> outputs;
+  for (size_t i = 0; i < shapes.size(); ++i) {
+    outputs.emplace_back(std::move(shapes[i]), dtypes[i], primitive, inputs);
+  }
+  // For each node in |outputs|, its siblings are the other nodes.
+  for (size_t i = 0; i < outputs.size(); ++i) {
+    auto siblings = outputs;
+    siblings.erase(siblings.begin() + i);
+    outputs[i].set_siblings(std::move(siblings), i);
+  }
+  return outputs;
+}
+
+array array::unsafe_weak_copy(const array& other) {
+  auto cpy = array(other.shape(), other.dtype(), nullptr, {});
+  cpy.set_data(
+      other.buffer(),
+      other.data_size(),
+      other.strides(),
+      other.flags(),
+      [](auto) {});
+  cpy.array_desc_->offset = other.array_desc_->offset;
+  return cpy;
+}
+
+array::array(std::initializer_list<float> data)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              float32)) {
+  init(data.begin());
+}
+
+array::array(std::initializer_list<int> data, Dtype dtype)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              dtype)) {
+  init(data.begin());
+}
+
+array::array(
+    void* data,
+    Shape shape,
+    Dtype dtype,
+    const std::function<void(void*)>& deleter)
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  auto buffer = allocator::make_buffer(data, nbytes());
+  if (buffer.ptr() == nullptr) {
+    set_data(allocator::malloc(nbytes()));
+    auto ptr = static_cast<char*>(data);
+    std::copy(ptr, ptr + nbytes(), this->data<char>());
+    deleter(data);
+  } else {
+    auto wrapped_deleter = [deleter](allocator::Buffer buffer) {
+      auto ptr = buffer.raw_ptr();
+      allocator::release(buffer);
+      return deleter(ptr);
+    };
+    set_data(buffer, std::move(wrapped_deleter));
+  }
+}
+
+/* Build an array from a shared buffer */
+array::array(allocator::Buffer data, Shape shape, Dtype dtype, Deleter deleter)
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  set_data(data, deleter);
+}
+
+void array::detach() {
+  array_desc_->primitive = nullptr;
+  for (auto& s : array_desc_->siblings) {
+    s.array_desc_->primitive = nullptr;
+  }
+  for (auto& s : array_desc_->siblings) {
+    s.array_desc_->inputs.clear();
+    s.array_desc_->siblings.clear();
+    s.array_desc_->position = 0;
+  }
+  array_desc_->inputs.clear();
+  array_desc_->siblings.clear();
+  array_desc_->position = 0;
+}
+
+bool array::is_available() const {
+  if (status() == Status::available) {
+    return true;
+  } else if (
+      status() == Status::evaluated &&
+      (!event().valid() || event().is_signaled())) {
+    set_status(Status::available);
+    return true;
+  }
+  return false;
+}
+
+void array::wait() {
+  if (!is_available()) {
+    if (event().valid()) {
+      event().wait();
+      detach_event();
+    }
+    set_status(Status::available);
+  }
+}
+
+void array::eval() {
+  // Ensure the array is ready to be read
+  if (status() == Status::unscheduled) {
+    mlx::core::eval({*this});
+  } else {
+    wait();
+  }
+}
+
+bool array::is_tracer() const {
+  return (array_desc_->is_tracer && detail::in_tracing()) ||
+      detail::retain_graph();
+}
+
+void array::set_data(allocator::Buffer buffer, Deleter d) {
+  array_desc_->data = std::make_shared<Data>(buffer, d);
+  array_desc_->offset = 0;
+  array_desc_->data_size = size();
+  array_desc_->flags.contiguous = true;
+  array_desc_->flags.row_contiguous = true;
+  auto max_dim = std::max_element(shape().begin(), shape().end());
+  array_desc_->flags.col_contiguous = size() <= 1 || size() == *max_dim;
+}
+
+void array::set_data(
+    allocator::Buffer buffer,
+    size_t data_size,
+    Strides strides,
+    Flags flags,
+    Deleter d) {
+  array_desc_->data = std::make_shared<Data>(buffer, d);
+  array_desc_->offset = 0;
+  array_desc_->data_size = data_size;
+  array_desc_->strides = std::move(strides);
+  array_desc_->flags = flags;
+}
+
+void array::copy_shared_buffer(
+    const array& other,
+    const Strides& strides,
+    Flags flags,
+    size_t data_size,
+    int64_t offset /* = 0 */) {
+  array_desc_->data = other.array_desc_->data;
+  array_desc_->strides = strides;
+  array_desc_->flags = flags;
+  array_desc_->data_size = data_size;
+  array_desc_->offset =
+      sizeof(char) * itemsize() * offset + other.array_desc_->offset;
+}
+
+void array::copy_shared_buffer(const array& other) {
+  copy_shared_buffer(other, other.strides(), other.flags(), other.data_size());
+}
+
+array::~array() {
+  if (array_desc_ == nullptr) {
+    return;
+  }
+
+  // Detached/detaching
+  if (array_desc_->primitive == nullptr) {
+    return;
+  }
+
+  // Break circular reference for non-detached arrays with siblings
+  if (auto n = siblings().size(); n > 0) {
+    bool do_detach = true;
+    // If all siblings have siblings.size() references except
+    // the one we are currently destroying (which has siblings.size() + 1)
+    // then there are no more external references
+    do_detach &= (array_desc_.use_count() == (n + 1));
+    for (auto& s : siblings()) {
+      do_detach &= (s.array_desc_.use_count() == n);
+      if (!do_detach) {
+        break;
+      }
+    }
+    if (do_detach) {
+      for (auto& s : siblings()) {
+        for (auto& ss : s.siblings()) {
+          // Set to null here to avoid descending into array destructor
+          // for siblings
+          ss.array_desc_ = nullptr;
+        }
+        s.array_desc_->siblings.clear();
+      }
+    }
+  }
+}
+
+void array::ArrayDesc::init() {
+  strides.resize(shape.size());
+  size = 1;
+  for (int i = shape.size() - 1; i >= 0; --i) {
+    strides[i] = size;
+    size *= shape[i];
+  }
+  for (const auto& in : inputs) {
+    is_tracer |= in.is_tracer();
+  }
+}
+
+array::ArrayDesc::ArrayDesc(Shape shape, Dtype dtype)
+    : shape(std::move(shape)), dtype(dtype), status(Status::available) {
+  init();
+}
+
+array::ArrayDesc::ArrayDesc(
+    Shape shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array> inputs)
+    : shape(std::move(shape)),
+      dtype(dtype),
+      primitive(std::move(primitive)),
+      status(Status::unscheduled),
+      inputs(std::move(inputs)) {
+  init();
+}
+
+array::ArrayDesc::~ArrayDesc() {
+  // When an array description is destroyed it will delete a bunch of arrays
+  // that may also destroy their corresponding descriptions and so on and so
+  // forth.
+  //
+  // This calls recursively the destructor and can result in stack overflow, we
+  // instead put them in a vector and destroy them one at a time resulting in a
+  // max stack depth of 2.
+  if (inputs.empty()) {
+    return;
+  }
+
+  std::vector<std::shared_ptr<ArrayDesc>> for_deletion;
+
+  auto append_deletable_inputs = [&for_deletion](ArrayDesc& ad) {
+    std::unordered_map<std::uintptr_t, array> input_map;
+    for (array& a : ad.inputs) {
+      if (a.array_desc_) {
+        input_map.insert({a.id(), a});
+        for (auto& s : a.siblings()) {
+          input_map.insert({s.id(), s});
+        }
+      }
+    }
+    ad.inputs.clear();
+    for (auto& [_, a] : input_map) {
+      bool is_deletable =
+          (a.array_desc_.use_count() <= a.siblings().size() + 1);
+      // An array with siblings is deletable only if all of its siblings
+      // are deletable
+      for (auto& s : a.siblings()) {
+        if (!is_deletable) {
+          break;
+        }
+        int is_input = (input_map.find(s.id()) != input_map.end());
+        is_deletable &=
+            s.array_desc_.use_count() <= a.siblings().size() + is_input;
+      }
+      if (is_deletable) {
+        for_deletion.push_back(std::move(a.array_desc_));
+      }
+    }
+  };
+
+  append_deletable_inputs(*this);
+
+  while (!for_deletion.empty()) {
+    // top is going to be deleted at the end of the block *after* the arrays
+    // with inputs have been moved into the vector
+    auto top = std::move(for_deletion.back());
+    for_deletion.pop_back();
+    append_deletable_inputs(*top);
+
+    // Clear out possible siblings to break circular references
+    for (auto& s : top->siblings) {
+      // Set to null here to avoid descending into top-level
+      // array destructor for siblings
+      s.array_desc_ = nullptr;
+    }
+    top->siblings.clear();
+  }
+}
+
+array::ArrayIterator::ArrayIterator(const array& arr, int idx)
+    : arr(arr), idx(idx) {
+  if (arr.ndim() == 0) {
+    throw std::invalid_argument("Cannot iterate over 0-d array.");
+  }
+}
+
+array::ArrayIterator::reference array::ArrayIterator::operator*() const {
+  auto start = Shape(arr.ndim(), 0);
+  auto end = arr.shape();
+  auto shape = arr.shape();
+  shape.erase(shape.begin());
+  start[0] = idx;
+  end[0] = idx + 1;
+  return reshape(slice(arr, start, end), shape);
+};
+
+} // namespace mlx::core