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

Add synchronize function (#1006)

Browse Source 
* add synchronize function

* fix linux

* fix linux

* fix and fix docs

* fix test

* try synchronize in stream destroy

* synchronize works for both cpu and gpu
This commit is contained in:
Awni Hannun
2024-04-22 08:25:46 -07:00
committed by GitHub
parent b0012cdd0f
commit 3d405fb3b1
14 changed files with 95 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_metal.py
View File

@@ -24,14 +24,16 @@ class TestMetal(mlx_tests.MLXTestCase):
self.assertTrue(mx.metal.set_memory_limit(old_limit), old_limit)
# Query active and peak memory
a = mx.zeros((4096,), stream=mx.cpu)
a = mx.zeros((4096,))
mx.eval(a)
mx.synchronize()
active_mem = mx.metal.get_active_memory()
self.assertTrue(active_mem >= 4096 * 4)
b = mx.zeros((4096,), stream=mx.cpu)
b = mx.zeros((4096,))
mx.eval(b)
del b
mx.synchronize()
new_active_mem = mx.metal.get_active_memory()
self.assertEqual(new_active_mem, active_mem)