mlx 0.30.7

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

@@ -0,0 +1,41 @@
+// Copyright @ 2023 - 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/reduce.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/stream.h"
+
+namespace mlx::core {
+
+using metal::CommandEncoder;
+
+void all_reduce_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s);
+
+void row_reduce_general_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    const ReductionPlan& plan,
+    const std::vector<int>& axes,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s);
+
+void strided_reduce_general_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    const ReductionPlan& plan,
+    const std::vector<int>& axes,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s);
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/resident.cpp

@@ -0,0 +1,100 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/resident.h"
+
+namespace mlx::core::metal {
+
+ResidencySet::ResidencySet(MTL::Device* d) {
+  if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
+    return;
+  } else if (__builtin_available(macOS 15, iOS 18, *)) {
+    auto pool = new_scoped_memory_pool();
+    auto desc = MTL::ResidencySetDescriptor::alloc()->init();
+    NS::Error* error;
+    wired_set_ = d->newResidencySet(desc, &error);
+    desc->release();
+    if (!wired_set_) {
+      std::ostringstream msg;
+      msg << "[metal::Device] Unable to construct residency set.\n";
+      if (error) {
+        msg << error->localizedDescription()->utf8String() << "\n";
+      }
+      throw std::runtime_error(msg.str());
+    }
+    wired_set_->requestResidency();
+  }
+}
+
+void ResidencySet::insert(MTL::Allocation* buf) {
+  if (!wired_set_) {
+    return;
+  }
+  if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) {
+    wired_set_->addAllocation(buf);
+    wired_set_->commit();
+  } else {
+    unwired_set_.insert(buf);
+  }
+}
+
+void ResidencySet::erase(MTL::Allocation* buf) {
+  if (!wired_set_) {
+    return;
+  }
+  if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) {
+    unwired_set_.erase(it);
+  } else {
+    wired_set_->removeAllocation(buf);
+    wired_set_->commit();
+  }
+}
+
+void ResidencySet::resize(size_t size) {
+  if (!wired_set_) {
+    return;
+  }
+
+  if (capacity_ == size) {
+    return;
+  }
+  capacity_ = size;
+
+  size_t current_size = wired_set_->allocatedSize();
+
+  if (current_size < size) {
+    auto pool = new_scoped_memory_pool();
+    // Add unwired allocations to the set
+    for (auto it = unwired_set_.begin(); it != unwired_set_.end();) {
+      auto buf_size = (*it)->allocatedSize();
+      if (current_size + buf_size > size) {
+        it++;
+      } else {
+        current_size += buf_size;
+        wired_set_->addAllocation(*it);
+        unwired_set_.erase(it++);
+      }
+    }
+    wired_set_->commit();
+  } else if (current_size > size) {
+    auto pool = new_scoped_memory_pool();
+    // Remove wired allocations until under capacity
+    auto allocations = wired_set_->allAllocations();
+    auto num_allocations = wired_set_->allocationCount();
+    for (int i = 0; i < num_allocations && current_size > size; ++i) {
+      auto buf = static_cast<const MTL::Allocation*>(allocations->object(i));
+      wired_set_->removeAllocation(buf);
+      current_size -= buf->allocatedSize();
+      unwired_set_.insert(buf);
+    }
+    wired_set_->commit();
+  }
+}
+
+ResidencySet::~ResidencySet() {
+  if (wired_set_) {
+    auto pool = new_scoped_memory_pool();
+    wired_set_->release();
+  }
+}
+
+} // namespace mlx::core::metal

data/mlx/mlx/backend/metal/resident.h

@@ -0,0 +1,32 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/device.h"
+
+namespace mlx::core::metal {
+
+class ResidencySet {
+ public:
+  ResidencySet(MTL::Device* d);
+  ~ResidencySet();
+
+  ResidencySet(const ResidencySet&) = delete;
+  ResidencySet& operator=(const ResidencySet&) = delete;
+
+  const MTL::ResidencySet* mtl_residency_set() {
+    return wired_set_;
+  }
+
+  void insert(MTL::Allocation* buf);
+  void erase(MTL::Allocation* buf);
+
+  void resize(size_t size);
+
+ private:
+  MTL::ResidencySet* wired_set_{nullptr};
+  std::unordered_set<const MTL::Allocation*> unwired_set_;
+  size_t capacity_{0};
+};
+
+} // namespace mlx::core::metal

data/mlx/mlx/backend/metal/rope.cpp

@@ -0,0 +1,165 @@
+// Copyright © 2023-2024 Apple Inc.
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core::fast {
+
+constexpr int n_per_thread = 4;
+
+bool RoPE::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void RoPE::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(outputs.size() == 1);
+  auto& in = inputs[0];
+  auto& out = outputs[0];
+
+  auto& s = out.primitive().stream();
+  auto& d = metal::device(s.device);
+
+  int64_t strides[3];
+  int64_t out_strides[3];
+  bool donated = false;
+  int ndim = in.ndim();
+  int B = in.shape(0);
+  int T = in.shape(-2);
+  int D = in.shape(-1);
+  size_t mat_size = T * D;
+  bool large = in.data_size() > INT32_MAX || in.size() > INT32_MAX;
+
+  int dispatch_ndim = ndim;
+  while (in.shape(-dispatch_ndim) == 1 && dispatch_ndim > 3) {
+    dispatch_ndim--;
+  }
+
+  int N = 1;
+  for (int i = 1; i < (ndim - 2); ++i) {
+    N *= in.shape(i);
+  }
+
+  bool head_seq_transpose = false;
+
+  if (dims_ < D) {
+    donated = true;
+    auto ctype =
+        (in.flags().row_contiguous) ? CopyType::Vector : CopyType::General;
+    copy_gpu(in, out, ctype, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  } else if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      donated = true;
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(allocator::malloc(out.nbytes()));
+    }
+    strides[0] = mat_size;
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else if (dispatch_ndim == 3) {
+    // Handle non-contiguous 3D inputs
+    out.set_data(allocator::malloc(out.nbytes()));
+    strides[0] = in.strides()[ndim - 3];
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else if (
+      ndim == 4 &&
+      // batch dim is regularly strided
+      in.strides()[0] == T * N * D &&
+      // sequence and head dimensions are transposed
+      in.strides()[1] == D && in.strides()[2] == N * D) {
+    head_seq_transpose = true;
+    out.set_data(allocator::malloc(out.nbytes()));
+    strides[0] = in.strides()[1];
+    strides[1] = in.strides()[2];
+    strides[2] = in.strides()[3];
+  } else {
+    // Copy non-contiguous > 3D inputs into the output and treat
+    // input as donated
+    donated = true;
+    copy_gpu(in, out, CopyType::General, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  }
+  out_strides[0] = mat_size;
+  out_strides[1] = out.strides()[ndim - 2];
+  out_strides[2] = out.strides()[ndim - 1];
+
+  // Special case for inference (single time step, contiguous, one offset)
+  auto& offset = inputs[1];
+  bool single = in.flags().row_contiguous && T == 1 && offset.size() == 1;
+
+  bool with_freqs = inputs.size() == 3;
+  std::string kname;
+  concatenate(
+      kname,
+      "rope_",
+      single ? "single_" : "",
+      (with_freqs) ? "freqs_" : "",
+      large ? "large_" : "",
+      type_to_name(in));
+  std::string hash_name;
+  concatenate(
+      hash_name,
+      kname,
+      "_",
+      forward_ ? "" : "vjp_",
+      traditional_ ? "traditional_" : "",
+      head_seq_transpose ? "transpose" : "");
+  metal::MTLFCList func_consts = {
+      {&forward_, MTL::DataType::DataTypeBool, 1},
+      {&traditional_, MTL::DataType::DataTypeBool, 2},
+      {&head_seq_transpose, MTL::DataType::DataTypeBool, 3}};
+
+  auto kernel = d.get_kernel(kname, hash_name, func_consts);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  float base = std::log2(base_);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(donated ? out : in, 0);
+  compute_encoder.set_output_array(out, 1);
+
+  compute_encoder.set_input_array(offset, 2);
+  compute_encoder.set_bytes(scale_, 3);
+
+  MTL::Size group_dims;
+  MTL::Size grid_dims;
+  if (single) {
+    compute_encoder.set_bytes(out_strides, 1, 4);
+    uint32_t dim0 = dims_ / 2;
+    group_dims = get_block_dims(dim0, N, 1);
+    grid_dims = MTL::Size(dim0, N, 1);
+  } else {
+    compute_encoder.set_bytes(strides, 3, 4);
+    compute_encoder.set_bytes(out_strides, 3, 5);
+    int64_t offset_stride = 0;
+    if (offset.ndim() > 0) {
+      offset_stride = offset.strides()[0];
+    }
+    compute_encoder.set_bytes(offset_stride, 6);
+    compute_encoder.set_bytes(N, 7);
+    uint32_t dim0 = dims_ / 2;
+    uint32_t dim1 = T;
+    uint32_t dim2 = B * ((N + n_per_thread - 1) / n_per_thread);
+    group_dims = get_block_dims(dim0, dim1, dim2);
+    grid_dims = MTL::Size(dim0, dim1, dim2);
+  }
+
+  if (with_freqs) {
+    auto& freqs = inputs[2];
+    compute_encoder.set_input_array(freqs, 10);
+    auto freq_stride = freqs.strides()[0];
+    compute_encoder.set_bytes(freq_stride, 11);
+  } else {
+    compute_encoder.set_bytes(base, 10);
+  }
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+} // namespace mlx::core::fast