Implemented Cholesky on CPU (#1119)

mlx/backend/accelerate/primitives.cpp

@@ -80,6 +80,7 @@ DEFAULT(StopGradient)
 DEFAULT_MULTI(SVD)
 DEFAULT(Transpose)
 DEFAULT(Inverse)
+DEFAULT(Cholesky)
 
 void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);

mlx/backend/common/CMakeLists.txt

@@ -56,6 +56,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/qrf.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/svd.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/inverse.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/cholesky.cpp
   ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
 )
 

mlx/backend/common/cholesky.cpp

@@ -0,0 +1,109 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/copy.h"
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+
+#ifdef ACCELERATE_NEW_LAPACK
+#include <Accelerate/Accelerate.h>
+#else
+#include <lapack.h>
+#endif
+
+namespace mlx::core {
+
+namespace {
+
+// Delegate to the Cholesky factorization taking into account differences in
+// LAPACK implementations (basically how to pass the 'uplo' string to fortran).
+int spotrf_wrapper(char uplo, float* matrix, int N) {
+  int info;
+
+#ifdef LAPACK_FORTRAN_STRLEN_END
+  spotrf_(
+      /* uplo = */ &uplo,
+      /* n = */ &N,
+      /* a = */ matrix,
+      /* lda = */ &N,
+      /* info = */ &info,
+      /* uplo_len = */ static_cast<size_t>(1));
+#else
+  spotrf_(
+      /* uplo = */ &uplo,
+      /* n = */ &N,
+      /* a = */ matrix,
+      /* lda = */ &N,
+      /* info = */ &info);
+#endif
+
+  return info;
+}
+
+} // namespace
+
+void cholesky_impl(const array& a, array& factor, bool upper) {
+  // Lapack uses the column-major convention. We take advantage of the fact that
+  // the matrix should be symmetric:
+  //   (A)ᵀ = A
+  // and that a column-major lower triangular matrix is a row-major upper
+  // triangular matrix, so uplo is the opposite of what we would expect from
+  // upper
+
+  char uplo = (upper) ? 'L' : 'U';
+
+  // The decomposition is computed in place, so just copy the input to the
+  // output.
+  copy(
+      a,
+      factor,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+
+  const int N = a.shape(-1);
+  const size_t num_matrices = a.size() / (N * N);
+
+  float* matrix = factor.data<float>();
+
+  for (int i = 0; i < num_matrices; i++) {
+    // Compute Cholesky factorization.
+    int info = spotrf_wrapper(uplo, matrix, N);
+
+    // TODO: We do nothing when the matrix is not positive semi-definite
+    // because throwing an error would result in a crash. If we figure out how
+    // to catch errors from the implementation we should throw.
+    if (info < 0) {
+      std::stringstream msg;
+      msg << "[cholesky] Cholesky decomposition failed with error code "
+          << info;
+      throw std::runtime_error(msg.str());
+    }
+
+    // Zero out the upper/lower triangle while advancing the pointer to the
+    // next matrix at the same time.
+    for (int row = 0; row < N; row++) {
+      if (upper) {
+        std::fill(matrix, matrix + row, 0);
+      } else {
+        std::fill(matrix + row + 1, matrix + N, 0);
+      }
+      matrix += N;
+    }
+  }
+}
+
+void Cholesky::eval(const std::vector<array>& inputs, array& output) {
+  if (inputs[0].dtype() != float32) {
+    throw std::runtime_error("[Cholesky::eval] only supports float32.");
+  }
+  cholesky_impl(inputs[0], output, upper_);
+}
+
+std::pair<std::vector<array>, std::vector<int>> Cholesky::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0] >= 0 ? 0 : -1;
+  auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  return {{linalg::cholesky(a, upper_, stream())}, {ax}};
+}
+
+} // namespace mlx::core

mlx/backend/common/default_primitives.cpp

@@ -113,6 +113,7 @@ DEFAULT(Tan)
 DEFAULT(Tanh)
 DEFAULT(Transpose)
 DEFAULT(Inverse)
+DEFAULT(Cholesky)
 
 namespace {
 

mlx/backend/metal/primitives.cpp

@@ -1012,4 +1012,9 @@ void Inverse::eval_gpu(const std::vector<array>& inputs, array& output) {
   throw std::runtime_error("[Inverse::eval_gpu] Metal inversion NYI.");
 }
 
+void Cholesky::eval_gpu(const std::vector<array>& inputs, array& out) {
+  throw std::runtime_error(
+      "[Cholesky::eval_gpu] Metal Cholesky decomposition NYI.");
+}
+
 } // namespace mlx::core

mlx/backend/no_metal/primitives.cpp

@@ -107,6 +107,7 @@ NO_GPU(Tan)
 NO_GPU(Tanh)
 NO_GPU(Transpose)
 NO_GPU(Inverse)
+NO_GPU(Cholesky)
 
 namespace fast {
 NO_GPU_MULTI(LayerNorm)

mlx/linalg.cpp

@@ -261,4 +261,35 @@ array inv(const array& a, StreamOrDevice s /* = {} */) {
       a.shape(), a.dtype(), std::make_shared<Inverse>(to_stream(s)), {a});
 }
 
+array cholesky(
+    const array& a,
+    bool upper /* = false */,
+    StreamOrDevice s /* = {} */) {
+  if (a.dtype() != float32) {
+    std::ostringstream msg;
+    msg << "[linalg::cholesky] Arrays must type float32. Received array "
+        << "with type " << a.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::cholesky] Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(
+        "[linalg::cholesky] Cholesky decomposition is only defined for square "
+        "matrices.");
+  }
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Cholesky>(to_stream(s), upper),
+      {a});
+}
+
 } // namespace mlx::core::linalg

mlx/linalg.h

@@ -66,4 +66,6 @@ std::vector<array> svd(const array& a, StreamOrDevice s = {});
 
 array inv(const array& a, StreamOrDevice s = {});
 
+array cholesky(const array& a, bool upper = false, StreamOrDevice s = {});
+
 } // namespace mlx::core::linalg

mlx/primitives.h

@@ -2093,4 +2093,20 @@ class Inverse : public UnaryPrimitive {
   void eval(const std::vector<array>& inputs, array& output);
 };
 
+class Cholesky : public UnaryPrimitive {
+ public:
+  explicit Cholesky(Stream stream, bool upper)
+      : UnaryPrimitive(stream), upper_(upper) {};
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_PRINT(Cholesky)
+
+ private:
+  void eval(const std::vector<array>& inputs, array& output);
+  bool upper_;
+};
+
 } // namespace mlx::core