mlx-cpu 0.30.1__py3-none-manylinux_2_35_x86_64.whl

mlx/include/mlx/backend/metal/kernels/hadamard.h

@@ -0,0 +1,182 @@
+// Copyright © 2024 Apple Inc.
+#include <metal_common>
+#include <metal_compute>
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+using namespace metal;
+
+// Thread local Hadamard transform for 2^R
+template <short R>
+METAL_FUNC void radix_func(thread float* x) {
+  constexpr short logR = __builtin_ctz(R);
+  short h = 1;
+  STEEL_PRAGMA_UNROLL
+  for (short s = 0; s < logR; s++) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < R / 2; i++) {
+      short k = i & (h - 1);
+      short j = ((i - k) << 1) + k;
+      float a = x[j];
+      float b = x[j + h];
+      x[j] = a + b;
+      x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+
+template <typename T, int N, int max_radix, int read_width, int stride = 1>
+[[kernel]] void hadamard_n(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const float& scale,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Compute a Hadamard transform of size N = 2^k
+  //
+  // Equivalent to:
+  //    from scipy.linalg import hadamard
+  //    y = hadamard(len(x)) @ x
+
+  constexpr short num_threads = N / max_radix;
+  constexpr short logN = __builtin_ctz(N);
+  constexpr short logR = __builtin_ctz(max_radix);
+  constexpr short num_steps = logN / logR;
+  constexpr short logFinal = logN % logR;
+  constexpr short final_radix = 1 << (logFinal);
+
+  int batch_idx = elem.y * N * stride + elem.z;
+  short i = elem.x;
+
+  threadgroup T buf[N];
+
+  // Read values from device
+  if (stride == 1) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix / read_width; j++) {
+      short index = j * read_width * num_threads + i * read_width;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < read_width; r++) {
+        buf[index + r] = in[batch_idx + index + r];
+      }
+    }
+  } else {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix; j++) {
+      buf[j * num_threads + i] = in[batch_idx + (j * num_threads + i) * stride];
+    }
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  float x[max_radix];
+  short h = 1;
+
+  STEEL_PRAGMA_UNROLL
+  for (short s = 0; s < num_steps; s++) {
+    short k = i & (h - 1);
+    short j = ((i - k) << logR) + k;
+
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < max_radix; r++) {
+      x[r] = buf[j + h * r];
+    }
+
+    radix_func<max_radix>(x);
+
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < max_radix; r++) {
+      buf[j + h * r] = T(x[r]);
+    }
+
+    h <<= logR;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // Do the final radix
+  // e.g. max_radix = 16
+  //      N = 1024 = 16 * 16 * 4
+  if (final_radix > 1) {
+    // Each thread does multiple butterflies
+    STEEL_PRAGMA_UNROLL
+    for (int t = 0; t < max_radix / final_radix; t++) {
+      short index = i + t * num_threads;
+      short k = index & (h - 1);
+      short j = ((index - k) << logFinal) + k;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < final_radix; r++) {
+        x[r] = buf[j + h * r];
+      }
+
+      radix_func<final_radix>(x);
+
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < final_radix; r++) {
+        buf[j + h * r] = T(x[r]);
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // Write values to device
+  if (stride == 1) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix / read_width; j++) {
+      short index = j * read_width * num_threads + i * read_width;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < read_width; r++) {
+        out[batch_idx + index + r] = T(buf[index + r] * scale);
+      }
+    }
+  } else {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix; j++) {
+      out[batch_idx + (j * num_threads + i) * stride] =
+          buf[j * num_threads + i];
+    }
+  }
+}
+
+template <typename T, int N, int M, int read_width>
+[[kernel]] void hadamard_m(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const float& scale,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Compute a Hadamard transform of size M
+  // using a naive O(M^2) codelet.
+  //
+  // This kernel is the second stage in the computation
+  // of a Hadamard transform of size M*N where N = 2^k.
+
+  int index = elem.x * grid.y + elem.y;
+  short i = index % (N / read_width);
+  int batch_idx = index / (N / read_width) * M * N;
+
+  float x[read_width][M];
+  STEEL_PRAGMA_UNROLL
+  for (short c = 0; c < M; c++) {
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      x[r][c] = in[batch_idx + c * N + i * read_width + r];
+    }
+  }
+
+  STEEL_PRAGMA_UNROLL
+  for (short r = 0; r < read_width; r++) {
+    // This function is JIT compiled for M
+    // using the Hadamard matrix strings in `metal/hadamard.cpp`
+    hadamard_radix_m(x[r]);
+  }
+
+  // Write back to device
+  STEEL_PRAGMA_UNROLL
+  for (short c = 0; c < M; c++) {
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      out[batch_idx + c * N + i * read_width + r] = T(x[r][c] * scale);
+    }
+  }
+}

mlx/include/mlx/backend/metal/kernels/indexing/gather.h

@@ -0,0 +1,51 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <typename T, typename IdxT, int NIDX, int IDX_NDIM, typename LocT>
+METAL_FUNC void gather_impl(
+    const device T* src [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const constant int* src_shape [[buffer(2)]],
+    const constant int64_t* src_strides [[buffer(3)]],
+    const constant size_t& src_ndim [[buffer(4)]],
+    const constant int* slice_sizes [[buffer(5)]],
+    const constant int* axes [[buffer(6)]],
+    const thread Indices<IdxT, NIDX>& indices,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  LocT src_idx = 0;
+  for (int i = 0; i < NIDX; ++i) {
+    LocT idx_loc;
+    if (IDX_NDIM == 0) {
+      idx_loc = 0;
+    } else if (IDX_NDIM == 1) {
+      idx_loc = index.x * static_cast<LocT>(indices.strides[indices.ndim * i]);
+    } else {
+      idx_loc = index.x * static_cast<LocT>(indices.strides[indices.ndim * i]);
+      idx_loc += indices.row_contiguous[i]
+          ? index.y
+          : elem_to_loc<LocT>(
+                index.y,
+                &indices.shapes[indices.ndim * i + 1],
+                &indices.strides[indices.ndim * i + 1],
+                indices.ndim - 1);
+    }
+    auto ax = axes[i];
+    auto idx_val = offset_neg_idx(indices.buffers[i][idx_loc], src_shape[ax]);
+    src_idx += static_cast<LocT>(idx_val) * static_cast<LocT>(src_strides[ax]);
+  }
+
+  auto src_offset =
+      elem_to_loc<LocT>(index.z, slice_sizes, src_strides, src_ndim);
+
+  LocT out_idx = index.z;
+  if (IDX_NDIM == 1) {
+    out_idx += static_cast<LocT>(grid_dim.z) * index.x;
+  } else if (IDX_NDIM >= 2) {
+    out_idx += grid_dim.z * (index.x * static_cast<LocT>(grid_dim.y) + index.y);
+  }
+  out[out_idx] = src[src_offset + src_idx];
+}

mlx/include/mlx/backend/metal/kernels/indexing/gather_axis.h

@@ -0,0 +1,44 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+template <typename T, typename IdxT, typename LocT, bool SrcC, bool IdxC>
+[[kernel]] void gather_axis(
+    const device T* src [[buffer(0)]],
+    const device IdxT* indices [[buffer(1)]],
+    device T* out [[buffer(2)]],
+    const constant int* shape [[buffer(3)]],
+    const constant int64_t* src_strides [[buffer(4)]],
+    const constant int64_t* idx_strides [[buffer(5)]],
+    const constant size_t& ndim [[buffer(6)]],
+    const constant int& axis [[buffer(7)]],
+    const constant int& axis_size [[buffer(8)]],
+    const constant size_t& src_ax_stride [[buffer(9)]],
+    const constant size_t& idx_ax_stride [[buffer(10)]],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  LocT elem_idx = index.z * static_cast<LocT>(grid_dim.x);
+  LocT out_idx = elem_idx * grid_dim.y + index.x;
+
+  LocT idx_loc = index.y * static_cast<LocT>(idx_ax_stride);
+  if (IdxC) {
+    idx_loc += out_idx;
+  } else {
+    idx_loc += elem_to_loc<LocT>(elem_idx + index.x, shape, idx_strides, ndim);
+  }
+
+  auto idx_val = indices[idx_loc];
+  if (is_signed_v<IdxT>) {
+    idx_val = (idx_val < 0) ? idx_val + axis_size : idx_val;
+  }
+
+  LocT src_idx = idx_val * static_cast<LocT>(src_ax_stride);
+  if (SrcC) {
+    src_idx += elem_idx * axis_size + index.x;
+  } else {
+    src_idx += elem_to_loc<LocT>(elem_idx + index.x, shape, src_strides, ndim);
+  }
+
+  out_idx += index.y * static_cast<LocT>(grid_dim.x);
+  out[out_idx] = src[src_idx];
+}

mlx/include/mlx/backend/metal/kernels/indexing/gather_front.h

@@ -0,0 +1,24 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <typename T, typename IdxT, typename LocT, int N>
+[[kernel]] void gather_front(
+    const device T* src,
+    const device IdxT* indices,
+    device T* out,
+    const constant int64_t& stride,
+    const constant int& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto idx = offset_neg_idx(indices[index.y], size);
+  LocT src_idx = static_cast<LocT>(stride) * idx;
+  LocT out_idx = static_cast<LocT>(stride) * index.y;
+
+  int s_idx = N * index.x;
+  for (int i = 0; i < N && s_idx < stride; ++i, ++s_idx) {
+    out[out_idx + s_idx] = src[src_idx + s_idx];
+  }
+}

mlx/include/mlx/backend/metal/kernels/indexing/indexing.h

@@ -0,0 +1,23 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+template <typename IdxT, int NIDX>
+struct Indices {
+  const array<const device IdxT*, NIDX> buffers;
+  const constant int* shapes;
+  const constant int64_t* strides;
+  const constant bool* row_contiguous;
+  const int ndim;
+};
+
+template <typename IdxT>
+METAL_FUNC size_t offset_neg_idx(IdxT idx, int size) {
+  if (is_unsigned_v<IdxT>) {
+    return idx;
+  } else {
+    return (idx < 0) ? idx + size : idx;
+  }
+}

mlx/include/mlx/backend/metal/kernels/indexing/masked_scatter.h

@@ -0,0 +1,38 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+template <typename T, bool src_contiguous>
+[[kernel]] void masked_assign_impl(
+    const device bool* mask [[buffer(0)]],
+    const device uint* scatter_offsets [[buffer(1)]],
+    const device T* src [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    const constant int* src_shapes [[buffer(4)]],
+    const constant int64_t* src_strides [[buffer(5)]],
+    const constant int& src_ndim [[buffer(6)]],
+    const constant int64_t& src_batch_size [[buffer(7)]],
+    const constant int64_t& mask_batch_size [[buffer(8)]],
+    uint idx [[thread_position_in_grid]]) {
+  const bool mask_value = mask[idx];
+  if (!mask_value) {
+    return;
+  }
+
+  const uint src_index = scatter_offsets[idx];
+  if (src_index >= src_batch_size) {
+    return;
+  }
+
+  const uint batch_idx = idx / mask_batch_size;
+
+  if (src_contiguous) {
+    out[idx] = src[batch_idx * src_batch_size + src_index];
+  } else {
+    out[idx] = src[elem_to_loc<uint>(
+        batch_idx * src_batch_size + src_index,
+        src_shapes,
+        src_strides,
+        src_ndim)];
+  }
+}

mlx/include/mlx/backend/metal/kernels/indexing/scatter.h

@@ -0,0 +1,59 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <
+    typename T,
+    typename IdxT,
+    typename Op,
+    int NIDX,
+    bool UPD_ROW_CONTIG,
+    int NWORK,
+    typename LocT>
+METAL_FUNC void scatter_impl(
+    const device T* updates,
+    device mlx_atomic<T>* out,
+    const constant int* upd_shape,
+    const constant int64_t* upd_strides,
+    const constant size_t& upd_ndim,
+    const constant size_t& upd_size,
+    const constant int* out_shape,
+    const constant int64_t* out_strides,
+    const constant size_t& out_ndim,
+    const constant int* axes,
+    const constant size_t& idx_size,
+    const thread Indices<IdxT, NIDX>& indices,
+    uint2 gid [[thread_position_in_grid]]) {
+  Op op;
+
+  auto ind_idx = gid.y * NWORK;
+  LocT out_offset = 0;
+  if (upd_size > 1) {
+    out_offset = elem_to_loc<LocT>(
+        gid.x, upd_shape + indices.ndim, out_strides, out_ndim);
+  }
+
+  for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) {
+    LocT out_idx = out_offset;
+    for (int i = 0; i < NIDX; ++i) {
+      auto idx_loc = indices.row_contiguous[i]
+          ? ind_idx
+          : elem_to_loc<LocT>(
+                ind_idx,
+                &indices.shapes[indices.ndim * i],
+                &indices.strides[indices.ndim * i],
+                indices.ndim);
+      auto ax = axes[i];
+      auto idx_val = offset_neg_idx(indices.buffers[i][idx_loc], out_shape[ax]);
+      out_idx +=
+          static_cast<LocT>(idx_val) * static_cast<LocT>(out_strides[ax]);
+    }
+    auto upd_idx = ind_idx * static_cast<LocT>(upd_size) + gid.x;
+    if constexpr (!UPD_ROW_CONTIG) {
+      upd_idx = elem_to_loc<LocT>(upd_idx, upd_shape, upd_strides, upd_ndim);
+    }
+    op.atomic_update(out, updates[upd_idx], out_idx);
+  }
+}

mlx/include/mlx/backend/metal/kernels/indexing/scatter_axis.h

@@ -0,0 +1,52 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+template <
+    typename T,
+    typename IdxT,
+    typename LocT,
+    typename Op,
+    bool UpdC,
+    bool IdxC>
+[[kernel]] void scatter_axis(
+    const device T* upd [[buffer(0)]],
+    const device IdxT* indices [[buffer(1)]],
+    device mlx_atomic<T>* out [[buffer(2)]],
+    const constant int* shape [[buffer(3)]],
+    const constant int64_t* upd_strides [[buffer(4)]],
+    const constant int64_t* idx_strides [[buffer(5)]],
+    const constant size_t& ndim [[buffer(6)]],
+    const constant int& axis [[buffer(7)]],
+    const constant int& out_axis_size [[buffer(8)]],
+    const constant size_t& upd_ax_stride [[buffer(9)]],
+    const constant size_t& idx_ax_stride [[buffer(10)]],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  Op op;
+
+  LocT elem_idx = index.z * static_cast<LocT>(grid_dim.x);
+
+  LocT idx_loc = index.y * static_cast<LocT>(idx_ax_stride);
+  if (IdxC) {
+    idx_loc += elem_idx * grid_dim.y + index.x;
+  } else {
+    idx_loc += elem_to_loc<LocT>(elem_idx + index.x, shape, idx_strides, ndim);
+  }
+
+  auto idx_val = indices[idx_loc];
+  if (is_signed_v<IdxT>) {
+    idx_val = (idx_val < 0) ? idx_val + out_axis_size : idx_val;
+  }
+
+  LocT upd_idx = index.y * static_cast<LocT>(upd_ax_stride);
+  if (UpdC) {
+    upd_idx += elem_idx * grid_dim.y + index.x;
+  } else {
+    upd_idx += elem_to_loc<LocT>(elem_idx + index.x, shape, upd_strides, ndim);
+  }
+
+  LocT out_idx = elem_idx * static_cast<LocT>(out_axis_size) +
+      idx_val * grid_dim.x + index.x;
+  op.atomic_update(out, upd[upd_idx], out_idx);
+}

mlx/include/mlx/backend/metal/kernels/logsumexp.h

@@ -0,0 +1,140 @@
+// Copyright © 2025 Apple Inc.
+
+template <typename T, typename AccT = float, int N_READS = 4>
+[[kernel]] void logsumexp(
+    const device T* in,
+    device T* out,
+    constant int& axis_size,
+    uint gid [[threadgroup_position_in_grid]],
+    uint _lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  int lid = _lid;
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup AccT local_max[SIMD_SIZE];
+  threadgroup AccT local_normalizer[SIMD_SIZE];
+
+  AccT ld[N_READS];
+
+  in += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      ld[i] = AccT(in[i]);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      ld[i] =
+          ((lid * N_READS + i) < axis_size) ? AccT(in[i]) : Limits<AccT>::min;
+    }
+  }
+  if (simd_group_id == 0) {
+    local_max[simd_lane_id] = Limits<AccT>::min;
+    local_normalizer[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Get the max
+  AccT maxval = Limits<AccT>::finite_min;
+  for (int i = 0; i < N_READS; i++) {
+    maxval = (maxval < ld[i]) ? ld[i] : maxval;
+  }
+  maxval = simd_max(maxval);
+  if (simd_lane_id == 0) {
+    local_max[simd_group_id] = maxval;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    maxval = simd_max(local_max[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_max[0] = maxval;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  maxval = local_max[0];
+
+  // Compute exp(x_i - maxval) and store the partial sums in local_normalizer
+  AccT normalizer = 0;
+  for (int i = 0; i < N_READS; i++) {
+    normalizer += fast::exp(ld[i] - maxval);
+  }
+  normalizer = simd_sum(normalizer);
+  if (simd_lane_id == 0) {
+    local_normalizer[simd_group_id] = normalizer;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    normalizer = simd_sum(local_normalizer[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
+    }
+  }
+}
+
+template <typename T, typename AccT = float, int N_READS = 4>
+[[kernel]] void logsumexp_looped(
+    const device T* in,
+    device T* out,
+    constant int& axis_size,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  in += gid * size_t(axis_size);
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup AccT local_max[SIMD_SIZE];
+  threadgroup AccT local_normalizer[SIMD_SIZE];
+
+  // Get the max and the normalizer in one go
+  AccT prevmax;
+  AccT maxval = Limits<AccT>::finite_min;
+  AccT normalizer = 0;
+  for (int r = 0; r < static_cast<int>(ceildiv(axis_size, N_READS * lsize));
+       r++) {
+    int offset = r * lsize * N_READS + lid * N_READS;
+    AccT vals[N_READS];
+    if (offset + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] = AccT(in[offset + i]);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] =
+            (offset + i < axis_size) ? AccT(in[offset + i]) : Limits<AccT>::min;
+      }
+    }
+    prevmax = maxval;
+    for (int i = 0; i < N_READS; i++) {
+      maxval = (maxval < vals[i]) ? vals[i] : maxval;
+    }
+    normalizer *= fast::exp(prevmax - maxval);
+    for (int i = 0; i < N_READS; i++) {
+      normalizer += fast::exp(vals[i] - maxval);
+    }
+  }
+  prevmax = maxval;
+  maxval = simd_max(maxval);
+  normalizer *= fast::exp(prevmax - maxval);
+  normalizer = simd_sum(normalizer);
+
+  prevmax = maxval;
+  if (simd_lane_id == 0) {
+    local_max[simd_group_id] = maxval;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  maxval = simd_max(local_max[simd_lane_id]);
+  normalizer *= fast::exp(prevmax - maxval);
+  if (simd_lane_id == 0) {
+    local_normalizer[simd_group_id] = normalizer;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  normalizer = simd_sum(local_normalizer[simd_lane_id]);
+
+  if (lid == 0) {
+    out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
+  }
+}