mlx 1.0.0

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

@@ -0,0 +1,375 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <fmt/format.h>
+
+#include "mlx/array.h"
+#include "mlx/backend/metal/device.h"
+
+namespace mlx::core {
+
+MTL::ComputePipelineState* get_arange_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out);
+
+MTL::ComputePipelineState* get_unary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op);
+
+MTL::ComputePipelineState* get_binary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op);
+
+MTL::ComputePipelineState* get_binary_two_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op);
+
+MTL::ComputePipelineState* get_ternary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype type,
+    const char* op);
+
+MTL::ComputePipelineState* get_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out);
+
+MTL::ComputePipelineState* get_dynamic_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out);
+
+MTL::ComputePipelineState* get_softmax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool precise,
+    const array& out);
+
+MTL::ComputePipelineState* get_logsumexp_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out);
+
+MTL::ComputePipelineState* get_scan_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool reverse,
+    bool inclusive,
+    const std::string& reduce_type,
+    const array& in,
+    const array& out);
+
+MTL::ComputePipelineState* get_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    int bn,
+    int tn);
+
+MTL::ComputePipelineState* get_mb_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& idx,
+    int bn,
+    int tn);
+
+MTL::ComputePipelineState* get_reduce_init_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& func_name,
+    const std::string& op_name,
+    const Dtype& out_type);
+
+MTL::ComputePipelineState* get_reduce_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& func_name,
+    const std::string& op_name,
+    const Dtype& in_type,
+    const Dtype& out_type,
+    const std::string& idx_t,
+    int ndim = -1,
+    int bm = -1,
+    int bn = -1);
+
+MTL::ComputePipelineState* get_steel_gemm_fused_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool mn_aligned,
+    bool k_aligned);
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_accum_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    bool axbpy);
+
+MTL::ComputePipelineState* get_steel_gemm_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool mn_aligned,
+    bool k_aligned);
+
+MTL::ComputePipelineState* get_steel_gemm_gather_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs);
+
+MTL::ComputePipelineState* get_steel_gemm_segmented_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_conv_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    int n_channel_specialization,
+    bool small_filter);
+
+MTL::ComputePipelineState* get_gemv_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_mat,
+    int bm,
+    int bn,
+    int sm,
+    int sn,
+    int tm,
+    int tn,
+    bool contiguous);
+
+MTL::ComputePipelineState* get_steel_conv_general_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_fft_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const std::string& template_def);
+
+MTL::ComputePipelineState* get_quantized_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode);
+
+MTL::ComputePipelineState* get_gather_qmm_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose);
+
+MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs);
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode);
+
+MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose);
+
+MTL::ComputePipelineState* get_steel_attention_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m);
+
+MTL::ComputePipelineState* get_steel_attention_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m);
+
+// Create a GPU kernel template definition for JIT compilation
+template <typename... Args>
+std::string get_template_definition(
+    std::string_view name,
+    std::string_view func,
+    Args... args) {
+  std::ostringstream s;
+  s << func << "<";
+  bool first = true;
+  auto add_arg = [&s, &first](const auto& arg) {
+    if (!first) {
+      s << ", ";
+    }
+    first = false;
+    s << arg;
+  };
+  (add_arg(args), ...);
+  s << ">";
+  return fmt::format(
+      "\ntemplate [[host_name(\"{0}\")]] [[kernel]] decltype({1}) {1};\n",
+      name,
+      s.str());
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/logsumexp.cpp

@@ -0,0 +1,95 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+constexpr int LOGSUMEXP_LOOPED_LIMIT = 4096;
+
+void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  if (!issubdtype(out.dtype(), floating)) {
+    throw std::runtime_error(
+        "[logsumexp] Does not support non-floating point types.");
+  }
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Make sure that the last dimension is contiguous
+  auto ensure_contiguous = [&s, &d](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      d.add_temporary(x_copy, s.index);
+      return x_copy;
+    }
+  };
+
+  auto in = ensure_contiguous(inputs[0]);
+  if (in.flags().row_contiguous) {
+    out.set_data(allocator::malloc(out.nbytes()));
+  } else {
+    auto n = in.shape(-1);
+    auto flags = in.flags();
+    auto strides = in.strides();
+    for (auto& s : strides) {
+      s /= n;
+    }
+    bool col_contig = strides[0] == 1;
+    for (int i = 1; col_contig && i < strides.size(); ++i) {
+      col_contig &=
+          (out.shape(i) == 1 || strides[i - 1] == out.shape(i) * strides[i]);
+    }
+    flags.col_contiguous = col_contig;
+    out.set_data(
+        allocator::malloc(in.nbytes() / n),
+        in.data_size() / n,
+        std::move(strides),
+        flags);
+  }
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = 4;
+  const int looped_limit = LOGSUMEXP_LOOPED_LIMIT;
+
+  std::string kernel_name = (axis_size > looped_limit) ? "looped_" : "block_";
+  kernel_name += "logsumexp_";
+  kernel_name += type_to_name(out);
+
+  auto kernel = get_logsumexp_kernel(d, kernel_name, out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(axis_size, 2);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/make_compiled_preamble.sh

@@ -0,0 +1,120 @@
+#!/bin/bash
+#
+# This script generates a C++ function that provides the Metal source code
+# at runtime for use with kernel generation.
+#
+# The steps executed are as follows 
+# - Take as input a metal header file in the mlx metal backend 
+# - Use the metal compiler to expand the dependency headers 
+# - Sort the headers in order of inclusion 
+# - Expand the headers in order of inclusion 
+# - Export the generated source code content as a C++ function
+#
+# Doing the expansion this way allows us to retain macros, comments, and 
+# formatting in the expanded source. This adds user readibility, and also 
+# enables use of the metal macros in the source code which can then be 
+# handled by the metal runtime compiler
+#
+# Copyright © 2023-25 Apple Inc.
+
+OUTPUT_DIR=$1
+CC=$2
+SRC_DIR=$3
+SRC_FILE=$4
+CFLAGS=$5
+SRC_NAME=$(basename -- "${SRC_FILE}")
+JIT_INCLUDES=${SRC_DIR}/mlx/backend/metal/kernels/jit
+INPUT_FILE=${SRC_DIR}/mlx/backend/metal/kernels/${SRC_FILE}.h
+OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
+
+# Prepare output
+mkdir -p "$OUTPUT_DIR"
+
+# Use the metal compiler to get a list of headers (with depth)
+CCC="xcrun -sdk macosx metal -x metal"
+HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null )
+
+# Remove any included system frameworks (for MetalPerformancePrimitive headers)
+HDRS=$(echo "$HDRS" | grep -v "Xcode")
+
+# Use the header depth to sort the files in order of inclusion
+declare -a HDRS_LIST=($HDRS)
+declare -a HDRS_STACK=()
+declare -a HDRS_SORTED=()
+
+length=${#HDRS_LIST[@]}
+
+HDRS_LIST+=(".")
+
+for ((i=0; i<${length}; i+=2));
+do 
+
+  header="${HDRS_LIST[$i+1]#$SRC_DIR/}"
+
+  str_this="${HDRS_LIST[$i]}"
+  str_next="${HDRS_LIST[$i + 2]}"
+
+  depth_this=${#str_this}
+  depth_next=${#str_next}
+
+  # If we have a dependency then we stack it
+  if [ $depth_next -gt $depth_this ]; then 
+    HDRS_STACK=($header ${HDRS_STACK[@]})
+
+  # If we are done with this level 
+  else 
+    # We add the header to out list
+    HDRS_SORTED+=($header) 
+
+    # Pop the stacked up dependencies
+    pop_len=$((depth_this - depth_next))
+    for popped_header in "${HDRS_STACK[@]:0:$pop_len}"
+    do 
+      HDRS_SORTED+=($popped_header)
+    done 
+
+    HDRS_STACK=(${HDRS_STACK[@]:$pop_len})
+  fi  
+
+done
+
+# Make sure the given metal header is also expanded in the source content
+HDRS_SORTED+=("${INPUT_FILE#$SRC_DIR/}")
+
+# Expand the headers in order of inclusion 
+CONTENT=$(
+echo "// Copyright © 2025 Apple Inc."
+echo "" 
+echo "// Auto generated source for ${INPUT_FILE#$SRC_DIR/}"
+echo ""
+
+for header in "${HDRS_SORTED[@]}"
+do 
+  echo "///////////////////////////////////////////////////////////////////////////////"
+  echo "// Contents from \"${header}\""
+  echo "///////////////////////////////////////////////////////////////////////////////"
+  echo ""
+
+  echo "#line 1 \"${header}\""
+
+  grep -h -v -G -e "#include \".*.h\"" -e "#pragma once" "${SRC_DIR}/${header}" 
+  
+  echo ""
+  
+done
+
+echo "///////////////////////////////////////////////////////////////////////////////"
+)
+
+# Export the generated source code content as a C++ function
+cat << EOF > "$OUTPUT_FILE"
+namespace mlx::core::metal {
+
+const char* $SRC_NAME() {
+  return R"preamble(
+$CONTENT
+)preamble";
+}
+
+} // namespace mlx::core::metal
+EOF