Skip using Residency sets in VMs (#1537)

* 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

View File

@ -5,11 +5,10 @@
namespace mlx::core::metal { namespace mlx::core::metal {
// TODO maybe worth including tvos / visionos
#define supported __builtin_available(macOS 15, iOS 18, *)
ResidencySet::ResidencySet(MTL::Device* d) { 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 pool = new_scoped_memory_pool();
auto desc = MTL::ResidencySetDescriptor::alloc()->init(); auto desc = MTL::ResidencySetDescriptor::alloc()->init();
NS::Error* error; NS::Error* error;
@ -27,7 +26,9 @@ ResidencySet::ResidencySet(MTL::Device* d) {
} }
void ResidencySet::insert(MTL::Allocation* buf) { void ResidencySet::insert(MTL::Allocation* buf) {
if (supported) { if (!wired_set_) {
return;
}
if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) { if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) {
wired_set_->addAllocation(buf); wired_set_->addAllocation(buf);
wired_set_->commit(); wired_set_->commit();
@ -35,22 +36,25 @@ void ResidencySet::insert(MTL::Allocation* buf) {
} else { } else {
unwired_set_.insert(buf); unwired_set_.insert(buf);
} }
}
} }
void ResidencySet::erase(MTL::Allocation* buf) { void ResidencySet::erase(MTL::Allocation* buf) {
if (supported) { if (!wired_set_) {
return;
}
if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) { if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) {
unwired_set_.erase(it); unwired_set_.erase(it);
} else { } else {
wired_set_->removeAllocation(buf); wired_set_->removeAllocation(buf);
wired_set_->commit(); wired_set_->commit();
} }
}
} }
void ResidencySet::resize(size_t size) { void ResidencySet::resize(size_t size) {
if (supported) { if (!wired_set_) {
return;
}
if (capacity_ == size) { if (capacity_ == size) {
return; return;
} }
@ -84,11 +88,10 @@ void ResidencySet::resize(size_t size) {
} }
wired_set_->commit(); wired_set_->commit();
} }
}
} }
ResidencySet::~ResidencySet() { ResidencySet::~ResidencySet() {
if (supported) { if (wired_set_) {
wired_set_->release(); wired_set_->release();
} }
} }