usbgpu: copyin size is 16k (#10240)

* usbgpu: copyin size is 16k

* ush

test/external/external_test_usb_asm24.py

@@ -2,6 +2,7 @@ import unittest, time
 from tinygrad.runtime.support.usb import ASM24Controller
 from tinygrad.helpers import Timing
 from tinygrad import Tensor, Device
+import numpy as np
 
 class TestASMController(unittest.TestCase):
   @classmethod
@@ -58,5 +59,15 @@ class TestDevCopySpeeds(unittest.TestCase):
       with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
         t.to('CPU').realize()
 
+  def testValidateCopies(self):
+    t = Tensor.randn(self.sz, self.sz, device="CPU").contiguous().realize()
+    x = t.to(Device.DEFAULT).realize()
+    Device[Device.DEFAULT].synchronize()
+
+    y = x.to('CPU').realize()
+
+    np.testing.assert_equal(t.numpy(), y.numpy())
+    del x, y, t
+
 if __name__ == "__main__":
   unittest.main()

tinygrad/runtime/ops_amd.py

@@ -466,7 +466,8 @@ class AMDProgram(HCQProgram):
     if hasattr(self, 'lib_gpu'): self.dev.allocator.free(self.lib_gpu, self.lib_gpu.size, BufferSpec(cpu_access=True, nolru=True))
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
-  def __init__(self, dev:AMDDevice): super().__init__(dev, copy_bufs=getattr(dev.dev_iface, 'copy_bufs', None))
+  def __init__(self, dev:AMDDevice):
+    super().__init__(dev, copy_bufs=getattr(dev.dev_iface, 'copy_bufs', None), max_copyout_size=0x1000 if dev.is_usb() else None)
 
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.dev_iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access)
@@ -819,7 +820,7 @@ class USBIface(PCIIface):
     self.usb._pci_cacheable += [self.bars[2]] # doorbell region is cacheable
 
     # special regions
-    self.copy_bufs = [self._new_dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x1000)]
+    self.copy_bufs = [self._new_dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x4000)]
     self.sys_buf, self.sys_next_off = self._new_dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0
 
   def _new_dma_region(self, ctrl_addr, sys_addr, size):

tinygrad/runtime/support/hcq.py

@@ -436,10 +436,10 @@ class HCQAllocatorBase(LRUAllocator, Generic[DeviceType]):
   This class implements basic copy operations following the HCQ API, utilizing both types of `HWQueue`.
   """
 
-  def __init__(self, dev:DeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None):
+  def __init__(self, dev:DeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, max_copyout_size:int|None=None):
     self.dev:DeviceType = dev
     self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True)) for _ in range(batch_cnt)]
-    self.b_timeline, self.b_next = [0] * len(self.b), 0
+    self.b_timeline, self.b_next, self.max_copyout_size = [0] * len(self.b), 0, max_copyout_size
     super().__init__()
 
   def map(self, buf:HCQBuffer): pass
@@ -481,9 +481,9 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]):
 
     assert self.dev.hw_copy_queue_t is not None
     with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> CPU", enabled=PROFILE):
-      for i in range(0, dest.nbytes, self.b[0].size):
+      for i in range(0, dest.nbytes, cp_size:=(self.max_copyout_size or self.b[0].size)):
         self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
-                                  .copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(self.b[0].size, dest.nbytes-i)) \
+                                  .copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(cp_size, dest.nbytes-i)) \
                                   .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
         self.dev.timeline_signal.wait(self.dev.timeline_value - 1)
         dest[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]

tinygrad/runtime/support/usb.py

@@ -1,7 +1,7 @@
 import ctypes, struct, dataclasses, array, itertools
 from typing import Sequence
 from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, to_mv
+from tinygrad.helpers import DEBUG, to_mv, round_up
 from tinygrad.runtime.support.hcq import MMIOInterface
 
 class USB3:
@@ -143,7 +143,9 @@ class ASM24Controller:
         addr = (op.addr & 0x1FFFF) | 0x500000
         _add_req(struct.pack('>BBBHB', 0xE4, op.size, addr >> 16, addr & 0xFFFF, 0), op.size, None)
         for i in range(op.size): self._cache[addr + i] = None
-      elif isinstance(op, ScsiWriteOp): _add_req(struct.pack('>BBQIBB', 0x8A, 0, op.lba, 4096//512, 0, 0), 0, op.data+b'\x00'*(4096-len(op.data)))
+      elif isinstance(op, ScsiWriteOp):
+        sectors = round_up(len(op.data), 512) // 512
+        _add_req(struct.pack('>BBQIBB', 0x8A, 0, op.lba, sectors, 0, 0), 0, op.data+b'\x00'*((sectors*512)-len(op.data)))
 
     return self.usb.send_batch(cdbs, idata, odata)