nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

Merge branch 'master' into dsp_search

Browse Source
This commit is contained in:
George Hotz
2025-03-21 17:13:10 +08:00
committed by GitHub
parent f66b03f0a6 8e555c586c
commit c715c25420
3 changed files with 92 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_quantize_onnx.py
+35 -34
View File
@@ -1,12 +1,13 @@
# ruff: noqa: E501
import numpy as np
import unittest
from dataclasses import replace
from tinygrad import Tensor, Context, Device, dtypes
from tinygrad.ops import Ops, UOp
from tinygrad.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import
from tinygrad.codegen.kernel import Kernel, Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item
from tinygrad.engine.search import bufs_from_lin
from tinygrad.shape.shapetracker import ShapeTracker, View
from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import
N = 512
@@ -297,45 +298,45 @@ class TestDSPCache(unittest.TestCase):
k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
for opt in opts: k.apply_opt(opt)
prg = k.to_program()
print(prg.src)
#print(prg.src)
new_src = """
typedef int int32 __attribute__((aligned(128),vector_size(128)));
typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128)));
typedef unsigned char unsigned_char8 __attribute__((aligned(8),vector_size(8)));
typedef int int32 __attribute__((aligned(128),vector_size(128)));
typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128)));
typedef unsigned char unsigned_char8 __attribute__((aligned(8),vector_size(8)));
typedef unsigned char unsigned_char4 __attribute__((aligned(4),vector_size(4)));
typedef unsigned char unsigned_char128 __attribute__((aligned(128),vector_size(128)));
__attribute__((noinline)) void r_196_24_8_32_4(unsigned char* restrict __attribute__((align_value(128))) data0, unsigned char* restrict __attribute__((align_value(128))) data1, signed char* restrict __attribute__((align_value(
128))) data2, int* restrict __attribute__((align_value(128))) data3) {
int32 cast0 = (int32){0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};
int32 val0 = *((int32*)((data3+0)));
for (int ridx0 = 0; ridx0 < 196; ridx0++) {
int32 acc0 = cast0;
int32 acc1 = cast0;
int32 acc2 = cast0;
int32 acc3 = cast0;
int32 val0 = *((int32*)((data3+0)));
for (int ridx0 = 0; ridx0 < 196; ridx0++) {
int32 acc0 = cast0;
int32 acc1 = cast0;
int32 acc2 = cast0;
int32 acc3 = cast0;
__builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768);
__builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768+192);
__builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768+384);
__builtin_HEXAGON_Y2_dcfetch(data1+ridx0*768+576);
for (int ridx1 = 0; ridx1 < 24; ridx1++) {
signed_char128 val1 = *((signed_char128*)((data2+(ridx1<<8))));
for (int ridx1 = 0; ridx1 < 24; ridx1++) {
signed_char128 val1 = *((signed_char128*)((data2+(ridx1<<8))));
signed_char128 val2 = *((signed_char128*)((data2+((1+(ridx1<<1))<<7))));
int alu0 = ((ridx0*768)+(ridx1<<3));
int alu0 = ((ridx0*768)+(ridx1<<3));
unsigned_char8 val3 = *((unsigned_char8*)((data1+alu0)));
__builtin_HEXAGON_Y2_dcfetch(((data1+alu0)+16));
unsigned_char8 val4 = *((unsigned_char8*)((data1+(alu0+192))));
__builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+192))+16));
unsigned_char8 val5 = *((unsigned_char8*)((data1+(alu0+384))));
__builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+384))+16));
unsigned_char8 val6 = *((unsigned_char8*)((data1+(alu0+576))));
__builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+576))+16));
unsigned_char8 val3 = *((unsigned_char8*)((data1+alu0)));
__builtin_HEXAGON_Y2_dcfetch(((data1+alu0)+16));
unsigned_char8 val4 = *((unsigned_char8*)((data1+(alu0+192))));
__builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+192))+16));
unsigned_char8 val5 = *((unsigned_char8*)((data1+(alu0+384))));
__builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+384))+16));
unsigned_char8 val6 = *((unsigned_char8*)((data1+(alu0+576))));
__builtin_HEXAGON_Y2_dcfetch(((data1+(alu0+576))+16));
unsigned_char4 alu5 = __builtin_shufflevector(val3, val3, 0, 1, 2, 3);
unsigned_char4 alu6 = __builtin_shufflevector(val4, val4, 0, 1, 2, 3);
unsigned_char4 alu7 = __builtin_shufflevector(val5, val5, 0, 1, 2, 3);
unsigned_char4 alu7 = __builtin_shufflevector(val5, val5, 0, 1, 2, 3);
unsigned_char4 alu8 = __builtin_shufflevector(val6, val6, 0, 1, 2, 3);
acc0 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc0, val1, (*((unsigned int*)&alu5)));
acc1 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc1, val1, (*((unsigned int*)&alu6)));
@@ -346,19 +347,19 @@ __attribute__((noinline)) void r_196_24_8_32_4(unsigned char* restrict __attribu
unsigned_char4 alu10 = __builtin_shufflevector(val4, val4, 4, 5, 6, 7);
unsigned_char4 alu11 = __builtin_shufflevector(val5, val5, 4, 5, 6, 7);
unsigned_char4 alu12 = __builtin_shufflevector(val6, val6, 4, 5, 6, 7);
acc0 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc0, val2, (*((unsigned int*)&alu9)));
acc1 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc1, val2, (*((unsigned int*)&alu10)));
acc2 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc2, val2, (*((unsigned int*)&alu11)));
acc3 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc3, val2, (*((unsigned int*)&alu12)));
}
unsigned_char128 alu18 = __builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B((((((acc3+val0)*203)+32767)/65536)+136), (((((acc2+val0)*203)+32767)/65536)+136)), __builtin_HEXAGON_V6_vpackwh_sat_128B((((((acc1+val0)*203)+32767)/65536)+136), (((((acc0+val0)*203)+32767)/65536)+136)));
*((unsigned_char128*)((data0+(ridx0<<7)))) = alu18;
}
}
acc0 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc0, val2, (*((unsigned int*)&alu9)));
acc1 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc1, val2, (*((unsigned int*)&alu10)));
acc2 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc2, val2, (*((unsigned int*)&alu11)));
acc3 = __builtin_HEXAGON_V6_vrmpybus_acc_128B(acc3, val2, (*((unsigned int*)&alu12)));
}
unsigned_char128 alu18 = __builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B((((((acc3+val0)*203)+32767)/65536)+136), (((((acc2+val0)*203)+32767)/65536)+136)), __builtin_HEXAGON_V6_vpackwh_sat_128B((((((acc1+val0)*203)+32767)/65536)+136), (((((acc0+val0)*203)+32767)/65536)+136)));
*((unsigned_char128*)((data0+(ridx0<<7)))) = alu18;
}
}
"""
prg = replace(prg, src=new_src+prg.src.split("/* DSP boilerplate */ ")[1])
rt = CompiledRunner(prg)
Device.default.compiler.disassemble(rt.lib)
#Device.default.compiler.disassemble(rt.lib)
ei = ExecItem(rt, bufs_from_lin(k))
tm = ei.run(wait=True)
print(f"final time {tm*1e6:.2f} us")
tinygrad/runtime/support/am/amdev.py
+14 -8
View File
@@ -220,13 +220,15 @@ class AMMemoryManager:
if contigous: paddrs = [(self.palloc(size, zero=True), size)]
else:
# Traverse the PT to find the largest contiguous sizes we need to allocate. Try to allocate the longest segment to reduce TLB pressure.
paddrs = []
ctx = AMPageTableTraverseContext(self.adev, self.root_page_table, va, create_pts=True)
for off, _, _, seg_cnt, seg_size in ctx.next(size):
while seg_cnt > 0:
rem_len = seg_cnt * seg_size
while rem_len > 0:
# Try to allocate as long segment (power of 2) as possible
cont_seg_sz, paddr = 1 << (self._frag_size(ctx.vaddr+off, seg_cnt*seg_size) + 12), None
while cont_seg_sz >= seg_size:
cont_seg_sz, paddr = 1 << (self._frag_size(ctx.vaddr+off, rem_len) + 12), None
while cont_seg_sz >= 0x1000:
try: paddr = self.palloc(cont_seg_sz, zero=True)
except MemoryError: cont_seg_sz //= 2
else: break
@@ -234,8 +236,8 @@ class AMMemoryManager:
if paddr is not None: paddrs += [(paddr, cont_seg_sz)]
else:
for paddr, _ in paddrs: self.pa_allocator.free(paddr)
raise MemoryError(f"Failed to allocate contigous {cont_seg_sz=:#x} bytes (size={size:#x})")
seg_cnt, off = seg_cnt - cont_seg_sz // seg_size, off + cont_seg_sz
raise MemoryError(f"Failed to allocate contigous a page. (allocation size={size:#x})")
rem_len, off = rem_len - cont_seg_sz, off + cont_seg_sz
return self.map_range(va, size, paddrs, uncached=uncached)
@@ -298,10 +300,14 @@ class AMDev:
if DEBUG >= 2: print(f"am {self.devfmt}: MEC is active. Issue a full reset.")
self.partial_boot = False
# Init sw for all IP blocks
for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu, self.gfx, self.sdma]: ip.init_sw()
# Init hw for IP blocks where it is needed
if not self.partial_boot:
if self.psp.is_sos_alive() and self.smu.is_smu_alive(): self.smu.mode1_reset()
for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]:
ip.init()
ip.init_hw()
if DEBUG >= 2: print(f"am {self.devfmt}: {ip.__class__.__name__} initialized")
# Booting done
@@ -309,7 +315,7 @@ class AMDev:
# Re-initialize main blocks
for ip in [self.gfx, self.sdma]:
ip.init()
ip.init_hw()
if DEBUG >= 2: print(f"am {self.devfmt}: {ip.__class__.__name__} initialized")
self.smu.set_clocks(level=-1) # last level, max perf.
@@ -319,7 +325,7 @@ class AMDev:
def fini(self):
if DEBUG >= 2: print(f"am {self.devfmt}: Finalizing")
for ip in [self.sdma, self.gfx]: ip.fini()
for ip in [self.sdma, self.gfx]: ip.fini_hw()
self.smu.set_clocks(level=0)
self.ih.interrupt_handler()
tinygrad/runtime/support/am/ip.py
+43 -47
View File
@@ -5,12 +5,13 @@ from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG
class AM_IP:
def __init__(self, adev): self.adev = adev
def init(self): raise NotImplementedError("IP block init must be implemeted")
def fini(self): pass
def set_clockgating_state(self): pass
def init_sw(self): pass # Prepare sw/allocations for this IP
def init_hw(self): pass # Initialize hw for this IP
def fini_hw(self): pass # Finalize hw for this IP
def set_clockgating_state(self): pass # Set clockgating state for this IP
class AM_SOC(AM_IP):
def init(self):
def init_hw(self):
self.adev.regRCC_DEV0_EPF2_STRAP2.update(strap_no_soft_reset_dev0_f2=0x0)
self.adev.regRCC_DEV0_EPF0_RCC_DOORBELL_APER_EN.write(0x1)
def set_clockgating_state(self): self.adev.regHDP_MEM_POWER_CTRL.update(atomic_mem_power_ctrl_en=1, atomic_mem_power_ds_en=1)
@@ -21,9 +22,7 @@ class AM_SOC(AM_IP):
f"s2a_doorbell_port{port}_range_size": size})
class AM_GMC(AM_IP):
def __init__(self, adev):
super().__init__(adev)
def init_sw(self):
# Memory controller aperture
self.mc_base = (self.adev.regMMMC_VM_FB_LOCATION_BASE.read() & 0xFFFFFF) << 24
self.mc_end = self.mc_base + self.adev.mm.vram_size - 1
@@ -39,7 +38,7 @@ class AM_GMC(AM_IP):
self.dummy_page_paddr = self.adev.mm.palloc(0x1000, zero=not self.adev.partial_boot, boot=True)
self.hub_initted = {"MM": False, "GC": False}
def init(self): self.init_hub("MM")
def init_hw(self): self.init_hub("MM")
def flush_hdp(self): self.adev.wreg(self.adev.reg("regBIF_BX0_REMAP_HDP_MEM_FLUSH_CNTL").read() // 4, 0x0)
def flush_tlb(self, ip:Literal["MM", "GC"], vmid, flush_type=0):
@@ -116,12 +115,11 @@ class AM_GMC(AM_IP):
if self.adev.reg(f"reg{ip}VM_L2_PROTECTION_FAULT_STATUS").read(): raise RuntimeError(f"{ip}VM_L2_PROTECTION_FAULT_STATUS: {st:#x} {va:#x}")
class AM_SMU(AM_IP):
def __init__(self, adev):
super().__init__(adev)
def init_sw(self):
self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP, prever_prefix='v')
self.driver_table_paddr = self.adev.mm.palloc(0x4000, zero=not self.adev.partial_boot, boot=True)
def init(self):
def init_hw(self):
self._send_msg(self.smu_mod.PPSMC_MSG_SetDriverDramAddrHigh, hi32(self.adev.paddr2mc(self.driver_table_paddr)), poll=True)
self._send_msg(self.smu_mod.PPSMC_MSG_SetDriverDramAddrLow, lo32(self.adev.paddr2mc(self.driver_table_paddr)), poll=True)
self._send_msg(self.smu_mod.PPSMC_MSG_EnableAllSmuFeatures, 0, poll=True)
@@ -165,7 +163,7 @@ class AM_SMU(AM_IP):
return self.adev.mmMP1_SMN_C2PMSG_82.read() if read_back_arg else None
class AM_GFX(AM_IP):
def init(self):
def init_hw(self):
# Wait for RLC autoload to complete
while self.adev.regCP_STAT.read() != 0 and self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read(bootload_complete=1) != 0: pass
@@ -202,7 +200,7 @@ class AM_GFX(AM_IP):
# NOTE: Wait for MEC to be ready. The kernel does udelay here as well.
time.sleep(0.05)
def fini(self):
def fini_hw(self):
self._grbm_select(me=1, pipe=0, queue=0)
self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
self.adev.regSPI_COMPUTE_QUEUE_RESET.write(1)
@@ -274,24 +272,13 @@ class AM_GFX(AM_IP):
_config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=4, me=1)
class AM_IH(AM_IP):
def __init__(self, adev):
super().__init__(adev)
def init_sw(self):
self.ring_size = 512 << 10
def _alloc_ring(size): return (self.adev.mm.palloc(size, zero=not self.adev.partial_boot, boot=True),
self.adev.mm.palloc(0x1000, zero=not self.adev.partial_boot, boot=True))
self.rings = [(*_alloc_ring(self.ring_size), "", 0), (*_alloc_ring(self.ring_size), "_RING1", 1)]
def interrupt_handler(self):
_, rwptr_vm, suf, _ = self.rings[0]
wptr = to_mv(self.adev.paddr2cpu(rwptr_vm), 8).cast('Q')[0]
if self.adev.reg(f"regIH_RB_WPTR{suf}").read(rb_overflow=1):
self.adev.reg(f"regIH_RB_WPTR{suf}").update(rb_overflow=0)
self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=1)
self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=0)
self.adev.regIH_RB_RPTR.write(wptr % self.ring_size)
def init(self):
def init_hw(self):
for ring_vm, rwptr_vm, suf, ring_id in self.rings:
self.adev.wreg_pair("regIH_RB_BASE", suf, f"_HI{suf}", self.adev.paddr2mc(ring_vm) >> 8)
@@ -315,7 +302,33 @@ class AM_IH(AM_IP):
self.adev.soc.doorbell_enable(port=1, awid=0x0, awaddr_31_28_value=0x0, offset=am.AMDGPU_NAVI10_DOORBELL_IH*2, size=2)
def interrupt_handler(self):
_, rwptr_vm, suf, _ = self.rings[0]
wptr = to_mv(self.adev.paddr2cpu(rwptr_vm), 8).cast('Q')[0]
if self.adev.reg(f"regIH_RB_WPTR{suf}").read(rb_overflow=1):
self.adev.reg(f"regIH_RB_WPTR{suf}").update(rb_overflow=0)
self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=1)
self.adev.reg(f"regIH_RB_CNTL{suf}").update(wptr_overflow_clear=0)
self.adev.regIH_RB_RPTR.write(wptr % self.ring_size)
class AM_SDMA(AM_IP):
def init_hw(self):
for pipe in range(2):
self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100) # 10s, 100ms per unit
self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9)
self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=0x2, wr_l2_policy=0x3, llc_noalloc=1) # rd=noa, wr=bypass
self.adev.reg(f"regSDMA{pipe}_F32_CNTL").update(halt=0, th1_reset=0)
self.adev.reg(f"regSDMA{pipe}_CNTL").update(ctxempty_int_enable=1, trap_enable=1)
self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x3, offset=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2, size=4)
def fini_hw(self):
self.adev.regSDMA0_QUEUE0_RB_CNTL.update(rb_enable=0)
self.adev.regSDMA0_QUEUE0_IB_CNTL.update(ib_enable=0)
self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
time.sleep(0.01)
self.adev.regGRBM_SOFT_RESET.write(0x0)
def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, doorbell:int, pipe:int, queue:int):
# Setup the ring
self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_MINOR_PTR_UPDATE").write(0x1)
@@ -331,26 +344,8 @@ class AM_SDMA(AM_IP):
f32_wptr_poll_enable=1, rb_size=(ring_size//4).bit_length()-1, rb_enable=1, rb_priv=1)
self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_IB_CNTL").update(ib_enable=1)
def init(self):
for pipe in range(2):
self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100) # 10s, 100ms per unit
self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9)
self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=0x2, wr_l2_policy=0x3, llc_noalloc=1) # rd=noa, wr=bypass
self.adev.reg(f"regSDMA{pipe}_F32_CNTL").update(halt=0, th1_reset=0)
self.adev.reg(f"regSDMA{pipe}_CNTL").update(ctxempty_int_enable=1, trap_enable=1)
self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x3, offset=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2, size=4)
def fini(self):
self.adev.regSDMA0_QUEUE0_RB_CNTL.update(rb_enable=0)
self.adev.regSDMA0_QUEUE0_IB_CNTL.update(ib_enable=0)
self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
time.sleep(0.01)
self.adev.regGRBM_SOFT_RESET.write(0x0)
class AM_PSP(AM_IP):
def __init__(self, adev):
super().__init__(adev)
def init_sw(self):
self.msg1_paddr = self.adev.mm.palloc(am.PSP_1_MEG, align=am.PSP_1_MEG, zero=not self.adev.partial_boot, boot=True)
self.cmd_paddr = self.adev.mm.palloc(am.PSP_CMD_BUFFER_SIZE, zero=not self.adev.partial_boot, boot=True)
self.fence_paddr = self.adev.mm.palloc(am.PSP_FENCE_BUFFER_SIZE, zero=not self.adev.partial_boot, boot=True)
@@ -361,8 +356,7 @@ class AM_PSP(AM_IP):
self.max_tmr_size = 0x1300000
self.tmr_paddr = self.adev.mm.palloc(self.max_tmr_size, align=am.PSP_TMR_ALIGNMENT, zero=not self.adev.partial_boot, boot=True)
def is_sos_alive(self): return self.adev.regMP0_SMN_C2PMSG_81.read() != 0x0
def init(self):
def init_hw(self):
sos_components_load_order = [
(am.PSP_FW_TYPE_PSP_KDB, am.PSP_BL__LOAD_KEY_DATABASE), (am.PSP_FW_TYPE_PSP_KDB, am.PSP_BL__LOAD_TOS_SPL_TABLE),
(am.PSP_FW_TYPE_PSP_SYS_DRV, am.PSP_BL__LOAD_SYSDRV), (am.PSP_FW_TYPE_PSP_SOC_DRV, am.PSP_BL__LOAD_SOCDRV),
@@ -383,6 +377,8 @@ class AM_PSP(AM_IP):
for psp_desc in self.adev.fw.descs: self._load_ip_fw_cmd(*psp_desc)
self._rlc_autoload_cmd()
def is_sos_alive(self): return self.adev.regMP0_SMN_C2PMSG_81.read() != 0x0
def _wait_for_bootloader(self): self.adev.wait_reg(self.adev.regMP0_SMN_C2PMSG_35, mask=0xFFFFFFFF, value=0x80000000)
def _prep_msg1(self, data):