tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-30 00:06:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

usb amd: single-part chunks, stock submit path

Browse Source
This commit is contained in:
tiny
2026-08-21 05:26:51 +00:00
parent 0cdb48d65f
commit d0f1afe156
2 changed files with 26 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/ops_amd.py
+24 -53
View File
@@ -652,81 +652,52 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
def _copyin(self, dest:HCQBuffer, src:memoryview):
if not self.dev.is_usb(): return super()._copyin(dest, src)
from tinygrad.runtime.support.usb import alloc_cbuffer
# Pipelined USB copyin: 240KB chunks stream over EP2 into two alternating SRAM bounce buffers, every 120KB part
# followed by a 512B sentinel sector carrying a unique value. The prebuilt SDMA ring polls a part's sentinel
# before copying it to VRAM (GPU blocks until the engine lands it) and bumps the buffer's drain fence after the
# chunk; the host waits on that fence before re-arming the buffer (USB blocks until the GPU drains). 1 doorbell.
# Pipelined USB copyin: 240KB chunks stream over EP2 into two alternating SRAM bounce buffers, each chunk
# ending in a 512B sentinel sector with a unique value. The prebuilt SDMA ring polls a chunk's sentinel before
# copying it to VRAM and bumps the buffer's drain fence after; the host waits on that fence before re-arming.
dev, usb = self.dev, self.dev.iface.pci_dev.usb
sdq, ts, sdma = dev.sdma_queue(0), dev.timeline_signal, dev.sdma
ts, sdma = dev.timeline_signal, dev.sdma
with hcq_profile(self.dev, queue_type=dev.hw_copy_queue_t, desc=TracingKey(f"TINY -> {dev.device}", ret=src.nbytes), enabled=PROFILE,
dev_suff="SDMA:0"):
CP, PART = 0x3C000, 0x1E000 # 240KB chunks of two 120KB parts; sentinels ride the buffer's last SRAM slot
src_mv, nchunks = src.cast('B'), ceildiv(src.nbytes, CP)
CP, src_mv = 0x3C000, src.cast('B')
nchunks = ceildiv(src.nbytes, CP)
if nchunks == 0: return
slots = [(self.b[bi].cpu_view().addr - 0xf000) >> 14 for bi in range(2)] # 0xF2 slot base per buffer
if not hasattr(self, '_usb_seq'):
self._usb_seq, self._usb_prev_tail, self._usb_buf_ctr = 0, (0, 0), [0, 0]
self._usb_stage = [alloc_cbuffer(0x40000) for _ in range(2)] # wire-image staging, one per buffer
self._usb_seq, self._usb_buf_ctr = 0, [0, 0]
self._usb_stage = [alloc_cbuffer(0x40000) for _ in range(2)] # wire-image staging, one per bounce buffer
for bi in range(2): # clear drain fences (PCIe 0x820800 = xdata 0xA800) and zero buffers (sentinel EQ needs it)
usb.write(0xA800 + bi * 8, bytes(8))
usb.scsi_write(bytes(0x40000), slot_start=slots[bi])
seq, ctr0 = self._usb_seq, list(self._usb_buf_ctr)
POLL_EQ = sdma.SDMA_OP_POLL_REGMEM | sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(3) | sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
POLL_CFG = sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)
def wait_full(bi, ctr): # spin until the GPU has fully drained chunk ctr of buffer bi
while int.from_bytes(usb.read(0xA800 + bi * 8, 8), 'little') < ctr: pass
# the whole ring: one wait on pre-copyin GPU work, then per chunk, per part: [POLL sentinel][copy][signal][fence]
# the whole ring: per chunk: [POLL sentinel][copy][signal][fence write]
q = dev.hw_copy_queue_t().wait(ts, dev.timeline_value - 1)
for c in range(nchunks):
bi, lsize = c % 2, min(CP, src.nbytes - c * CP)
for j in range(2 if lsize > PART else 1):
plen = min(PART, lsize - j * PART)
q.q(POLL_EQ, *data64_le(self.b[bi].va_addr + j * (PART + 512) + round_up(plen, 512)),
0x51000000 | (seq + 2 * c + j & 0xFFFFFF), 0xFFFFFFFF, POLL_CFG)
q.copy(dest.offset(c * CP + j * PART), self.b[bi].offset(j * (PART + 512)), plen)
q.q(POLL_EQ, *data64_le(self.b[bi].va_addr + round_up(lsize, 512)), 0x51000000 | (self._usb_seq + c & 0xFFFFFF), 0xFFFFFFFF,
sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
q.copy(dest.offset(c * CP), self.b[bi], lsize)
q.signal(ts, dev.next_timeline())
q.write(dev.iface.sys_buf.offset(0x800 + bi * 8, 8), ctr0[bi] + c // 2 + 1, b64=True)
wait_full(*self._usb_prev_tail) # the previous copyin's ring entries must be consumed before overwriting
# stage the ring in one EP2 write while the engine is disarmed (the GPU ring is circular, dword 0 is a NOP)
cmds = array.array('I', q._q).tobytes()
rb, pv_start = sdq.ring.nbytes, sdq.put_value
off = pv_start % rb
if off + len(cmds) > rb: # zero-fill to the ring end and continue at offset 0
sdq.ring.view(off, rb - off, fmt='B')[:] = bytes(rb - off)
pv_start, off = pv_start + rb - off, 0
sdq.ring.view(off, len(cmds), fmt='B')[:] = cmds
sdq.put_value = pv_start + len(cmds)
sdq.write_ptr[0] = sdq.put_value
sdq.doorbell[0] = sdq.put_value
def push(c):
q.write(dev.iface.sys_buf.offset(0x800 + bi * 8, 8), self._usb_buf_ctr[bi] + c // 2 + 1, b64=True)
q.submit(dev) # one EP2 write for the ring while the engine is disarmed, then a single doorbell
def push(c): # arm chunk c and queue its staged wire image [payload][sentinel]; it streams once c-1 completes
bi, lsize = c % 2, min(CP, src.nbytes - c * CP)
sbuf, smv = self._usb_stage[c % 2]
sptr = ctypes.addressof(from_mv(src_mv[c * CP:c * CP + lsize]))
for j in range(2 if lsize > PART else 1): # stage the wire image [part][sentinel]...
plen = min(PART, lsize - j * PART)
ctypes.memmove(ctypes.addressof(sbuf) + j * (PART + 512), sptr + j * PART, plen)
struct.pack_into('<I', sbuf, j * (PART + 512) + round_up(plen, 512), 0x51000000 | (seq + 2 * c + j & 0xFFFFFF))
wait_full(bi, ctr0[bi] + c // 2) # the previous chunk on this buffer must be fully drained before re-arming
usb.scsi_write_arm(round_up(min(PART, lsize), 512) + 512 + (round_up(lsize - PART, 512) + 512 if lsize > PART else 0),
slot_start=slots[bi])
return usb.usb.bulk_write_async(smv[:round_up(lsize, 512) + 512 * (2 if lsize > PART else 1)])
ctypes.memmove(ctypes.addressof(sbuf), ctypes.addressof(from_mv(src_mv[c * CP:c * CP + lsize])), lsize)
struct.pack_into('<I', sbuf, round_up(lsize, 512), 0x51000000 | (self._usb_seq + c & 0xFFFFFF))
wait_full(bi, self._usb_buf_ctr[bi] + c // 2) # the previous chunk on this buffer must be fully drained first
usb.scsi_write_arm(round_up(lsize, 512) + 512, slot_start=slots[bi])
return usb.usb.bulk_write_async(smv[:round_up(lsize, 512) + 512])
pending = [push(0)] + ([push(1)] if nchunks > 1 else [])
for c in range(1, nchunks):
usb.usb.bulk_wait(pending.pop(0)) # chunk c-1 has landed in SRAM
if c + 1 < nchunks: pending.append(push(c + 1)) # arm+queue the next chunk while chunk c streams
if c + 1 < nchunks: pending.append(push(c + 1))
usb.usb.bulk_wait(pending.pop(0))
tbi = (nchunks - 1) % 2
self._usb_prev_tail = (tbi, ctr0[tbi] + (nchunks - 1) // 2 + 1)
wait_full(*self._usb_prev_tail) # the last chunk must be in VRAM before returning
self._usb_seq = (seq + 2 * (nchunks - 1) + (2 if lsize > PART else 1)) & 0xFFFFFF
self._usb_buf_ctr = [ctr0[0] + (nchunks + 1) // 2, ctr0[1] + nchunks // 2]
self._usb_seq, self._usb_buf_ctr = (self._usb_seq + nchunks) & 0xFFFFFF, \
[self._usb_buf_ctr[0] + (nchunks + 1) // 2, self._usb_buf_ctr[1] + nchunks // 2]
wait_full((nchunks - 1) % 2, self._usb_buf_ctr[(nchunks - 1) % 2]) # the last chunk must be in VRAM before returning
def _copyout(self, dest:memoryview, src:HCQBuffer):
if not self.dev.is_usb(): return super()._copyout(dest, src)
self.dev.synchronize()
tinygrad/runtime/support/usb.py
+2 -5
View File
@@ -78,14 +78,13 @@ class USB3:
(self.handle, 0x02, self._bulk_buf, len(payload), self._transferred, timeout)
assert self._transferred.value == len(payload), f"bulk OUT short write: {self._transferred.value}/{len(payload)} bytes"
def _on_bulk_done(self, xfer):
# NOTE: runs inside libusb event handling; exceptions here are unraisable, so latch errors for bulk_wait.
def _on_bulk_done(self, xfer): # runs inside libusb event handling; exceptions are unraisable, so latch errors
if xfer.contents.status != 0 or xfer.contents.actual_length != xfer.contents.length: self._async_err = xfer.contents.status or -1
self._async_pending.pop(int(xfer.contents.user_data or 0), None)
self._async_pool.append(xfer)
def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:
"""Zero-copy async bulk OUT on EP 0x02. Returns a tag for bulk_wait. The payload must stay alive until bulk_wait."""
# zero-copy async bulk OUT on EP 0x02; the payload must stay alive until bulk_wait(tag)
assert payload.contiguous, "bulk_write_async requires a contiguous buffer"
tr = self._async_pool.pop() if self._async_pool else libusb.libusb_alloc_transfer(0)
tr.contents.dev_handle, tr.contents.endpoint, tr.contents.type = self.handle, 0x02, libusb.LIBUSB_TRANSFER_TYPE_BULK
@@ -199,8 +198,6 @@ class CustomASM24Controller:
num_slots = ceildiv(nbytes, 0x4000) # 16KB per slot
self.usb.control_write(0xF2, value=sectors, index=(slot_start & 0xFF) | ((num_slots & 0xFF) << 8))
def bulk_wait(self, tag:int): self.usb.bulk_wait(tag)
def scsi_read_arm(self, size:int):
windex = (ceildiv(size, 0x4000) & 0xFF) << 8
self.usb.control_write(0xF2, value=(ceildiv(size, 512) & 0x7FFF) | 0x8000, index=windex)