mlx 0.30.7

data/mlx/mlx/backend/metal/copy.cpp

@@ -0,0 +1,235 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+
+namespace mlx::core {
+
+constexpr int MAX_COPY_SPECIALIZED_DIMS = 3;
+
+void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
+  bool donated = set_copy_output_data(in, out, ctype);
+  if (donated && in.dtype() == out.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
+  }
+  if (ctype == CopyType::GeneralGeneral) {
+    ctype = CopyType::General;
+  }
+  copy_gpu_inplace(in, out, ctype, s);
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Shape& data_shape,
+    const Strides& strides_in_pre,
+    const Strides& strides_out_pre,
+    int64_t inp_offset,
+    int64_t out_offset,
+    CopyType ctype,
+    const Stream& s,
+    std::optional<array> dynamic_i_offset /* = std::nullopt */,
+    std::optional<array> dynamic_o_offset /* = std::nullopt */) {
+  if (out.size() == 0) {
+    return;
+  }
+
+  // Try to collapse contiguous dims
+  auto maybe_collapse =
+      [ctype, &data_shape, &strides_in_pre, &strides_out_pre]() {
+        if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+          auto [shape, strides] = collapse_contiguous_dims(
+              data_shape,
+              std::vector{strides_in_pre, strides_out_pre},
+              /* size_cap = */ INT32_MAX);
+          return std::make_tuple(shape, strides[0], strides[1]);
+        } else {
+          Strides e{};
+          return std::make_tuple(Shape{}, e, e);
+        }
+      };
+  auto [shape, strides_in_, strides_out_] = maybe_collapse();
+  int ndim = shape.size();
+  bool large;
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    // Allow for negative strides
+    large = in.data_size() > INT32_MAX || out.data_size() > INT32_MAX;
+  } else {
+    large = out.data_size() > UINT32_MAX;
+  }
+  bool dynamic = dynamic_i_offset || dynamic_o_offset;
+  auto& d = metal::device(s.device);
+  int work_per_thread = 1;
+  std::string kernel_name;
+  switch (ctype) {
+    case CopyType::Scalar:
+      kernel_name = large ? "s2" : "s";
+      break;
+    case CopyType::Vector:
+      kernel_name = large ? "v2" : "v";
+      break;
+    case CopyType::General:
+      kernel_name = "g";
+      break;
+    case CopyType::GeneralGeneral:
+      kernel_name = "gg";
+      break;
+  }
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    if (shape.size() <= MAX_COPY_SPECIALIZED_DIMS) {
+      kernel_name += std::to_string(shape.size());
+    } else {
+      work_per_thread = large ? 4 : 2;
+      concatenate(kernel_name, "n", std::to_string(work_per_thread));
+    }
+    if (large) {
+      kernel_name += "large";
+    }
+    if (dynamic) {
+      kernel_name += "_dynamic";
+      if (ctype != CopyType::GeneralGeneral) {
+        throw std::runtime_error(
+            "[Copy::eval_gpu] Dynamic output offset requires GeneralGeneral copy");
+      }
+    }
+  } else {
+    work_per_thread = get_work_per_thread(out.dtype(), out.data_size());
+    if (!large && work_per_thread > 1) {
+      kernel_name += "n";
+    }
+  }
+  concatenate(kernel_name, "_copy", type_to_name(in), type_to_name(out));
+  auto kernel = dynamic ? get_dynamic_copy_kernel(d, kernel_name, in, out)
+                        : get_copy_kernel(d, kernel_name, in, out);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  inp_offset *= size_of(in.dtype());
+  out_offset *= size_of(out.dtype());
+
+  compute_encoder.set_input_array(in, 0, inp_offset);
+  compute_encoder.set_output_array(out, 1, out_offset);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    Strides strides_in{strides_in_.begin(), strides_in_.end()};
+    Strides strides_out{strides_out_.begin(), strides_out_.end()};
+    if (ndim > 3) {
+      compute_encoder.set_vector_bytes(shape, ndim, 2);
+    }
+    compute_encoder.set_vector_bytes(strides_in, ndim, 3);
+    if (ctype == CopyType::GeneralGeneral) {
+      compute_encoder.set_vector_bytes(strides_out, ndim, 4);
+    }
+
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+
+    size_t data_size = 1;
+    for (auto& s : shape)
+      data_size *= s;
+    size_t rest = data_size / (dim0 * dim1);
+
+    if (ndim > MAX_COPY_SPECIALIZED_DIMS) {
+      compute_encoder.set_bytes(ndim, 5);
+      dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    }
+    if (dynamic) {
+      if (dynamic_i_offset) {
+        compute_encoder.set_input_array(*dynamic_i_offset, 6);
+      } else {
+        compute_encoder.set_bytes(0ll, 6);
+      }
+      if (dynamic_o_offset) {
+        compute_encoder.set_input_array(*dynamic_o_offset, 7);
+      } else {
+        compute_encoder.set_bytes(0ll, 7);
+      }
+    }
+
+    // NB assuming thread_group_size is a power of 2 larger than 32 x 32
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::copy] Must use 1024 sized block");
+    }
+
+    auto group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(out.data_size(), 2);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(out.data_size(), 2);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void fill_gpu(const array& val, array& out, const Stream& s) {
+  if (out.size() == 0) {
+    return;
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+  bool large = out.data_size() > UINT32_MAX;
+  int work_per_thread = get_work_per_thread(out.dtype(), out.data_size());
+  auto& d = metal::device(s.device);
+  std::string kernel_name = large ? "s2" : (work_per_thread > 1 ? "sn" : "s");
+  concatenate(kernel_name, "_copy", type_to_name(val), type_to_name(out));
+  auto kernel = get_copy_kernel(d, kernel_name, val, out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(val, 0);
+  compute_encoder.set_output_array(out, 1);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+  if (thread_group_size > nthreads) {
+    thread_group_size = nthreads;
+  }
+  MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+  MTL::Size grid_dims;
+  if (large) {
+    compute_encoder.set_bytes<int64_t>(out.data_size(), 2);
+    grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+  } else {
+    compute_encoder.set_bytes<int>(out.data_size(), 2);
+    grid_dims = MTL::Size(nthreads, 1, 1);
+  }
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void reshape_gpu(const array& in, array& out, Stream s) {
+  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
+  if (copy_necessary) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    copy_gpu_inplace(
+        in,
+        out,
+        in.shape(),
+        in.strides(),
+        make_contiguous_strides(in.shape()),
+        0,
+        0,
+        CopyType::General,
+        s);
+  } else {
+    shared_buffer_reshape(in, out_strides, out);
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/custom_kernel.cpp

@@ -0,0 +1,430 @@
+// Copyright © 2024 Apple Inc.
+
+#include <iostream>
+#include <regex>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fast {
+
+struct CustomKernelCache {
+  std::unordered_map<std::string, std::string> libraries;
+};
+
+static CustomKernelCache& cache() {
+  static CustomKernelCache cache_;
+  return cache_;
+};
+
+std::string write_signature(
+    std::string func_name,
+    const std::string& header,
+    const std::string& source,
+    const std::vector<std::string>& input_names,
+    const std::vector<array>& inputs,
+    const std::vector<std::string>& output_names,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args,
+    const std::vector<std::string>& attributes,
+    const std::vector<std::tuple<bool, bool, bool>>& shape_infos,
+    bool atomic_outputs) {
+  std::string kernel_source;
+  kernel_source.reserve(header.size() + source.size() + 16384);
+  kernel_source += header;
+  // Auto-generate a function signature based on `template_args`
+  // and the dtype/shape of the arrays passed as `inputs`.
+  if (!template_args.empty()) {
+    kernel_source += "template <";
+    int i = 0;
+    for (const auto& [name, arg] : template_args) {
+      std::string param_type;
+      if (std::holds_alternative<int>(arg)) {
+        param_type = "int";
+      } else if (std::holds_alternative<bool>(arg)) {
+        param_type = "bool";
+      } else if (std::holds_alternative<Dtype>(arg)) {
+        param_type = "typename";
+      }
+      if (i > 0) {
+        kernel_source += ", ";
+      }
+      kernel_source += param_type;
+      kernel_source += " ";
+      kernel_source += name;
+      i++;
+    }
+    kernel_source += ">\n";
+  }
+  kernel_source += "[[kernel]] void ";
+  kernel_source += func_name;
+  kernel_source += "(\n";
+
+  int index = 0;
+  constexpr int max_constant_array_size = 8;
+  // Add inputs
+  for (int i = 0; i < inputs.size(); ++i) {
+    const auto& name = input_names[i];
+    const auto& arr = inputs[i];
+    auto dtype = get_type_string(arr.dtype());
+    std::string location =
+        arr.size() < max_constant_array_size ? "constant" : "device";
+    std::string ref = arr.ndim() == 0 ? "&" : "*";
+    kernel_source += "  const ";
+    kernel_source += location;
+    kernel_source += " ";
+    kernel_source += dtype;
+    kernel_source += ref;
+    kernel_source += " ";
+    kernel_source += name;
+    kernel_source += " [[buffer(";
+    kernel_source += std::to_string(index);
+    kernel_source += ")]],\n";
+    index++;
+    // Add input shape, strides and ndim if present in the source
+    if (arr.ndim() > 0) {
+      if (std::get<0>(shape_infos[i])) {
+        kernel_source +=
+            ("  const constant int* " + name + "_shape [[buffer(" +
+             std::to_string(index) + ")]],\n");
+        index++;
+      }
+      if (std::get<1>(shape_infos[i])) {
+        kernel_source +=
+            ("  const constant int64_t* " + name + "_strides [[buffer(" +
+             std::to_string(index) + ")]],\n");
+        index++;
+      }
+      if (std::get<2>(shape_infos[i])) {
+        kernel_source +=
+            ("  const constant int& " + name + "_ndim [[buffer(" +
+             std::to_string(index) + ")]],\n");
+        index++;
+      }
+    }
+  }
+  // Add outputs
+  for (int i = 0; i < output_names.size(); ++i) {
+    const auto& name = output_names[i];
+    const auto& dtype = output_dtypes[i];
+    kernel_source += "  device ";
+    auto type_string = get_type_string(dtype);
+    if (atomic_outputs) {
+      kernel_source += "atomic<";
+    }
+    kernel_source += type_string;
+    if (atomic_outputs) {
+      kernel_source += ">";
+    }
+    kernel_source += "* ";
+    kernel_source += name;
+    kernel_source += " [[buffer(";
+    kernel_source += std::to_string(index);
+    kernel_source += ")]]";
+    if (index < inputs.size() + output_names.size() - 1 ||
+        attributes.size() > 0) {
+      kernel_source += ",\n";
+    } else {
+      kernel_source += ") {\n";
+    }
+    index++;
+  }
+
+  index = 0;
+  for (const auto& attr : attributes) {
+    kernel_source += attr;
+    if (index < attributes.size() - 1) {
+      kernel_source += ",\n";
+    } else {
+      kernel_source += ") {\n";
+    }
+    index++;
+  }
+  kernel_source += source;
+  kernel_source += "\n}\n";
+  return kernel_source;
+}
+
+std::string write_template(
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args) {
+  std::ostringstream template_def;
+  template_def << "<";
+  int i = 0;
+  for (const auto& [name, arg] : template_args) {
+    if (i > 0) {
+      template_def << ", ";
+    }
+    if (std::holds_alternative<int>(arg)) {
+      template_def << std::get<int>(arg);
+    } else if (std::holds_alternative<bool>(arg)) {
+      template_def << std::get<bool>(arg);
+    } else if (std::holds_alternative<Dtype>(arg)) {
+      template_def << get_type_string(std::get<Dtype>(arg));
+    }
+    i++;
+  }
+  template_def << ">";
+  return template_def.str();
+}
+
+CustomKernelFunction metal_kernel(
+    const std::string& name,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::string>& output_names,
+    const std::string& source,
+    const std::string& header /* = "" */,
+    bool ensure_row_contiguous /* = true */,
+    bool atomic_outputs /* = false */) {
+  if (output_names.empty()) {
+    throw std::invalid_argument(
+        "[metal_kernel] Must specify at least one output.");
+  }
+  std::vector<std::tuple<bool, bool, bool>> shape_infos;
+  for (auto& n : input_names) {
+    std::tuple<bool, bool, bool> shape_info;
+    std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos;
+    std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos;
+    std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos;
+    shape_infos.push_back(shape_info);
+  }
+  const std::vector<std::pair<std::string, std::string>> metal_attributes = {
+      {"dispatch_quadgroups_per_threadgroup", "uint"},
+      {"dispatch_simdgroups_per_threadgroup", "uint"},
+      {"dispatch_threads_per_threadgroup", "uint3"},
+      {"grid_origin", "uint3"},
+      {"grid_size", "uint3"},
+      {"quadgroup_index_in_threadgroup", "uint"},
+      {"quadgroups_per_threadgroup", "uint"},
+      {"simdgroup_index_in_threadgroup", "uint"},
+      {"simdgroups_per_threadgroup", "uint"},
+      {"thread_execution_width", "uint"},
+      {"thread_index_in_quadgroup", "uint"},
+      {"thread_index_in_simdgroup", "uint"},
+      {"thread_index_in_threadgroup", "uint"},
+      {"thread_position_in_grid", "uint3"},
+      {"thread_position_in_threadgroup", "uint3"},
+      {"threadgroup_position_in_grid", "uint3"},
+      {"threadgroups_per_grid", "uint3"},
+      {"threads_per_grid", "uint3"},
+      {"threads_per_simdgroup", "uint"},
+      {"threads_per_threadgroup", "uint3"},
+  };
+
+  std::vector<std::string> attributes;
+  for (const auto& [attr, dtype] : metal_attributes) {
+    if (source.find(attr) != std::string::npos) {
+      attributes.push_back("  " + dtype + " " + attr + " [[" + attr + "]]");
+    }
+  }
+
+  return [=,
+          shape_infos = std::move(shape_infos),
+          attributes = std::move(attributes)](
+             const std::vector<array>& inputs,
+             const std::vector<Shape>& output_shapes,
+             const std::vector<Dtype>& output_dtypes,
+             std::tuple<int, int, int> grid,
+             std::tuple<int, int, int> threadgroup,
+             const std::vector<std::pair<std::string, TemplateArg>>&
+                 template_args = {},
+             std::optional<float> init_value = std::nullopt,
+             bool verbose = false,
+             StreamOrDevice s_ = {}) {
+    if (inputs.size() != input_names.size()) {
+      std::ostringstream msg;
+      msg << "[metal_kernel] Expected `inputs` to have size "
+          << input_names.size() << " but got size " << inputs.size() << "."
+          << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_shapes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[metal_kernel] Expected `output_shapes` to have size "
+          << output_names.size() << " but got size " << output_shapes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_dtypes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[metal_kernel] Expected `output_dtypes` to have size "
+          << output_names.size() << " but got size " << output_dtypes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+
+    auto s = to_stream(s_);
+    if (s.device != Device::gpu) {
+      throw std::invalid_argument("[metal_kernel] Only supports the GPU.");
+    }
+
+    std::string kernel_name = "custom_kernel_" + name;
+    std::string template_def = "";
+    if (!template_args.empty()) {
+      std::regex disallowed_chars("\\<|\\>|(, )");
+      template_def = write_template(template_args);
+      auto template_hash =
+          std::regex_replace(template_def, disallowed_chars, "_");
+      template_hash.pop_back();
+      kernel_name += "_";
+      kernel_name += template_hash;
+    }
+
+    std::string kernel_source = write_signature(
+        kernel_name,
+        header,
+        source,
+        input_names,
+        inputs,
+        output_names,
+        output_dtypes,
+        template_args,
+        attributes,
+        shape_infos,
+        atomic_outputs);
+
+    if (!template_args.empty()) {
+      template_def = kernel_name + template_def;
+      kernel_source += "\ntemplate [[host_name(\"";
+      kernel_source += kernel_name;
+      kernel_source += "\")]] [[kernel]] decltype(";
+      kernel_source += template_def;
+      kernel_source += ") ";
+      kernel_source += template_def;
+      kernel_source += ";\n";
+    }
+
+    if (verbose) {
+      std::cout << "Generated source code for `" << name << "`:" << std::endl
+                << "```" << std::endl
+                << kernel_source << std::endl
+                << "```" << std::endl;
+    }
+
+    return array::make_arrays(
+        std::move(output_shapes),
+        std::move(output_dtypes),
+        std::make_shared<CustomKernel>(
+            s,
+            std::move(kernel_name),
+            std::move(kernel_source),
+            grid,
+            threadgroup,
+            shape_infos,
+            ensure_row_contiguous,
+            init_value,
+            std::vector<ScalarArg>{},
+            false,
+            0),
+        std::move(inputs));
+  };
+}
+
+void CustomKernel::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  // silence some warnings
+  (void)is_precompiled_;
+  (void)shared_memory_;
+
+  auto& s = stream();
+
+  std::vector<array> copies;
+
+  for (auto& out : outputs) {
+    if (init_value_) {
+      copies.emplace_back(init_value_.value(), out.dtype());
+      fill_gpu(copies.back(), out, s);
+    } else {
+      out.set_data(allocator::malloc(out.nbytes()));
+    }
+  }
+
+  auto check_input = [&copies, &s, this](const array& x) -> const array {
+    bool no_copy = x.flags().row_contiguous;
+    if (!ensure_row_contiguous_ || no_copy) {
+      return x;
+    } else {
+      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
+      copy_gpu(x, copies.back(), CopyType::General, s);
+      return copies.back();
+    }
+  };
+  std::vector<array> checked_inputs;
+  for (const array& in : inputs) {
+    checked_inputs.push_back(check_input(in));
+  }
+
+  auto& d = metal::device(s.device);
+
+  {
+    // Clear kernels from the device library cache if needed
+    auto& kernel_cache = cache();
+    if (auto it = kernel_cache.libraries.find(name_);
+        it != kernel_cache.libraries.end()) {
+      if (it->second != source_) {
+        auto& d = metal::device(s.device);
+        d.clear_library(name_);
+        it->second = source_;
+      }
+    } else {
+      kernel_cache.libraries.emplace(name_, source_);
+    }
+  }
+
+  auto lib = d.get_library(name_, [this] { return metal::utils() + source_; });
+  auto kernel = d.get_kernel(name_, lib);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  int index = 0;
+  for (int i = 0; i < checked_inputs.size(); i++) {
+    const array& in = checked_inputs[i];
+    auto& shape_info = shape_infos_[i];
+    compute_encoder.set_input_array(in, index);
+    index++;
+    if (in.ndim() > 0) {
+      int ndim = in.ndim();
+      if (std::get<0>(shape_info)) {
+        compute_encoder.set_vector_bytes(in.shape(), ndim, index);
+        index++;
+      }
+      if (std::get<1>(shape_info)) {
+        compute_encoder.set_vector_bytes(in.strides(), ndim, index);
+        index++;
+      }
+      if (std::get<2>(shape_info)) {
+        compute_encoder.set_bytes(ndim, index);
+        index++;
+      }
+    }
+  }
+  for (auto& out : outputs) {
+    compute_encoder.set_output_array(out, index);
+    index++;
+  }
+
+  const auto [tx, ty, tz] = threadgroup_;
+  auto tg_size = tx * ty * tz;
+  auto max_tg_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (tg_size > max_tg_size) {
+    std::ostringstream msg;
+    msg << "Thread group size (" << tg_size << ") is greater than "
+        << " the maximum allowed threads per threadgroup (" << max_tg_size
+        << ").";
+    throw std::invalid_argument(msg.str());
+  }
+
+  const auto [gx, gy, gz] = grid_;
+  MTL::Size group_dims =
+      MTL::Size(std::min(tx, gx), std::min(ty, gy), std::min(tz, gz));
+  MTL::Size grid_dims = MTL::Size(gx, gy, gz);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+} // namespace mlx::core::fast