|
|
|
|
|
|
|
|
import unittest |
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx_tests |
|
|
|
|
|
|
|
|
class TestMemory(mlx_tests.MLXTestCase): |
|
|
def test_memory_info(self): |
|
|
old_limit = mx.set_cache_limit(0) |
|
|
|
|
|
a = mx.zeros((4096,)) |
|
|
mx.eval(a) |
|
|
del a |
|
|
self.assertEqual(mx.get_cache_memory(), 0) |
|
|
self.assertEqual(mx.set_cache_limit(old_limit), 0) |
|
|
self.assertEqual(mx.set_cache_limit(old_limit), old_limit) |
|
|
|
|
|
old_limit = mx.set_memory_limit(10) |
|
|
self.assertTrue(mx.set_memory_limit(old_limit), 10) |
|
|
self.assertTrue(mx.set_memory_limit(old_limit), old_limit) |
|
|
|
|
|
|
|
|
a = mx.zeros((4096,)) |
|
|
mx.eval(a) |
|
|
mx.synchronize() |
|
|
active_mem = mx.get_active_memory() |
|
|
self.assertTrue(active_mem >= 4096 * 4) |
|
|
|
|
|
b = mx.zeros((4096,)) |
|
|
mx.eval(b) |
|
|
del b |
|
|
mx.synchronize() |
|
|
|
|
|
new_active_mem = mx.get_active_memory() |
|
|
self.assertEqual(new_active_mem, active_mem) |
|
|
peak_mem = mx.get_peak_memory() |
|
|
self.assertTrue(peak_mem >= 4096 * 8) |
|
|
|
|
|
if mx.metal.is_available(): |
|
|
cache_mem = mx.get_cache_memory() |
|
|
self.assertTrue(cache_mem >= 4096 * 4) |
|
|
|
|
|
mx.clear_cache() |
|
|
self.assertEqual(mx.get_cache_memory(), 0) |
|
|
|
|
|
mx.reset_peak_memory() |
|
|
self.assertEqual(mx.get_peak_memory(), 0) |
|
|
|
|
|
@unittest.skipIf(not mx.metal.is_available(), "Metal is not available") |
|
|
def test_wired_memory(self): |
|
|
old_limit = mx.set_wired_limit(1000) |
|
|
old_limit = mx.set_wired_limit(0) |
|
|
self.assertEqual(old_limit, 1000) |
|
|
|
|
|
max_size = mx.metal.device_info()["max_recommended_working_set_size"] |
|
|
with self.assertRaises(ValueError): |
|
|
mx.set_wired_limit(max_size + 10) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
mlx_tests.MLXTestRunner() |
|
|
|
