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

Skip using Residency sets in VMs (#1537)

Browse Source 
* Skip using Residency sets in VMs

Attempting to use residency sets in a VM throws[^1]

    libc++abi: terminating due to uncaught exception of type std::runtime_error: [metal::Device] Unable to construct residency set.

Not quite sure if this is the best fix, but it does make the error go
away.

Note that it was previously possible to run simple programs that used
mlx in a VM prior to 0eb56d5be0. See
related discussion at Homebrew/homebrew-core#195627.

[^1]: https://github.com/Homebrew/homebrew-core/actions/runs/11525831492/job/32105148462#step:3:56

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>

* change residency check

---------

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>
Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Carlo Cabrera 2024-10-30 10:37:23 +08:00 committed by GitHub
parent d2ff04a4f2
commit 1a992e31e8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 53 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

103
mlx/backend/metal/resident.cpp
View File

@ -5,11 +5,10 @@
namespace mlx::core::metal {
// TODO maybe worth including tvos / visionos
#define supported __builtin_available(macOS 15, iOS 18, *)
ResidencySet::ResidencySet(MTL::Device* d) {
if (supported) {
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;
@ -27,68 +26,72 @@ ResidencySet::ResidencySet(MTL::Device* d) {
}
void ResidencySet::insert(MTL::Allocation* buf) {
if (supported) {
if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) {
wired_set_->addAllocation(buf);
wired_set_->commit();
wired_set_->requestResidency();
} else {
unwired_set_.insert(buf);
}
if (!wired_set_) {
return;
}
if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) {
wired_set_->addAllocation(buf);
wired_set_->commit();
wired_set_->requestResidency();
} else {
unwired_set_.insert(buf);
}
}
void ResidencySet::erase(MTL::Allocation* buf) {
if (supported) {
if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) {
unwired_set_.erase(it);
} else {
wired_set_->removeAllocation(buf);
wired_set_->commit();
}
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 (supported) {
if (capacity_ == size) {
return;
}
capacity_ = size;
if (!wired_set_) {
return;
}
size_t current_size = wired_set_->allocatedSize();
if (capacity_ == size) {
return;
}
capacity_ = size;
if (current_size < size) {
// 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++);
}
size_t current_size = wired_set_->allocatedSize();
if (current_size < size) {
// 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();
wired_set_->requestResidency();
} else if (current_size > size) {
// 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();
}
wired_set_->commit();
wired_set_->requestResidency();
} else if (current_size > size) {
// 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 (supported) {
if (wired_set_) {
wired_set_->release();
}
}