mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-25 01:41:17 +08:00
8a0677d56d
* more async eval * fix rebase * try correct async eval * fix async * more tests for async eval * use shared events for synchronization * comment + cleanup * with autorelease pool * fix no metal build * fix compile * fix patch * don't eval if asyn evale'd * don't use is_evaled * comments * more multi stream tests * try and cleanup use of is_evaled * use a status flag
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
# Copyright © 2023-2024 Apple Inc.
|
|
|
|
import unittest
|
|
|
|
import mlx.core as mx
|
|
import mlx_tests
|
|
|
|
|
|
class TestMetal(mlx_tests.MLXTestCase):
|
|
|
|
@unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
|
|
def test_memory_info(self):
|
|
old_limit = mx.metal.set_cache_limit(0)
|
|
|
|
a = mx.zeros((4096,))
|
|
mx.eval(a)
|
|
del a
|
|
self.assertEqual(mx.metal.get_cache_memory(), 0)
|
|
self.assertEqual(mx.metal.set_cache_limit(old_limit), 0)
|
|
self.assertEqual(mx.metal.set_cache_limit(old_limit), old_limit)
|
|
|
|
old_limit = mx.metal.set_memory_limit(10)
|
|
self.assertTrue(mx.metal.set_memory_limit(old_limit), 10)
|
|
self.assertTrue(mx.metal.set_memory_limit(old_limit), old_limit)
|
|
|
|
# Query active and peak memory
|
|
a = mx.zeros((4096,), stream=mx.cpu)
|
|
mx.eval(a)
|
|
active_mem = mx.metal.get_active_memory()
|
|
self.assertTrue(active_mem >= 4096 * 4)
|
|
|
|
b = mx.zeros((4096,), stream=mx.cpu)
|
|
mx.eval(b)
|
|
del b
|
|
|
|
new_active_mem = mx.metal.get_active_memory()
|
|
self.assertEqual(new_active_mem, active_mem)
|
|
peak_mem = mx.metal.get_peak_memory()
|
|
self.assertTrue(peak_mem >= 4096 * 8)
|
|
cache_mem = mx.metal.get_cache_memory()
|
|
self.assertTrue(cache_mem >= 4096 * 4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|