mlx 0.30.7

data/mlx/mlx/backend/metal/kernels/quantized_nax.metal

@@ -0,0 +1,106 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/loader.h"
+#include "mlx/backend/metal/kernels/quantized_nax.h"
+
+#define instantiate_quantized(name, type, group_size, bits, bm, bn, bk, wm, wn)  \
+  instantiate_kernel(                                                    \
+      #name "_" #type "_gs_" #group_size "_b_" #bits,                    \
+      name,                                                              \
+      type,                                                              \
+      group_size,                                                        \
+      bits, bm, bk, bn, wm, wn)
+
+#define instantiate_quantized_batched(name, type, group_size, bits, bm, bn, bk, wm, wn, batched)     \
+  instantiate_kernel(                                                    \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \
+      name,                                                              \
+      type,                                                              \
+      group_size,                                                        \
+      bits,                                                              \
+      batched, bm, bk, bn, wm, wn)
+
+#define instantiate_quantized_aligned(name, type, group_size, bits, bm, bn, bk, wm, wn, aligned)     \
+  instantiate_kernel(                                                                     \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \
+      name,                                                                  \
+      type,                                                                  \
+      group_size,                                                            \
+      bits,                                                                  \
+      aligned, bm, bk, bn, wm, wn)
+
+#define instantiate_quantized_aligned_batched(name, type, group_size, bits, bm, bn, bk, wm, wn, aligned, batched)     \
+  instantiate_kernel(                                                                     \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \
+      name,                                                                  \
+      type,                                                                  \
+      group_size,                                                            \
+      bits,                                                                  \
+      aligned,                                                               \
+      batched, bm, bk, bn, wm, wn)
+
+#define instantiate_gather_qmm_rhs(func, name, type, group_size, bits, bm, bn, bk, wm, wn, transpose)        \
+  instantiate_kernel(                                                                                        \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
+      func,                                                         \
+      type,                                                         \
+      group_size,                                                   \
+      bits,                                                         \
+      bm,                                                           \
+      bn,                                                           \
+      bk,                                                           \
+      wm,                                                           \
+      wn,                                                           \
+      transpose)
+
+#define instantiate_quantized_batched_wrap(name, type, group_size, bits) \
+  instantiate_quantized_batched(name, type, group_size, bits, 64, 64, 64, 2, 2, 1)      \
+  instantiate_quantized_batched(name, type, group_size, bits, 64, 64, 64, 2, 2, 0)
+
+#define instantiate_quantized_all_batched(type, group_size, bits) \
+  instantiate_quantized_batched_wrap(affine_qmm_n_nax, type, group_size, bits)
+
+
+#define instantiate_quantized_all_single(type, group_size, bits) \
+  instantiate_quantized(affine_gather_qmm_n_nax, type, group_size, bits, 64, 64, 64, 2, 2)
+
+#define instantiate_quantized_all_aligned(type, group_size, bits)   \
+  instantiate_quantized_aligned(affine_gather_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true) \
+  instantiate_quantized_aligned(affine_gather_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 1) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 0) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0)
+
+#define instantiate_quantized_all_rhs(type, group_size, bits) \
+  instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \
+  instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false)
+
+#define instantiate_quantized_funcs(type, group_size, bits) \
+  instantiate_quantized_all_batched(type, group_size, bits) \
+  instantiate_quantized_all_aligned(type, group_size, bits) \
+  instantiate_quantized_all_rhs(type, group_size, bits)
+
+#define instantiate_quantized_types(group_size, bits)       \
+  instantiate_quantized_funcs(float, group_size, bits)      \
+  instantiate_quantized_funcs(float16_t, group_size, bits)  \
+  instantiate_quantized_funcs(bfloat16_t, group_size, bits)  
+
+#define instantiate_quantized_groups(bits) \
+  instantiate_quantized_types(128, bits)   \
+  instantiate_quantized_types(64, bits)    \
+  instantiate_quantized_types(32, bits)
+
+#define instantiate_quantized_all() \
+  instantiate_quantized_groups(2) \
+  instantiate_quantized_groups(3) \
+  instantiate_quantized_groups(4) \
+  instantiate_quantized_groups(5) \
+  instantiate_quantized_groups(6) \
+  instantiate_quantized_groups(8)
+
+instantiate_quantized_all() // clang-format on

data/mlx/mlx/backend/metal/kernels/quantized_utils.h

@@ -0,0 +1,90 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+template <typename T, typename mma_t, typename loader_a_t, typename loader_b_t>
+METAL_FUNC void gemm_loop_aligned(
+    threadgroup T* As,
+    threadgroup T* Bs,
+    thread mma_t& mma_op,
+    thread loader_a_t& loader_a,
+    thread loader_b_t& loader_b,
+    const int k_iterations) {
+  for (int k = 0; k < k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load elements into threadgroup memory
+    loader_a.load_unsafe();
+    loader_b.load_unsafe();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+}
+
+template <
+    bool rows_aligned,
+    bool cols_aligned,
+    bool transpose,
+    typename T,
+    typename mma_t,
+    typename loader_a_t,
+    typename loader_b_t>
+METAL_FUNC void gemm_loop_unaligned(
+    threadgroup T* As,
+    threadgroup T* Bs,
+    thread mma_t& mma_op,
+    thread loader_a_t& loader_a,
+    thread loader_b_t& loader_b,
+    const int k_iterations,
+    const short tgp_bm,
+    const short tgp_bn,
+    const short tgp_bk) {
+  for (int k = 0; k < k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load elements into threadgroup memory
+    if (rows_aligned) {
+      loader_a.load_unsafe();
+    } else {
+      loader_a.load_safe(short2(tgp_bk, tgp_bm));
+    }
+    if (cols_aligned) {
+      loader_b.load_unsafe();
+    } else {
+      loader_b.load_safe(
+          transpose ? short2(tgp_bk, tgp_bn) : short2(tgp_bn, tgp_bk));
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+}
+
+template <typename T, typename mma_t, typename loader_a_t, typename loader_b_t>
+METAL_FUNC void gemm_loop_finalize(
+    threadgroup T* As,
+    threadgroup T* Bs,
+    thread mma_t& mma_op,
+    thread loader_a_t& loader_a,
+    thread loader_b_t& loader_b,
+    const short2 tile_a,
+    const short2 tile_b) {
+  loader_a.load_safe(tile_a);
+  loader_b.load_safe(tile_b);
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  mma_op.mma(As, Bs);
+}

data/mlx/mlx/backend/metal/kernels/random.metal

@@ -0,0 +1,103 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+static constexpr constant uint32_t rotations[2][4] = {
+    {13, 15, 26, 6},
+    {17, 29, 16, 24}};
+
+union rbits {
+  uint2 val;
+  uchar4 bytes[2];
+};
+
+rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
+  uint4 ks = {key.x, key.y, key.x ^ key.y ^ 0x1BD11BDA};
+
+  rbits v;
+  v.val.x = count.x + ks[0];
+  v.val.y = count.y + ks[1];
+
+  for (int i = 0; i < 5; ++i) {
+    for (auto r : rotations[i % 2]) {
+      v.val.x += v.val.y;
+      v.val.y = (v.val.y << r) | (v.val.y >> (32 - r));
+      v.val.y ^= v.val.x;
+    }
+    v.val.x += ks[(i + 1) % 3];
+    v.val.y += ks[(i + 2) % 3] + i + 1;
+  }
+
+  return v;
+}
+
+[[kernel]] void rbitsc(
+    device const uint32_t* keys,
+    device char* out,
+    constant const bool& odd,
+    constant const uint& bytes_per_key,
+    uint2 grid_dim [[threads_per_grid]],
+    uint2 index [[thread_position_in_grid]]) {
+  auto kidx = 2 * index.x;
+  auto key = uint2(keys[kidx], keys[kidx + 1]);
+  auto half_size = grid_dim.y - odd;
+  out += index.x * bytes_per_key;
+  bool drop_last = odd && (index.y == half_size);
+  auto bits = threefry2x32_hash(
+      key, uint2(index.y, drop_last ? 0 : index.y + grid_dim.y));
+  size_t idx = size_t(index.y) << 2;
+  for (int i = 0; i < 4; ++i) {
+    out[idx + i] = bits.bytes[0][i];
+  }
+  if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index.y) + grid_dim.y) << 2;
+    if ((index.y + 1) == half_size && (bytes_per_key % 4) > 0) {
+      int edge_bytes = (bytes_per_key % 4);
+      for (int i = 0; i < edge_bytes; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    }
+  }
+}
+
+[[kernel]] void rbits(
+    device const uint32_t* keys,
+    device char* out,
+    constant const bool& odd,
+    constant const uint& bytes_per_key,
+    constant const int& ndim,
+    constant const int* key_shape,
+    constant const int64_t* key_strides,
+    uint2 grid_dim [[threads_per_grid]],
+    uint2 index [[thread_position_in_grid]]) {
+  auto kidx = 2 * index.x;
+  auto k1_elem = elem_to_loc(kidx, key_shape, key_strides, ndim);
+  auto k2_elem = elem_to_loc(kidx + 1, key_shape, key_strides, ndim);
+  auto key = uint2(keys[k1_elem], keys[k2_elem]);
+  auto half_size = grid_dim.y - odd;
+  out += size_t(index.x) * bytes_per_key;
+  bool drop_last = odd && (index.y == half_size);
+  auto bits = threefry2x32_hash(
+      key, uint2(index.y, drop_last ? 0 : index.y + grid_dim.y));
+  size_t idx = size_t(index.y) << 2;
+  for (int i = 0; i < 4; ++i) {
+    out[idx + i] = bits.bytes[0][i];
+  }
+  if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index.y) + grid_dim.y) << 2;
+    if ((index.y + 1) == half_size && (bytes_per_key % 4) > 0) {
+      int edge_bytes = (bytes_per_key % 4);
+      for (int i = 0; i < edge_bytes; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    }
+  }
+}

data/mlx/mlx/backend/metal/kernels/reduce.h

@@ -0,0 +1,5 @@
+#pragma once
+#include "mlx/backend/metal/kernels/reduction/reduce_all.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_col.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_init.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_row.h"

data/mlx/mlx/backend/metal/kernels/reduce.metal

@@ -0,0 +1,169 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_atomic>
+#include <metal_simdgroup>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/atomic.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+#include "mlx/backend/metal/kernels/reduce.h"
+
+#define instantiate_init_reduce(name, tname, type, op) \
+  instantiate_kernel("init_reduce_" #name #tname, init_reduce, type, op<type>)
+
+instantiate_init_reduce(and, bool_, bool, And)
+instantiate_init_reduce(or, bool_, bool, Or)
+
+#define instantiate_init_sum_prod(name, op)                 \
+  instantiate_init_reduce(name, int32, int32_t, op)         \
+  instantiate_init_reduce(name, int64, int64_t, op)         \
+  instantiate_init_reduce(name, float16, float16_t, op)     \
+  instantiate_init_reduce(name, bfloat16, bfloat16_t, op)   \
+  instantiate_init_reduce(name, float32, float, op)         \
+  instantiate_init_reduce(name, complex64, complex64_t, op)
+
+instantiate_init_sum_prod(sum, Sum)
+instantiate_init_sum_prod(prod, Prod)
+
+#define instantiate_init_min_max(name, op)                   \
+  instantiate_init_reduce(name, bool_, bool, op)             \
+  instantiate_init_reduce(name, int8, int8_t, op)            \
+  instantiate_init_reduce(name, int16, int16_t, op)          \
+  instantiate_init_reduce(name, int32, int32_t, op)          \
+  instantiate_init_reduce(name, int64, int64_t, op)          \
+  instantiate_init_reduce(name, uint8, uint8_t, op)          \
+  instantiate_init_reduce(name, uint16, uint16_t, op)        \
+  instantiate_init_reduce(name, uint32, uint32_t, op)        \
+  instantiate_init_reduce(name, uint64, uint64_t, op)        \
+  instantiate_init_reduce(name, float16, float16_t, op)      \
+  instantiate_init_reduce(name, bfloat16, bfloat16_t, op)    \
+  instantiate_init_reduce(name, float32, float, op)          \
+  instantiate_init_reduce(name, complex64, complex64_t, op)
+
+instantiate_init_min_max(min, Min)
+instantiate_init_min_max(max, Max)
+
+#define instantiate_all_reduce(name, itype, otype, op) \
+  instantiate_kernel("all_reduce_" #name,              \
+                     all_reduce,                       \
+                     itype, otype, op)
+
+#define instantiate_col_reduce_small(name, itype, otype, op, dim)          \
+  instantiate_kernel("col_reduce_small_" #dim "_reduce_" #name,            \
+                     col_reduce_small,                                     \
+                     itype, otype, op, int, dim)                           \
+  instantiate_kernel("col_reduce_longcolumn_" #dim "_reduce_" #name,       \
+                     col_reduce_longcolumn,                                \
+                     itype, otype, op, int, dim)                           \
+  instantiate_kernel("col_reduce_small_large_" #dim "_reduce_" #name,      \
+                     col_reduce_small,                                     \
+                     itype, otype, op, int64_t, dim)                       \
+  instantiate_kernel("col_reduce_longcolumn_large_" #dim "_reduce_" #name, \
+                     col_reduce_longcolumn,                                \
+                     itype, otype, op, int64_t, dim)
+
+#define instantiate_col_reduce_looped_tile(name, itype, otype, op, dim, bm, bn)        \
+  instantiate_kernel("col_reduce_looped_" #dim "_" #bm "_" #bn "_reduce_" #name,       \
+                     col_reduce_looped,                                                \
+                     itype, otype, op, int, dim, bm, bn)                               \
+  instantiate_kernel("col_reduce_looped_large_" #dim "_" #bm "_" #bn "_reduce_" #name, \
+                     col_reduce_looped,                                                \
+                     itype, otype, op, int64_t, dim, bm, bn)
+
+#define instantiate_col_reduce_2pass_tile(name, itype, otype, op, dim, bm, bn)        \
+  instantiate_kernel("col_reduce_2pass_" #dim "_" #bm "_" #bn "_reduce_" #name,       \
+                     col_reduce_2pass,                                                \
+                     itype, otype, op, int, dim, bm, bn)                              \
+  instantiate_kernel("col_reduce_2pass_large_" #dim "_" #bm "_" #bn "_reduce_" #name, \
+                     col_reduce_2pass,                                                \
+                     itype, otype, op, int64_t, dim, bm, bn)
+
+#define instantiate_col_reduce_looped(name, itype, otype, op, dim)        \
+  instantiate_col_reduce_looped_tile(name, itype, otype, op, dim, 32, 32) \
+  instantiate_col_reduce_2pass_tile(name, itype, otype, op, dim, 32, 32)
+
+#define instantiate_col_reduce_general(name, itype, otype, op) \
+  instantiate_col_reduce_small(name, itype, otype, op, 1)      \
+  instantiate_col_reduce_small(name, itype, otype, op, 2)      \
+  instantiate_col_reduce_small(name, itype, otype, op, 5)      \
+  instantiate_col_reduce_looped(name, itype, otype, op, 1)     \
+  instantiate_col_reduce_looped(name, itype, otype, op, 2)     \
+  instantiate_col_reduce_looped(name, itype, otype, op, 5)
+
+#define instantiate_row_reduce_small(name, itype, otype, op, dim)     \
+  instantiate_kernel("row_reduce_small_" #dim "_reduce_" #name,       \
+                     row_reduce_small,                                \
+                     itype, otype, op, int, dim)                      \
+  instantiate_kernel("row_reduce_small_large_" #dim "_reduce_" #name, \
+                     row_reduce_small,                                \
+                     itype, otype, op, int64_t, dim)
+
+#define instantiate_row_reduce_looped(name, itype, otype, op, dim)       \
+  instantiate_kernel("row_reduce_looped_" #dim "_reduce_" #name,         \
+                     row_reduce_looped,                                  \
+                     itype, otype, op, int, dim)                         \
+  instantiate_kernel("row_reduce_looped_large_" #dim "_reduce_" #name,   \
+                     row_reduce_looped,                                  \
+                     itype, otype, op, int64_t, dim)
+
+#define instantiate_row_reduce_general(name, itype, otype, op) \
+  instantiate_row_reduce_small(name, itype, otype, op, 1)      \
+  instantiate_row_reduce_small(name, itype, otype, op, 2)      \
+  instantiate_row_reduce_small(name, itype, otype, op, 5)      \
+  instantiate_row_reduce_looped(name, itype, otype, op, 1)     \
+  instantiate_row_reduce_looped(name, itype, otype, op, 2)     \
+  instantiate_row_reduce_looped(name, itype, otype, op, 5)     \
+  instantiate_kernel("row_reduce_simple_" #name,               \
+                     row_reduce_simple,                        \
+                     itype, otype, op)
+
+#define instantiate_reduce_functions(name, tname, itype, otype, op)    \
+  instantiate_all_reduce(name##tname, itype, otype, op<otype>)         \
+  instantiate_row_reduce_general(name##tname, itype, otype, op<otype>) \
+  instantiate_col_reduce_general(name##tname, itype, otype, op<otype>)
+
+#define instantiate_and_or(name, op)                           \
+  instantiate_reduce_functions(name, bool_, bool, bool, op)    \
+  instantiate_reduce_functions(name, int16, int16_t, bool, op) \
+  instantiate_reduce_functions(name, int32, int32_t, bool, op) \
+  instantiate_reduce_functions(name, int64, int64_t, bool, op)
+
+instantiate_and_or(and, And)
+instantiate_and_or(or, Or)
+
+#define instantiate_sum_prod(name, op)                                       \
+  instantiate_reduce_functions(name, uint8, uint8_t, int32_t, op)            \
+  instantiate_reduce_functions(name, uint16, uint16_t, uint32_t, op)         \
+  instantiate_reduce_functions(name, uint32, uint32_t, uint32_t, op)         \
+  instantiate_reduce_functions(name, uint64, uint64_t, uint64_t, op)         \
+  instantiate_reduce_functions(name, int8, int8_t, int32_t, op)              \
+  instantiate_reduce_functions(name, int16, int16_t, int32_t, op)            \
+  instantiate_reduce_functions(name, int32, int32_t, int32_t, op)            \
+  instantiate_reduce_functions(name, int64, int64_t, int64_t, op)            \
+  instantiate_reduce_functions(name, float16, float16_t, float16_t, op)      \
+  instantiate_reduce_functions(name, bfloat16, bfloat16_t, bfloat16_t, op)   \
+  instantiate_reduce_functions(name, float32, float, float, op)              \
+  instantiate_reduce_functions(name, complex64, complex64_t, complex64_t, op)
+
+instantiate_sum_prod(sum, Sum)
+instantiate_sum_prod(prod, Prod)
+
+#define instantiate_min_max(name, op)                                        \
+  instantiate_reduce_functions(name, int8, int8_t, int8_t, op)               \
+  instantiate_reduce_functions(name, int16, int16_t, int16_t, op)            \
+  instantiate_reduce_functions(name, int32, int32_t, int32_t, op)            \
+  instantiate_reduce_functions(name, int64, int64_t, int64_t, op)            \
+  instantiate_reduce_functions(name, uint8, uint8_t, uint8_t, op)            \
+  instantiate_reduce_functions(name, uint16, uint16_t, uint16_t, op)         \
+  instantiate_reduce_functions(name, uint32, uint32_t, uint32_t, op)         \
+  instantiate_reduce_functions(name, uint64, uint64_t, uint64_t, op)         \
+  instantiate_reduce_functions(name, float16, float16_t, float16_t, op)      \
+  instantiate_reduce_functions(name, bfloat16, bfloat16_t, bfloat16_t, op)   \
+  instantiate_reduce_functions(name, float32, float, float, op)              \
+  instantiate_reduce_functions(name, complex64, complex64_t, complex64_t, op)
+
+instantiate_min_max(min, Min)
+instantiate_min_max(max, Max)
+    // clang-format on

data/mlx/mlx/backend/metal/kernels/reduce_utils.h

@@ -0,0 +1,6 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/atomic.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"