mlx 0.30.7

data/mlx/mlx/backend/cpu/matmul.cpp

@@ -0,0 +1,166 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstring>
+#include "mlx/array.h"
+#include "mlx/backend/cpu/binary.h"
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void matmul_dispatch(
+    const array& a,
+    const array& b,
+    array& out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    float alpha,
+    float beta,
+    Stream stream) {
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  T* out_ptr = out.data<T>();
+  size_t ldc = out.shape(-1);
+  size_t batch_size = a.size() / (a.shape(-2) * a.shape(-1));
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a_ptr,
+                    b_ptr,
+                    out_ptr,
+                    a_transposed,
+                    b_transposed,
+                    lda,
+                    ldb,
+                    ldc,
+                    alpha,
+                    beta,
+                    batch_size,
+                    a_shape = a.shape(),
+                    a_strides = a.strides(),
+                    b_shape = b.shape(),
+                    b_strides = b.strides()]() {
+    matmul<T>(
+        a_ptr,
+        b_ptr,
+        out_ptr,
+        a_transposed,
+        b_transposed,
+        lda,
+        ldb,
+        ldc,
+        alpha,
+        beta,
+        batch_size,
+        a_shape,
+        a_strides,
+        b_shape,
+        b_strides);
+  });
+}
+
+void matmul_general(
+    const array& a_pre,
+    const array& b_pre,
+    array& out,
+    Stream stream,
+    float alpha = 1.0f,
+    float beta = 0.0f) {
+  std::vector<array> temps;
+  auto check_transpose = [stream, &temps](const array& arr) {
+    auto stx = arr.strides()[arr.ndim() - 2];
+    auto sty = arr.strides()[arr.ndim() - 1];
+    if (stx == arr.shape(-1) && sty == 1) {
+      return std::make_tuple(false, stx, arr);
+    } else if (stx == 1 && sty == arr.shape(-2)) {
+      return std::make_tuple(true, sty, arr);
+    } else {
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy_cpu(arr, temps.back(), CopyType::General, stream);
+      stx = arr.shape(-1);
+      return std::make_tuple(false, stx, temps.back());
+    }
+  };
+
+  auto [a_transposed, lda, a] = check_transpose(a_pre);
+  auto [b_transposed, ldb, b] = check_transpose(b_pre);
+  size_t M = a.shape(-2);
+  size_t N = b.shape(-1);
+  if (M == 0 || N == 0) {
+    return;
+  }
+
+  if (out.dtype() == float32) {
+    matmul_dispatch<float>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == float16) {
+    matmul_dispatch<float16_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == bfloat16) {
+    matmul_dispatch<bfloat16_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == float64) {
+    matmul_dispatch<double>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == complex64) {
+    matmul_dispatch<complex64_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else {
+    throw std::runtime_error("[Matmul::eval_cpu] Invalid type.");
+  }
+  cpu::get_command_encoder(stream).add_temporaries(std::move(temps));
+}
+
+void Matmul::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+  if (inputs[0].shape(-1) == 0) {
+    auto& encoder = cpu::get_command_encoder(stream());
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<void>(), nbytes = out.nbytes()]() {
+      std::memset(out_ptr, 0, nbytes);
+    });
+    return;
+  }
+  matmul_general(inputs[0], inputs[1], out, stream());
+}
+
+void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    return;
+  }
+
+  // Handle empty matrix case (K=0)
+  if (inputs[0].shape(-1) == 0) {
+    auto& c = inputs[2];
+    if (beta_ == 1.0f) {
+      CopyType ctype = c.data_size() == 1
+          ? CopyType::Scalar
+          : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+      copy_cpu(c, out, ctype, stream());
+    } else {
+      array beta_scalar = array(beta_, c.dtype());
+      auto& encoder = cpu::get_command_encoder(stream());
+      binary_float_op_cpu(c, beta_scalar, out, detail::Multiply(), stream());
+      encoder.add_temporary(std::move(beta_scalar));
+    }
+    return;
+  }
+
+  // Fill output with C
+  auto& c = inputs[2];
+  CopyType ctype = c.data_size() == 1
+      ? CopyType::Scalar
+      : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+  copy_cpu(c, out, ctype, stream());
+  matmul_general(inputs[0], inputs[1], out, stream(), alpha_, beta_);
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/primitives.cpp

@@ -0,0 +1,478 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <numeric>
+#include <sstream>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/slicing.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/arange.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/threefry.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+void reshape(const array& in, array& out) {
+  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
+  if (copy_necessary) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    copy_cpu_inplace(in, out, CopyType::General, out.primitive().stream());
+  } else {
+    shared_buffer_reshape(in, out_strides, out);
+  }
+}
+
+static std::pair<array, bool> compute_dynamic_offset(
+    const array& indices,
+    const Strides& strides,
+    const std::vector<int>& axes,
+    Stream stream) {
+  array offset({1}, int64, nullptr, {});
+  bool donate = indices.is_donatable() &&
+      (indices.data_size() * indices.itemsize()) >= offset.itemsize();
+  if (donate) {
+    offset.copy_shared_buffer(indices);
+  } else {
+    offset.set_data(allocator::malloc(offset.itemsize()));
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(indices);
+  encoder.set_output_array(offset);
+  auto compute_offset =
+      [strides, axes, offset = offset.data<int64_t>()](const auto* indices) {
+        int64_t offset_ = 0;
+        for (int i = 0; i < axes.size(); ++i) {
+          offset_ += indices[i] * strides[axes[i]];
+        }
+        offset[0] = offset_;
+      };
+  switch (indices.dtype()) {
+    case int8:
+    case uint8:
+      encoder.dispatch(compute_offset, indices.data<uint8_t>());
+      break;
+    case int16:
+    case uint16:
+      encoder.dispatch(compute_offset, indices.data<uint16_t>());
+      break;
+    case int32:
+    case uint32:
+      encoder.dispatch(compute_offset, indices.data<uint32_t>());
+      break;
+    case int64:
+    case uint64:
+      encoder.dispatch(compute_offset, indices.data<uint64_t>());
+      break;
+    default:
+      throw std::runtime_error("Invalid indices type.");
+  }
+  return {offset, donate};
+}
+
+void AsStrided::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Broadcast::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void BroadcastAxes::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Copy::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void CustomTransforms::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+void Depends::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+void ExpandDims::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void NumberOfElements::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Slice::eval_cpu(const std::vector<array>& inputs, array& out) {
+  slice(inputs[0], out, start_indices_, strides_);
+}
+void Split::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+void Squeeze::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void StopGradient::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Transpose::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+
+void Arange::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 0);
+  out.set_data(allocator::malloc(out.nbytes()));
+  switch (out.dtype()) {
+    case bool_:
+      throw std::runtime_error("Bool type unsupported for arange.");
+      break;
+    case uint8:
+      arange<uint8_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint16:
+      arange<uint16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint32:
+      arange<uint32_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint64:
+      arange<uint64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int8:
+      arange<int8_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int16:
+      arange<int16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int32:
+      arange<int32_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int64:
+      arange<int64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float16:
+      arange<float16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float32:
+      arange<float>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float64:
+      arange<double>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case bfloat16:
+      arange<bfloat16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case complex64:
+      arange<complex64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+  }
+}
+
+void AsType::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(in, out, ctype, stream());
+}
+
+void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  std::vector<int> sizes;
+  sizes.push_back(0);
+  for (auto& p : inputs) {
+    sizes.push_back(p.shape(axis_));
+  }
+  std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto strides = out.strides();
+  auto flags = out.flags();
+  flags.row_contiguous = false;
+  flags.col_contiguous = false;
+  flags.contiguous = false;
+  for (int i = 0; i < inputs.size(); i++) {
+    array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
+    size_t data_offset = strides[axis_] * sizes[i];
+    out_slice.copy_shared_buffer(
+        out, strides, flags, out_slice.size(), data_offset);
+    copy_cpu_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, stream());
+  }
+}
+
+void Contiguous::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  constexpr size_t extra_bytes = 16384;
+  if (in.buffer_size() <= out.nbytes() + extra_bytes &&
+      (in.flags().row_contiguous ||
+       (allow_col_major_ && in.flags().col_contiguous))) {
+    out.copy_shared_buffer(in);
+  } else {
+    copy_cpu(in, out, CopyType::General, stream());
+  }
+}
+
+void Flatten::eval_cpu(const std::vector<array>& inputs, array& out) {
+  reshape(inputs[0], out);
+}
+
+void Unflatten::eval_cpu(const std::vector<array>& inputs, array& out) {
+  reshape(inputs[0], out);
+}
+
+void Full::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  assert(in.dtype() == out.dtype());
+  CopyType ctype;
+  if (in.data_size() == 1) {
+    ctype = CopyType::Scalar;
+  } else if (in.flags().contiguous) {
+    ctype = CopyType::Vector;
+  } else {
+    ctype = CopyType::General;
+  }
+  copy_cpu(in, out, ctype, stream());
+}
+
+void Pad::eval_cpu(const std::vector<array>& inputs, array& out) {
+  // Inputs must be base input array and scalar val array
+  assert(inputs.size() == 2);
+  auto& in = inputs[0];
+  auto& val = inputs[1];
+
+  // Padding value must be a scalar
+  assert(val.size() == 1);
+
+  // Padding value, input and output must be of the same type
+  assert(val.dtype() == in.dtype() && in.dtype() == out.dtype());
+
+  // Fill output with val
+  copy_cpu(val, out, CopyType::Scalar, stream());
+
+  // Find offset for start of input values
+  size_t data_offset = 0;
+  for (int i = 0; i < axes_.size(); i++) {
+    auto ax = axes_[i] < 0 ? out.ndim() + axes_[i] : axes_[i];
+    data_offset += out.strides()[ax] * low_pad_size_[i];
+  }
+
+  // Extract slice from output where input will be pasted
+  array out_slice(in.shape(), out.dtype(), nullptr, {});
+  out_slice.copy_shared_buffer(
+      out, out.strides(), out.flags(), out_slice.size(), data_offset);
+
+  // Copy input values into the slice
+  copy_cpu_inplace(in, out_slice, CopyType::GeneralGeneral, stream());
+}
+
+void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  // keys has shape (N1, ..., NK, 2)
+  // out has shape (N1, ..., NK, M1, M2, ...)
+  auto& keys = inputs[0];
+  size_t num_keys = keys.size() / 2;
+
+  size_t elems_per_key = out.size() / num_keys;
+  size_t bytes_per_key = out.itemsize() * elems_per_key;
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto kptr = inputs[0].data<uint32_t>();
+  auto cptr = out.data<char>();
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(inputs[0]);
+  encoder.set_output_array(out);
+  encoder.dispatch([kptr,
+                    cptr,
+                    bytes_per_key,
+                    num_keys,
+                    kshape = keys.shape(),
+                    kstrides = keys.strides()]() mutable {
+    auto copy_remaining = [&](char* cptr, size_t loc, uint32_t v) {
+      if (4 * loc + 4 <= bytes_per_key) {
+        reinterpret_cast<uint32_t*>(cptr)[loc] = v;
+      } else {
+        std::copy(
+            reinterpret_cast<char*>(&v),
+            reinterpret_cast<char*>(&v) + bytes_per_key - 4 * loc,
+            cptr + 4 * loc);
+      }
+    };
+
+    size_t out_skip = (bytes_per_key + 4 - 1) / 4;
+    auto half_size = out_skip / 2;
+    bool even = out_skip % 2 == 0;
+    for (int i = 0; i < num_keys; ++i, cptr += bytes_per_key) {
+      auto ptr = reinterpret_cast<uint32_t*>(cptr);
+      // Get ith key
+      auto kidx = 2 * i;
+      auto k1_elem = elem_to_loc(kidx, kshape, kstrides);
+      auto k2_elem = elem_to_loc(kidx + 1, kshape, kstrides);
+      auto key = std::make_pair(kptr[k1_elem], kptr[k2_elem]);
+
+      std::pair<uintptr_t, uintptr_t> count{0, half_size + !even};
+      for (; count.first + 1 < half_size; count.first++, count.second++) {
+        std::tie(ptr[count.first], ptr[count.second]) =
+            random::threefry2x32_hash(key, count);
+      }
+      if (count.first < half_size) {
+        auto rb = random::threefry2x32_hash(key, count);
+        ptr[count.first++] = rb.first;
+        copy_remaining(cptr, count.second, rb.second);
+      }
+      if (!even) {
+        count.second = 0;
+        copy_remaining(
+            cptr, half_size, random::threefry2x32_hash(key, count).first);
+      }
+    }
+  });
+}
+
+void Reshape::eval_cpu(const std::vector<array>& inputs, array& out) {
+  reshape(inputs[0], out);
+}
+
+void DynamicSlice::eval_cpu(const std::vector<array>& inputs, array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+  auto& in = inputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto [in_offset, donated] =
+      compute_dynamic_offset(inputs[1], in.strides(), axes_, stream());
+  copy_cpu_inplace(
+      /* const array& src = */ in,
+      /* array& dst = */ out,
+      /* const Shape& data_shape = */ out.shape(),
+      /* const Strides& i_strides = */ in.strides(),
+      /* const Strides& o_strides = */ out.strides(),
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream(),
+      /* const std::optional<array>& dynamic_i_offset = */ in_offset,
+      /* const std::optional<array>& dynamic_o_offset = */ std::nullopt);
+  if (!donated) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(in_offset));
+  }
+}
+
+void DynamicSliceUpdate::eval_cpu(
+    const std::vector<array>& inputs,
+    array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  // Copy or move src to dst
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  auto [out_offset, donated] =
+      compute_dynamic_offset(inputs[2], out.strides(), axes_, stream());
+  copy_cpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd.strides(),
+      /* const std::vector<stride_t>& o_strides = */ out.strides(),
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream(),
+      /* const std::optional<array>& dynamic_i_offset = */ std::nullopt,
+      /* const std::optional<array>& dynamic_o_offset = */ out_offset);
+  if (!donated) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(out_offset));
+  }
+}
+
+void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Check if materialization is needed
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [data_offset, out_strides] =
+      prepare_slice(out, start_indices_, strides_);
+
+  // Do copy
+  copy_cpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd.strides(),
+      /* const std::vector<stride_t>& o_strides = */ out_strides,
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ data_offset,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream());
+}
+
+void View::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  auto ibytes = size_of(in.dtype());
+  auto obytes = size_of(out.dtype());
+  // Conditions for buffer copying (disjunction):
+  // - type size is the same
+  // - type size is smaller and the last axis is contiguous
+  // - the entire array is row contiguous
+  if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) ||
+      in.flags().row_contiguous) {
+    auto strides = in.strides();
+    for (int i = 0; i < static_cast<int>(strides.size()) - 1; ++i) {
+      strides[i] *= ibytes;
+      strides[i] /= obytes;
+    }
+    out.copy_shared_buffer(
+        in, strides, in.flags(), in.data_size() * ibytes / obytes);
+  } else {
+    auto tmp = array(
+        in.shape(), in.dtype() == bool_ ? uint8 : in.dtype(), nullptr, {});
+    tmp.set_data(allocator::malloc(tmp.nbytes()));
+    if (in.dtype() == bool_) {
+      auto in_tmp = array(in.shape(), uint8, nullptr, {});
+      in_tmp.copy_shared_buffer(in);
+      copy_cpu_inplace(in_tmp, tmp, CopyType::General, stream());
+    } else {
+      copy_cpu_inplace(in, tmp, CopyType::General, stream());
+    }
+
+    auto flags = out.flags();
+    flags.contiguous = true;
+    flags.row_contiguous = true;
+    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
+    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
+    out.copy_shared_buffer(tmp, out.strides(), flags, out.size());
+  }
+}
+
+} // namespace mlx::core