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

hcq: use mmio iface in copies (#10111)

Browse Source 
* hcq: use mmio iface in copies

* linter

* fix_am

* am
This commit is contained in:
nimlgen
2025-04-30 11:05:13 +03:00
committed by GitHub
parent 5c7d004da5
commit b4c9a3d8f4
3 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/support/am/amdev.py
+2 -2
View File
@@ -1,5 +1,5 @@
from __future__ import annotations
import ctypes, collections, time, dataclasses, functools, fcntl, os, hashlib, array
import ctypes, collections, time, dataclasses, functools, fcntl, os, hashlib
from tinygrad.helpers import mv_address, getenv, round_up, DEBUG, temp, fetch
from tinygrad.runtime.autogen.am import am, mp_11_0
from tinygrad.runtime.support.hcq import MMIOInterface
@@ -358,7 +358,7 @@ class AMDev:
mmRCC_CONFIG_MEMSIZE = 0xde3
self.vram_size = self.rreg(mmRCC_CONFIG_MEMSIZE) << 20
self.bhdr = am.struct_binary_header.from_buffer(array.array('B', self.vram.view(self.vram_size - (64 << 10), (10 << 10))[:]))
self.bhdr = am.struct_binary_header.from_buffer(bytearray(self.vram.view(self.vram_size - (64 << 10), (10 << 10))[:]))
ihdr = am.struct_ip_discovery_header.from_address(ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.IP_DISCOVERY].offset)
assert ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE and not ihdr.base_addr_64_bit, f"0x{ihdr.signature:X} != 0x{am.DISCOVERY_TABLE_SIGNATURE:X}"
tinygrad/runtime/support/am/ip.py
+3 -3
View File
@@ -1,4 +1,4 @@
import ctypes, time, contextlib, importlib, array
import ctypes, time, contextlib, importlib
from typing import Literal
from tinygrad.runtime.autogen.am import am
from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG
@@ -153,7 +153,7 @@ class AM_SMU(AM_IP):
def read_table(self, table_t, cmd):
self._send_msg(self.smu_mod.PPSMC_MSG_TransferTableSmu2Dram, cmd)
return table_t.from_buffer(array.array('B', self.adev.vram.view(self.driver_table_paddr, ctypes.sizeof(table_t))[:]))
return table_t.from_buffer(bytearray(self.adev.vram.view(self.driver_table_paddr, ctypes.sizeof(table_t))[:]))
def read_metrics(self): return self.read_table(self.smu_mod.SmuMetricsExternal_t, self.smu_mod.TABLE_SMU_METRICS)
def set_clocks(self, level):
@@ -459,7 +459,7 @@ class AM_PSP(AM_IP):
while self.adev.vram.view(self.fence_paddr, 4, 'I')[0] != prev_wptr: pass
time.sleep(0.005)
resp = type(cmd).from_buffer(array.array('B', self.adev.vram.view(self.cmd_paddr, ctypes.sizeof(cmd))[:]))
resp = type(cmd).from_buffer(bytearray(self.adev.vram.view(self.cmd_paddr, ctypes.sizeof(cmd))[:]))
if resp.resp.status != 0: raise RuntimeError(f"PSP command failed {resp.cmd_id} {resp.resp.status}")
return resp
tinygrad/runtime/support/hcq.py
+6 -4
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
from typing import cast, Callable, Type, TypeVar, Generic, Any, ClassVar
import contextlib, decimal, statistics, time, ctypes, array, os, fcntl, struct
from tinygrad.helpers import PROFILE, from_mv, getenv, to_mv, round_up
from tinygrad.helpers import PROFILE, getenv, to_mv, round_up
from tinygrad.renderer import Renderer
from tinygrad.device import BufferSpec, Compiler, Compiled, LRUAllocator, ProfileRangeEvent, ProfileDeviceEvent, ProfileProgramEvent
from tinygrad.ops import sym_infer, sint, Variable, UOp
@@ -10,7 +10,7 @@ from tinygrad.runtime.autogen import libc
class MMIOInterface:
def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
def __len__(self): return self.nbytes // struct.calcsize(self.fmt)
def __getitem__(self, k): return self.mv[k].tolist() if isinstance(k, slice) else self.mv[k]
def __getitem__(self, k): return (bytes(self.mv[k]) if self.fmt == 'B' else self.mv[k].tolist()) if isinstance(k, slice) else self.mv[k]
def __setitem__(self, k, v): self.mv[k] = v
def view(self, offset:int=0, size:int|None=None, fmt=None) -> MMIOInterface:
return MMIOInterface(self.addr+offset, size or (self.nbytes - offset), fmt=fmt or self.fmt)
@@ -451,7 +451,9 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]):
for i in range(0, src.nbytes, self.b[0].size):
self.b_next = (self.b_next + 1) % len(self.b)
self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
ctypes.memmove(self.b[self.b_next].va_addr, from_mv(src[i:]), lsize:=min(self.b[self.b_next].size, src.nbytes-i))
lsize = min(self.b[self.b_next].size, src.nbytes - i)
self.b[self.b_next].cpu_view().view(size=lsize, fmt='B')[:] = src[i:i+lsize]
self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
.copy(dest.va_addr+i, self.b[self.b_next].va_addr, lsize) \
.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
@@ -483,7 +485,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]):
.copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(self.b[0].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)
ctypes.memmove(from_mv(dest[i:]), self.b[0].va_addr, lsize)
dest[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]
def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:DeviceType, dest_dev:DeviceType):
cast(HCQAllocator, src_dev.allocator).map(dest)