mlx-cpu 0.30.1__py3-none-manylinux_2_35_x86_64.whl

mlx/include/mlx/fast_primitives.h

@@ -0,0 +1,427 @@
+// Copyright © 2024 Apple Inc.
+
+#include <optional>
+#include <variant>
+
+#include "mlx/primitives.h"
+
+namespace mlx::core::fast {
+
+// Custom primitive accepts a fallback function which it uses for
+// transformations. Transformations are virtual so that derived classes may
+// override the default behavior.
+class Custom : public Primitive {
+ public:
+  explicit Custom(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback)
+      : Primitive(stream), fallback_(std::move(fallback)) {}
+
+  virtual std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  virtual std::vector<array> jvp(
+      const std::vector<array>& primals,
+      const std::vector<array>& tangents,
+      const std::vector<int>& argnums) override;
+
+  virtual std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+ protected:
+  std::function<std::vector<array>(std::vector<array>)> fallback_;
+};
+
+class RMSNorm : public Custom {
+ public:
+  RMSNorm(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  static bool use_fallback(Stream stream);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(RMSNorm)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class RMSNormVJP : public Custom {
+ public:
+  RMSNormVJP(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(RMSNormVJP)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class LayerNorm : public Custom {
+ public:
+  LayerNorm(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  static bool use_fallback(Stream s);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(LayerNorm)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class LayerNormVJP : public Custom {
+ public:
+  LayerNormVJP(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(LayerNormVJP)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class RoPE : public Custom {
+ public:
+  RoPE(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      int dims,
+      bool traditional,
+      float base,
+      float scale,
+      bool forward)
+      : Custom(stream, std::move(fallback)),
+        dims_(dims),
+        traditional_(traditional),
+        base_(base),
+        scale_(scale),
+        forward_(forward) {}
+
+  static bool use_fallback(Stream s);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(RoPE)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_tuple(
+        nullptr, dims_, traditional_, base_, scale_, forward_);
+  }
+
+ private:
+  int dims_;
+  bool traditional_;
+  float base_;
+  float scale_;
+  bool forward_;
+};
+
+class ScaledDotProductAttention : public Custom {
+ public:
+  ScaledDotProductAttention(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float scale,
+      bool do_causal,
+      bool has_sinks,
+      bool output_logsumexp)
+      : Custom(stream, std::move(fallback)),
+        scale_(scale),
+        do_causal_(do_causal),
+        has_sinks_(has_sinks),
+        output_logsumexp_(output_logsumexp) {}
+
+  static bool use_fallback(
+      const array& q,
+      const array& k,
+      const array& v,
+      bool has_mask,
+      bool has_arr_mask,
+      bool do_causal,
+      bool is_training,
+      bool output_logsumexp,
+      Stream s);
+  static bool supports_bool_mask();
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  bool is_equivalent(const Primitive& other) const override;
+
+  DEFINE_NAME(ScaledDotProductAttention);
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_tuple(
+        nullptr, scale_, do_causal_, has_sinks_, output_logsumexp_);
+  }
+
+ private:
+  float scale_;
+  bool do_causal_;
+  bool has_sinks_;
+  bool output_logsumexp_;
+};
+
+class ScaledDotProductAttentionVJP : public Custom {
+ public:
+  ScaledDotProductAttentionVJP(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float scale,
+      bool do_causal,
+      bool has_sinks)
+      : Custom(stream, std::move(fallback)),
+        scale_(scale),
+        do_causal_(do_causal),
+        has_sinks_(has_sinks) {}
+
+  static bool use_fallback(const array& q, Stream s);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(ScaledDotProductAttentionVJP);
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(nullptr, scale_, do_causal_, has_sinks_);
+  }
+
+ private:
+  float scale_;
+  bool do_causal_;
+  bool has_sinks_;
+};
+
+class ConvertFP8 : public Primitive {
+ public:
+  explicit ConvertFP8(Stream stream, bool to_fp8)
+      : Primitive(stream), to_fp8_(to_fp8) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  const char* name() const override {
+    if (to_fp8_) {
+      return "ToFP8";
+    } else {
+      return "FromFP8";
+    }
+  }
+  bool state() const {
+    return to_fp8_;
+  };
+
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE();
+
+ private:
+  bool to_fp8_;
+};
+
+class Quantize : public Custom {
+ public:
+  explicit Quantize(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      int group_size,
+      int bits,
+      QuantizationMode mode,
+      bool dequantize)
+      : Custom(stream, std::move(fallback)),
+        group_size_(group_size),
+        bits_(bits),
+        mode_(mode),
+        dequantize_(dequantize) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(Quantize);
+
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return std::make_tuple(nullptr, group_size_, bits_, mode_, dequantize_);
+  }
+
+ private:
+  int group_size_;
+  int bits_;
+  QuantizationMode mode_;
+  bool dequantize_;
+};
+
+using ScalarArg = std::variant<bool, int, float>;
+
+class CustomKernel : public Primitive {
+ public:
+  CustomKernel(
+      Stream stream,
+      std::string name,
+      std::string source,
+      std::tuple<int, int, int> grid,
+      std::tuple<int, int, int> threadgroup,
+      std::vector<std::tuple<bool, bool, bool>> shape_infos,
+      bool ensure_row_contiguous,
+      std::optional<float> init_value,
+      std::vector<ScalarArg> scalar_arguments,
+      bool is_precompiled,
+      int shared_memory)
+      : Primitive(stream),
+        name_(std::move(name)),
+        source_(std::move(source)),
+        grid_(grid),
+        threadgroup_(threadgroup),
+        shape_infos_(std::move(shape_infos)),
+        ensure_row_contiguous_(ensure_row_contiguous),
+        init_value_(init_value),
+        scalar_arguments_(std::move(scalar_arguments)),
+        is_precompiled_(is_precompiled),
+        shared_memory_(shared_memory) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("Custom kernels only run on GPU.");
+  }
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(CustomKernel);
+  auto state() const {
+    return std::make_tuple(
+        name_,
+        source_,
+        grid_,
+        threadgroup_,
+        shape_infos_,
+        ensure_row_contiguous_,
+        init_value_,
+        scalar_arguments_,
+        is_precompiled_,
+        shared_memory_);
+  }
+
+ private:
+  std::string name_;
+  std::string source_;
+  std::tuple<int, int, int> grid_;
+  std::tuple<int, int, int> threadgroup_;
+  std::vector<std::tuple<bool, bool, bool>> shape_infos_;
+  bool ensure_row_contiguous_;
+  std::optional<float> init_value_;
+  std::vector<ScalarArg> scalar_arguments_;
+  bool is_precompiled_;
+  int shared_memory_;
+};
+
+} // namespace mlx::core::fast

mlx/include/mlx/fence.h

@@ -0,0 +1,39 @@
+// Copyright © 2024 Apple Inc.
+
+#include <vector>
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+/* A fence to be used for synchronizing work between streams.
+ *
+ * Calls to `wait` wait in the given stream until all previous calls to update
+ * are complete on their given stream.
+ *
+ * The array passed to `update` is computed and visible after the call to
+ * `wait` returns. The array passed to `wait` will not be read until all
+ * previous calls to `update` have completed.
+ *
+ * Note, calls to `update` should always be from the same thread or explicitly
+ * synchronized so that they occur in sequence. Calls to `wait` can be on any
+ * thread.
+ *
+ * For the Metal back-end the fence supports slow (default) and fast mode.
+ * Fast mode requires setting the environment variable
+ * `MLX_METAL_FAST_SYNCH=1`. Fast mode also requires Metal 3.2+ (macOS 15+,
+ * iOS 18+).
+ */
+class Fence {
+ public:
+  Fence() {};
+  explicit Fence(Stream stream);
+
+  void update(Stream stream, const array& x, bool cross_device);
+  void wait(Stream stream, const array& x);
+
+ private:
+  std::shared_ptr<void> fence_{nullptr};
+};
+
+} // namespace mlx::core

mlx/include/mlx/fft.h

@@ -0,0 +1,167 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <variant>
+
+#include "array.h"
+#include "device.h"
+#include "utils.h"
+
+namespace mlx::core::fft {
+
+/** Compute the n-dimensional Fourier Transform. */
+array fftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array fftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+array fftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the n-dimensional inverse Fourier Transform. */
+array ifftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array ifftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array ifftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the one-dimensional Fourier Transform. */
+inline array fft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return fftn(a, {n}, {axis}, s);
+}
+inline array fft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return fftn(a, {axis}, s);
+}
+
+/** Compute the one-dimensional inverse Fourier Transform. */
+inline array ifft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return ifftn(a, {n}, {axis}, s);
+}
+inline array ifft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return ifftn(a, {axis}, s);
+}
+
+/** Compute the two-dimensional Fourier Transform. */
+inline array fft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return fftn(a, n, axes, s);
+}
+inline array fft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return fftn(a, axes, s);
+}
+
+/** Compute the two-dimensional inverse Fourier Transform. */
+inline array ifft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return ifftn(a, n, axes, s);
+}
+inline array ifft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return ifftn(a, axes, s);
+}
+
+/** Compute the n-dimensional Fourier Transform on a real input. */
+array rfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array rfftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array rfftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the n-dimensional inverse of `rfftn`. */
+array irfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array irfftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+array irfftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the one-dimensional Fourier Transform on a real input. */
+inline array rfft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return rfftn(a, {n}, {axis}, s);
+}
+inline array rfft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return rfftn(a, {axis}, s);
+}
+/** Compute the one-dimensional inverse of `rfft`. */
+inline array irfft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return irfftn(a, {n}, {axis}, s);
+}
+inline array irfft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return irfftn(a, {axis}, s);
+}
+
+/** Compute the two-dimensional Fourier Transform on a real input. */
+inline array rfft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return rfftn(a, n, axes, s);
+}
+inline array rfft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return rfftn(a, axes, s);
+}
+
+/** Compute the two-dimensional inverse of `rfft2`. */
+inline array irfft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return irfftn(a, n, axes, s);
+}
+inline array irfft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return irfftn(a, axes, s);
+}
+/** Shift the zero-frequency component to the center of the spectrum. */
+array fftshift(const array& a, StreamOrDevice s = {});
+
+/** Shift the zero-frequency component to the center of the spectrum along
+ * specified axes. */
+array fftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+
+/** The inverse of fftshift. */
+array ifftshift(const array& a, StreamOrDevice s = {});
+
+/** The inverse of fftshift along specified axes. */
+array ifftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+
+} // namespace mlx::core::fft

mlx/include/mlx/graph_utils.h

@@ -0,0 +1,66 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+struct NodeNamer {
+  std::unordered_map<std::uintptr_t, std::string> names;
+
+  const std::string& get_name(const array& x);
+  void set_name(const array& x, std::string n);
+};
+
+void print_graph(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs);
+
+inline void print_graph(std::ostream& os, const std::vector<array>& outputs) {
+  print_graph(os, NodeNamer{}, outputs);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void print_graph(std::ostream& os, Arrays&&... outputs) {
+  print_graph(
+      os, NodeNamer{}, std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void
+print_graph(std::ostream& os, NodeNamer namer, Arrays&&... outputs) {
+  print_graph(
+      os,
+      std::move(namer),
+      std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+void export_to_dot(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs);
+
+inline void export_to_dot(std::ostream& os, const std::vector<array>& outputs) {
+  export_to_dot(os, NodeNamer{}, outputs);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void export_to_dot(std::ostream& os, Arrays&&... outputs) {
+  export_to_dot(
+      os, NodeNamer{}, std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void
+export_to_dot(std::ostream& os, NodeNamer namer, Arrays&&... outputs) {
+  export_to_dot(
+      os,
+      std::move(namer),
+      std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+} // namespace mlx::core