mlx 1.0.0

data/mlx/mlx/backend/cpu/svd.cpp

@@ -0,0 +1,289 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T, class Enable = void>
+struct SVDWork {};
+
+template <typename T>
+struct SVDWork<
+    T,
+    typename std::enable_if<std::is_floating_point<T>::value>::type> {
+  using R = T;
+
+  int N;
+  int M;
+  int K;
+  int lda;
+  int ldu;
+  int ldvt;
+  char jobz;
+  std::vector<array::Data> buffers;
+  int lwork;
+
+  SVDWork(int N, int M, int K, char jobz)
+      : N(N), M(M), K(K), lda(N), ldu(N), ldvt(M), jobz(jobz) {
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    buffers.emplace_back(allocator::malloc(sizeof(int) * 8 * K));
+
+    int lwork_query = -1;
+    int info;
+
+    // Compute workspace size.
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ nullptr,
+        /* lda = */ &lda,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
+        /* ldu = */ &ldu,
+        /* vt = */ nullptr,
+        /* ldvt = */ &ldvt,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    lwork = workspace_dimension;
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, R* s, T* u, T* vt) {
+    int info;
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ a,
+        /* lda = */ &lda,
+        /* s = */ s,
+        // According to the identity above, lapack will write Vᵀᵀ as U.
+        /* u = */ u,
+        /* ldu = */ &ldu,
+        // According to the identity above, lapack will write Uᵀ as Vᵀ.
+        /* vt = */ vt,
+        /* ldvt = */ &ldvt,
+        /* work = */ static_cast<T*>(buffers[1].buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+};
+
+template <>
+struct SVDWork<std::complex<float>> {
+  using T = std::complex<float>;
+  using R = float;
+
+  int N;
+  int M;
+  int K;
+  int lda;
+  int ldu;
+  int ldvt;
+  char jobz;
+  std::vector<array::Data> buffers;
+  int lwork;
+
+  SVDWork(int N, int M, int K, char jobz)
+      : N(N), M(M), K(K), lda(N), ldu(N), ldvt(M), jobz(jobz) {
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    buffers.emplace_back(allocator::malloc(sizeof(int) * 8 * K));
+
+    const int lrwork =
+        jobz == 'A' ? std::max(1, 5 * K * K + 5 * K) : std::max(1, 7 * K);
+    buffers.emplace_back(allocator::malloc(sizeof(float) * lrwork));
+
+    int lwork_query = -1;
+    int work_query = -1;
+    int info;
+
+    // Compute workspace size.
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ nullptr,
+        /* lda = */ &lda,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
+        /* ldu = */ &ldu,
+        /* vt = */ nullptr,
+        /* ldvt = */ &ldvt,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
+        /* rwork = */ static_cast<float*>(buffers[1].buffer.raw_ptr()),
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    lwork = workspace_dimension.real();
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, R* s, T* u, T* vt) {
+    int info;
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ a,
+        /* lda = */ &lda,
+        /* s = */ s,
+        // According to the identity above, lapack will write Vᵀᵀ as U.
+        /* u = */ u,
+        /* ldu = */ &ldu,
+        // According to the identity above, lapack will write Uᵀ as Vᵀ.
+        /* vt = */ vt,
+        /* ldvt = */ &ldvt,
+        /* work = */ static_cast<T*>(buffers[2].buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* rwork = */ static_cast<float*>(buffers[1].buffer.raw_ptr()),
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+};
+
+template <typename T>
+void svd_impl(
+    const array& a,
+    std::vector<array>& outputs,
+    bool compute_uv,
+    Stream stream) {
+  // Lapack uses the column-major convention. To avoid having to transpose
+  // the input and then transpose the outputs, we swap the indices/sizes of the
+  // matrices and take advantage of the following identity (see
+  // https://math.stackexchange.com/a/30077)
+  //    A = UΣVᵀ
+  //    Aᵀ = VΣUᵀ
+  // As a result some of the indices/sizes are swapped as noted above.
+
+  // Rows and cols of the original matrix in row-major order.
+  const int M = a.shape(-2);
+  const int N = a.shape(-1);
+  const int K = std::min(M, N);
+
+  using R = typename SVDWork<T>::R;
+
+  size_t num_matrices = a.size() / (M * N);
+
+  // lapack clobbers the input, so we have to make a copy.
+  array in(a.shape(), a.dtype(), nullptr, {});
+  copy_cpu(
+      a,
+      in,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
+
+  // Allocate outputs.
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  auto in_ptr = in.data<T>();
+  T* u_ptr;
+  R* s_ptr;
+  T* vt_ptr;
+
+  if (compute_uv) {
+    array& u = outputs[0];
+    array& s = outputs[1];
+    array& vt = outputs[2];
+
+    u.set_data(allocator::malloc(u.nbytes()));
+    s.set_data(allocator::malloc(s.nbytes()));
+    vt.set_data(allocator::malloc(vt.nbytes()));
+
+    encoder.set_output_array(u);
+    encoder.set_output_array(s);
+    encoder.set_output_array(vt);
+
+    s_ptr = s.data<R>();
+    u_ptr = u.data<T>();
+    vt_ptr = vt.data<T>();
+  } else {
+    array& s = outputs[0];
+
+    s.set_data(allocator::malloc(s.nbytes()));
+
+    encoder.set_output_array(s);
+
+    s_ptr = s.data<R>();
+    u_ptr = nullptr;
+    vt_ptr = nullptr;
+  }
+
+  encoder.dispatch([in_ptr, u_ptr, s_ptr, vt_ptr, M, N, K, num_matrices]() {
+    auto jobz = (u_ptr) ? 'A' : 'N';
+    SVDWork<T> svd_work(N, M, K, jobz);
+    // Loop over matrices.
+    for (int i = 0; i < num_matrices; i++) {
+      svd_work.run(
+          in_ptr + M * N * i,
+          s_ptr + K * i,
+          vt_ptr ? vt_ptr + N * N * i : nullptr,
+          u_ptr ? u_ptr + M * M * i : nullptr);
+    }
+  });
+  encoder.add_temporary(in);
+}
+
+void SVD::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      svd_impl<float>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    case float64:
+      svd_impl<double>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    case complex64:
+      svd_impl<std::complex<float>>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[SVD::eval_cpu] only supports float32, float64, or complex64.");
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/ternary.h

@@ -0,0 +1,154 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include "mlx/array.h"
+#include "mlx/backend/common/ternary.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core {
+
+template <typename T1, typename T2, typename T3, typename U, typename Op, int D>
+void ternary_op_dims(
+    const T1* a,
+    const T2* b,
+    const T3* c,
+    U* out,
+    Op op,
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& c_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_c = c_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      ternary_op_dims<T1, T2, T3, U, Op, D - 1>(
+          a,
+          b,
+          c,
+          out,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          c_strides,
+          out_strides,
+          axis + 1);
+    } else {
+      *out = op(*a, *b, *c);
+    }
+    a += stride_a;
+    b += stride_b;
+    c += stride_c;
+    out += stride_out;
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dispatch_dims(
+    const T1* a_ptr,
+    const T2* b_ptr,
+    const T3* c_ptr,
+    U* out_ptr,
+    Op op,
+    size_t size,
+    Shape& shape,
+    std::vector<Strides>& strides) {
+  const auto& a_strides = strides[0];
+  const auto& b_strides = strides[1];
+  const auto& c_strides = strides[2];
+  const auto& out_strides = strides[3];
+  int ndim = shape.size();
+  switch (ndim) {
+    case 1:
+      ternary_op_dims<T1, T2, T3, U, Op, 1>(
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          c_strides,
+          out_strides,
+          0);
+      return;
+    case 2:
+      ternary_op_dims<T1, T2, T3, U, Op, 2>(
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          c_strides,
+          out_strides,
+          0);
+      return;
+  }
+
+  ContiguousIterator a_it(shape, a_strides, ndim - 2);
+  ContiguousIterator b_it(shape, b_strides, ndim - 2);
+  ContiguousIterator c_it(shape, c_strides, ndim - 2);
+  auto stride = out_strides[ndim - 3];
+  for (size_t elem = 0; elem < size; elem += stride) {
+    ternary_op_dims<T1, T2, T3, U, Op, 2>(
+        a_ptr + a_it.loc,
+        b_ptr + b_it.loc,
+        c_ptr + c_it.loc,
+        out_ptr + elem,
+        op,
+        shape,
+        a_strides,
+        b_strides,
+        c_strides,
+        out_strides,
+        ndim - 2);
+    a_it.step();
+    b_it.step();
+    c_it.step();
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op,
+    TernaryOpType topt) {
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+  U* out_ptr = out.data<U>();
+
+  if (topt == TernaryOpType::ScalarScalarScalar) {
+    *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
+  } else if (topt == TernaryOpType::VectorVectorVector) {
+    for (size_t i = 0; i < out.size(); ++i) {
+      *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
+      a_ptr++;
+      b_ptr++;
+      c_ptr++;
+      out_ptr++;
+    }
+  } else {
+    auto [shape, strides] = collapse_contiguous_dims(
+        a.shape(), {a.strides(), b.strides(), c.strides(), out.strides()});
+    ternary_op_dispatch_dims<T1, T2, T3, U>(
+        a_ptr, b_ptr, c_ptr, out_ptr, op, out.size(), shape, strides);
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/threefry.cpp

@@ -0,0 +1,31 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/backend/cpu/threefry.h"
+
+namespace mlx::core::random {
+
+std::pair<uint32_t, uint32_t> threefry2x32_hash(
+    const std::pair<uint32_t, uint32_t>& key,
+    std::pair<uint32_t, uint32_t> count) {
+  constexpr static uint32_t rotations[2][4] = {
+      {13, 15, 26, 6}, {17, 29, 16, 24}};
+
+  uint32_t ks[3] = {key.first, key.second, key.first ^ key.second ^ 0x1BD11BDA};
+
+  count.first += ks[0];
+  count.second += ks[1];
+
+  for (int i = 0; i < 5; ++i) {
+    for (auto r : rotations[i % 2]) {
+      count.first += count.second;
+      count.second = (count.second << r) | (count.second >> (32 - r));
+      count.second ^= count.first;
+    }
+    count.first += ks[(i + 1) % 3];
+    count.second += ks[(i + 2) % 3] + i + 1;
+  }
+
+  return count;
+}
+
+} // namespace mlx::core::random

data/mlx/mlx/backend/cpu/threefry.h

@@ -0,0 +1,21 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <cstdint>
+#include <utility>
+
+namespace mlx::core::random {
+
+/** Applies the Threefry 2x32 hash function.
+ * This code is based on the Jax counter-based and splittable PRNG
+ * https://github.com/google/jax/blob/main/docs/jep/263-prng.md
+ *
+ * Original Threefry reference:
+ * http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
+ */
+std::pair<uint32_t, uint32_t> threefry2x32_hash(
+    const std::pair<uint32_t, uint32_t>& key,
+    std::pair<uint32_t, uint32_t> count);
+
+} // namespace mlx::core::random