mlx 0.30.7

data/mlx/mlx/backend/common/utils.h

@@ -0,0 +1,205 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <filesystem>
+#include <tuple>
+#include <vector>
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+// Return the directory that contains current shared library.
+std::filesystem::path current_binary_dir();
+
+inline int64_t
+elem_to_loc(int elem, const Shape& shape, const Strides& strides) {
+  int64_t loc = 0;
+  for (int i = shape.size() - 1; i >= 0; --i) {
+    auto q_and_r = ldiv(elem, shape[i]);
+    loc += q_and_r.rem * strides[i];
+    elem = q_and_r.quot;
+  }
+  return loc;
+}
+
+inline int64_t elem_to_loc(int elem, const array& a) {
+  if (a.flags().row_contiguous) {
+    return elem;
+  }
+  return elem_to_loc(elem, a.shape(), a.strides());
+}
+
+inline Strides make_contiguous_strides(const Shape& shape) {
+  Strides strides(shape.size(), 1);
+  for (int i = shape.size() - 1; i > 0; i--) {
+    strides[i - 1] = strides[i] * shape[i];
+  }
+  return strides;
+}
+
+// Collapse dims that are contiguous to possibly route to a better kernel
+// e.g. for x = transpose(array({0, 1, 2, 3, 4, 5, 6, 7}, {2, 2, 2}), {2, 0, 1})
+// should return {{2, 4}, {{1, 2}}}.
+//
+// When multiple arrays are passed they should all have the same shape. The
+// collapsed axes are also the same so one shape is returned.
+std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
+    const Shape& shape,
+    const std::vector<Strides>& strides,
+    int64_t size_cap = std::numeric_limits<int32_t>::max());
+
+inline std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
+    const std::vector<array>& xs,
+    size_t size_cap = std::numeric_limits<int32_t>::max()) {
+  std::vector<Strides> strides;
+  for (auto& x : xs) {
+    strides.emplace_back(x.strides());
+  }
+  return collapse_contiguous_dims(xs[0].shape(), strides, size_cap);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline auto collapse_contiguous_dims(Arrays&&... xs) {
+  return collapse_contiguous_dims(
+      std::vector<array>{std::forward<Arrays>(xs)...});
+}
+
+// The single array version of the above.
+std::pair<Shape, Strides> collapse_contiguous_dims(
+    const Shape& shape,
+    const Strides& strides,
+    int64_t size_cap = std::numeric_limits<int32_t>::max());
+std::pair<Shape, Strides> collapse_contiguous_dims(
+    const array& a,
+    int64_t size_cap = std::numeric_limits<int32_t>::max());
+
+// Compute the thread block dimensions which fit the given
+// input dimensions.
+// - The thread block dimensions will be powers of two
+// - The thread block size will be less than 2^pow2
+using Dims = std::tuple<uint32_t, uint32_t, uint32_t>;
+Dims get_block_dims_common(int dim0, int dim1, int dim2, int pow2 = 10);
+
+// Computes a 2D grid where each element is < UINT_MAX
+// Assumes:
+// - overall size (product of non-broadcasted dimensions) is < UINT_MAX^2
+// - shape and strides correspond to a contiguous (no holes) but
+//   possibly broadcasted array
+Dims get_2d_grid_dims_common(const Shape& shape, const Strides& strides);
+
+// Same as above but we do an implicit division with divisor.
+// Basically, equivalent to factorizing
+//    Prod(s \forall s in shape if strides[s] > 0) / divisor.
+Dims get_2d_grid_dims_common(
+    const Shape& shape,
+    const Strides& strides,
+    size_t divisor);
+
+// Get both the block and a grid of blocks that covers dim0, dim1 and dim2.
+std::pair<Dims, Dims> get_grid_and_block_common(int dim0, int dim1, int dim2);
+
+struct ContiguousIterator {
+  inline void step() {
+    int dims = shape_.size();
+    if (dims == 0) {
+      return;
+    }
+    int i = dims - 1;
+    while (pos_[i] == (shape_[i] - 1) && i > 0) {
+      pos_[i] = 0;
+      loc -= (shape_[i] - 1) * strides_[i];
+      i--;
+    }
+    pos_[i]++;
+    loc += strides_[i];
+  }
+
+  void seek(int64_t n) {
+    loc = 0;
+    for (int i = shape_.size() - 1; i >= 0; --i) {
+      auto q_and_r = ldiv(n, shape_[i]);
+      loc += q_and_r.rem * strides_[i];
+      pos_[i] = q_and_r.rem;
+      n = q_and_r.quot;
+    }
+  }
+
+  void reset() {
+    loc = 0;
+    std::fill(pos_.begin(), pos_.end(), 0);
+  }
+
+  ContiguousIterator() {};
+
+  explicit ContiguousIterator(const array& a)
+      : shape_(a.shape()), strides_(a.strides()) {
+    if (!shape_.empty()) {
+      std::tie(shape_, strides_) = collapse_contiguous_dims(shape_, strides_);
+      pos_ = Shape(shape_.size(), 0);
+    }
+  }
+
+  explicit ContiguousIterator(
+      const Shape& shape,
+      const Strides& strides,
+      int dims)
+      : shape_(shape.begin(), shape.begin() + dims),
+        strides_(strides.begin(), strides.begin() + dims) {
+    if (!shape_.empty()) {
+      std::tie(shape_, strides_) = collapse_contiguous_dims(shape_, strides_);
+      pos_ = Shape(shape_.size(), 0);
+    }
+  }
+
+  int64_t loc{0};
+
+ private:
+  Shape shape_;
+  Strides strides_;
+  Shape pos_;
+};
+
+inline auto check_contiguity(const Shape& shape, const Strides& strides) {
+  size_t no_broadcast_data_size = 1;
+  int64_t f_stride = 1;
+  int64_t b_stride = 1;
+  bool is_row_contiguous = true;
+  bool is_col_contiguous = true;
+
+  for (int i = 0, ri = shape.size() - 1; ri >= 0; i++, ri--) {
+    is_col_contiguous &= strides[i] == f_stride || shape[i] == 1;
+    is_row_contiguous &= strides[ri] == b_stride || shape[ri] == 1;
+    f_stride *= shape[i];
+    b_stride *= shape[ri];
+    if (strides[i] > 0) {
+      no_broadcast_data_size *= shape[i];
+    }
+  }
+
+  return std::make_tuple(
+      no_broadcast_data_size, is_row_contiguous, is_col_contiguous);
+}
+
+inline bool is_donatable(const array& in, const array& out) {
+  constexpr size_t donation_extra = 16384;
+
+  return in.is_donatable() && in.itemsize() == out.itemsize() &&
+      in.buffer_size() <= out.nbytes() + donation_extra;
+}
+
+std::pair<bool, Strides> prepare_reshape(const array& in, const array& out);
+
+void shared_buffer_reshape(
+    const array& in,
+    const Strides& out_strides,
+    array& out);
+
+template <typename T>
+inline SmallVector<T> remove_index(SmallVector<T> vec, size_t index) {
+  vec.erase(std::next(vec.begin(), index));
+  return vec;
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/CMakeLists.txt

@@ -0,0 +1,88 @@
+if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+  set(COMPILER ${CMAKE_C_COMPILER})
+  set(CLANG TRUE)
+else()
+  set(COMPILER ${CMAKE_CXX_COMPILER})
+endif()
+
+set(COMPILE_DEPS
+    ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
+    ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
+    ${PROJECT_SOURCE_DIR}/mlx/types/bf16.h
+    ${PROJECT_SOURCE_DIR}/mlx/types/complex.h
+    simd/simd.h
+    simd/base_simd.h
+    simd/math.h
+    simd/type.h
+    unary_ops.h
+    binary_ops.h)
+
+if(MSVC)
+  set(SHELL_EXT ps1)
+  set(SHELL_CMD powershell -ExecutionPolicy Bypass -File)
+else()
+  set(SHELL_EXT sh)
+  set(SHELL_CMD bash)
+endif()
+
+add_custom_command(
+  OUTPUT compiled_preamble.cpp
+  COMMAND
+    ${SHELL_CMD} ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.${SHELL_EXT}
+    ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp ${COMPILER}
+    ${PROJECT_SOURCE_DIR} ${CLANG} ${CMAKE_SYSTEM_PROCESSOR}
+  DEPENDS make_compiled_preamble.${SHELL_EXT} compiled_preamble.h
+          ${COMPILE_DEPS})
+
+add_custom_target(cpu_compiled_preamble DEPENDS compiled_preamble.cpp)
+
+add_dependencies(mlx cpu_compiled_preamble)
+
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eig.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eigh.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/encoder.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cblas.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/masked_mm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/luf.cpp
+          ${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_SOURCE_DIR}/unary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
+          ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp)
+
+if(MLX_BUILD_ACCELERATE)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/bnns.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/simd_fp16.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/gemms/simd_bf16.cpp)
+endif()
+
+if(IOS)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../no_cpu/compiled.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/jit_compiler.cpp)
+endif()

data/mlx/mlx/backend/cpu/arange.h

@@ -0,0 +1,28 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T>
+void arange(T start, T next, array& out, size_t size, Stream stream) {
+  auto ptr = out.data<T>();
+  auto step_size = next - start;
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([ptr, start, step_size, size]() mutable {
+    for (int i = 0; i < size; ++i) {
+      ptr[i] = start;
+      start += step_size;
+    }
+  });
+}
+
+} // namespace
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/arg_reduce.cpp

@@ -0,0 +1,124 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename InT, typename OpT>
+void arg_reduce(const array& in, array& out, const OpT& op, int axis) {
+  auto axis_size = in.shape()[axis];
+  auto axis_stride = in.strides()[axis];
+  Strides strides = remove_index(in.strides(), axis);
+  Shape shape = remove_index(in.shape(), axis);
+  auto in_ptr = in.data<InT>();
+  auto out_ptr = out.data<uint32_t>();
+
+  for (uint32_t i = 0; i < out.size(); ++i) {
+    auto loc = elem_to_loc(i, shape, strides);
+    auto local_in_ptr = in_ptr + loc;
+    uint32_t ind_v = 0;
+    InT v = (*local_in_ptr);
+    for (uint32_t j = 0; j < axis_size; ++j, local_in_ptr += axis_stride) {
+      op(j, (*local_in_ptr), &ind_v, &v);
+    }
+    out_ptr[i] = ind_v;
+  }
+}
+
+template <typename InT>
+void arg_reduce_dispatch(
+    const array& in,
+    array& out,
+    ArgReduce::ReduceType rtype,
+    int axis) {
+  switch (rtype) {
+    case ArgReduce::ArgMin: {
+      auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
+        if (x < (*y)) {
+          (*y) = x;
+          (*ind_y) = ind_x;
+        }
+      };
+      arg_reduce<InT>(in, out, op, axis);
+      break;
+    }
+    case ArgReduce::ArgMax: {
+      auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
+        if (x > (*y)) {
+          (*y) = x;
+          (*ind_y) = ind_x;
+        }
+      };
+      arg_reduce<InT>(in, out, op, axis);
+      break;
+    }
+  }
+}
+
+} // namespace
+
+void ArgReduce::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    reduce_type_ = reduce_type_,
+                    axis_ = axis_]() mutable {
+    switch (in.dtype()) {
+      case bool_:
+        arg_reduce_dispatch<bool>(in, out, reduce_type_, axis_);
+        break;
+      case uint8:
+        arg_reduce_dispatch<uint8_t>(in, out, reduce_type_, axis_);
+        break;
+      case uint16:
+        arg_reduce_dispatch<uint16_t>(in, out, reduce_type_, axis_);
+        break;
+      case uint32:
+        arg_reduce_dispatch<uint32_t>(in, out, reduce_type_, axis_);
+        break;
+      case uint64:
+        arg_reduce_dispatch<uint64_t>(in, out, reduce_type_, axis_);
+        break;
+      case int8:
+        arg_reduce_dispatch<int8_t>(in, out, reduce_type_, axis_);
+        break;
+      case int16:
+        arg_reduce_dispatch<int16_t>(in, out, reduce_type_, axis_);
+        break;
+      case int32:
+        arg_reduce_dispatch<int32_t>(in, out, reduce_type_, axis_);
+        break;
+      case int64:
+        arg_reduce_dispatch<int64_t>(in, out, reduce_type_, axis_);
+        break;
+      case float16:
+        arg_reduce_dispatch<float16_t>(in, out, reduce_type_, axis_);
+        break;
+      case float32:
+        arg_reduce_dispatch<float>(in, out, reduce_type_, axis_);
+        break;
+      case bfloat16:
+        arg_reduce_dispatch<bfloat16_t>(in, out, reduce_type_, axis_);
+        break;
+      case float64:
+        arg_reduce_dispatch<double>(in, out, reduce_type_, axis_);
+        break;
+      case complex64:
+        arg_reduce_dispatch<complex64_t>(in, out, reduce_type_, axis_);
+        break;
+    }
+  });
+}
+
+} // namespace mlx::core