zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-21 12:06:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

remove cache for now

Browse Source
This commit is contained in:
Awni Hannun 2023-12-25 13:45:05 -08:00
parent c67a48be48
commit 8fbd760b90
1 changed files with 21 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
mlx/backend/metal/allocator.cpp
View File

@ -1,6 +1,5 @@
// Copyright © 2023 Apple Inc. // Copyright © 2023 Apple Inc.
#include <iostream>
#include "mlx/backend/metal/allocator.h" #include "mlx/backend/metal/allocator.h"
#include "mlx/backend/metal/metal.h" #include "mlx/backend/metal/metal.h"
@ -31,33 +30,34 @@ MetalAllocator::MetalAllocator()
Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) { Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
// Align up memory // Align up memory
///if (size > vm_page_size) { /// if (size > vm_page_size) {
// size = vm_page_size * ((size + vm_page_size - 1) / vm_page_size); // size = vm_page_size * ((size + vm_page_size - 1) / vm_page_size);
//} //}
// MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size); // MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size);
// Prepare to allocate new memory as needed // Prepare to allocate new memory as needed
// if (!buf) { // if (!buf) {
// If we have memory pressure, first check if we can reclaim some memory // If we have memory pressure, first check if we can reclaim some memory
// from the cache // from the cache
// if (auto new_size = device_->currentAllocatedSize() + size; new_size >= block_limit_) { // if (auto new_size = device_->currentAllocatedSize() + size; new_size >=
// buffer_cache_.clear(); // block_limit_) {
// buffer_cache_.release_cached_buffers( // buffer_cache_.clear();
// std::max(new_size - block_limit_, size)); // buffer_cache_.release_cached_buffers(
// } // std::max(new_size - block_limit_, size));
// }
// If there is still too much memory pressure, fail (likely causes a wait). // If there is still too much memory pressure, fail (likely causes a wait).
// size + allocated (to avoid going over the limit) // size + allocated (to avoid going over the limit)
if (!allow_swap && device_->currentAllocatedSize() + size >= block_limit_) { if (!allow_swap && device_->currentAllocatedSize() + size >= block_limit_) {
return Buffer{nullptr}; return Buffer{nullptr};
} }
// } // }
// Allocate new buffer if needed // Allocate new buffer if needed
size_t res_opt = MTL::ResourceStorageModeShared; size_t res_opt = MTL::ResourceStorageModeShared;
res_opt |= MTL::ResourceHazardTrackingModeTracked; res_opt |= MTL::ResourceHazardTrackingModeTracked;
auto buf = device_->newBuffer(size, res_opt); auto buf = device_->newBuffer(size, res_opt);
peak_allocated_size_ = peak_allocated_size_ =
std::max(peak_allocated_size_, device_->currentAllocatedSize()); std::max(peak_allocated_size_, device_->currentAllocatedSize());