mlx 0.30.7

data/mlx/mlx/random.h

@@ -0,0 +1,283 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <chrono>
+#include <optional>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/stream.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::random {
+
+class MLX_API KeySequence {
+ public:
+  explicit KeySequence(uint64_t seed);
+
+  void seed(uint64_t seed);
+  array next();
+
+  // static default
+  static KeySequence& default_() {
+    static KeySequence ks(get_current_time_seed());
+    return ks;
+  }
+
+ private:
+  array key_;
+  static uint64_t get_current_time_seed() {
+    auto now = std::chrono::system_clock::now();
+    return std::chrono::duration_cast<std::chrono::milliseconds>(
+               now.time_since_epoch())
+        .count();
+  }
+};
+
+/** Get a PRNG key from a seed. */
+MLX_API array key(uint64_t seed);
+
+/** Seed the default PRNG key. */
+MLX_API void seed(uint64_t seed);
+
+/** Generate an array with type uint32 filled with random bits. */
+MLX_API array bits(
+    const Shape& shape,
+    int width,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+inline array bits(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return bits(shape, 4, key, s);
+}
+
+/** Split the rng key into a pair of keys. */
+MLX_API std::pair<array, array> split(const array& key, StreamOrDevice s = {});
+
+/** Split the rng key into `num` keys. */
+MLX_API array split(const array& key, int num, StreamOrDevice s = {});
+
+/** Generate uniform random numbers between low and high. */
+MLX_API array uniform(
+    const array& low,
+    const array& high,
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+template <typename T, typename U>
+array uniform(
+    T low,
+    U high,
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return uniform(array(low), array(high), shape, dtype, key, to_stream(s));
+}
+
+/** Generate uniform random numbers between 0 and 1. */
+MLX_API array uniform(
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+inline array uniform(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return uniform(shape, float32, key, s);
+}
+
+/** Generate samples from the standard normal distribution. */
+MLX_API array normal(
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& loc,
+    const std::optional<array>& scale,
+    const std::optional<array>& key,
+    StreamOrDevice s = {});
+inline array normal(
+    const Shape& shape,
+    Dtype dtype,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  auto loc_ = loc == 0 ? std::nullopt : std::make_optional(array(loc, dtype));
+  auto scale_ =
+      scale == 1 ? std::nullopt : std::make_optional(array(scale, dtype));
+  return normal(shape, dtype, loc_, scale_, key, s);
+}
+inline array normal(
+    const Shape& shape,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return normal(shape, float32, loc, scale, key, s);
+}
+inline array normal(
+    const Shape& shape,
+    const Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return normal(shape, dtype, std::nullopt, std::nullopt, key, s);
+}
+inline array normal(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return normal(shape, float32, std::nullopt, std::nullopt, key, s);
+}
+
+/** Generate samples from a multivariate normal distribution. **/
+MLX_API array multivariate_normal(
+    const array& mean,
+    const array& cov,
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Generate integer samples uniformly at random */
+MLX_API array randint(
+    const array& low,
+    const array& high,
+    const Shape& shape,
+    Dtype dtype = int32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+template <typename T, typename U>
+array randint(
+    T low,
+    U high,
+    const Shape& shape,
+    Dtype dtype = int32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return randint(array(low), array(high), shape, dtype, key, to_stream(s));
+}
+
+/** Generate binary variables with probability to be true equal to p */
+MLX_API array bernoulli(
+    const array& p,
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+MLX_API array bernoulli(
+    const array& p,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+template <typename T>
+array bernoulli(
+    T p,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return bernoulli(array(p), key, s);
+}
+
+template <typename T>
+array bernoulli(
+    T p,
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return bernoulli(array(p), shape, key, s);
+}
+
+MLX_API array bernoulli(
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array truncated_normal(
+    const array& lower,
+    const array& upper,
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array truncated_normal(
+    const array& lower,
+    const array& upper,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array gumbel(
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array categorical(
+    const array& logits,
+    int axis,
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array categorical(
+    const array& logits_,
+    int axis,
+    int num_samples,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array categorical(
+    const array& logits,
+    int axis = -1,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Generate samples from the laplace distribution. */
+MLX_API array laplace(
+    const Shape& shape,
+    Dtype dtype,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+inline array laplace(
+    const Shape& shape,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return laplace(shape, float32, loc, scale, key, s);
+}
+inline array laplace(
+    const Shape& shape,
+    const Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return laplace(shape, dtype, 0.0, 1.0, key, s);
+}
+inline array laplace(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return laplace(shape, float32, 0.0, 1.0, key, s);
+}
+
+/* Randomly permute the elements of x along the given axis. */
+MLX_API array permutation(
+    const array& x,
+    int axis = 0,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+/* A random permutation of `arange(x)` */
+MLX_API array permutation(
+    int x,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+} // namespace mlx::core::random

data/mlx/mlx/scheduler.cpp

@@ -0,0 +1,73 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/scheduler.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/gpu/eval.h"
+
+namespace mlx::core {
+
+Stream default_stream(Device d) {
+  if (!gpu::is_available() && d == Device::gpu) {
+    throw std::invalid_argument(
+        "[default_stream] Cannot get gpu stream without gpu backend.");
+  }
+  return scheduler::scheduler().get_default_stream(d);
+}
+
+void set_default_stream(Stream s) {
+  if (!gpu::is_available() && s.device == Device::gpu) {
+    throw std::invalid_argument(
+        "[set_default_stream] Cannot set gpu stream without gpu backend.");
+  }
+  return scheduler::scheduler().set_default_stream(s);
+}
+
+Stream get_stream(int index) {
+  return scheduler::scheduler().get_stream(index);
+}
+
+Stream new_stream(Device d) {
+  if (!gpu::is_available() && d == Device::gpu) {
+    throw std::invalid_argument(
+        "[new_stream] Cannot make gpu stream without gpu backend.");
+  }
+  return scheduler::scheduler().new_stream(d);
+}
+
+Stream new_stream() {
+  return scheduler::scheduler().new_stream(default_device());
+}
+
+void synchronize(Stream s) {
+  if (s.device == mlx::core::Device::cpu) {
+    auto p = std::make_shared<std::promise<void>>();
+    std::future<void> f = p->get_future();
+    scheduler::enqueue(s, [p = std::move(p)]() { p->set_value(); });
+    f.wait();
+  } else {
+    gpu::synchronize(s);
+  }
+}
+
+void synchronize() {
+  synchronize(default_stream(default_device()));
+}
+
+namespace scheduler {
+
+/** A singleton scheduler to manage devices, streams, and task execution. */
+Scheduler& scheduler() {
+  // Leak the scheduler on Windows to avoid joining threads on exit, can be
+  // removed after Visual Studio fixes bug:
+  // https://developercommunity.visualstudio.com/t/1654756
+#ifdef _WIN32
+  static Scheduler* scheduler = new Scheduler;
+  return *scheduler;
+#else
+  static Scheduler scheduler;
+  return scheduler;
+#endif
+}
+
+} // namespace scheduler
+} // namespace mlx::core

data/mlx/mlx/scheduler.h

@@ -0,0 +1,189 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <atomic>
+#include <future>
+#include <queue>
+#include <thread>
+#include <unordered_map>
+
+#include "mlx/api.h"
+#include "mlx/backend/gpu/eval.h"
+#include "mlx/device.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::scheduler {
+
+struct StreamThread {
+  std::mutex mtx;
+  std::queue<std::function<void()>> q;
+  std::condition_variable cond;
+  bool stop;
+  std::thread thread;
+
+  StreamThread() : stop(false), thread(&StreamThread::thread_fn, this) {}
+
+  ~StreamThread() {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      stop = true;
+    }
+    cond.notify_one();
+    thread.join();
+  }
+
+  void thread_fn() {
+    while (true) {
+      std::function<void()> task;
+      {
+        std::unique_lock<std::mutex> lk(mtx);
+        cond.wait(lk, [this] { return !this->q.empty() || this->stop; });
+        if (q.empty() && stop) {
+          return;
+        }
+        task = std::move(q.front());
+        q.pop();
+      }
+
+      task();
+    }
+  }
+
+  template <typename F>
+  void enqueue(F&& f) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      if (stop) {
+        throw std::runtime_error(
+            "Cannot enqueue work after stream is stopped.");
+      }
+      q.emplace(std::forward<F>(f));
+    }
+    cond.notify_one();
+  }
+};
+
+class Scheduler {
+ public:
+  Scheduler() : n_active_tasks_(0) {
+    if (is_available(Device::gpu)) {
+      default_streams_.insert({Device::gpu, new_stream(Device::gpu)});
+    }
+    default_streams_.insert({Device::cpu, new_stream(Device::cpu)});
+  }
+
+  // Not copyable or moveable
+  Scheduler(const Scheduler&) = delete;
+  Scheduler(Scheduler&&) = delete;
+  Scheduler& operator=(const Scheduler&) = delete;
+  Scheduler& operator=(Scheduler&&) = delete;
+
+  Stream new_stream(const Device& d) {
+    streams_.emplace_back(streams_.size(), d);
+    if (d == Device::gpu) {
+      threads_.push_back(nullptr);
+      gpu::new_stream(streams_.back());
+    } else {
+      threads_.push_back(new StreamThread{});
+    }
+    return streams_.back();
+  }
+
+  template <typename F>
+  void enqueue(const Stream& stream, F&& f);
+
+  Stream get_default_stream(const Device& d) const {
+    return default_streams_.at(d.type);
+  }
+  Stream get_stream(int index) const {
+    return streams_.at(index);
+  }
+
+  void set_default_stream(const Stream& s) {
+    default_streams_.at(s.device.type) = s;
+  }
+
+  void notify_new_task(const Stream& stream) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      n_active_tasks_++;
+    }
+    completion_cv.notify_all();
+  }
+
+  void notify_task_completion(const Stream& stream) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      n_active_tasks_--;
+    }
+    completion_cv.notify_all();
+  }
+
+  int n_active_tasks() const {
+    return n_active_tasks_;
+  }
+
+  void wait_for_one() {
+    std::unique_lock<std::mutex> lk(mtx);
+    int n_tasks_old = n_active_tasks();
+    if (n_tasks_old > 1) {
+      completion_cv.wait(lk, [this, n_tasks_old] {
+        return this->n_active_tasks() < n_tasks_old;
+      });
+    }
+  }
+
+  ~Scheduler() {
+    for (auto s : streams_) {
+      try {
+        synchronize(s);
+      } catch (const std::runtime_error&) {
+        // ignore errors if synch fails
+      }
+    }
+    for (auto t : threads_) {
+      if (t != nullptr) {
+        delete t;
+      }
+    }
+  }
+
+ private:
+  int n_active_tasks_;
+  std::vector<StreamThread*> threads_;
+  std::vector<Stream> streams_;
+  std::unordered_map<Device::DeviceType, Stream> default_streams_;
+  std::condition_variable completion_cv;
+  std::mutex mtx;
+};
+
+template <typename F>
+void Scheduler::enqueue(const Stream& stream, F&& f) {
+  threads_[stream.index]->enqueue(std::forward<F>(f));
+}
+
+MLX_API Scheduler& scheduler();
+
+template <typename F>
+void enqueue(const Stream& stream, F&& f) {
+  scheduler().enqueue(stream, std::forward<F>(f));
+}
+
+inline int n_active_tasks() {
+  return scheduler().n_active_tasks();
+}
+
+inline void notify_new_task(const Stream& stream) {
+  scheduler().notify_new_task(stream);
+}
+
+inline void notify_task_completion(const Stream& stream) {
+  scheduler().notify_task_completion(stream);
+}
+
+inline void wait_for_one() {
+  scheduler().wait_for_one();
+}
+
+} // namespace mlx::core::scheduler