tinygrad 0.10.0__py3-none-any.whl → 0.10.2__py3-none-any.whl

tinygrad/runtime/ops_cloud.py

@@ -5,7 +5,7 @@
 # it should be a secure (example: no use of pickle) boundary. HTTP is used for RPC
 
 from __future__ import annotations
-from typing import Tuple, Optional, Dict, Any, DefaultDict, List
+from typing import Optional, Any
 from collections import defaultdict
 from dataclasses import dataclass, field
 import multiprocessing, functools, http.client, hashlib, json, time, os, binascii, struct, ast, contextlib
@@ -13,14 +13,14 @@ from http.server import HTTPServer, BaseHTTPRequestHandler
 from tinygrad.renderer import Renderer
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, Timing
-from tinygrad.device import Compiled, Allocator, Compiler, Device, BufferOptions
+from tinygrad.device import Compiled, Allocator, Compiler, Device, BufferSpec
 
 # ***** API *****
 
 class CloudRequest: pass
 
 @dataclass(frozen=True)
-class BufferAlloc(CloudRequest): buffer_num: int; size: int; options: BufferOptions # noqa: E702
+class BufferAlloc(CloudRequest): buffer_num: int; size: int; options: BufferSpec # noqa: E702
 
 @dataclass(frozen=True)
 class BufferFree(CloudRequest): buffer_num: int # noqa: E702
@@ -39,11 +39,11 @@ class ProgramFree(CloudRequest): name: str; datahash: str # noqa: E702
 
 @dataclass(frozen=True)
 class ProgramExec(CloudRequest):
-  name: str; datahash: str; bufs: Tuple[int, ...]; vals: Tuple[int, ...] # noqa: E702
-  global_size: Optional[Tuple[int, ...]]; local_size: Optional[Tuple[int, ...]]; wait: bool # noqa: E702
+  name: str; datahash: str; bufs: tuple[int, ...]; vals: tuple[int, ...] # noqa: E702
+  global_size: Optional[tuple[int, ...]]; local_size: Optional[tuple[int, ...]]; wait: bool # noqa: E702
 
 # for safe deserialization
-whitelist = {x.__name__:x for x in [BufferAlloc, BufferFree, CopyIn, CopyOut, ProgramAlloc, ProgramFree, ProgramExec, BufferOptions]}
+whitelist = {x.__name__:x for x in [BufferAlloc, BufferFree, CopyIn, CopyOut, ProgramAlloc, ProgramFree, ProgramExec, BufferSpec]}
 eval_fxns = {ast.Constant: lambda x: x.value, ast.Tuple: lambda x: tuple(map(safe_eval, x.elts)), ast.List: lambda x: list(map(safe_eval, x.elts)),
   ast.Call: lambda x: safe_eval(x.func)(*[safe_eval(arg) for arg in x.args], **{kwarg.arg: safe_eval(kwarg.value) for kwarg in x.keywords}),
   ast.Name: lambda x: whitelist[x.id], ast.Attribute: lambda x: {"imagef": dtypes.imagef, "imageh": dtypes.imageh}[x.attr]}
@@ -51,8 +51,8 @@ def safe_eval(node): return eval_fxns[node.__class__](node)
 
 class BatchRequest:
   def __init__(self):
-    self._q: List[CloudRequest] = []
-    self._h: Dict[str, bytes] = {}
+    self._q: list[CloudRequest] = []
+    self._h: dict[str, bytes] = {}
   def h(self, d:bytes) -> str:
     binhash = hashlib.sha256(d).digest()
     self._h[datahash:=binascii.hexlify(binhash).decode()] = binhash+struct.pack("<Q", len(d))+d
@@ -74,14 +74,14 @@ class BatchRequest:
 
 @dataclass
 class CloudSession:
-  programs: Dict[Tuple[str, str], Any] = field(default_factory=dict)
+  programs: dict[tuple[str, str], Any] = field(default_factory=dict)
   # TODO: the buffer should track this internally
-  buffers: Dict[int, Tuple[Any, int, Optional[BufferOptions]]] = field(default_factory=dict)
+  buffers: dict[int, tuple[Any, int, Optional[BufferSpec]]] = field(default_factory=dict)
 
 class CloudHandler(BaseHTTPRequestHandler):
   protocol_version = 'HTTP/1.1'
-  dname: str
-  sessions: DefaultDict[str, CloudSession] = defaultdict(CloudSession)
+  device: str
+  sessions: defaultdict[str, CloudSession] = defaultdict(CloudSession)
 
   def setup(self):
     super().setup()
@@ -99,18 +99,18 @@ class CloudHandler(BaseHTTPRequestHandler):
         match c:
           case BufferAlloc():
             assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already allocated"
-            session.buffers[c.buffer_num] = (Device[CloudHandler.dname].allocator.alloc(c.size, c.options), c.size, c.options)
+            session.buffers[c.buffer_num] = (Device[CloudHandler.device].allocator.alloc(c.size, c.options), c.size, c.options)
           case BufferFree():
             buf,sz,buffer_options = session.buffers[c.buffer_num]
-            Device[CloudHandler.dname].allocator.free(buf,sz,buffer_options)
+            Device[CloudHandler.device].allocator.free(buf,sz,buffer_options)
             del session.buffers[c.buffer_num]
-          case CopyIn(): Device[CloudHandler.dname].allocator.copyin(session.buffers[c.buffer_num][0], memoryview(bytearray(req._h[c.datahash])))
+          case CopyIn(): Device[CloudHandler.device].allocator._copyin(session.buffers[c.buffer_num][0], memoryview(bytearray(req._h[c.datahash])))
           case CopyOut():
             buf,sz,_ = session.buffers[c.buffer_num]
-            Device[CloudHandler.dname].allocator.copyout(memoryview(ret:=bytearray(sz)), buf)
+            Device[CloudHandler.device].allocator._copyout(memoryview(ret:=bytearray(sz)), buf)
           case ProgramAlloc():
-            lib = Device[CloudHandler.dname].compiler.compile_cached(req._h[c.datahash].decode())
-            session.programs[(c.name, c.datahash)] = Device[CloudHandler.dname].runtime(c.name, lib)
+            lib = Device[CloudHandler.device].compiler.compile_cached(req._h[c.datahash].decode())
+            session.programs[(c.name, c.datahash)] = Device[CloudHandler.device].runtime(c.name, lib)
           case ProgramFree(): del session.programs[(c.name, c.datahash)]
           case ProgramExec():
             bufs = [session.buffers[x][0] for x in c.bufs]
@@ -118,7 +118,7 @@ class CloudHandler(BaseHTTPRequestHandler):
             r = session.programs[(c.name, c.datahash)](*bufs, vals=c.vals, wait=c.wait, **extra_args)
             if r is not None: ret = str(r).encode()
     elif self.path == "/renderer" and method == "GET":
-      cls, args = Device[CloudHandler.dname].renderer.__reduce__()
+      cls, args = Device[CloudHandler.device].renderer.__reduce__()
       ret = json.dumps((cls.__module__, cls.__name__, args)).encode()
     else: status_code = 404
     self.send_response(status_code)
@@ -131,42 +131,42 @@ class CloudHandler(BaseHTTPRequestHandler):
 
 def cloud_server(port:int):
   multiprocessing.current_process().name = "MainProcess"
-  CloudHandler.dname = getenv("CLOUDDEV", "METAL") if Device.DEFAULT == "CLOUD" else Device.DEFAULT
-  print(f"start cloud server on {port} with device {CloudHandler.dname}")
+  CloudHandler.device = getenv("CLOUDDEV", "METAL") if Device.DEFAULT == "CLOUD" else Device.DEFAULT
+  print(f"start cloud server on {port} with device {CloudHandler.device}")
   server = HTTPServer(('', port), CloudHandler)
   server.serve_forever()
 
 # ***** frontend *****
 
 class CloudAllocator(Allocator):
-  def __init__(self, device:CloudDevice):
-    self.device = device
+  def __init__(self, dev:CloudDevice):
+    self.device = dev
     super().__init__()
   # TODO: ideally we shouldn't have to deal with images here
-  def _alloc(self, size:int, options:BufferOptions) -> int:
+  def _alloc(self, size:int, options:BufferSpec) -> int:
     self.device.buffer_num += 1
     self.device.req.q(BufferAlloc(self.device.buffer_num, size, options))
     return self.device.buffer_num
   # TODO: options should not be here in any Allocator
   def _free(self, opaque:int, options): self.device.req.q(BufferFree(opaque))
-  def copyin(self, dest:int, src:memoryview): self.device.req.q(CopyIn(dest, self.device.req.h(bytes(src))))
-  def copyout(self, dest:memoryview, src:int):
+  def _copyin(self, dest:int, src:memoryview): self.device.req.q(CopyIn(dest, self.device.req.h(bytes(src))))
+  def _copyout(self, dest:memoryview, src:int):
     self.device.req.q(CopyOut(src))
     resp = self.device.batch_submit()
     assert len(resp) == len(dest), f"buffer length mismatch {len(resp)} != {len(dest)}"
     dest[:] = resp
 
 class CloudProgram:
-  def __init__(self, device:CloudDevice, name:str, lib:bytes):
-    self.device, self.name = device, name
-    self.datahash = self.device.req.h(lib)
-    self.device.req.q(ProgramAlloc(self.name, self.datahash))
+  def __init__(self, dev:CloudDevice, name:str, lib:bytes):
+    self.dev, self.name = dev, name
+    self.datahash = self.dev.req.h(lib)
+    self.dev.req.q(ProgramAlloc(self.name, self.datahash))
     super().__init__()
-  def __del__(self): self.device.req.q(ProgramFree(self.name, self.datahash))
+  def __del__(self): self.dev.req.q(ProgramFree(self.name, self.datahash))
 
-  def __call__(self, *bufs, global_size=None, local_size=None, vals:Tuple[int, ...]=(), wait=False):
-    self.device.req.q(ProgramExec(self.name, self.datahash, bufs, vals, global_size, local_size, wait))
-    if wait: return float(self.device.batch_submit())
+  def __call__(self, *bufs, global_size=None, local_size=None, vals:tuple[int, ...]=(), wait=False):
+    self.dev.req.q(ProgramExec(self.name, self.datahash, bufs, vals, global_size, local_size, wait))
+    if wait: return float(self.dev.batch_submit())
 
 class CloudDevice(Compiled):
   def __init__(self, device:str):

tinygrad/runtime/ops_cpu.py

@@ -0,0 +1,24 @@
+import platform, subprocess, sys
+from tinygrad.helpers import capstone_flatdump, getenv
+from tinygrad.device import Compiled, Compiler, MallocAllocator, CPUProgram
+from tinygrad.runtime.support.elf import jit_loader
+from tinygrad.renderer.cstyle import ClangRenderer
+
+class ClangJITCompiler(Compiler):
+  def __init__(self, cachekey="compile_clang_jit"): super().__init__(cachekey)
+
+  def compile(self, src:str) -> bytes:
+    # -fno-math-errno is required for __builtin_sqrt to become an instruction instead of a function call
+    # x18 is a reserved platform register. It is clobbered on context switch in macos and is used to store TEB pointer in windows on arm, don't use it
+    target = 'x86_64' if sys.platform == 'win32' else platform.machine()
+    args = ['-march=native', f'--target={target}-none-unknown-elf', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib']
+    arch_args = ['-ffixed-x18'] if target == 'arm64' else []
+    obj = subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', *args, *arch_args, '-', '-o', '-'], input=src.encode('utf-8'))
+    return jit_loader(obj)
+
+  def disassemble(self, lib:bytes): return capstone_flatdump(lib)
+
+class ClangDevice(Compiled):
+  def __init__(self, device:str): super().__init__(device, MallocAllocator, ClangRenderer(), ClangJITCompiler(), CPUProgram)
+
+CPUDevice = ClangDevice

tinygrad/runtime/ops_cuda.py

@@ -1,25 +1,25 @@
 from __future__ import annotations
 import ctypes, ctypes.util, functools
-from typing import Tuple, Optional, List
 from tinygrad.helpers import DEBUG, getenv, from_mv, init_c_var, init_c_struct_t
-from tinygrad.device import Compiled, BufferOptions, LRUAllocator
+from tinygrad.device import Compiled, BufferSpec, LRUAllocator
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
 from tinygrad.runtime.support.compiler_cuda import cuda_disassemble, pretty_ptx, CUDACompiler, PTXCompiler, PTX
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint: disable=unused-import
+if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported
 
 def check(status):
   if status != 0: raise RuntimeError(f"CUDA Error {status}, {ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode()}")  # noqa: E501
 
-def encode_args(args, vals) -> Tuple[ctypes.Structure, ctypes.Array]:
+def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]:
   c_args = init_c_struct_t(tuple([(f'f{i}', cuda.CUdeviceptr_v2) for i in range(len(args))] +
                                  [(f'v{i}', ctypes.c_int) for i in range(len(vals))]))(*args, *vals)
   vargs = (ctypes.c_void_p * 5)(ctypes.c_void_p(1), ctypes.cast(ctypes.byref(c_args), ctypes.c_void_p), ctypes.c_void_p(2),
                                 ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(c_args))), ctypes.c_void_p), ctypes.c_void_p(0))
   return c_args, vargs
 
-def cu_time_execution(cb, enable=False) -> Optional[float]:
+def cu_time_execution(cb, enable=False) -> float|None:
   if not enable: return cb()
   evs = [init_c_var(cuda.CUevent(), lambda x: cuda.cuEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
   cuda.cuEventRecord(evs[0], None)
@@ -31,17 +31,17 @@ def cu_time_execution(cb, enable=False) -> Optional[float]:
   return ret.value * 1e-3
 
 class CUDAProgram:
-  def __init__(self, device:CUDADevice, name:str, lib:bytes, smem:int=0):
-    self.device, self.name, self.lib, self.smem = device, name, lib, smem
+  def __init__(self, dev:CUDADevice, name:str, lib:bytes, smem:int=0):
+    self.dev, self.name, self.lib, self.smem = dev, name, lib, smem
     if DEBUG >= 5: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(lib.decode('utf-8')).split("\n"))]))
-    if DEBUG >= 6: cuda_disassemble(lib, device.arch)
+    if DEBUG >= 6: cuda_disassemble(lib, dev.arch)
 
-    check(cuda.cuCtxSetCurrent(self.device.context))
+    check(cuda.cuCtxSetCurrent(self.dev.context))
     self.module = cuda.CUmodule()
     status = cuda.cuModuleLoadData(ctypes.byref(self.module), lib)
     if status != 0:
       del self.module
-      cuda_disassemble(lib, device.arch)
+      cuda_disassemble(lib, dev.arch)
       raise RuntimeError(f"module load failed with status code {status}: {cuda.cudaError_enum__enumvalues[status]}")
     check(cuda.cuModuleGetFunction(ctypes.byref(prg := cuda.CUfunction()), self.module, name.encode("utf-8")))
     self.prg = prg
@@ -50,47 +50,50 @@ class CUDAProgram:
   def __del__(self):
     if hasattr(self, 'module'): check(cuda.cuModuleUnload(self.module))
 
-  def __call__(self, *args, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False):
-    check(cuda.cuCtxSetCurrent(self.device.context))
+  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+    check(cuda.cuCtxSetCurrent(self.dev.context))
     if not hasattr(self, "vargs"):
       self.c_args, self.vargs = encode_args(args, vals)
+
+      # HACK: For MOCKGPU send the args struct itself.
+      if MOCKGPU: self.vargs = self.c_args # type: ignore[assignment]
     else:
       for i in range(len(args)): self.c_args.__setattr__(f'f{i}', args[i])
       for i in range(len(vals)): self.c_args.__setattr__(f'v{i}', vals[i])
     return cu_time_execution(lambda: check(cuda.cuLaunchKernel(self.prg, *global_size, *local_size, self.smem, None, None, self.vargs)), enable=wait)
 
 class CUDAAllocator(LRUAllocator):
-  def __init__(self, device:CUDADevice):
-    self.device = device
+  def __init__(self, dev:CUDADevice):
+    self.dev = dev
     super().__init__()
-  def _alloc(self, size, options:BufferOptions):
-    check(cuda.cuCtxSetCurrent(self.device.context))
+  def _alloc(self, size, options:BufferSpec):
+    check(cuda.cuCtxSetCurrent(self.dev.context))
     if options.host: return init_c_var(ctypes.c_void_p(), lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0x01)))
     return init_c_var(cuda.CUdeviceptr(), lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size)))
-  def _free(self, opaque, options:BufferOptions):
+  def _free(self, opaque, options:BufferSpec):
     if options.host: check(cuda.cuMemFreeHost(opaque))
     else: check(cuda.cuMemFree_v2(opaque))
-  def copyin(self, dest, src:memoryview):
-    check(cuda.cuCtxSetCurrent(self.device.context))
-    host_mem = self.alloc(len(src), BufferOptions(host=True))
-    self.device.pending_copyin.append((host_mem, len(src), BufferOptions(host=True)))
+  def _copyin(self, dest, src:memoryview):
+    check(cuda.cuCtxSetCurrent(self.dev.context))
+    host_mem = self.alloc(len(src), BufferSpec(host=True))
+    self.dev.pending_copyin.append((host_mem, len(src), BufferSpec(host=True)))
     ctypes.memmove(host_mem, from_mv(src), len(src))
     check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem, len(src), None))
-  def copyout(self, dest:memoryview, src):
+  def _copyout(self, dest:memoryview, src):
     CUDADevice.synchronize_system()
-    check(cuda.cuCtxSetCurrent(self.device.context))
+    check(cuda.cuCtxSetCurrent(self.dev.context))
     check(cuda.cuMemcpyDtoH_v2(from_mv(dest), src, len(dest)))
-  def transfer(self, dest, src, sz:int, src_dev, dest_dev):
+  def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
     check(cuda.cuCtxSetCurrent(src_dev.context))
     check(cuda.cuEventCreate(ctypes.byref(sync_event := cuda.CUevent()), 0))
     check(cuda.cuMemcpyDtoDAsync_v2(dest, src, sz, None))
     check(cuda.cuEventRecord(sync_event, None))
     check(cuda.cuCtxSetCurrent(dest_dev.context))
     check(cuda.cuStreamWaitEvent(None, sync_event, 0)) # sync the default stream on the dest dev
-  def offset(self, buf, size:int, offset:int): return cuda.CUdeviceptr_v2(buf.value + offset)
+  def _offset(self, buf, size:int, offset:int): return cuda.CUdeviceptr_v2(buf.value + offset)
 
 class CUDADevice(Compiled):
-  devices: List[CUDADevice] = []
+  devices: list[CUDADevice] = []
   peer_access = False
 
   def __init__(self, device:str):
@@ -110,12 +113,12 @@ class CUDADevice(Compiled):
       CUDADevice.peer_access = True
 
     self.arch = f"sm_{major.value}{minor.value}"
-    self.pending_copyin: List[Tuple[int, int, Optional[BufferOptions]]] = []
+    self.pending_copyin: list[tuple[int, int, BufferSpec|None]] = []
     CUDADevice.devices.append(self)
 
     from tinygrad.runtime.graph.cuda import CUDAGraph
     super().__init__(device, CUDAAllocator(self), PTXRenderer(self.arch) if PTX else CUDARenderer(self.arch),
-                     PTXCompiler(self.arch) if PTX else CUDACompiler(self.arch), functools.partial(CUDAProgram, self), graph=CUDAGraph)
+                     PTXCompiler(self.arch) if PTX else CUDACompiler(self.arch), functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)
 
   def synchronize(self):
     check(cuda.cuCtxSetCurrent(self.context))

tinygrad/runtime/ops_disk.py

@@ -1,72 +1,11 @@
-from __future__ import annotations
 import os, sys, mmap, io, ctypes, ctypes.util, contextlib
-from typing import Optional, Generator, Tuple, Callable, List
+from typing import Optional, Generator, Callable
 from tinygrad.helpers import OSX, round_up
 from tinygrad.device import Compiled, Allocator
 with contextlib.suppress(ImportError):
   import _posixshmem
   from tinygrad.runtime.autogen import io_uring, libc
 
-class DiskBuffer:
-  def __init__(self, device:DiskDevice, size:int, offset=0):
-    self.device, self.size, self.offset = device, size, offset
-  def __repr__(self): return f"<DiskBuffer size={self.size} offset={self.offset}>"
-  def _buf(self) -> memoryview:
-    assert self.device.mem is not None, "DiskBuffer wasn't opened"
-    return memoryview(self.device.mem)[self.offset:self.offset+self.size]
-
-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:DiskDevice): self.device = device
-  def _alloc(self, size:int, options):
-    self.device._might_open(size)
-    return DiskBuffer(self.device, size)
-  def _free(self, opaque, options): self.device._might_close()
-  def as_buffer(self, src:DiskBuffer): return src._buf()
-  def copyin(self, dest:DiskBuffer, src:memoryview): dest._buf()[:] = src
-  def copyout(self, dest:memoryview, src:DiskBuffer):
-    if OSX and self.device.fd is not None:
-      # OSX doesn't seem great at mmap, this is faster
-      with io.FileIO(self.device.fd, "a+b", closefd=False) as fo:
-        fo.seek(src.offset)
-        fo.readinto(dest)
-    else:
-      dest[:] = src._buf()
-
-  def _copyout_sharded(self, src:DiskBuffer, size:int, _get_free_buf:Callable, seg_len:int) -> Generator[Tuple[int, int, int, int], None, None]:
-    assert hasattr(DiskDevice, 'io_uring'), "function requires io uring support"
-
-    fd_offset = src.offset - (minor_offset := src.offset % mmap.PAGESIZE)
-    processed_reqs_cnt, copied_in, next_read_offset, total_copy_size = 0, 0, 0, round_up(size + minor_offset, mmap.PAGESIZE)
-    reqs: List[Tuple[int, int, int, int]] = []
-
-    while next_read_offset < total_copy_size or len(reqs) != processed_reqs_cnt:
-      if next_read_offset < total_copy_size and (copy_batch := _get_free_buf()) is not None:
-        # Prepare sqe
-        sqe_index = (tail:=DiskDevice.io_uring.sq.ktail[0]) & DiskDevice.io_uring.sq.kring_mask[0]
-        sqe = DiskDevice.io_uring.sq.sqes[sqe_index]
-        sqe.opcode, sqe.fd, sqe.off = io_uring.IORING_OP_READ, self.device.fd, fd_offset + next_read_offset
-        sqe.addr, sqe.len, sqe.user_data = copy_batch[0], min(seg_len, total_copy_size - next_read_offset), len(reqs)
-
-        # Send sqe
-        DiskDevice.io_uring.sq.array[sqe_index] = sqe_index
-        DiskDevice.io_uring.sq.ktail[0] = tail + 1
-        libc.syscall(io_uring.NR_io_uring_enter, DiskDevice.io_uring.ring_fd, 1, 1, io_uring.IORING_ENTER_GETEVENTS)
-
-        reqs.append((copy_batch, copied_in, minor_offset, real_copy_size:=min(sqe.len - minor_offset, size - copied_in)))
-        next_read_offset += sqe.len
-        copied_in += real_copy_size
-        minor_offset = 0
-
-      if (head:=DiskDevice.io_uring.cq.khead[0]) != DiskDevice.io_uring.cq.ktail[0]:
-        cqe = DiskDevice.io_uring.cq.cqes[head & DiskDevice.io_uring.cq.kring_mask[0]]
-        assert cqe.res >= 0, f"read from disk failed, err: {cqe.res}"
-        yield reqs[cqe.user_data]
-        DiskDevice.io_uring.cq.khead[0] = head + 1 # advance
-        processed_reqs_cnt += 1
-
-  def offset(self, buf:DiskBuffer, size:int, offset:int): return DiskBuffer(buf.device, size, offset)
-
 class DiskDevice(Compiled):
   _tried_io_uring_init = False
 
@@ -81,7 +20,7 @@ class DiskDevice(Compiled):
     self.count += 1
     assert self.size is None or size <= self.size, f"can't reopen Disk tensor with larger size, opened with {self.size}, tried to open with {size}"
     if self.size is not None: return
-    filename = self.dname[len("disk:"):]
+    filename = self.device[len("disk:"):]
     self.size = size
 
     if sys.platform != "win32" and filename.startswith("shm:"):
@@ -122,3 +61,63 @@ class DiskDevice(Compiled):
         kring_mask=u32ptr(sq_ptr+p.cq_off.ring_mask), cqes=ctypes.cast(cq_ptr+p.cq_off.cqes, ctypes.POINTER(io_uring.struct_io_uring_cqe)))
 
       DiskDevice.io_uring = io_uring.struct_io_uring(ring_fd=fd, sq=sqdesc, cq=cqdesc) # type: ignore
+
+class DiskBuffer:
+  def __init__(self, device:DiskDevice, size:int, offset=0):
+    self.device, self.size, self.offset = device, size, offset
+  def __repr__(self): return f"<DiskBuffer size={self.size} offset={self.offset}>"
+  def _buf(self) -> memoryview:
+    assert hasattr(self.device, "mem"), f"DiskBuffer wasn't opened: {self.device.device}"
+    return memoryview(self.device.mem)[self.offset:self.offset+self.size]
+
+MAP_LOCKED, MAP_POPULATE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000)
+class DiskAllocator(Allocator):
+  def __init__(self, dev:DiskDevice): self.dev = dev
+  def _alloc(self, size:int, options):
+    self.dev._might_open(size)
+    return DiskBuffer(self.dev, size)
+  def _free(self, opaque, options): self.dev._might_close()
+  def _as_buffer(self, src:DiskBuffer): return src._buf()
+  def _copyin(self, dest:DiskBuffer, src:memoryview): dest._buf()[:] = src
+  def _copyout(self, dest:memoryview, src:DiskBuffer):
+    if OSX and self.dev.fd is not None:
+      # OSX doesn't seem great at mmap, this is faster
+      with io.FileIO(self.dev.fd, "a+b", closefd=False) as fo:
+        fo.seek(src.offset)
+        fo.readinto(dest)
+    else:
+      dest[:] = src._buf()
+
+  def _copyout_sharded(self, src:DiskBuffer, size:int, _get_free_buf:Callable, seg_len:int) -> Generator[tuple[int, int, int, int], None, None]:
+    assert hasattr(DiskDevice, 'io_uring'), "function requires io uring support"
+
+    fd_offset = src.offset - (minor_offset := src.offset % mmap.PAGESIZE)
+    processed_reqs_cnt, copied_in, next_read_offset, total_copy_size = 0, 0, 0, round_up(size + minor_offset, mmap.PAGESIZE)
+    reqs: list[tuple[int, int, int, int]] = []
+
+    while next_read_offset < total_copy_size or len(reqs) != processed_reqs_cnt:
+      if next_read_offset < total_copy_size and (copy_batch := _get_free_buf()) is not None:
+        # Prepare sqe
+        sqe_index = (tail:=DiskDevice.io_uring.sq.ktail[0]) & DiskDevice.io_uring.sq.kring_mask[0]
+        sqe = DiskDevice.io_uring.sq.sqes[sqe_index]
+        sqe.opcode, sqe.fd, sqe.off = io_uring.IORING_OP_READ, self.dev.fd, fd_offset + next_read_offset
+        sqe.addr, sqe.len, sqe.user_data = copy_batch[0], min(seg_len, total_copy_size - next_read_offset), len(reqs)
+
+        # Send sqe
+        DiskDevice.io_uring.sq.array[sqe_index] = sqe_index
+        DiskDevice.io_uring.sq.ktail[0] = tail + 1
+        libc.syscall(io_uring.NR_io_uring_enter, DiskDevice.io_uring.ring_fd, 1, 1, io_uring.IORING_ENTER_GETEVENTS)
+
+        reqs.append((copy_batch, copied_in, minor_offset, real_copy_size:=min(sqe.len - minor_offset, size - copied_in)))
+        next_read_offset += sqe.len
+        copied_in += real_copy_size
+        minor_offset = 0
+
+      if (head:=DiskDevice.io_uring.cq.khead[0]) != DiskDevice.io_uring.cq.ktail[0]:
+        cqe = DiskDevice.io_uring.cq.cqes[head & DiskDevice.io_uring.cq.kring_mask[0]]
+        assert cqe.res >= 0, f"read from disk failed, err: {cqe.res}"
+        yield reqs[cqe.user_data]
+        DiskDevice.io_uring.cq.khead[0] = head + 1 # advance
+        processed_reqs_cnt += 1
+
+  def _offset(self, buf:DiskBuffer, size:int, offset:int): return DiskBuffer(buf.device, size, offset)