mlx 0.30.7

data/mlx/mlx/backend/cpu/inverse.cpp

@@ -0,0 +1,160 @@
+// Copyright © 2023-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>
+void general_inv(T* inv, int N) {
+  int info;
+  auto ipiv = array::Data{allocator::malloc(sizeof(int) * N)};
+  // Compute LU factorization.
+  getrf<T>(
+      /* m = */ &N,
+      /* n = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] LU factorization failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+
+  static const int lwork_query = -1;
+  T workspace_size = 0;
+
+  // Compute workspace size.
+  getri<T>(
+      /* m = */ &N,
+      /* a = */ nullptr,
+      /* lda = */ &N,
+      /* ipiv = */ nullptr,
+      /* work = */ &workspace_size,
+      /* lwork = */ &lwork_query,
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] LU workspace calculation failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+
+  const int lwork = workspace_size;
+  auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)};
+
+  // Compute inverse.
+  getri<T>(
+      /* m = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+      /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
+      /* lwork = */ &lwork,
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] inversion failed with error code " << info;
+    throw std::runtime_error(ss.str());
+  }
+}
+
+template <typename T>
+void tri_inv(T* inv, int N, bool upper) {
+  const char uplo = upper ? 'L' : 'U';
+  const char diag = 'N';
+  int info;
+  trtri<T>(
+      /* uplo = */ &uplo,
+      /* diag = */ &diag,
+      /* N = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* info = */ &info);
+
+  // zero out the other triangle
+  if (upper) {
+    for (int i = 0; i < N; i++) {
+      std::fill(inv, inv + i, 0.0f);
+      inv += N;
+    }
+  } else {
+    for (int i = 0; i < N; i++) {
+      std::fill(inv + i + 1, inv + N, 0.0f);
+      inv += N;
+    }
+  }
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] triangular inversion failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+}
+
+template <typename T>
+void inverse_impl(
+    const array& a,
+    array& inv,
+    bool tri,
+    bool upper,
+    Stream stream) {
+  // Lapack uses the column-major convention. We take advantage of the following
+  // identity to avoid transposing (see
+  // https://math.stackexchange.com/a/340234):
+  //   (A⁻¹)ᵀ = (Aᵀ)⁻¹
+
+  // The inverse is computed in place, so just copy the input to the output.
+  copy_cpu(
+      a,
+      inv,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
+
+  const int N = a.shape(-1);
+  const size_t num_matrices = a.size() / (N * N);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(inv);
+
+  auto inv_ptr = inv.data<T>();
+  if (tri) {
+    encoder.dispatch([inv_ptr, N, num_matrices, upper]() {
+      for (int i = 0; i < num_matrices; i++) {
+        tri_inv<T>(inv_ptr + N * N * i, N, upper);
+      }
+    });
+  } else {
+    encoder.dispatch([inv_ptr, N, num_matrices]() {
+      for (int i = 0; i < num_matrices; i++) {
+        general_inv<T>(inv_ptr + N * N * i, N);
+      }
+    });
+  }
+}
+
+void Inverse::eval_cpu(const std::vector<array>& inputs, array& output) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      inverse_impl<float>(inputs[0], output, tri_, upper_, stream());
+      break;
+    case float64:
+      inverse_impl<double>(inputs[0], output, tri_, upper_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Inverse::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/jit_compiler.cpp

@@ -0,0 +1,166 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/cpu/jit_compiler.h"
+
+#include <algorithm>
+#include <sstream>
+#include <vector>
+
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+#ifdef _MSC_VER
+
+namespace {
+
+// Split string into array.
+std::vector<std::string> str_split(const std::string& str, char delimiter) {
+  std::vector<std::string> tokens;
+  std::string token;
+  std::istringstream tokenStream(str);
+  while (std::getline(tokenStream, token, delimiter)) {
+    tokens.push_back(token);
+  }
+  return tokens;
+}
+
+// Get path information about MSVC.
+struct VisualStudioInfo {
+  VisualStudioInfo() {
+#ifdef _M_ARM64
+    arch = "arm64";
+#else
+    arch = "x64";
+#endif
+    // Get path of Visual Studio.
+    // Use -latest to get only the most recent installation when multiple
+    // versions are installed, avoiding path concatenation issues.
+    std::string vs_path = JitCompiler::exec(
+        fmt::format(
+            "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
+            " -latest -property installationPath",
+            std::getenv("ProgramFiles(x86)")));
+    if (vs_path.empty()) {
+      throw std::runtime_error("Can not find Visual Studio.");
+    }
+    // Trim any trailing whitespace/newlines from the path
+    vs_path.erase(
+        std::find_if(
+            vs_path.rbegin(),
+            vs_path.rend(),
+            [](unsigned char ch) { return !std::isspace(ch); })
+            .base(),
+        vs_path.end());
+    // Read the envs from vcvarsall.
+    std::string envs = JitCompiler::exec(
+        fmt::format(
+            "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
+            vs_path,
+            arch));
+    for (const std::string& line : str_split(envs, '\n')) {
+      // Each line is in the format "ENV_NAME=values".
+      auto pos = line.find_first_of('=');
+      if (pos == std::string::npos || pos == 0 || pos == line.size() - 1)
+        continue;
+      std::string name = line.substr(0, pos);
+      std::string value = line.substr(pos + 1);
+      if (name == "LIB") {
+        libpaths = str_split(value, ';');
+      } else if (name == "VCToolsInstallDir" || name == "VCTOOLSINSTALLDIR") {
+        cl_exe = fmt::format("{0}\\bin\\Host{1}\\{1}\\cl.exe", value, arch);
+      }
+    }
+  }
+  std::string arch;
+  std::string cl_exe;
+  std::vector<std::string> libpaths;
+};
+
+const VisualStudioInfo& GetVisualStudioInfo() {
+  static VisualStudioInfo info;
+  return info;
+}
+
+} // namespace
+
+#endif // _MSC_VER
+
+std::string JitCompiler::build_command(
+    const std::filesystem::path& dir,
+    const std::string& source_file_name,
+    const std::string& shared_lib_name) {
+#ifdef _MSC_VER
+  const VisualStudioInfo& info = GetVisualStudioInfo();
+  std::string libpaths;
+  for (const std::string& lib : info.libpaths) {
+    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
+  }
+  return fmt::format(
+      "\""
+      "cd /D \"{0}\" && "
+      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
+      "/link /out:\"{3}\" {4} 2>&1"
+      "\"",
+      dir.string(),
+      info.cl_exe,
+      source_file_name,
+      shared_lib_name,
+      libpaths);
+#else
+  return fmt::format(
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared \"{0}\" -o \"{1}\" 2>&1",
+      (dir / source_file_name).string(),
+      (dir / shared_lib_name).string());
+#endif
+}
+
+std::string JitCompiler::exec(const std::string& cmd) {
+#ifdef _MSC_VER
+  FILE* pipe = _popen(cmd.c_str(), "r");
+#else
+  FILE* pipe = popen(cmd.c_str(), "r");
+#endif
+  if (!pipe) {
+    throw std::runtime_error("popen() failed.");
+  }
+  char buffer[128];
+  std::string ret;
+  while (fgets(buffer, sizeof(buffer), pipe)) {
+    ret += buffer;
+  }
+  // Trim trailing spaces.
+  ret.erase(
+      std::find_if(
+          ret.rbegin(),
+          ret.rend(),
+          [](unsigned char ch) { return !std::isspace(ch); })
+          .base(),
+      ret.end());
+
+#ifdef _MSC_VER
+  int status = _pclose(pipe);
+#else
+  int status = pclose(pipe);
+#endif
+  if (status == -1) {
+    throw std::runtime_error("pclose() failed.");
+  }
+#if defined(_WIN32) || defined(__FreeBSD__)
+  int code = status;
+#else
+  int code = WEXITSTATUS(status);
+#endif
+  if (code != 0) {
+    throw std::runtime_error(
+        fmt::format(
+            "Failed to execute command with return code {0}: \"{1}\", "
+            "the output is: {2}",
+            code,
+            cmd,
+            ret));
+  }
+  return ret;
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/jit_compiler.h

@@ -0,0 +1,20 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <filesystem>
+
+namespace mlx::core {
+
+class JitCompiler {
+ public:
+  // Build a shell command that compiles a source code file to a shared library.
+  static std::string build_command(
+      const std::filesystem::path& dir,
+      const std::string& source_file_name,
+      const std::string& shared_lib_name);
+
+  // Run a command and get its output.
+  static std::string exec(const std::string& cmd);
+};
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/lapack.h

@@ -0,0 +1,80 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <complex>
+#define LAPACK_COMPLEX_CUSTOM
+#define lapack_complex_float std::complex<float>
+#define lapack_complex_double std::complex<double>
+#define lapack_complex_float_real(z) ((z).real())
+#define lapack_complex_float_imag(z) ((z).imag())
+#define lapack_complex_double_real(z) ((z).real())
+#define lapack_complex_double_imag(z) ((z).imag())
+
+#ifdef MLX_USE_ACCELERATE
+#include <Accelerate/Accelerate.h>
+#else
+#include <cblas.h>
+#include <lapack.h>
+#endif
+
+#if defined(LAPACK_GLOBAL) || defined(LAPACK_NAME)
+
+// This is to work around a change in the function signatures of lapack >= 3.9.1
+// where functions taking char* also include a strlen argument, see a similar
+// change in OpenCV:
+// https://github.com/opencv/opencv/blob/1eb061f89de0fb85c4c75a2deeb0f61a961a63ad/cmake/OpenCVFindLAPACK.cmake#L57
+#define MLX_LAPACK_FUNC(f) LAPACK_##f
+
+#else
+
+#define MLX_LAPACK_FUNC(f) f##_
+
+#endif
+
+#define INSTANTIATE_LAPACK_REAL(FUNC)                        \
+  template <typename T, typename... Args>                    \
+  void FUNC(Args... args) {                                  \
+    if constexpr (std::is_same_v<T, float>) {                \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...); \
+    } else if constexpr (std::is_same_v<T, double>) {        \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...); \
+    }                                                        \
+  }
+
+INSTANTIATE_LAPACK_REAL(geqrf)
+INSTANTIATE_LAPACK_REAL(orgqr)
+INSTANTIATE_LAPACK_REAL(syevd)
+INSTANTIATE_LAPACK_REAL(potrf)
+INSTANTIATE_LAPACK_REAL(getrf)
+INSTANTIATE_LAPACK_REAL(getri)
+INSTANTIATE_LAPACK_REAL(trtri)
+
+#define INSTANTIATE_LAPACK_COMPLEX(FUNC)                            \
+  template <typename T, typename... Args>                           \
+  void FUNC(Args... args) {                                         \
+    if constexpr (std::is_same_v<T, std::complex<float>>) {         \
+      MLX_LAPACK_FUNC(c##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<double>>) { \
+      MLX_LAPACK_FUNC(z##FUNC)(std::forward<Args>(args)...);        \
+    }                                                               \
+  }
+
+INSTANTIATE_LAPACK_COMPLEX(heevd)
+
+#define INSTANTIATE_LAPACK_ALL(FUNC)                                \
+  template <typename T, typename... Args>                           \
+  void FUNC(Args... args) {                                         \
+    if constexpr (std::is_same_v<T, float>) {                       \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, double>) {               \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<float>>) {  \
+      MLX_LAPACK_FUNC(c##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<double>>) { \
+      MLX_LAPACK_FUNC(z##FUNC)(std::forward<Args>(args)...);        \
+    }                                                               \
+  }
+
+INSTANTIATE_LAPACK_ALL(geev)
+INSTANTIATE_LAPACK_ALL(gesdd)

data/mlx/mlx/backend/cpu/logsumexp.cpp

@@ -0,0 +1,139 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cassert>
+#include <cmath>
+
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
+#include "mlx/types/limits.h"
+
+namespace mlx::core {
+
+namespace {
+
+using namespace mlx::core::simd;
+
+template <typename T, typename AccT>
+void logsumexp(const array& in, array& out, Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
+  const T* in_ptr = in.data<T>();
+  T* out_ptr = out.data<T>();
+
+  int M = in.shape().back();
+  int L = in.data_size() / M;
+
+  encoder.dispatch([in_ptr, out_ptr, M, L]() mutable {
+    constexpr int N = std::min(max_size<AccT>, max_size<T>);
+
+    const T* current_in_ptr;
+
+    for (int i = 0; i < L; i++, in_ptr += M, out_ptr += 1) {
+      // Find the maximum
+      current_in_ptr = in_ptr;
+      Simd<AccT, N> vmaximum(-numeric_limits<AccT>::infinity());
+      size_t s = M;
+      while (s >= N) {
+        Simd<AccT, N> vals = load<T, N>(current_in_ptr);
+        vmaximum = maximum(vals, vmaximum);
+        current_in_ptr += N;
+        s -= N;
+      }
+
+      AccT maximum = max(vmaximum);
+      while (s-- > 0) {
+        maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
+        current_in_ptr++;
+      }
+
+      // Compute the normalizer and the exponentials
+      Simd<AccT, N> vnormalizer(0.0);
+      current_in_ptr = in_ptr;
+      s = M;
+      while (s >= N) {
+        Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
+        vexp = exp(vexp - maximum);
+        vnormalizer = vnormalizer + vexp;
+        current_in_ptr += N;
+        s -= N;
+      }
+      AccT normalizer = sum(vnormalizer);
+      while (s-- > 0) {
+        AccT _exp = std::exp(*current_in_ptr - maximum);
+        normalizer += _exp;
+        current_in_ptr++;
+      }
+      // Normalize
+      *out_ptr = std::isinf(maximum)
+          ? static_cast<T>(maximum)
+          : static_cast<T>(std::log(normalizer) + maximum);
+    }
+  });
+}
+
+} // namespace
+
+void LogSumExp::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  // Make sure that the last dimension is contiguous
+  auto s = stream();
+  auto& encoder = cpu::get_command_encoder(s);
+  auto ensure_contiguous = [&s, &encoder](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_cpu(x, s);
+      encoder.add_temporary(x_copy);
+      return x_copy;
+    }
+  };
+
+  auto in = ensure_contiguous(inputs[0]);
+  if (in.flags().row_contiguous) {
+    out.set_data(allocator::malloc(out.nbytes()));
+  } else {
+    auto n = in.shape(-1);
+    auto flags = in.flags();
+    auto strides = in.strides();
+    for (auto& s : strides) {
+      s /= n;
+    }
+    bool col_contig = strides[0] == 1;
+    for (int i = 1; col_contig && i < strides.size(); ++i) {
+      col_contig &=
+          (out.shape(i) == 1 || strides[i - 1] == out.shape(i) * strides[i]);
+    }
+    flags.col_contiguous = col_contig;
+    out.set_data(
+        allocator::malloc(in.nbytes() / n),
+        in.data_size() / n,
+        std::move(strides),
+        flags);
+  }
+
+  switch (in.dtype()) {
+    case float32:
+      logsumexp<float, float>(in, out, stream());
+      break;
+    case float16:
+      logsumexp<float16_t, float>(in, out, stream());
+      break;
+    case bfloat16:
+      logsumexp<bfloat16_t, float>(in, out, stream());
+      break;
+    case float64:
+      logsumexp<double, double>(in, out, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[logsumexp] only supports floating point types");
+      break;
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/luf.cpp

@@ -0,0 +1,120 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#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>
+void luf_impl(
+    const array& a,
+    array& lu,
+    array& pivots,
+    array& row_indices,
+    Stream stream) {
+  int M = a.shape(-2);
+  int N = a.shape(-1);
+  int K = std::min(M, N);
+
+  // Copy a into lu and make it col contiguous
+  auto ndim = lu.ndim();
+  auto flags = lu.flags();
+  flags.col_contiguous = ndim == 2;
+  flags.row_contiguous = false;
+  flags.contiguous = true;
+  auto strides = lu.strides();
+  strides[ndim - 1] = M;
+  strides[ndim - 2] = 1;
+  lu.set_data(allocator::malloc(lu.nbytes()), lu.nbytes(), strides, flags);
+  copy_cpu_inplace(
+      a,
+      lu,
+      a.shape(),
+      a.strides(),
+      strides,
+      0,
+      0,
+      CopyType::GeneralGeneral,
+      stream);
+
+  auto a_ptr = lu.data<T>();
+  pivots.set_data(allocator::malloc(pivots.nbytes()));
+  row_indices.set_data(allocator::malloc(row_indices.nbytes()));
+  auto pivots_ptr = pivots.data<uint32_t>();
+  auto row_indices_ptr = row_indices.data<uint32_t>();
+  size_t num_matrices = a.size() / (M * N);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(lu);
+  encoder.set_output_array(pivots);
+  encoder.set_output_array(row_indices);
+
+  encoder.dispatch(
+      [a_ptr, pivots_ptr, row_indices_ptr, num_matrices, M, N, K]() mutable {
+        int info;
+        for (size_t i = 0; i < num_matrices; ++i) {
+          // Compute LU factorization of A
+          getrf<T>(
+              /* m */ &M,
+              /* n */ &N,
+              /* a */ a_ptr,
+              /* lda */ &M,
+              /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
+              /* info */ &info);
+
+          if (info != 0) {
+            std::stringstream ss;
+            ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
+               << ((info > 0) ? " because matrix is singular"
+                              : " because argument had an illegal value");
+            throw std::runtime_error(ss.str());
+          }
+
+          // Subtract 1 to get 0-based index
+          int j = 0;
+          for (; j < K; ++j) {
+            pivots_ptr[j]--;
+            row_indices_ptr[j] = j;
+          }
+          for (; j < M; ++j) {
+            row_indices_ptr[j] = j;
+          }
+          for (int j = K - 1; j >= 0; --j) {
+            auto piv = pivots_ptr[j];
+            auto t1 = row_indices_ptr[piv];
+            auto t2 = row_indices_ptr[j];
+            row_indices_ptr[j] = t1;
+            row_indices_ptr[piv] = t2;
+          }
+
+          // Advance pointers to the next matrix
+          a_ptr += M * N;
+          pivots_ptr += K;
+          row_indices_ptr += M;
+        }
+      });
+}
+
+void LUF::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  switch (inputs[0].dtype()) {
+    case float32:
+      luf_impl<float>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      break;
+    case float64:
+      luf_impl<double>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[LUF::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core