mlx 1.0.0

data/mlx/mlx/backend/cpu/binary_two.h

@@ -0,0 +1,166 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/binary.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, typename U, typename Op, int D>
+void binary_op_dims(
+    const T* a,
+    const T* b,
+    U* out_a,
+    U* out_b,
+    Op op,
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      binary_op_dims<T, U, Op, D - 1>(
+          a,
+          b,
+          out_a,
+          out_b,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          axis + 1);
+    } else {
+      std::tie(*out_a, *out_b) = op(*a, *b);
+    }
+    a += stride_a;
+    b += stride_b;
+    out_a += stride_out;
+    out_b += stride_out;
+  }
+}
+
+template <typename T, typename U, typename Op>
+void binary_op_dispatch_dims(
+    const array& a,
+    const array& b,
+    array& out_a,
+    array& out_b,
+    Op op) {
+  auto [shape, strides] = collapse_contiguous_dims(
+      a.shape(), {a.strides(), b.strides(), out_a.strides()});
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  U* out_a_ptr = out_a.data<U>();
+  U* out_b_ptr = out_b.data<U>();
+
+  const auto& a_strides = strides[0];
+  const auto& b_strides = strides[1];
+  const auto& out_strides = strides[2];
+  int ndim = shape.size();
+  switch (ndim) {
+    case 1:
+      binary_op_dims<T, U, Op, 1>(
+          a_ptr,
+          b_ptr,
+          out_a_ptr,
+          out_b_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
+      return;
+    case 2:
+      binary_op_dims<T, U, Op, 2>(
+          a_ptr,
+          b_ptr,
+          out_a_ptr,
+          out_b_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
+      return;
+  }
+
+  ContiguousIterator a_it(shape, a_strides, ndim - 2);
+  ContiguousIterator b_it(shape, b_strides, ndim - 2);
+  auto stride = out_strides[ndim - 3];
+  for (size_t elem = 0; elem < a.size(); elem += stride) {
+    binary_op_dims<T, U, Op, 2>(
+        a_ptr + a_it.loc,
+        b_ptr + b_it.loc,
+        out_a_ptr + elem,
+        out_b_ptr + elem,
+        op,
+        shape,
+        a_strides,
+        b_strides,
+        out_strides,
+        ndim - 2);
+    a_it.step();
+    b_it.step();
+  }
+}
+
+template <typename T, typename U = T, typename Op>
+void binary_op(
+    const array& a,
+    const array& b,
+    array& out_a,
+    array& out_b,
+    Op op,
+    BinaryOpType bopt) {
+  // The full computation is scalar scalar so call the base op once
+  if (bopt == BinaryOpType::General) {
+    binary_op_dispatch_dims<T, U, Op>(a, b, out_a, out_b, op);
+    return;
+  }
+
+  auto a_ptr = a.data<T>();
+  auto b_ptr = b.data<T>();
+  auto out_a_ptr = out_a.data<U>();
+  auto out_b_ptr = out_b.data<U>();
+  if (bopt == BinaryOpType::ScalarScalar) {
+    std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+  } else if (bopt == BinaryOpType::ScalarVector) {
+    for (size_t i = 0; i < b.data_size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      b_ptr++;
+    }
+  } else if (bopt == BinaryOpType::VectorScalar) {
+    for (size_t i = 0; i < a.data_size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      a_ptr++;
+    }
+  } else { // VectorVector
+    for (size_t i = 0; i < a.size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      a_ptr++;
+      b_ptr++;
+    }
+  }
+}
+
+} // namespace
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/cholesky.cpp

@@ -0,0 +1,85 @@
+// 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/linalg.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) {
+  // 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
+
+  // The decomposition is computed in place, so just copy the input to the
+  // output.
+  copy_cpu(
+      a,
+      factor,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(factor);
+  encoder.dispatch([matrix = factor.data<T>(),
+                    upper,
+                    N = a.shape(-1),
+                    size = a.size()]() mutable {
+    char uplo = (upper) ? 'L' : 'U';
+    size_t num_matrices = size / (N * N);
+    for (int i = 0; i < num_matrices; i++) {
+      // Compute Cholesky factorization.
+      int info;
+      potrf<T>(
+          /* uplo = */ &uplo,
+          /* n = */ &N,
+          /* a = */ matrix,
+          /* lda = */ &N,
+          /* info = */ &info);
+
+      // 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::eval_cpu] 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_cpu(const std::vector<array>& inputs, array& output) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      cholesky_impl<float>(inputs[0], output, upper_, stream());
+      break;
+    case float64:
+      cholesky_impl<double>(inputs[0], output, upper_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Cholesky::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/compiled.cpp

@@ -0,0 +1,357 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <dlfcn.h>
+#include <filesystem>
+#include <fstream>
+#include <list>
+#include <mutex>
+#include <shared_mutex>
+
+#include <fmt/format.h>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cpu/compiled_preamble.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/jit_compiler.h"
+#include "mlx/device.h"
+#include "mlx/graph_utils.h"
+#include "mlx/version.h"
+
+namespace mlx::core {
+
+struct CompilerCache {
+  struct DLib {
+    DLib(const std::string& libname) {
+      lib = dlopen(libname.c_str(), RTLD_NOW);
+      if (!lib) {
+        std::ostringstream msg;
+        msg << "Could not load C++ shared library " << dlerror();
+        throw std::runtime_error(msg.str());
+      }
+    }
+
+    ~DLib() {
+      dlclose(lib);
+    }
+    void* lib;
+  };
+  // Statics to cache compiled libraries and functions
+  std::list<DLib> libs;
+  std::unordered_map<std::string, void*> kernels;
+  std::shared_mutex mtx;
+};
+
+static CompilerCache& cache() {
+  static CompilerCache cache_;
+  return cache_;
+};
+
+// GPU compile is always available if the GPU is available and since we are in
+// this file CPU compile is also available.
+namespace detail {
+bool compile_available_for_device(const Device& device) {
+  return true;
+}
+
+} // namespace detail
+
+// Return a pointer to a compiled function
+void* compile(
+    const std::string& kernel_name,
+    const std::function<std::string(void)>& source_builder) {
+  {
+    std::shared_lock lock(cache().mtx);
+    if (auto it = cache().kernels.find(kernel_name);
+        it != cache().kernels.end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock lock(cache().mtx);
+  if (auto it = cache().kernels.find(kernel_name);
+      it != cache().kernels.end()) {
+    return it->second;
+  }
+  std::string source_code = source_builder();
+  std::string kernel_file_name;
+
+  // Deal with long kernel names. Maximum length for filename on macOS is 255
+  // characters, and on Windows the maximum length for whole path is 260. Clip
+  // file name with a little extra room and append a 16 character hash.
+#ifdef _WIN32
+  constexpr int max_file_name_length = 140;
+#else
+  constexpr int max_file_name_length = 245;
+#endif
+  if (kernel_name.size() > max_file_name_length) {
+    std::ostringstream file_name;
+    file_name
+        << std::string_view(kernel_name).substr(0, max_file_name_length - 16);
+    auto file_id =
+        std::hash<std::string>{}(kernel_name.substr(max_file_name_length - 16));
+    file_name << "_" << std::hex << std::setw(16) << file_id << std::dec;
+    kernel_file_name = file_name.str();
+  } else {
+    kernel_file_name = kernel_name;
+  }
+
+  auto output_dir =
+      std::filesystem::temp_directory_path() / "mlx" / version() / "cpu";
+  if (!std::filesystem::exists(output_dir)) {
+    std::filesystem::create_directories(output_dir);
+  }
+
+  std::string shared_lib_name = "lib" + kernel_file_name + ".so";
+  auto shared_lib_path = (output_dir / shared_lib_name).string();
+  bool lib_exists = false;
+  {
+    std::ifstream f(shared_lib_path.c_str());
+    lib_exists = f.good();
+  }
+
+  if (!lib_exists) {
+    // Open source file and write source code to it
+    std::string source_file_name = kernel_file_name + ".cpp";
+    auto source_file_path = (output_dir / source_file_name).string();
+
+    std::ofstream source_file(source_file_path);
+    source_file << source_code;
+    source_file.close();
+
+    try {
+      JitCompiler::exec(
+          JitCompiler::build_command(
+              output_dir, source_file_name, shared_lib_name));
+    } catch (const std::exception& error) {
+      throw std::runtime_error(
+          fmt::format(
+              "[Compile::eval_cpu] Failed to compile function {0}: {1}",
+              kernel_name,
+              error.what()));
+    }
+  }
+
+  // load library
+  cache().libs.emplace_back(shared_lib_path);
+
+  // Load function
+  void* fun = dlsym(cache().libs.back().lib, kernel_name.c_str());
+  if (!fun) {
+    std::ostringstream msg;
+    msg << "[Compile::eval_cpu] Failed to load compiled function "
+        << kernel_name << std::endl
+        << dlerror();
+    throw std::runtime_error(msg.str());
+  }
+  cache().kernels.insert({kernel_name, fun});
+  return fun;
+}
+
+inline void build_kernel(
+    std::ostream& os,
+    const std::string& kernel_name,
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::function<bool(size_t)>& is_constant,
+    bool contiguous,
+    int ndim) {
+  NodeNamer namer;
+
+#ifdef _MSC_VER
+  // Export the symbol
+  os << "__declspec(dllexport) ";
+#endif
+
+  // Start the kernel
+  os << "void " << kernel_name
+     << "(int* shape, int64_t** strides, void** args) {" << std::endl;
+
+  // Add the input arguments
+  int cnt = 0;
+  int strides_index = 1;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    // Skip constants from the input list
+    if (is_constant(i)) {
+      continue;
+    }
+
+    const auto& x = inputs[i];
+    auto& xname = namer.get_name(x);
+
+    auto tstr = get_type_string(x.dtype());
+    os << "  " << tstr << "* " << xname << " = (" << tstr << "*)args[" << cnt++
+       << "];" << std::endl;
+    // Scalars and contiguous need no strides
+    if (!is_scalar(x) && !contiguous) {
+      os << "  const int64_t* " << xname << "_strides = strides["
+         << strides_index++ << "];" << std::endl;
+    }
+  }
+
+  // Add the output arguments
+  for (auto& x : outputs) {
+    auto tstr = get_type_string(x.dtype());
+    os << "  " << tstr << "* " << namer.get_name(x) << " = (" << tstr
+       << "*)args[" << cnt++ << "];" << std::endl;
+  }
+  // Add output size
+  if (contiguous) {
+    os << "  const size_t size = (size_t)args[" << cnt++ << "];" << std::endl;
+  }
+
+  if (contiguous) {
+    os << "  for (size_t i = 0; i < size; ++i) {" << std::endl;
+  } else {
+    for (int d = 0; d < ndim; ++d) {
+      os << "  for (int i" << d << " = 0; i" << d << " < shape[" << d
+         << "]; ++i" << d << ") {" << std::endl;
+    }
+  }
+
+  // Read the inputs in tmps
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    const auto& x = inputs[i];
+    auto& xname = namer.get_name(x);
+
+    if (is_constant(i)) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = ";
+      print_constant(os, x);
+      os << ";" << std::endl;
+    } else if (is_scalar(x)) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = "
+         << xname << "[0];" << std::endl;
+    } else if (contiguous) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = "
+         << xname << "[i];" << std::endl;
+    } else {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = *"
+         << xname << ";" << std::endl;
+    }
+  }
+
+  // Actually write the computation
+  for (auto& x : tape) {
+    os << "  " << get_type_string(x.dtype()) << " tmp_" << namer.get_name(x)
+       << " = ";
+    if (is_static_cast(x.primitive())) {
+      os << "static_cast<" << get_type_string(x.dtype()) << ">(tmp_"
+         << namer.get_name(x.inputs()[0]) << ");" << std::endl;
+    } else {
+      os << x.primitive().name();
+      os << "()(";
+      for (int i = 0; i < x.inputs().size() - 1; i++) {
+        os << "tmp_" << namer.get_name(x.inputs()[i]) << ", ";
+      }
+      os << "tmp_" << namer.get_name(x.inputs().back()) << ");" << std::endl;
+    }
+  }
+
+  // Write the outputs from tmps
+  for (auto& x : outputs) {
+    if (contiguous) {
+      os << "  " << namer.get_name(x) << "[i] = tmp_" << namer.get_name(x)
+         << ";" << std::endl;
+    } else {
+      os << "  *" << namer.get_name(x) << "++ = tmp_" << namer.get_name(x)
+         << ";" << std::endl;
+    }
+  }
+
+  // Close loops
+  if (contiguous) {
+    os << "  }" << std::endl;
+  } else {
+    for (int d = ndim - 1; d >= 0; --d) {
+      // Update pointers
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        if (is_constant(i) || is_scalar(x)) {
+          continue;
+        }
+        auto& xname = namer.get_name(x);
+        os << "  " << xname << " += " << xname << "_strides[" << d << "];"
+           << std::endl;
+        if (d < ndim - 1) {
+          os << "  " << xname << " -= " << xname << "_strides[" << d + 1 << "]"
+             << " * shape[" << d + 1 << "];" << std::endl;
+        }
+      }
+      os << "  }" << std::endl;
+    }
+  }
+
+  // Finish the kernel
+  os << "}" << std::endl;
+}
+
+void Compiled::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  // Collapse contiguous dims to route to a faster kernel if possible. Also
+  // handle all broadcasting.
+  auto [contiguous, shape, strides] =
+      compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
+
+  // Collect function input arguments.
+  std::vector<void*> args;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    if (is_constant_(i)) {
+      continue;
+    }
+    const auto& x = inputs[i];
+    encoder.set_input_array(x);
+    args.push_back((void*)x.data<void>());
+  }
+
+  // Get the kernel name from the lib
+  int ndim = shape.size();
+  auto kernel_name = kernel_lib_ + (contiguous ? "_contiguous" : "_strided_");
+  if (!contiguous) {
+    kernel_name += std::to_string(ndim);
+  }
+
+  // Get the function
+  auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() {
+    std::ostringstream kernel;
+    kernel << get_kernel_preamble() << std::endl;
+    kernel << "extern \"C\"  {" << std::endl;
+    build_kernel(
+        kernel,
+        kernel_name,
+        inputs_,
+        outputs_,
+        tape_,
+        is_constant_,
+        contiguous,
+        ndim);
+    // Close extern "C"
+    kernel << "}" << std::endl;
+    return kernel.str();
+  });
+
+  compiled_allocate_outputs(inputs, outputs, is_constant_, contiguous);
+
+  for (auto& x : outputs) {
+    args.push_back(x.data<void>());
+    encoder.set_output_array(x);
+  }
+  if (contiguous) {
+    args.push_back((void*)outputs[0].data_size());
+  }
+  auto fun = reinterpret_cast<void (*)(int*, int64_t**, void**)>(fn_ptr);
+  encoder.dispatch([fun,
+                    args = std::move(args),
+                    strides = std::move(strides),
+                    shape = std::move(shape)]() mutable {
+    SmallVector<int64_t*> strides_ptrs;
+    for (auto& s : strides) {
+      strides_ptrs.push_back(s.data());
+    }
+    fun(shape.data(), strides_ptrs.data(), args.data());
+  });
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/compiled_preamble.h

@@ -0,0 +1,12 @@
+// Copyright © 2023-24 Apple Inc.
+
+#pragma once
+
+// clang-format off
+#include "mlx/types/half_types.h"
+#include "mlx/types/complex.h"
+#include "mlx/backend/cpu/unary_ops.h"
+#include "mlx/backend/cpu/binary_ops.h"
+// clang-format on
+
+const char* get_kernel_preamble();