copyout sharded w/o ioring (#15562)

* copyout sharded w/o ioring

* x

* x

* f

tinygrad/engine/realize.py

@@ -71,9 +71,8 @@ class BufferCopy(Runner):
     name = f"{type(self).__name__[6:].lower()} {sz}, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
     super().__init__(colored(name, "yellow"), dest_device, Estimates(lds=total_sz, mem=total_sz))
   def copy(self, dest, src):
-    disk_supports_fast_copyout = src.device.startswith("DISK") and hasattr(src.allocator.dev, 'io_uring') and \
-      getattr(src.allocator.dev, 'fd', None) is not None and dest.allocator.supports_copy_from_disk
-    if disk_supports_fast_copyout and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096:
+    disk_supports_fast_copyout = src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None
+    if disk_supports_fast_copyout and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
       dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
     elif isinstance(src.device, str) and src.device.startswith(("DISK", "TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
       # fast(ish) path, uses readinto in diskbuffers

tinygrad/runtime/ops_amd.py

@@ -635,8 +635,7 @@ class AMDProgram(HCQProgram):
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
     super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None), max_copyout_size=0x1000 if dev.is_usb() else None,
-                     supports_copy_from_disk=(not dev.is_am() or dev.iface.is_local()) and not dev.is_usb() and dev.has_sdma_queue,
-                     supports_transfer=dev.has_sdma_queue)
+                     supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue)
 
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)

tinygrad/runtime/ops_disk.py

@@ -95,20 +95,31 @@ class DiskAllocator(Allocator):
     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"
-
+  def _copyout_sharded(self, src:DiskBuffer, size:int, _get_free_buf:Callable, seg_len:int,
+                       use_ioring:bool=True) -> Generator[tuple[int, int, int, int], None, None]:
     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]] = []
 
+    if not hasattr(DiskDevice, 'io_uring') or not use_ioring:
+      local_buf = memoryview(bytearray(seg_len))
+      for off in range(0, total_copy_size, seg_len):
+        while (copy_batch := _get_free_buf()) is None: pass
+        read_size = min(seg_len, total_copy_size - off, src.device.size - fd_offset - off)
+        self._copyout(local_buf[:read_size], DiskBuffer(src.device, read_size, fd_offset + off))
+        copy_batch[0].view(size=read_size)[:] = local_buf[:read_size]
+        real_copy_size = min(read_size - minor_offset, size - copied_in)
+        yield (copy_batch, copied_in, minor_offset, real_copy_size)
+        copied_in, minor_offset = copied_in + real_copy_size, 0
+      return
+
+    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)
+        sqe.addr, sqe.len, sqe.user_data = copy_batch[0].addr, min(seg_len, total_copy_size - next_read_offset), len(reqs)
 
         # Send sqe
         DiskDevice.io_uring.sq.array[sqe_index] = sqe_index

tinygrad/runtime/ops_nv.py

@@ -326,8 +326,6 @@ class NVProgram(HCQProgram):
     return res
 
 class NVAllocator(HCQAllocator['NVDevice']):
-  def __init__(self, dev:'NVDevice'): super().__init__(dev, supports_copy_from_disk=not dev.is_nvd() or dev.iface.is_local())
-
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host)
 

tinygrad/runtime/support/hcq.py

@@ -544,15 +544,16 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
       # Check if the next buffer is safe to be used (its signal has passed) and reserve it.
       if self.b_timeline[(self.b_next + 1) % len(self.b)] <= self.dev.timeline_signal.value:
         self.b_timeline[(self.b_next + 1) % len(self.b)], self.b_next = (1 << 64), (self.b_next + 1) % len(self.b)
-        return (self.b[self.b_next].va_addr, self.b_next)
+        return (self.b[self.b_next].cpu_view(), self.b_next)
       return None
 
     assert self.dev.hw_copy_queue_t is not None
     with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"DISK -> {self.dev.device}", ret=size), enabled=PROFILE,
                      dev_suff="SDMA:0"):
-      for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size):
+      for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size,
+                                                                                             use_ioring=type(self.b[0].cpu_view()) is MMIOInterface):
         self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
-                                  .copy(dest.va_addr + dst_off, batch_info[0] + src_off, copy_size) \
+                                  .copy(dest.va_addr + dst_off, self.b[batch_info[1]].va_addr + src_off, copy_size) \
                                   .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
         self.b_timeline[batch_info[1]] = self.dev.timeline_value - 1