mlx-cpu 0.30.1__py3-none-manylinux_2_35_x86_64.whl

mlx/include/mlx/small_vector.h

@@ -0,0 +1,540 @@
+// Copyright © 2025 Apple Inc.
+// Copyright © 2018 the V8 project authors.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//     * Redistributions of source code must retain the above copyright
+//       notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above
+//       copyright notice, this list of conditions and the following
+//       disclaimer in the documentation and/or other materials provided
+//       with the distribution.
+//     * Neither the name of Google Inc. nor the names of its
+//       contributors may be used to endorse or promote products derived
+//       from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <type_traits>
+#include <utility>
+
+namespace mlx::core {
+
+#if defined(__has_builtin)
+#define MLX_HAS_BUILTIN(x) __has_builtin(x)
+#else
+#define MLX_HAS_BUILTIN(x) 0
+#endif
+
+#if defined(__has_attribute)
+#define MLX_HAS_ATTRIBUTE(x) __has_attribute(x)
+#else
+#define MLX_HAS_ATTRIBUTE(x) 0
+#endif
+
+#if MLX_HAS_BUILTIN(__builtin_expect)
+#define MLX_LIKELY(condition) (__builtin_expect(!!(condition), 1))
+#define MLX_UNLIKELY(condition) (__builtin_expect(!!(condition), 0))
+#else
+#define MLX_LIKELY(condition) (condition)
+#define MLX_UNLIKELY(condition) (condition)
+#endif
+
+#if MLX_HAS_ATTRIBUTE(noinline)
+#define MLX_NOINLINE __attribute__((noinline))
+#else
+#define MLX_NOINLINE
+#endif
+
+template <typename T, typename = void>
+struct is_iterator : std::false_type {};
+
+template <typename T>
+struct is_iterator<
+    T,
+    std::void_t<
+        typename std::iterator_traits<T>::difference_type,
+        typename std::iterator_traits<T>::iterator_category,
+        typename std::iterator_traits<T>::pointer,
+        typename std::iterator_traits<T>::reference,
+        typename std::iterator_traits<T>::value_type>> : std::true_type {};
+
+template <typename T>
+constexpr bool is_iterator_v = is_iterator<T>::value;
+
+// Minimal SmallVector implementation. Uses inline storage first, switches to
+// dynamic storage when it overflows.
+//
+// Notes:
+// * The default inline storage size is MAX_NDIM, as it is mainly used for
+//   shapes and strides, users should choose a better size for other cases.
+// * The data() returns real address even for empty vector.
+// * The pointer returned by data() will change after moving the vector as it
+//   points to the inline storage.
+// * For trivial elements the storage will not be default constructed,
+//   i.e. SmallVector<int>(10) will not be filled with 0 by default.
+template <typename T, size_t kSize = 10, typename Allocator = std::allocator<T>>
+class SmallVector {
+ public:
+  using value_type = T;
+  using reference = T&;
+  using const_reference = const T&;
+  using iterator = T*;
+  using const_iterator = const T*;
+  using difference_type = std::ptrdiff_t;
+  using size_type = std::size_t;
+
+  SmallVector() = default;
+
+  explicit SmallVector(const Allocator& allocator) : allocator_(allocator) {}
+
+  explicit SmallVector(size_t size, const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    resize(size);
+  }
+
+  SmallVector(
+      size_t size,
+      const T& initial_value,
+      const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    resize(size, initial_value);
+  }
+
+  SmallVector(
+      std::initializer_list<T> init,
+      const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    if (init.size() > capacity()) {
+      grow(init.size());
+    }
+    assert(capacity() >= init.size()); // sanity check
+    std::uninitialized_move(init.begin(), init.end(), begin_);
+    end_ = begin_ + init.size();
+  }
+
+  template <typename Iter, typename = std::enable_if_t<is_iterator_v<Iter>>>
+  SmallVector(Iter begin, Iter end, const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    size_t size = std::distance(begin, end);
+    if (size > capacity()) {
+      grow(size);
+    }
+    assert(capacity() >= size); // sanity check
+    std::uninitialized_copy(begin, end, begin_);
+    end_ = begin_ + size;
+  }
+
+  SmallVector(const SmallVector& other) : allocator_(other.allocator_) {
+    *this = other;
+  }
+  SmallVector(const SmallVector& other, const Allocator& allocator)
+      : allocator_(allocator) {
+    *this = other;
+  }
+  SmallVector(SmallVector&& other) : allocator_(std::move(other.allocator_)) {
+    *this = std::move(other);
+  }
+  SmallVector(SmallVector&& other, const Allocator& allocator)
+      : allocator_(allocator) {
+    *this = std::move(other);
+  }
+
+  ~SmallVector() {
+    free_storage();
+  }
+
+  SmallVector& operator=(const SmallVector& other) {
+    if (this == &other) {
+      return *this;
+    }
+    size_t other_size = other.size();
+    if (capacity() < other_size) {
+      // Create large-enough heap-allocated storage.
+      free_storage();
+      begin_ = allocator_.allocate(other_size);
+      end_of_storage_ = begin_ + other_size;
+      std::uninitialized_copy(other.begin_, other.end_, begin_);
+    } else if constexpr (kHasTrivialElement) {
+      std::copy(other.begin_, other.end_, begin_);
+    } else {
+      ptrdiff_t to_copy =
+          std::min(static_cast<ptrdiff_t>(other_size), end_ - begin_);
+      std::copy(other.begin_, other.begin_ + to_copy, begin_);
+      if (other.begin_ + to_copy < other.end_) {
+        std::uninitialized_copy(
+            other.begin_ + to_copy, other.end_, begin_ + to_copy);
+      } else {
+        std::destroy_n(begin_ + to_copy, size() - to_copy);
+      }
+    }
+    end_ = begin_ + other_size;
+    return *this;
+  }
+
+  SmallVector& operator=(SmallVector&& other) {
+    if (this == &other) {
+      return *this;
+    }
+    if (other.is_big()) {
+      free_storage();
+      begin_ = other.begin_;
+      end_ = other.end_;
+      end_of_storage_ = other.end_of_storage_;
+    } else {
+      assert(capacity() >= other.size()); // sanity check
+      size_t other_size = other.size();
+      if constexpr (kHasTrivialElement) {
+        std::move(other.begin_, other.end_, begin_);
+      } else {
+        ptrdiff_t to_move =
+            std::min(static_cast<ptrdiff_t>(other_size), end_ - begin_);
+        std::move(other.begin_, other.begin_ + to_move, begin_);
+        if (other.begin_ + to_move < other.end_) {
+          std::uninitialized_move(
+              other.begin_ + to_move, other.end_, begin_ + to_move);
+        } else {
+          std::destroy_n(begin_ + to_move, size() - to_move);
+        }
+      }
+      end_ = begin_ + other_size;
+    }
+    other.reset_to_inline_storage();
+    return *this;
+  }
+
+  bool operator==(const SmallVector& other) const {
+    if (size() != other.size()) {
+      return false;
+    }
+    return std::equal(begin_, end_, other.begin_);
+  }
+
+  bool operator!=(const SmallVector& other) const {
+    return !(*this == other);
+  }
+
+  T* data() {
+    return begin_;
+  }
+  const T* data() const {
+    return begin_;
+  }
+
+  iterator begin() {
+    return begin_;
+  }
+  const_iterator begin() const {
+    return begin_;
+  }
+
+  iterator end() {
+    return end_;
+  }
+  const_iterator end() const {
+    return end_;
+  }
+
+  const_iterator cbegin() const {
+    return begin_;
+  }
+
+  const_iterator cend() const {
+    return end_;
+  }
+
+  auto rbegin() {
+    return std::make_reverse_iterator(end_);
+  }
+  auto rbegin() const {
+    return std::make_reverse_iterator(end_);
+  }
+
+  auto rend() {
+    return std::make_reverse_iterator(begin_);
+  }
+  auto rend() const {
+    return std::make_reverse_iterator(begin_);
+  }
+
+  size_t size() const {
+    return end_ - begin_;
+  }
+  bool empty() const {
+    return end_ == begin_;
+  }
+  size_t capacity() const {
+    return end_of_storage_ - begin_;
+  }
+
+  T& front() {
+    assert(size() != 0);
+    return begin_[0];
+  }
+  const T& front() const {
+    assert(size() != 0);
+    return begin_[0];
+  }
+
+  T& back() {
+    assert(size() != 0);
+    return end_[-1];
+  }
+  const T& back() const {
+    assert(size() != 0);
+    return end_[-1];
+  }
+
+  T& at(size_t index) {
+    if (index >= size()) {
+      throw std::out_of_range("SmallVector out of range.");
+    }
+    return begin_[index];
+  }
+  const T& at(size_t index) const {
+    return const_cast<SmallVector*>(this)->at(index);
+  }
+
+  T& operator[](size_t index) {
+    assert(size() > index);
+    return begin_[index];
+  }
+  const T& operator[](size_t index) const {
+    return const_cast<SmallVector*>(this)->operator[](index);
+  }
+
+  template <typename... Args>
+  void emplace_back(Args&&... args) {
+    if (MLX_UNLIKELY(end_ == end_of_storage_)) {
+      grow();
+    }
+    void* storage = end_;
+    end_ += 1;
+    new (storage) T(std::forward<Args>(args)...);
+  }
+
+  void push_back(T x) {
+    emplace_back(std::move(x));
+  }
+
+  void pop_back(size_t count = 1) {
+    assert(size() >= count);
+    end_ -= count;
+    std::destroy_n(end_, count);
+  }
+
+  iterator insert(iterator pos, T value) {
+    return insert(pos, static_cast<size_t>(1), std::move(value));
+  }
+
+  iterator insert(iterator pos, size_t count, T value) {
+    assert(pos <= end_);
+    size_t offset = pos - begin_;
+    size_t old_size = size();
+    resize(old_size + count);
+    pos = begin_ + offset;
+    iterator old_end = begin_ + old_size;
+    assert(old_end <= end_);
+    std::move_backward(pos, old_end, end_);
+    if constexpr (kHasTrivialElement) {
+      std::fill_n(pos, count, value);
+    } else {
+      std::fill_n(pos + 1, count - 1, value);
+      *pos = std::move(value);
+    }
+    return pos;
+  }
+
+  template <typename Iter, typename = std::enable_if_t<is_iterator_v<Iter>>>
+  iterator insert(iterator pos, Iter begin, Iter end) {
+    if constexpr (std::is_same_v<std::decay_t<Iter>, iterator>) {
+      // The implementation can not take overlapping range.
+      assert(!(begin >= pos && begin < pos + std::distance(begin, end)));
+      assert(!(end > pos && end <= pos + std::distance(begin, end)));
+    }
+
+    assert(pos <= end_);
+    size_t offset = pos - begin_;
+    size_t count = std::distance(begin, end);
+    size_t old_size = size();
+    resize(old_size + count);
+    pos = begin_ + offset;
+    iterator old_end = begin_ + old_size;
+    assert(old_end <= end_);
+    std::move_backward(pos, old_end, end_);
+    std::copy(begin, end, pos);
+    return pos;
+  }
+
+  iterator insert(iterator pos, std::initializer_list<const T> values) {
+    return insert(pos, values.begin(), values.end());
+  }
+
+  iterator erase(iterator erase_start, iterator erase_end) {
+    assert(erase_start >= begin_);
+    assert(erase_start <= erase_end);
+    assert(erase_end <= end_);
+    iterator new_end = std::move(erase_end, end_, erase_start);
+    std::destroy_n(new_end, std::distance(new_end, end_));
+    end_ = new_end;
+    return erase_start;
+  }
+
+  iterator erase(iterator pos) {
+    return erase(pos, pos + 1);
+  }
+
+  void resize(size_t new_size) {
+    if (new_size > capacity()) {
+      grow(new_size);
+    }
+    T* new_end = begin_ + new_size;
+    if constexpr (!kHasTrivialElement) {
+      if (new_end > end_) {
+        std::uninitialized_default_construct(end_, new_end);
+      } else {
+        std::destroy_n(new_end, end_ - new_end);
+      }
+    }
+    end_ = new_end;
+  }
+
+  void resize(size_t new_size, const T& initial_value) {
+    if (new_size > capacity()) {
+      grow(new_size);
+    }
+    T* new_end = begin_ + new_size;
+    if (new_end > end_) {
+      std::uninitialized_fill(end_, new_end, initial_value);
+    } else {
+      std::destroy_n(new_end, end_ - new_end);
+    }
+    end_ = new_end;
+  }
+
+  void reserve(size_t new_capacity) {
+    if (new_capacity > capacity()) {
+      grow(new_capacity);
+    }
+  }
+
+  // Clear without reverting back to inline storage.
+  void clear() {
+    std::destroy_n(begin_, end_ - begin_);
+    end_ = begin_;
+  }
+
+ private:
+  // Grows the backing store by a factor of two, and at least to {min_capacity}.
+  // TODO: Move to private after removing external code using this method.
+  MLX_NOINLINE void grow(size_t min_capacity = 0) {
+    size_t new_capacity = std::max(min_capacity, 2 * capacity());
+    // Round up to power of 2.
+    new_capacity--;
+    new_capacity |= new_capacity >> 1;
+    new_capacity |= new_capacity >> 2;
+    new_capacity |= new_capacity >> 4;
+    new_capacity |= new_capacity >> 8;
+    new_capacity |= new_capacity >> 16;
+    if constexpr (sizeof(size_t) == sizeof(uint64_t)) {
+      new_capacity |= new_capacity >> 32;
+    }
+    new_capacity++;
+
+    T* new_storage = allocator_.allocate(new_capacity);
+    if (new_storage == nullptr) {
+      throw std::bad_alloc();
+    }
+
+    size_t in_use = end_ - begin_;
+    std::uninitialized_move(begin_, end_, new_storage);
+    free_storage();
+    begin_ = new_storage;
+    end_ = new_storage + in_use;
+    end_of_storage_ = new_storage + new_capacity;
+  }
+
+  MLX_NOINLINE void free_storage() {
+    std::destroy_n(begin_, end_ - begin_);
+    if (is_big()) {
+      allocator_.deallocate(begin_, end_of_storage_ - begin_);
+    }
+  }
+
+  // Clear and go back to inline storage. Dynamic storage is *not* freed. For
+  // internal use only.
+  void reset_to_inline_storage() {
+    if constexpr (!kHasTrivialElement) {
+      if (!is_big())
+        std::destroy_n(begin_, end_ - begin_);
+    }
+    begin_ = inline_storage_begin();
+    end_ = begin_;
+    end_of_storage_ = begin_ + kSize;
+  }
+
+  bool is_big() const {
+    return begin_ != inline_storage_begin();
+  }
+
+  T* inline_storage_begin() {
+    return reinterpret_cast<T*>(inline_storage_);
+  }
+  const T* inline_storage_begin() const {
+    return reinterpret_cast<const T*>(inline_storage_);
+  }
+
+  Allocator allocator_;
+
+  // Invariants:
+  // 1. The elements in the range between `begin_` (included) and `end_` (not
+  //    included) will be initialized at all times.
+  // 2. All other elements outside the range, both in the inline storage and in
+  //    the dynamic storage (if it exists), will be uninitialized at all times.
+
+  T* begin_ = inline_storage_begin();
+  T* end_ = begin_;
+  T* end_of_storage_ = begin_ + kSize;
+
+  alignas(T) char inline_storage_[sizeof(T) * kSize];
+
+  static constexpr bool kHasTrivialElement =
+      std::is_trivially_copyable<T>::value &&
+      std::is_trivially_destructible<T>::value;
+};
+
+template <typename>
+struct is_vector : std::false_type {};
+
+template <typename T, size_t Size, typename Allocator>
+struct is_vector<SmallVector<T, Size, Allocator>> : std::true_type {};
+
+template <typename T, typename Allocator>
+struct is_vector<std::vector<T, Allocator>> : std::true_type {};
+
+template <typename Vec>
+inline constexpr bool is_vector_v = is_vector<Vec>::value;
+
+#undef MLX_HAS_BUILTIN
+#undef MLX_HAS_ATTRIBUTE
+#undef MLX_LIKELY
+#undef MLX_UNLIKELY
+#undef MLX_NOINLINE
+
+} // namespace mlx::core

mlx/include/mlx/stream.h

@@ -0,0 +1,41 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/device.h"
+
+namespace mlx::core {
+
+struct Stream {
+  int index;
+  Device device;
+  explicit Stream(int index, Device device) : index(index), device(device) {}
+};
+
+/** Get the default stream for the given device. */
+Stream default_stream(Device d);
+
+/** Make the stream the default for its device. */
+void set_default_stream(Stream s);
+
+/** Make a new stream on the given device. */
+Stream new_stream(Device d);
+
+/** Get the stream with the given index. */
+Stream get_stream(int index);
+
+inline bool operator==(const Stream& lhs, const Stream& rhs) {
+  return lhs.index == rhs.index;
+}
+
+inline bool operator!=(const Stream& lhs, const Stream& rhs) {
+  return !(lhs == rhs);
+}
+
+/* Synchronize with the default stream. */
+void synchronize();
+
+/* Synchronize with the provided stream. */
+void synchronize(Stream);
+
+} // namespace mlx::core

mlx/include/mlx/threadpool.h

@@ -0,0 +1,133 @@
+// This code was modified from https://github.com/progschj/ThreadPool
+// The original License is copied below:
+//
+// Copyright (c) 2012 Jakob Progsch, Václav Zeman
+// This software is provided 'as-is', without any express or implied
+// warranty. In no event will the authors be held liable for any damages
+// arising from the use of this software.
+//
+// Permission is granted to anyone to use this software for any purpose,
+// including commercial applications, and to alter it and redistribute it
+// freely, subject to the following restrictions:
+//
+//    1. The origin of this software must not be misrepresented; you must not
+//    claim that you wrote the original software. If you use this software
+//    in a product, an acknowledgment in the product documentation would be
+//    appreciated but is not required.
+//
+//    2. Altered source versions must be plainly marked as such, and must not be
+//    misrepresented as being the original software.
+//
+//    3. This notice may not be removed or altered from any source
+//    distribution.
+#pragma once
+
+#include <condition_variable>
+#include <functional>
+#include <future>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <stdexcept>
+#include <thread>
+#include <vector>
+
+class ThreadPool {
+ public:
+  ThreadPool(size_t);
+  template <class F, class... Args>
+  auto enqueue(F&& f, Args&&... args)
+      -> std::future<typename std::invoke_result_t<F, Args...>>;
+  void resize(size_t);
+  ~ThreadPool();
+
+ private:
+  void stop_and_wait();
+  void start_threads(size_t);
+
+  std::vector<std::thread> workers;
+  std::queue<std::function<void()>> tasks;
+  std::mutex queue_mutex;
+  std::condition_variable condition;
+  bool stop;
+};
+
+inline ThreadPool::ThreadPool(size_t threads) : stop(false) {
+  start_threads(threads);
+}
+
+template <class F, class... Args>
+auto ThreadPool::enqueue(F&& f, Args&&... args)
+    -> std::future<typename std::invoke_result_t<F, Args...>> {
+  using return_type = typename std::invoke_result_t<F, Args...>;
+
+  auto task = std::make_shared<std::packaged_task<return_type()>>(
+      std::bind(std::forward<F>(f), std::forward<Args>(args)...));
+
+  std::future<return_type> res = task->get_future();
+  {
+    std::unique_lock<std::mutex> lock(queue_mutex);
+
+    if (stop) {
+      throw std::runtime_error(
+          "[ThreadPool::enqueue] Not allowed on stopped ThreadPool");
+    }
+
+    tasks.emplace([task]() { (*task)(); });
+  }
+  condition.notify_one();
+  return res;
+}
+
+inline void ThreadPool::resize(size_t threads) {
+  if (workers.size() == threads) {
+    return;
+  }
+
+  if (workers.size() > threads) {
+    stop_and_wait();
+  }
+  start_threads(threads - workers.size());
+}
+
+inline ThreadPool::~ThreadPool() {
+  stop_and_wait();
+}
+
+inline void ThreadPool::stop_and_wait() {
+  // Stop the current threads and wait until they finish
+  {
+    std::unique_lock<std::mutex> lock(queue_mutex);
+    stop = true;
+  }
+  condition.notify_all();
+  for (std::thread& worker : workers) {
+    worker.join();
+  }
+
+  // Reset the member variables so that the threadpool is reusable
+  stop = false;
+  workers.clear();
+}
+
+inline void ThreadPool::start_threads(size_t threads) {
+  for (size_t i = 0; i < threads; ++i) {
+    workers.emplace_back([this] {
+      for (;;) {
+        std::function<void()> task;
+
+        {
+          std::unique_lock<std::mutex> lock(this->queue_mutex);
+          this->condition.wait(
+              lock, [this] { return this->stop || !this->tasks.empty(); });
+          if (this->stop && this->tasks.empty())
+            return;
+          task = std::move(this->tasks.front());
+          this->tasks.pop();
+        }
+
+        task();
+      }
+    });
+  }
+}