mlx 0.30.7

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

@@ -0,0 +1,144 @@
+function(make_jit_source SRC_FILE)
+  # This function takes a metal header file, runs the C preprocessesor on it,
+  # and makes the processed contents available as a string in a C++ function
+  # mlx::core::metal::${SRC_NAME}()
+  #
+  # To use the function, declare it in jit/includes.h and include
+  # jit/includes.h.
+  #
+  # Additional arguments to this function are treated as dependencies in the
+  # Cmake build system.
+  get_filename_component(SRC_NAME ${SRC_FILE} NAME)
+  add_custom_command(
+    OUTPUT jit/${SRC_NAME}.cpp
+    COMMAND
+      bash ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
+      ${CMAKE_CURRENT_BINARY_DIR}/jit ${CMAKE_C_COMPILER} ${PROJECT_SOURCE_DIR}
+      ${SRC_FILE}
+    DEPENDS make_compiled_preamble.sh kernels/${SRC_FILE}.h ${ARGN})
+  add_custom_target(${SRC_NAME} DEPENDS jit/${SRC_NAME}.cpp)
+  add_dependencies(mlx ${SRC_NAME})
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/jit/${SRC_NAME}.cpp)
+endfunction(make_jit_source)
+
+make_jit_source(utils kernels/bf16.h kernels/bf16_math.h kernels/complex.h
+                kernels/defines.h kernels/logging.h)
+make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h kernels/fp8.h)
+make_jit_source(binary_ops)
+make_jit_source(ternary_ops)
+make_jit_source(reduce_utils kernels/atomic.h kernels/reduction/ops.h)
+make_jit_source(indexing/scatter kernels/indexing/indexing.h)
+make_jit_source(indexing/masked_scatter)
+make_jit_source(indexing/gather kernels/indexing/indexing.h)
+make_jit_source(indexing/gather_front kernels/indexing/indexing.h)
+make_jit_source(indexing/gather_axis)
+make_jit_source(indexing/scatter_axis)
+make_jit_source(hadamard)
+
+if(MLX_METAL_JIT)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jit_kernels.cpp)
+  make_jit_source(arange)
+  make_jit_source(copy)
+  make_jit_source(unary)
+  make_jit_source(binary)
+  make_jit_source(binary_two)
+  make_jit_source(fft kernels/fft/radix.h kernels/fft/readwrite.h)
+  make_jit_source(logsumexp)
+  make_jit_source(ternary)
+  make_jit_source(softmax)
+  make_jit_source(scan)
+  make_jit_source(sort)
+  make_jit_source(
+    reduce kernels/reduction/reduce_all.h kernels/reduction/reduce_col.h
+    kernels/reduction/reduce_row.h kernels/reduction/reduce_init.h)
+  make_jit_source(
+    steel/gemm/gemm kernels/steel/utils.h kernels/steel/gemm/loader.h
+    kernels/steel/gemm/mma.h kernels/steel/gemm/params.h
+    kernels/steel/gemm/transforms.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_fused)
+  make_jit_source(steel/gemm/kernels/steel_gemm_masked kernels/steel/defines.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_gather)
+  make_jit_source(steel/gemm/kernels/steel_gemm_splitk)
+  make_jit_source(steel/gemm/kernels/steel_gemm_segmented)
+  make_jit_source(
+    steel/conv/conv
+    kernels/steel/utils.h
+    kernels/steel/defines.h
+    kernels/steel/gemm/mma.h
+    kernels/steel/gemm/transforms.h
+    kernels/steel/conv/params.h
+    kernels/steel/conv/loader.h
+    kernels/steel/conv/loaders/loader_channel_l.h
+    kernels/steel/conv/loaders/loader_channel_n.h)
+  make_jit_source(steel/conv/kernels/steel_conv)
+  make_jit_source(steel/conv/kernels/steel_conv_general kernels/steel/defines.h
+                  kernels/steel/conv/loaders/loader_general.h)
+
+  make_jit_source(quantized_utils)
+  make_jit_source(quantized kernels/quantized_utils.h)
+  make_jit_source(fp_quantized kernels/quantized_utils.h kernels/fp8.h
+                  kernels/fp4.h)
+  make_jit_source(gemv_masked)
+
+  make_jit_source(steel/attn/kernels/steel_attention)
+
+  make_jit_source(
+    steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h
+    kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax)
+  make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax)
+  make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax)
+
+  make_jit_source(quantized_nax kernels/quantized_utils.h)
+  make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h
+                  kernels/fp4.h)
+
+  make_jit_source(steel/attn/kernels/steel_attention_nax)
+
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp)
+endif()
+
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/normalization.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/unary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/resident.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
+
+if(NOT MLX_METAL_PATH)
+  set(MLX_METAL_PATH ${CMAKE_CURRENT_BINARY_DIR}/kernels/)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels)
+
+target_compile_definitions(mlx
+                           PRIVATE METAL_PATH="${MLX_METAL_PATH}/mlx.metallib")

data/mlx/mlx/backend/metal/allocator.cpp

@@ -0,0 +1,279 @@
+// Copyright © 2023-2024 Apple Inc.
+#include "mlx/backend/metal/allocator.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/backend/metal/resident.h"
+#include "mlx/memory.h"
+
+#include <mach/vm_page_size.h>
+#include <unistd.h>
+#include <cstdlib>
+
+namespace mlx::core {
+
+constexpr size_t resource_options =
+    MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked;
+
+namespace allocator {
+
+Allocator& allocator() {
+  return metal::allocator();
+}
+
+void* Buffer::raw_ptr() {
+  if (!ptr_) {
+    return nullptr;
+  }
+  return static_cast<MTL::Buffer*>(ptr_)->contents();
+}
+
+} // namespace allocator
+
+namespace metal {
+
+MetalAllocator::MetalAllocator()
+    : device_(device(mlx::core::Device::gpu).mtl_device()),
+      buffer_cache_(
+          vm_page_size,
+          [](MTL::Buffer* buf) { return buf->length(); },
+          [this](MTL::Buffer* buf) {
+            if (!buf->heap()) {
+              residency_set_.erase(buf);
+            }
+            buf->release();
+          }),
+      residency_set_(device_) {
+  auto pool = metal::new_scoped_memory_pool();
+  const auto& info = gpu::device_info(0);
+  auto memsize = std::get<size_t>(info.at("memory_size"));
+  auto max_rec_size =
+      std::get<size_t>(info.at("max_recommended_working_set_size"));
+  resource_limit_ = std::get<size_t>(info.at("resource_limit"));
+  block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize);
+  gc_limit_ = std::min(static_cast<size_t>(0.95 * max_rec_size), block_limit_);
+  max_pool_size_ = block_limit_;
+  device(mlx::core::Device::gpu)
+      .set_residency_set(residency_set_.mtl_residency_set());
+  bool is_vm = std::get<std::string>(info.at("device_name")) ==
+      "Apple Paravirtual device";
+  if (is_vm) {
+    return;
+  }
+  auto heap_desc = MTL::HeapDescriptor::alloc()->init();
+  heap_desc->setResourceOptions(resource_options);
+  heap_desc->setSize(heap_size_);
+  heap_ = device_->newHeap(heap_desc);
+  heap_desc->release();
+  residency_set_.insert(heap_);
+}
+
+MetalAllocator::~MetalAllocator() {
+  auto pool = metal::new_scoped_memory_pool();
+  if (heap_) {
+    heap_->release();
+  }
+  buffer_cache_.clear();
+}
+
+size_t MetalAllocator::set_cache_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
+  std::swap(limit, max_pool_size_);
+  return limit;
+};
+
+size_t MetalAllocator::set_memory_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
+  std::swap(limit, block_limit_);
+  gc_limit_ = std::min(
+      block_limit_,
+      static_cast<size_t>(0.95 * device_->recommendedMaxWorkingSetSize()));
+  return limit;
+};
+
+size_t MetalAllocator::get_memory_limit() {
+  return block_limit_;
+}
+
+size_t MetalAllocator::set_wired_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
+  std::swap(limit, wired_limit_);
+  residency_set_.resize(wired_limit_);
+  return limit;
+};
+
+Buffer MetalAllocator::malloc(size_t size) {
+  // Metal doesn't like empty buffers
+  if (size == 0) {
+    return Buffer{nullptr};
+  }
+
+  // More helpful message if maximum buffer length is exceeded
+  if (size > device_->maxBufferLength()) {
+    std::ostringstream msg;
+    msg << "[metal::malloc] Attempting to allocate " << size
+        << " bytes which is greater than"
+        << " the maximum allowed buffer size of " << device_->maxBufferLength()
+        << " bytes.";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Align up memory
+  if (size > vm_page_size) {
+    size = vm_page_size * ((size + vm_page_size - 1) / vm_page_size);
+  }
+
+  // Try the cache
+  std::unique_lock lk(mutex_);
+  MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size);
+  if (!buf) {
+    size_t mem_required = get_active_memory() + get_cache_memory() + size;
+
+    auto pool = metal::new_scoped_memory_pool();
+
+    // If we have a lot of memory pressure try to reclaim memory from the cache
+    if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) {
+      num_resources_ -=
+          buffer_cache_.release_cached_buffers(mem_required - gc_limit_);
+    }
+
+    // Allocate new buffer if needed
+    if (num_resources_ >= resource_limit_) {
+      std::ostringstream msg;
+      msg << "[metal::malloc] Resource limit (" << resource_limit_
+          << ") exceeded.";
+      throw std::runtime_error(msg.str());
+    }
+    lk.unlock();
+    if (size < small_size_ && heap_) {
+      buf = heap_->newBuffer(size, resource_options);
+    }
+    if (!buf) {
+      buf = device_->newBuffer(size, resource_options);
+    }
+    if (!buf) {
+      std::ostringstream msg;
+      msg << "[malloc] Unable to allocate " << size << " bytes.";
+      throw std::runtime_error(msg.str());
+    }
+    lk.lock();
+    num_resources_++;
+    if (!buf->heap()) {
+      residency_set_.insert(buf);
+    }
+  }
+
+  active_memory_ += buf->length();
+  peak_memory_ = std::max(peak_memory_, active_memory_);
+
+  // Maintain the cache below the requested limit
+  if (get_cache_memory() > max_pool_size_) {
+    auto pool = metal::new_scoped_memory_pool();
+    num_resources_ -= buffer_cache_.release_cached_buffers(
+        get_cache_memory() - max_pool_size_);
+  }
+
+  return Buffer{static_cast<void*>(buf)};
+}
+
+void MetalAllocator::clear_cache() {
+  std::unique_lock lk(mutex_);
+  auto pool = metal::new_scoped_memory_pool();
+  num_resources_ -= buffer_cache_.clear();
+}
+
+void MetalAllocator::free(Buffer buffer) {
+  auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
+  if (buf == nullptr) {
+    return;
+  }
+  std::unique_lock lk(mutex_);
+  active_memory_ -= buf->length();
+  if (get_cache_memory() < max_pool_size_) {
+    buffer_cache_.recycle_to_cache(buf);
+  } else {
+    num_resources_--;
+    if (!buf->heap()) {
+      residency_set_.erase(buf);
+    }
+    lk.unlock();
+    auto pool = metal::new_scoped_memory_pool();
+    buf->release();
+  }
+}
+
+size_t MetalAllocator::size(Buffer buffer) const {
+  return static_cast<MTL::Buffer*>(buffer.ptr())->length();
+}
+
+Buffer MetalAllocator::make_buffer(void* ptr, size_t size) {
+  auto buf = device_->newBuffer(ptr, size, resource_options, nullptr);
+  if (!buf) {
+    return Buffer{nullptr};
+  }
+  std::unique_lock lk(mutex_);
+  residency_set_.insert(buf);
+  active_memory_ += buf->length();
+  peak_memory_ = std::max(peak_memory_, active_memory_);
+  num_resources_++;
+  return Buffer{static_cast<void*>(buf)};
+}
+
+void MetalAllocator::release(Buffer buffer) {
+  auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
+  if (buf == nullptr) {
+    return;
+  }
+  std::unique_lock lk(mutex_);
+  active_memory_ -= buf->length();
+  num_resources_--;
+  residency_set_.erase(buf);
+  lk.unlock();
+  auto pool = metal::new_scoped_memory_pool();
+  buf->release();
+}
+
+MetalAllocator& allocator() {
+  // By creating the |allocator_| on heap, the destructor of MetalAllocator
+  // will not be called on exit and buffers in the cache will be leaked. This
+  // can save some time at program exit.
+  static MetalAllocator* allocator_ = new MetalAllocator;
+  return *allocator_;
+}
+
+} // namespace metal
+
+size_t set_cache_limit(size_t limit) {
+  return metal::allocator().set_cache_limit(limit);
+}
+size_t set_memory_limit(size_t limit) {
+  return metal::allocator().set_memory_limit(limit);
+}
+size_t get_memory_limit() {
+  return metal::allocator().get_memory_limit();
+}
+size_t set_wired_limit(size_t limit) {
+  if (limit > std::get<size_t>(
+                  gpu::device_info(0).at("max_recommended_working_set_size"))) {
+    throw std::invalid_argument(
+        "[metal::set_wired_limit] Setting a wired limit larger than "
+        "the maximum working set size is not allowed.");
+  }
+  return metal::allocator().set_wired_limit(limit);
+}
+size_t get_active_memory() {
+  return metal::allocator().get_active_memory();
+}
+size_t get_peak_memory() {
+  return metal::allocator().get_peak_memory();
+}
+void reset_peak_memory() {
+  metal::allocator().reset_peak_memory();
+}
+size_t get_cache_memory() {
+  return metal::allocator().get_cache_memory();
+}
+void clear_cache() {
+  return metal::allocator().clear_cache();
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/allocator.h

@@ -0,0 +1,79 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <map>
+#include <mutex>
+#include <vector>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/buffer_cache.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/resident.h"
+
+namespace mlx::core::metal {
+
+using allocator::Buffer;
+
+class MetalAllocator : public allocator::Allocator {
+  /** Allocator for Metal GPUs. */
+ public:
+  virtual Buffer malloc(size_t size) override;
+  virtual void free(Buffer buffer) override;
+  virtual size_t size(Buffer buffer) const override;
+  virtual Buffer make_buffer(void* ptr, size_t size) override;
+  virtual void release(Buffer buffer) override;
+
+  size_t get_active_memory() {
+    return active_memory_;
+  };
+  size_t get_peak_memory() {
+    return peak_memory_;
+  };
+  void reset_peak_memory() {
+    std::unique_lock lk(mutex_);
+    peak_memory_ = 0;
+  };
+  size_t get_cache_memory() {
+    return buffer_cache_.cache_size();
+  };
+  size_t set_cache_limit(size_t limit);
+  size_t set_memory_limit(size_t limit);
+  size_t get_memory_limit();
+  size_t set_wired_limit(size_t limit);
+  void clear_cache();
+
+ private:
+  MTL::Device* device_;
+
+  // The size of allocations which go on the heap until it is full. This size
+  // is chosen because it is the actual minimum size of a buffer allocated from
+  // the heap, a heap can have at most heap.size() / 256 buffers.
+  static constexpr int small_size_ = 256;
+  static constexpr int heap_size_ = 1 << 20;
+  MTL::Heap* heap_;
+  MetalAllocator();
+  ~MetalAllocator();
+  friend MetalAllocator& allocator();
+
+  // Caching allocator
+  BufferCache<MTL::Buffer> buffer_cache_;
+
+  ResidencySet residency_set_;
+
+  // Allocation stats
+  size_t block_limit_;
+  size_t gc_limit_;
+  size_t active_memory_{0};
+  size_t peak_memory_{0};
+  size_t max_pool_size_;
+  size_t wired_limit_{0};
+  size_t num_resources_{0};
+  size_t resource_limit_{0};
+
+  std::mutex mutex_;
+};
+
+MetalAllocator& allocator();
+
+} // namespace mlx::core::metal