diff --git a/tinygrad/device.py b/tinygrad/device.py index 3ebebb07a7..eb60565fdc 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -236,10 +236,13 @@ class CPUProgram: PAGE_EXECUTE_READWRITE = 0x40 MEM_COMMIT = 0x1000 MEM_RESERVE = 0x2000 - ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_uint64 - ptr = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_int(0), ctypes.c_int(len(lib)), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE) - ctypes.memmove(ptr, lib, len(lib)) - self.fxn = ctypes.CFUNCTYPE(None)(ptr) + ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p + self.mem = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(lib)), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE) + ctypes.memmove(self.mem, lib, len(lib)) + ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p + proc = ctypes.windll.kernel32.GetCurrentProcess() + ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.mem), ctypes.c_size_t(len(lib))) + self.fxn = ctypes.CFUNCTYPE(None)(self.mem) else: from mmap import mmap, PROT_READ, PROT_WRITE, PROT_EXEC, MAP_ANON, MAP_PRIVATE # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/ @@ -268,6 +271,9 @@ class CPUProgram: if platform.machine() == "arm64" and OSX: args = args[:8] + [ctypes.c_int64(a) if isinstance(a, int) else a for a in args[8:]] return cpu_time_execution(lambda: self.fxn(*args), enable=wait) + def __del__(self): + if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.mem), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE + # **************** for Compiled Devices **************** class CompileError(Exception): pass