mirror of
https://github.com/ml-explore/mlx.git
synced 2025-07-29 22:01:17 +08:00
0eb56d5be0
* expose residency sets as wire/unwire * returns wired size * fix * runtime support check * fix os check * fix test * fix no metal build * docs * nit * nits in docs * nits
33 lines
652 B
C++
33 lines
652 B
C++
// 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
|