diff --git a/test/test_copy_speed.py b/test/test_copy_speed.py index bd27153574..8eac7b9c57 100644 --- a/test/test_copy_speed.py +++ b/test/test_copy_speed.py @@ -9,11 +9,10 @@ class TestCopySpeed(unittest.TestCase): @classmethod def setUpClass(cls): Device[Device.DEFAULT].synchronize() - @unittest.skipIf(OSX, "no shm on OSX") def testCopySHMtoDefault(self): s = shared_memory.SharedMemory(name="test_X", create=True, size=N*N*4) s.close() - if CI: + if CI and not OSX: t = Tensor.empty(N, N, device="disk:/dev/shm/test_X").realize() else: t = Tensor.empty(N, N, device="disk:shm:test_X").realize() diff --git a/tinygrad/runtime/ops_disk.py b/tinygrad/runtime/ops_disk.py index 41bf8230e2..4c6f341add 100644 --- a/tinygrad/runtime/ops_disk.py +++ b/tinygrad/runtime/ops_disk.py @@ -1,6 +1,4 @@ -import os, mmap -try: import _posixshmem -except Exception: pass +import os, mmap, _posixshmem from typing import Callable, Dict, Tuple from tinygrad.helpers import prod, DType, OSX, dtypes from tinygrad.device import Interpreted, Allocator @@ -27,21 +25,15 @@ class DiskBuffer: disk_fxn_for_op: Dict[Op, Callable] = { UnaryOps.CAST: DiskBuffer.cast, MovementOps.AS_STRIDED: DiskBuffer.as_strided } -MAP_LOCKED, MAP_POPULATE = 0x2000, 0x008000 +MAP_LOCKED, MAP_POPULATE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000) class DiskAllocator(Allocator): def __init__(self, device): self.device = device def _alloc(self, size): if str(self.device).startswith("shm:"): - if OSX: - with open(f"/tmp/shm_{self.device[4:]}", "w+b") as f: - f.truncate(size) - shm = mmap.mmap(f.fileno(), size, flags=mmap.MAP_SHARED) - else: - fd = _posixshmem.shm_open(self.device[4:], os.O_RDWR, 0o600) - # TODO: these flags are somewhat platform specific, but python doesn't expose the ones we need - shm = mmap.mmap(fd, size, flags=mmap.MAP_SHARED | MAP_LOCKED | MAP_POPULATE) - shm.madvise(mmap.MADV_HUGEPAGE) # type: ignore # not on OSX - os.close(fd) + fd = _posixshmem.shm_open("/"+self.device[4:].lstrip("/"), os.O_RDWR, 0o600) + shm = mmap.mmap(fd, size, flags=mmap.MAP_SHARED | MAP_POPULATE | MAP_LOCKED) + if (hp := getattr(mmap, "MADV_HUGEPAGE", None)) is not None: shm.madvise(hp) # type: ignore + os.close(fd) buf = UnderlyingDiskBuffer(None, shm) else: f = open(self.device, "a+b")