mlx 0.30.7

data/mlx/mlx/backend/cpu/scan.cpp

@@ -0,0 +1,338 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, typename U, typename Op>
+void contiguous_scan(
+    const T* input,
+    U* output,
+    int count,
+    int stride,
+    bool reverse,
+    bool inclusive,
+    const Op& op,
+    U init) {
+  if (!reverse) {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        *output = *input;
+        for (int j = 1; j < stride; j++) {
+          input++;
+          output++;
+          *output = op(*(output - 1), *input);
+        }
+        output++;
+        input++;
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        *output = init;
+        for (int j = 1; j < stride; j++) {
+          *(output + 1) = op(*output, *input);
+          input++;
+          output++;
+        }
+        output++;
+        input++;
+      }
+    }
+  } else {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        output += stride - 1;
+        input += stride - 1;
+        *output = *input;
+        for (int j = 1; j < stride; j++) {
+          input--;
+          output--;
+          *output = op(*(output + 1), *input);
+        }
+        output += stride;
+        input += stride;
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        output += stride - 1;
+        input += stride - 1;
+        *output = init;
+        for (int j = 1; j < stride; j++) {
+          *(output - 1) = op(*output, *input);
+          input--;
+          output--;
+        }
+        output += stride;
+        input += stride;
+      }
+    }
+  }
+};
+
+template <typename T, typename U, typename Op>
+void strided_scan(
+    const T* input,
+    U* output,
+    int count,
+    int size,
+    int stride,
+    bool reverse,
+    bool inclusive,
+    const Op& op,
+    U init) {
+  // TODO: Vectorize the following naive implementation
+  if (!reverse) {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        std::copy(input, input + stride, output);
+        output += stride;
+        input += stride;
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            *output = op(*(output - stride), *input);
+            output++;
+            input++;
+          }
+        }
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        std::fill(output, output + stride, init);
+        output += stride;
+        input += stride;
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            *output = op(*(output - stride), *(input - stride));
+            output++;
+            input++;
+          }
+        }
+      }
+    }
+  } else {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        output += (size - 1) * stride;
+        input += (size - 1) * stride;
+        std::copy(input, input + stride, output);
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            output--;
+            input--;
+            *output = op(*(output + stride), *input);
+          }
+        }
+        output += size * stride;
+        input += size * stride;
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        output += (size - 1) * stride;
+        input += (size - 1) * stride;
+        std::fill(output, output + stride, init);
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            output--;
+            input--;
+            *output = op(*(output + stride), *(input + stride));
+          }
+        }
+        output += size * stride;
+        input += size * stride;
+      }
+    }
+  }
+};
+
+template <typename T, typename U, typename Op>
+void scan_op(
+    const array& in,
+    array& out,
+    int axis,
+    bool reverse,
+    bool inclusive,
+    const Op& op,
+    U init) {
+  if (in.flags().row_contiguous) {
+    if (in.strides()[axis] == 1) {
+      contiguous_scan(
+          in.data<T>(),
+          out.data<U>(),
+          in.size() / in.shape(axis),
+          in.shape(axis),
+          reverse,
+          inclusive,
+          op,
+          init);
+    } else {
+      strided_scan(
+          in.data<T>(),
+          out.data<U>(),
+          in.size() / in.shape(axis) / in.strides()[axis],
+          in.shape(axis),
+          in.strides()[axis],
+          reverse,
+          inclusive,
+          op,
+          init);
+    }
+  } else {
+    throw std::runtime_error("Scan op supports only contiguous inputs");
+  }
+}
+
+template <typename T, typename U>
+void scan_dispatch(
+    Scan::ReduceType rtype,
+    const array& in,
+    array& out,
+    int axis,
+    bool reverse,
+    bool inclusive) {
+  switch (rtype) {
+    case Scan::Sum: {
+      auto op = [](U y, T x) { return y + x; };
+      auto init = static_cast<U>(0);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::Prod: {
+      auto op = [](U y, T x) { return y * x; };
+      auto init = static_cast<U>(1);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::Min: {
+      auto op = [](U y, T x) { return x < y ? x : y; };
+      auto init = (issubdtype(in.dtype(), floating))
+          ? static_cast<U>(std::numeric_limits<float>::infinity())
+          : std::numeric_limits<U>::max();
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::Max: {
+      auto op = [](U y, T x) { return x < y ? y : x; };
+      auto init = (issubdtype(in.dtype(), floating))
+          ? static_cast<U>(-std::numeric_limits<float>::infinity())
+          : std::numeric_limits<U>::min();
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::LogAddExp: {
+      auto op = [](U a, T b) {
+        return detail::LogAddExp{}(a, static_cast<U>(b));
+      };
+      auto init = (issubdtype(in.dtype(), floating))
+          ? static_cast<U>(-std::numeric_limits<float>::infinity())
+          : std::numeric_limits<U>::min();
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+  }
+}
+
+} // namespace
+
+void Scan::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  // Ensure contiguity
+  auto in = inputs[0];
+  if (!in.flags().row_contiguous) {
+    in = contiguous_copy_cpu(in, stream());
+    encoder.add_temporary(in);
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    axis_ = axis_,
+                    reduce_type_ = reduce_type_,
+                    reverse_ = reverse_,
+                    inclusive_ = inclusive_]() mutable {
+    switch (in.dtype()) {
+      case bool_: {
+        // We could do a full dtype x dtype switch but this is the only case
+        // where we accumulate in a different type, for now.
+        //
+        // TODO: If we add the option to accumulate floats in higher precision
+        //       floats perhaps we should add the full all-to-all dispatch.
+        if (reduce_type_ == Scan::Sum && out.dtype() == int32) {
+          scan_dispatch<bool, int32_t>(
+              reduce_type_, in, out, axis_, reverse_, inclusive_);
+        } else {
+          scan_dispatch<bool, bool>(
+              reduce_type_, in, out, axis_, reverse_, inclusive_);
+        }
+        break;
+      }
+      case uint8:
+        scan_dispatch<uint8_t, uint8_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case uint16:
+        scan_dispatch<uint16_t, uint16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case uint32:
+        scan_dispatch<uint32_t, uint32_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case uint64:
+        scan_dispatch<uint64_t, uint64_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int8:
+        scan_dispatch<int8_t, int8_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int16:
+        scan_dispatch<int16_t, int16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int32:
+        scan_dispatch<int32_t, int32_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int64:
+        scan_dispatch<int64_t, int64_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case float16:
+        scan_dispatch<float16_t, float16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case float32:
+        scan_dispatch<float, float>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case float64:
+        scan_dispatch<double, double>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case bfloat16:
+        scan_dispatch<bfloat16_t, bfloat16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case complex64:
+        scan_dispatch<complex64_t, complex64_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+    }
+  });
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/select.cpp

@@ -0,0 +1,95 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/ternary.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename Op>
+void select_op(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op,
+    Stream stream) {
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+  set_ternary_op_output_data(a, b, c, out, topt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    c = array::unsafe_weak_copy(c),
+                    out = array::unsafe_weak_copy(out),
+                    op,
+                    topt]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        ternary_op<bool, bool, bool, bool>(a, b, c, out, op, topt);
+        break;
+      case uint8:
+        ternary_op<bool, uint8_t, uint8_t, uint8_t>(a, b, c, out, op, topt);
+        break;
+      case uint16:
+        ternary_op<bool, uint16_t, uint16_t, uint16_t>(a, b, c, out, op, topt);
+        break;
+      case uint32:
+        ternary_op<bool, uint32_t, uint32_t, uint32_t>(a, b, c, out, op, topt);
+        break;
+      case uint64:
+        ternary_op<bool, uint64_t, uint64_t, uint64_t>(a, b, c, out, op, topt);
+        break;
+      case int8:
+        ternary_op<bool, int8_t, int8_t, int8_t>(a, b, c, out, op, topt);
+        break;
+      case int16:
+        ternary_op<bool, int16_t, int16_t, int16_t>(a, b, c, out, op, topt);
+        break;
+      case int32:
+        ternary_op<bool, int32_t, int32_t, int32_t>(a, b, c, out, op, topt);
+        break;
+      case int64:
+        ternary_op<bool, int64_t, int64_t, int64_t>(a, b, c, out, op, topt);
+        break;
+      case float16:
+        ternary_op<bool, float16_t, float16_t, float16_t>(
+            a, b, c, out, op, topt);
+        break;
+      case float32:
+        ternary_op<bool, float, float, float>(a, b, c, out, op, topt);
+        break;
+      case float64:
+        ternary_op<bool, double, double, double>(a, b, c, out, op, topt);
+        break;
+      case bfloat16:
+        ternary_op<bool, bfloat16_t, bfloat16_t, bfloat16_t>(
+            a, b, c, out, op, topt);
+        break;
+      case complex64:
+        ternary_op<bool, complex64_t, complex64_t, complex64_t>(
+            a, b, c, out, op, topt);
+        break;
+    }
+  });
+}
+
+} // namespace
+
+void Select::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+  const auto& condition = inputs[0];
+  const auto& a = inputs[1];
+  const auto& b = inputs[2];
+  select_op(condition, a, b, out, detail::Select(), stream());
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cpu/simd/accelerate_fp16_simd.h

@@ -0,0 +1,56 @@
+#pragma once
+
+#include "mlx/backend/cpu/simd/base_simd.h"
+
+#if MLX_SIMD_LIBRARY_VERSION < 6
+#include "mlx/backend/cpu/simd/neon_fp16_simd.h"
+#endif
+
+namespace mlx::core::simd {
+
+#if MLX_SIMD_LIBRARY_VERSION >= 6
+constexpr int N = 8;
+template <int N>
+struct ScalarT<float16_t, N> {
+  using v = _Float16;
+};
+#endif
+
+template <>
+inline constexpr int max_size<float16_t> = N;
+
+#define SIMD_FP16_DEFAULT_UNARY(op)                    \
+  template <>                                          \
+  inline Simd<float16_t, N> op(Simd<float16_t, N> v) { \
+    Simd<float, N> in = v;                             \
+    return op(in);                                     \
+  }
+
+SIMD_FP16_DEFAULT_UNARY(acos)
+SIMD_FP16_DEFAULT_UNARY(acosh)
+SIMD_FP16_DEFAULT_UNARY(asin)
+SIMD_FP16_DEFAULT_UNARY(asinh)
+SIMD_FP16_DEFAULT_UNARY(atan)
+SIMD_FP16_DEFAULT_UNARY(atanh)
+SIMD_FP16_DEFAULT_UNARY(cosh)
+SIMD_FP16_DEFAULT_UNARY(expm1)
+SIMD_FP16_DEFAULT_UNARY(log)
+SIMD_FP16_DEFAULT_UNARY(log2)
+SIMD_FP16_DEFAULT_UNARY(log10)
+SIMD_FP16_DEFAULT_UNARY(log1p)
+SIMD_FP16_DEFAULT_UNARY(sinh)
+SIMD_FP16_DEFAULT_UNARY(tan)
+SIMD_FP16_DEFAULT_UNARY(tanh)
+
+#define SIMD_FP16_DEFAULT_BINARY(op)                                         \
+  template <>                                                                \
+  inline Simd<float16_t, N> op(Simd<float16_t, N> x, Simd<float16_t, N> y) { \
+    Simd<float, N> a = x;                                                    \
+    Simd<float, N> b = y;                                                    \
+    return op(a, b);                                                         \
+  }
+SIMD_FP16_DEFAULT_BINARY(atan2)
+SIMD_FP16_DEFAULT_BINARY(remainder)
+SIMD_FP16_DEFAULT_BINARY(pow)
+
+} // namespace mlx::core::simd