mlx 0.30.7

data/mlx/mlx/backend/metal/kernels/CMakeLists.txt

@@ -0,0 +1,193 @@
+set(BASE_HEADERS
+    bf16.h
+    bf16_math.h
+    complex.h
+    defines.h
+    erf.h
+    expm1f.h
+    fp8.h
+    logging.h
+    utils.h)
+
+function(build_kernel_base TARGET SRCFILE DEPS)
+  set(METAL_FLAGS
+      -x
+      metal
+      -Wall
+      -Wextra
+      -fno-fast-math
+      -Wno-c++17-extensions
+      -Wno-c++20-extensions)
+  if(MLX_METAL_DEBUG)
+    set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
+  endif()
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug" AND MLX_METAL_VERSION GREATER_EQUAL 320)
+    set(METAL_FLAGS ${METAL_FLAGS} -fmetal-enable-logging)
+  endif()
+  if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
+    set(METAL_FLAGS ${METAL_FLAGS}
+                    "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+  endif()
+  add_custom_command(
+    COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE}
+            -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air
+    DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
+    OUTPUT ${TARGET}.air
+    COMMENT "Building ${TARGET}.air"
+    VERBATIM)
+endfunction(build_kernel_base)
+
+function(build_kernel KERNEL)
+  set(SRCFILE ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL}.metal)
+  cmake_path(GET KERNEL STEM TARGET)
+  build_kernel_base(${TARGET} ${SRCFILE} "${ARGN}")
+  set(KERNEL_AIR
+      ${TARGET}.air ${KERNEL_AIR}
+      PARENT_SCOPE)
+endfunction(build_kernel)
+
+build_kernel(arg_reduce)
+build_kernel(conv steel/conv/params.h)
+build_kernel(gemv steel/utils.h)
+build_kernel(layer_norm)
+build_kernel(random)
+build_kernel(rms_norm)
+build_kernel(rope)
+build_kernel(scaled_dot_product_attention sdpa_vector.h)
+if(MLX_METAL_VERSION GREATER_EQUAL 320)
+  build_kernel(fence)
+endif()
+
+set(STEEL_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/conv/conv.h
+    steel/conv/loader.h
+    steel/conv/loaders/loader_channel_l.h
+    steel/conv/loaders/loader_channel_n.h
+    steel/conv/loaders/loader_general.h
+    steel/conv/kernels/steel_conv.h
+    steel/conv/kernels/steel_conv_general.h
+    steel/gemm/gemm.h
+    steel/gemm/mma.h
+    steel/gemm/loader.h
+    steel/gemm/params.h
+    steel/gemm/transforms.h
+    steel/gemm/kernels/steel_gemm_fused.h
+    steel/gemm/kernels/steel_gemm_gather.h
+    steel/gemm/kernels/steel_gemm_masked.h
+    steel/gemm/kernels/steel_gemm_segmented.h
+    steel/gemm/kernels/steel_gemm_splitk.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h)
+
+set(STEEL_ATTN_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/gemm/gemm.h
+    steel/gemm/mma.h
+    steel/gemm/loader.h
+    steel/gemm/transforms.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/attn/attn.h
+    steel/attn/loader.h
+    steel/attn/mma.h
+    steel/attn/params.h
+    steel/attn/transforms.h
+    steel/attn/kernels/steel_attention.h)
+
+set(STEEL_NAX_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/gemm/params.h
+    steel/gemm/transforms.h
+    steel/gemm/nax.h
+    steel/gemm/gemm_nax.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/gemm/kernels/steel_gemm_fused_nax.h
+    steel/gemm/kernels/steel_gemm_gather_nax.h
+    steel/gemm/kernels/steel_gemm_splitk_nax.h)
+
+set(STEEL_NAX_ATTN_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/attn/nax.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/attn/params.h
+    steel/attn/kernels/steel_attention_nax.h)
+
+if(NOT MLX_METAL_JIT)
+  build_kernel(arange arange.h)
+  build_kernel(binary binary.h binary_ops.h)
+  build_kernel(binary_two binary_two.h)
+  build_kernel(copy copy.h)
+  build_kernel(fft fft.h fft/radix.h fft/readwrite.h)
+  build_kernel(
+    reduce
+    atomic.h
+    reduction/ops.h
+    reduction/reduce_init.h
+    reduction/reduce_all.h
+    reduction/reduce_col.h
+    reduction/reduce_row.h)
+  build_kernel(quantized quantized.h quantized_utils.h ${STEEL_HEADERS})
+  build_kernel(fp_quantized fp4.h fp8.h fp_quantized.h quantized_utils.h
+               ${STEEL_HEADERS})
+  build_kernel(scan scan.h)
+  build_kernel(softmax softmax.h)
+  build_kernel(logsumexp logsumexp.h)
+  build_kernel(sort sort.h)
+  build_kernel(ternary ternary.h ternary_ops.h)
+  build_kernel(unary unary.h unary_ops.h)
+  build_kernel(steel/conv/kernels/steel_conv ${STEEL_HEADERS})
+  build_kernel(steel/conv/kernels/steel_conv_general ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_fused ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_gather ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_masked ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_splitk ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_segmented ${STEEL_HEADERS})
+  build_kernel(gemv_masked steel/utils.h)
+  build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
+
+  if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL
+                                                26.2))
+
+    build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS})
+    build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS})
+    build_kernel(steel/gemm/kernels/steel_gemm_splitk_nax ${STEEL_NAX_HEADERS})
+
+    build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS})
+    build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h
+                 ${STEEL_NAX_HEADERS})
+
+    build_kernel(steel/attn/kernels/steel_attention_nax
+                 ${STEEL_NAX_ATTN_HEADERS})
+
+  else()
+    target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
+  endif()
+
+endif()
+
+add_custom_command(
+  OUTPUT ${MLX_METAL_PATH}/mlx.metallib
+  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
+          ${MLX_METAL_PATH}/mlx.metallib
+  DEPENDS ${KERNEL_AIR}
+  COMMENT "Building mlx.metallib"
+  VERBATIM)
+
+add_custom_target(mlx-metallib DEPENDS ${MLX_METAL_PATH}/mlx.metallib)
+
+add_dependencies(mlx mlx-metallib)
+
+# Install metallib
+include(GNUInstallDirs)
+
+install(
+  FILES ${MLX_METAL_PATH}/mlx.metallib
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  COMPONENT metallib)

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

@@ -0,0 +1,9 @@
+// Copyright © 2023-2024 Apple Inc.
+template <typename T>
+[[kernel]] void arange(
+    constant const T& start,
+    constant const T& step,
+    device T* out,
+    uint index [[thread_position_in_grid]]) {
+  out[index] = start + index * step;
+}

data/mlx/mlx/backend/metal/kernels/arange.metal

@@ -0,0 +1,20 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/arange.h"
+
+#define instantiate_arange(tname, type)                                 \
+  instantiate_kernel("arange" #tname, arange, type)
+
+instantiate_arange(uint8, uint8_t)
+instantiate_arange(uint16, uint16_t)
+instantiate_arange(uint32, uint32_t)
+instantiate_arange(uint64, uint64_t)
+instantiate_arange(int8, int8_t)
+instantiate_arange(int16, int16_t)
+instantiate_arange(int32, int32_t)
+instantiate_arange(int64, int64_t)
+instantiate_arange(float16, half)
+instantiate_arange(float32, float)
+instantiate_arange(bfloat16, bfloat16_t) // clang-format on

data/mlx/mlx/backend/metal/kernels/arg_reduce.metal

@@ -0,0 +1,182 @@
+// Copyright © 2023 Apple Inc.
+
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+template <typename U>
+struct IndexValPair {
+  uint32_t index;
+  U val;
+};
+
+template <typename U>
+struct ArgMin {
+  static constexpr constant U init = Limits<U>::max;
+
+  IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> current) {
+    if (best.val > current.val ||
+        (best.val == current.val && best.index > current.index)) {
+      return current;
+    } else {
+      return best;
+    }
+  }
+
+  template <int N>
+  IndexValPair<U>
+  reduce_many(IndexValPair<U> best, thread U* vals, uint32_t offset) {
+    for (int i = 0; i < N; i++) {
+      if (vals[i] < best.val) {
+        best.val = vals[i];
+        best.index = offset + i;
+      }
+    }
+    return best;
+  }
+};
+
+template <typename U>
+struct ArgMax {
+  static constexpr constant U init = Limits<U>::min;
+
+  IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> current) {
+    if (best.val < current.val ||
+        (best.val == current.val && best.index > current.index)) {
+      return current;
+    } else {
+      return best;
+    }
+  }
+
+  template <int N>
+  IndexValPair<U>
+  reduce_many(IndexValPair<U> best, thread U* vals, uint32_t offset) {
+    for (int i = 0; i < N; i++) {
+      if (vals[i] > best.val) {
+        best.val = vals[i];
+        best.index = offset + i;
+      }
+    }
+    return best;
+  }
+};
+
+template <typename U>
+IndexValPair<U> simd_shuffle_down(IndexValPair<U> data, uint16_t delta) {
+  return IndexValPair<U>{
+      simd_shuffle_down(data.index, delta), simd_shuffle_down(data.val, delta)};
+}
+
+template <typename T, typename Op, int N_READS = 4>
+[[kernel]] void arg_reduce_general(
+    const device T* in [[buffer(0)]],
+    device uint32_t* out [[buffer(1)]],
+    const constant int* shape [[buffer(2)]],
+    const constant int64_t* in_strides [[buffer(3)]],
+    const constant int64_t* out_strides [[buffer(4)]],
+    const constant size_t& ndim [[buffer(5)]],
+    const constant int64_t& axis_stride [[buffer(6)]],
+    const constant size_t& axis_size [[buffer(7)]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 gsize [[threads_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint simd_size [[threads_per_simdgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  // Shapes and strides *do not* contain the reduction axis. The reduction size
+  // and stride are provided in axis_stride and axis_size.
+  //
+  // Note: in shape == out shape with this convention.
+  //
+  // The sketch of the kernel is as follows.
+  //    1. Launch prod(shape) * thread_group_size threads.
+  //    2. Loop ceildiv(axis_size / lsize) times
+  //    3. Read input values
+  //    4. Reduce among them and go to 3
+  //    4. Reduce in each simd_group
+  //    6. Write in the thread local memory
+  //    6. Reduce them across thread group
+  //    7. Write the output without need for atomic
+  Op op;
+
+  // Compute the input/output index. There is one beginning and one output for
+  // the whole threadgroup.
+  int64_t row_idx = gid.y + static_cast<int64_t>(gsize.y) * gid.z;
+  auto in_idx = elem_to_loc(row_idx, shape, in_strides, ndim);
+  auto out_idx = elem_to_loc(row_idx, shape, out_strides, ndim);
+
+  IndexValPair<T> best{0, Op::init};
+
+  threadgroup IndexValPair<T> local_data[32];
+
+  // Loop over the reduction axis in lsize*N_READS buckets
+  for (uint r = 0; r < ceildiv(axis_size, N_READS * lsize.x); r++) {
+    // Read the current value
+    uint32_t current_index = r * lsize.x * N_READS + lid.x * N_READS;
+    uint32_t offset = current_index;
+    const device T* current_in = in + in_idx + current_index * axis_stride;
+    T vals[N_READS];
+    for (int i = 0; i < N_READS; i++) {
+      vals[i] = (current_index < axis_size) ? *current_in : T(Op::init);
+      current_index++;
+      current_in += axis_stride;
+    }
+    best = op.template reduce_many<N_READS>(best, vals, offset);
+  }
+  // At this point we have reduced the axis into thread group best values so we
+  // need to reduce across the thread group.
+
+  // First per simd reduction.
+  for (uint offset = simd_size / 2; offset > 0; offset /= 2) {
+    IndexValPair<T> neighbor = simd_shuffle_down(best, offset);
+    best = op.reduce(best, neighbor);
+  }
+
+  // Write to the threadgroup memory
+  if (simd_lane_id == 0) {
+    local_data[simd_group_id] = best;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id != 0) {
+    return;
+  }
+
+  // Read the appropriate value from local data and perform one simd reduction
+  uint simd_groups = ceildiv(lsize.x, simd_size);
+  if (simd_lane_id < simd_groups) {
+    best = local_data[simd_lane_id];
+  }
+  for (uint offset = simd_size / 2; offset > 0; offset /= 2) {
+    IndexValPair<T> neighbor = simd_shuffle_down(best, offset);
+    best = op.reduce(best, neighbor);
+  }
+
+  // Finally write the output
+  if (lid.x == 0) {
+    out[out_idx] = best.index;
+  }
+}
+
+// clang-format off
+#define instantiate_arg_reduce(name, itype)                      \
+  instantiate_kernel(                                            \
+      "argmin_" #name, arg_reduce_general, itype, ArgMin<itype>) \
+  instantiate_kernel(                                            \
+      "argmax_" #name, arg_reduce_general, itype, ArgMax<itype>)
+
+instantiate_arg_reduce(bool_, bool)
+instantiate_arg_reduce(uint8, uint8_t)
+instantiate_arg_reduce(uint16, uint16_t)
+instantiate_arg_reduce(uint32, uint32_t)
+instantiate_arg_reduce(uint64, uint64_t)
+instantiate_arg_reduce(int8, int8_t)
+instantiate_arg_reduce(int16, int16_t)
+instantiate_arg_reduce(int32, int32_t)
+instantiate_arg_reduce(int64, int64_t)
+instantiate_arg_reduce(float16, half)
+instantiate_arg_reduce(float32, float)
+instantiate_arg_reduce(bfloat16, bfloat16_t) // clang-format on