from __future__ import annotations from typing import cast, Any import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, itertools, collections, atexit from dataclasses import dataclass, replace from tinygrad.runtime.support.hcq2 import HWQueue, encode_cmdbuf, to_name, patch, unwrap_view, layout_args, pack_args, make_program, ccall from tinygrad.uop.ops import sint, UOp, ProgramInfo, Ops, UPat, PatternMatcher, KernelInfo, uopfunc from tinygrad.device import BufferStorage, BufferSpec, Buffer, Device, Allocator, Compiled from tinygrad.dtype import dtypes, DType from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, prod, colored from tinygrad.helpers import ceildiv, unwrap, pluralize, ContextVar, VIZ, DEV from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm, libc from tinygrad.runtime.autogen.am import am from tinygrad.runtime.support.elf import elf_loader from tinygrad.runtime.support.memory import MMIOInterface, AddrSpace from tinygrad.runtime.support.system import FileIOInterface, filter_visible_devices, PCIIfaceBase, USBPCIDevice, MAP_FIXED, MAP_NORESERVE from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc from tinygrad.runtime.support.usb import USB3, setup_usb_rules, usb_reset from tinygrad.engine.realize import get_call_kernel_args if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import SQTT = ContextVar("SQTT", abs(VIZ.value)>=2) SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE = \ ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("SQTT_LIMIT_SE", 0), ContextVar("SQTT_SIMD_SEL", 0), ContextVar("SQTT_TOKEN_EXCLUDE", 0) PMC = ContextVar("PMC", abs(VIZ.value)>=2) EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h WAIT_REG_MEM_FUNCTION_EQ = 3 # == WAIT_REG_MEM_FUNCTION_NEQ = 4 # != WAIT_REG_MEM_FUNCTION_GEQ = 5 # >= AQL_HDR = (1 << hsa.HSA_PACKET_HEADER_BARRIER) | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE) \ | (hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE) @dataclass(frozen=True) class ProfileSQTTEvent(ProfileEvent): device:str; kern:int; se:int; blob:bytes; itrace:bool; exec_tag:int # noqa: E702 @dataclass(frozen=True) class PMCSample: name:str; block:str; xcc:int; inst:int; se:int; sa:int; wgp:int; off:int; size:int; regsample:str # noqa: E702 @dataclass(frozen=True) class ProfilePMCEvent(ProfileEvent): device:str; kern:int; sched:list[PMCSample]; blob:bytes; exec_tag:int # noqa: E702 # ***************** # PM4 def _queue_args(hq:HWQueue, q) -> list[UOp]: # the ring and its pointers, named {device}_{name}_{queue}. put is the host's copy of the write pointer shapes = [("ring", (q.ring.nbytes // 4,), dtypes.uint32, hq.devs[0])] + [(n, (1,), dtypes.uint64, hq.devs[0]) for n in ("write_ptr", "doorbell")] shapes += [("put_value", (1,), dtypes.uint64, hq.dev.host)] return [UOp.alloc(s, dt, 0, device=d, name=to_name(hq.dev.device, n, hq.queue)) for n, s, dt, d in shapes] @uopfunc def amd_push(cmdbuf:UOp, words:UOp, ring:UOp, wptr:UOp, doorbell:UOp, put:UOp, unit:int=4, doorbell_lag:int=0) -> UOp: rs, n, p = ring.max_numel(), words.max_numel() // 4, put.index(0).load() # put counts units, the ring dwords first = (rs - (tail:=((p * (unit // 4)) % rs).cast(dtypes.int))).minimum(n) for dst, src, count in ((tail, 0, first), (0, first, n - first)): cmdbuf = ccall(libc.memcpy, ring.after(cmdbuf).index(dst), words.bitcast(dtypes.uint32).index(src), (count * 4).cast(dtypes.uint64)) w = wptr.after(cmdbuf).index(0).store(nxt:=p + words.max_numel() // unit).barrier() return doorbell.after(put.after(w).index(0).store(nxt)).index(0).store(nxt - doorbell_lag).sink() @uopfunc def amd_sdma_submit(cmdbuf:UOp, ring:UOp, wptr:UOp, doorbell:UOp, put:UOp) -> UOp: # sdma needs the cmdbuf contiguous in the ring: if it won't fit before the ring end, restart at 0 and zero the tail rs, size_dw = ring.max_numel(), cmdbuf.max_numel() // 4 if size_dw > rs: raise RuntimeError(f"SDMA command buffer ({size_dw*4} bytes) exceeds ring size ({rs*4} bytes)") put_b = put.index(0).load() tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int) fits = (size_dw <= rs - tail).cast(dtypes.int) start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail) zero_tail = ccall(libc.memset, ring.after(cmdbuf).index(tail), 0, (zero_amt * 4).cast(dtypes.uint64)) copy = ccall(libc.memcpy, ring.after(zero_tail).index(start_dw), cmdbuf.bitcast(dtypes.uint32).index(0), UOp.const(size_dw * 4, dtypes.uint64)) next_put = put_b + ((zero_amt + size_dw) * 4).cast(put_b.dtype) w = wptr.after(copy).index(0).store(next_put).barrier() return doorbell.after(put.after(w).index(0).store(next_put)).index(0).store(next_put).sink() def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals) def dispatch_packet(data:AMDProgramData, info:ProgramInfo, kernel_object:UOp=UOp.const(0, dtypes.uint64), kernarg_address:UOp=UOp.const(0, dtypes.uint64)) -> list: # as words: the grid may be symbolic pkt = bytes(hsa.hsa_kernel_dispatch_packet_t(header=AQL_HDR | (hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE), setup=3 << hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS, private_segment_size=data.private_segment_size, group_segment_size=data.group_segment_size, **{f"workgroup_size_{d}": l for d, l in zip("xyz", info.local_size)})) grid = [(g * l).cast(dtypes.uint32) if isinstance(g, UOp) else g * l for g, l in zip(cast(tuple[Any, ...], info.global_size), info.local_size)] return [UOp(Ops.BINARY, arg=pkt[:12]), *grid, UOp(Ops.BINARY, arg=pkt[24:32]), kernel_object, kernarg_address, UOp(Ops.BINARY, arg=pkt[48:])] class AMDComputeQueue(HWQueue): dev:AMDDevice def __init__(self, submit): super().__init__(submit) self.pm4, self.gc, self.soc, self.nbio, self.target = self.dev.pm4, self.dev.gc, self.dev.soc, self.dev.nbio, self.dev.target self.profiled:list[UOp] = [] if self.dev.pmc_enabled: self.pmc_start() def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, _dw(vals) - 1), *vals) def wreg(self, reg:AMDReg, *args:sint, **kwargs:int): if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified') if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_SH_REG_END: set_packet, set_packet_start = self.pm4.PACKET3_SET_SH_REG, self.pm4.PACKET3_SET_SH_REG_START elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1: set_packet, set_packet_start = self.pm4.PACKET3_SET_UCONFIG_REG, self.pm4.PACKET3_SET_UCONFIG_REG_START else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet') self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),))) @contextlib.contextmanager def pred_exec(self, xcc_mask:int): # the count fills in when the block closes if self.dev.xccs > 1: self.pkt3(self.pm4.PACKET3_PRED_EXEC, xcc_mask << 24) start = len(self.blob) yield if self.dev.xccs > 1: cnt, = struct.unpack("I", self.blob[start-4:start]) self.blob[start-4:start] = struct.pack("I", cnt | (len(self.blob) - start) // 4) def set_grbm(self, instance=None, se=None, sh=None, wgp=None): instance_val = (wgp << 2 | (instance or 0)) if wgp is not None else instance self.wreg(self.gc.regGRBM_GFX_INDEX, **{(f'{key}_broadcast_writes' if val is None else f'{key}_index'): (1 if val is None else val) for key, val in [('instance', instance_val), ('se', se), ('sh' if self.target[0] == 9 else 'sa', sh)]}) def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ): wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \ | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0) self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *((mem,) if mem is not None else (reg, reg_done)), value, mask, 4) def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1): if self.target[0] != 9: cache_flags_dw = self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \ | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \ | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \ | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \ | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2) return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw) cp_coher_cntl = self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \ self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \ self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \ self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \ self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2) return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A) def release_mem(self, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False): if self.target[0] != 9: cache_flags_dw = 0 if not cache_flush else (self.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \ | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \ | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | self.pm4.PACKET3_RELEASE_MEM_GCR_SEQ) event_dw = self.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \ | self.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe) memsel_dw = self.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | self.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \ | self.pm4.PACKET3_RELEASE_MEM_DST_SEL(0) else: cache_flags_dw = 0 if not cache_flush else (self.pm4.EOP_TC_WB_ACTION_EN | self.pm4.EOP_TC_NC_ACTION_EN) event_dw = self.pm4.EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | \ self.pm4.EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe) memsel_dw = self.pm4.DATA_SEL(data_sel) | self.pm4.INT_SEL(int_sel) ctxid = 0 addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64) val_w = UOp.const(value, dtypes.uint64) self.pkt3(self.pm4.PACKET3_RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid) def memory_barrier(self): (req, done), mask = self.dev.hdp_flush, self.dev.hdp_flush[0].fields_mask('cp2') self.wait_reg_mem(reg=req.addr[0], reg_done=done.addr[0], value=mask, mask=mask, op=WAIT_REG_MEM_FUNCTION_EQ) self.acquire_mem() def spi_config(self, tracing:bool): self.wreg(self.gc.regSPI_CONFIG_CNTL, ps_pkr_priority_cntl=3, exp_priority_order=3, gpr_write_priority=0x2c688, enable_sqg_bop_events=int(tracing), enable_sqg_top_events=int(tracing)) ### profiling: a kernel's slot holds its counters and trace until a synchronize reads them back def prof_buf(self, name:str, dtype:DType) -> UOp: return UOp.alloc((getattr(self.dev, name).nbytes // dtype.itemsize,), dtype, 0, device=self.devs[0], name=name) def prof_start(self, info:ProgramInfo, lib:UOp) -> UOp: log = self.prof_buf("prof_log", dtypes.uint64) slot = (log.index(0).load() + len(self.profiled)) % self.dev.prof_slots self.profiled.append(log.index(1 + slot.cast(dtypes.int)).store(UOp.const(unwrap_view(lib)[0].arg.slot, dtypes.uint64))) if self.dev.sqtt_enabled: self.sqtt_start(slot) self.sqtt_setup_exec(info) return slot def prof_stop(self, slot:UOp): self.exec_barrier() if self.dev.pmc_enabled: self.pmc_read(slot) if self.dev.sqtt_enabled: self.sqtt_stop(slot) def prof_bump(self, cmdbuf:UOp) -> UOp: if not self.profiled: return cmdbuf log = self.prof_buf("prof_log", dtypes.uint64) return cmdbuf.after(log.after(cmdbuf, *self.profiled).index(0).store(log.index(0).load() + len(self.profiled))) ### PMC def pmc_reset_counters(self, en=True): self.set_grbm() self.wreg(self.gc.regCP_PERFMON_CNTL if self.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=0) if en: self.wreg(self.gc.regCP_PERFMON_CNTL if self.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=1) def pmc_start(self): # every submit self.pmc_reset_counters(en=False) self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1, **({'vmid_mask':0xffff} if (gfx9:=self.target[0] == 9) else {})) if not gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff) end_off, sched = 0, [] block2pid:dict[str, itertools.count] = collections.defaultdict(lambda: itertools.count()) for name in self.dev.pmc_names: block, idx = self.dev.pmc_counters[name] # sq block on gfx11+ goes down to wgps inst_cnt, se_cnt, sa_cnt, wgp_cnt = {"GRBM": (1, 1, 1, 1), "GL2C": (32, 1, 1, 1), "TCC": (16, 1, 1, 1), "SQ": (1, self.dev.se_cnt) + ((1, 1) if gfx9 else (2, self.dev.iface.props['cu_per_simd_array'] // 2))}[block] end_off += (rec_size:=prod((self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt)) * 8) # gfx11+ and later require even-numbered SQ *_SELECT registers regsample = f'reg{block}_PERFCOUNTER{(pcid:=next(block2pid[block]))}' if (regsel:=getattr(self.gc, (f'reg{block}_PERFCOUNTER{(pcid*2) if not gfx9 and block=="SQ" else pcid}_SELECT'), None)) is None: raise RuntimeError(f'{block} is out of perfcounter registers: ({regsample} is not found)') self.wreg(regsel, perf_sel=idx, **({'simd_mask':0xf, 'sqc_bank_mask':0xf, 'sqc_client_mask':0xf} if gfx9 and block == "SQ" else {})) sched.append(PMCSample(name, block, self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt, end_off-rec_size, rec_size, regsample)) self.dev.pmc_sched = sched if gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_MASK, sh0_mask=0xffff, sh1_mask=0xffff) self.wreg(self.gc.regCOMPUTE_PERFCOUNT_ENABLE, 1) self.pmc_reset_counters(en=True) def pmc_read(self, slot:UOp): buf = self.prof_buf("pmc_buf", dtypes.uint8).getaddr(self.devs) + slot * self.dev.pmc_size self.set_grbm() self.wreg(self.gc.regCP_PERFMON_CNTL if self.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=1, perfmon_sample_enable=1) for smp in self.dev.pmc_sched: offset = itertools.count(smp.off, step=8) for xcc in range(smp.xcc): with self.pred_exec(xcc_mask=1 << xcc): for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(smp.inst), range(smp.se), range(smp.sa), range(smp.wgp)): loff = next(offset) if smp.wgp > 1 and not self.dev.iface.is_wgp_active(xcc, se_idx, sa_idx, wgp_idx): continue self.set_grbm(**({'instance':inst} if smp.inst > 1 else ({'se':se_idx}|({'sh':sa_idx, 'wgp':wgp_idx} if self.target[0] != 9 else {})))) # Copy counter to memory (src_sel = perf, dst_sel = tc_l2) lo, hi = getattr(self.gc, f'{smp.regsample}_LO'), getattr(self.gc, f'{smp.regsample}_HI', None) self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, lo.addr[0], 0, buf + loff) if hi is not None: self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, hi.addr[0], 0, buf + (loff + 4)) self.pmc_reset_counters(en=True) ### SQTT def sqtt_config(self, tracing:bool): trace_ctrl = {'rt_freq': self.soc.SQ_TT_RT_FREQ_4096_CLK} if self.target < (12,0,0) else {} self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1, util_timer=1, mode=int(tracing), **trace_ctrl) def sqtt_setup_exec(self, info:ProgramInfo): if SQTT_LIMIT_SE: # Calculate number of CUs per SE to enable based on blocks count. 4 is maximum simd per CU, but on rdna we can trace only 1. cu_per_se = prod([x if isinstance(x, int) else 1 for x in info.global_size]) // ((self.dev.cu_cnt // self.dev.se_cnt) * 4) for xcc in range(self.dev.xccs): with self.pred_exec(xcc_mask=1 << xcc): for i in range(8 if self.target[0] != 9 else 4): if SQTT_LIMIT_SE > 1: mask = 1 if SQTT_ITRACE_SE_MASK.value & (1 << i) else 0 # only run unmasked shader engines else: sa_mask = (1 << (self.dev.iface.props['cu_per_simd_array'] // 2)) - 1 cu_mask = (1 << (cu_per_se + (1 if i == 0 else 0))) - 1 mask = lo32((cu_mask & sa_mask) | (cu_mask & (sa_mask << 16)) << 16) self.wreg(getattr(self.gc, f'regCOMPUTE_STATIC_THREAD_MGMT_SE{i}'), mask) def sqtt_start(self, slot:UOp): self.memory_barrier() win, ses = self.dev.sqtt_win, self.dev.sqtt_ses base = self.prof_buf("sqtt_buf", dtypes.uint8).getaddr(self.devs) + slot * win if self.target[0] == 9: self.set_grbm() self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, simd_en=0xf, cu_sel=0, sq_stall_en=1, spi_stall_en=1, reg_stall_en=1, vm_id_mask=0) for se in range(ses): mask = (__SQTT_MISC:=1<<0) | (__SQTT_TIME:=1<<1) | (__SQTT_REG:=1<<2) | (__SQTT_WAVE_START:=1<<3) | (__SQTT_WAVE_END:=1<<6) \ | (__SQTT_USERDATA:=1<<12) | (__SQTT_REG_CS:=1<<5) | (__SQTT_REG_CS_PRIV:=1<<15) if (SQTT_ITRACE_SE_MASK.value >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13) buf0_lo, buf0_hi = [((base + se * self.dev.prof_slots * win) >> sh).cast(dtypes.uint32) for sh in (12, 44)] with self.pred_exec(xcc_mask=1<<(se // self.dev.se_cnt)): self.set_grbm(se=se % self.dev.se_cnt, sh=0) self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_mask=0xf, token_mask=mask) self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK2, inst_mask=0xffffffff) self.wreg(self.gc.regSQ_THREAD_TRACE_BASE, buf0_lo) self.wreg(self.gc.regSQ_THREAD_TRACE_BASE2, buf0_hi) self.wreg(self.gc.regSQ_THREAD_TRACE_SIZE, size=win >> 12) self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, reset_buffer=1) self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=1) else: self.spi_config(tracing=True) # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well for se in range(ses): self.set_grbm(se=se, sh=0) buf0_lo, buf0_hi = [((base + se * self.dev.prof_slots * win) >> sh).cast(dtypes.uint32) for sh in (12, 44)] if self.target >= (12,0,0): self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, size=win >> 12) self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_LO, buf0_lo) self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_HI, buf0_hi) else: self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, self.gc.regSQ_THREAD_TRACE_BUF0_SIZE.encode(size=win >> 12) | buf0_hi) self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, buf0_lo) # NOTE: SQTT can only trace instructions on one simd per se, this selects the simd in first wgp in first sa. # For RGP to display instruction trace it has to see it on first SE. Howerver ACE/MEC/whatever does the dispatching starting with second se, # and on amdgpu/non-AM it also does weird things with dispatch order inside se: around 7 times out of 10 it starts from the last cu, but # sometimes not, especially if the kernel has more than one wavefront which means that kernels with small global size might get unlucky and # be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the # CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE and trace even kernels that only have one wavefront. # Use SQTT_SIMD_SEL to select which SIMD to trace (0-3). Memory ops show different InstOp values (0x2x vs 0x5x) based on SIMD. cs_wtype = (1 << 6) if self.target >= (12,0,0) else self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=cs_wtype, simd_sel=SQTT_SIMD_SEL.value, wgp_sel=0, sa_sel=0) reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \ self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT token_exclude = SQTT_TOKEN_EXCLUDE.value | ((1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.target < (12,0,0) else 0) # disable instr tracing if not (SQTT_ITRACE_SE_MASK.value >> se) & 0b1: # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11. token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \ 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \ 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT) if self.target < (12,0,0) else 0x927 self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1, **({} if self.target < (12,0,0) else {'exclude_barrier_wait': 1})) self.sqtt_config(tracing=True) self.set_grbm() if self.target[0] != 9: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1) self.memory_barrier() # Magic values from src/amd/common/ac_sqtt.c:ac_sqtt_emit_stop and src/amd/common/ac_sqtt.c:ac_sqtt_emit_wait def sqtt_stop(self, slot:UOp): self.memory_barrier() self.set_grbm() ses = self.dev.sqtt_ses wptrs = self.prof_buf("sqtt_wptrs", dtypes.uint32).getaddr(self.devs) + slot * (ses * 4) # Start shutting everything down if self.target[0] == 9: self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=0) else: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 0) self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_FINISH) | self.pm4.EVENT_INDEX(0)) # For each SE wait for finish to complete and copy regSQ_THREAD_TRACE_WPTR to know where in the buffer trace data ends for se in range(ses): with self.pred_exec(xcc_mask=1<<(se // self.dev.se_cnt)): self.set_grbm(se=se % self.dev.se_cnt, sh=0) regstatus = self.gc.regSQ_THREAD_TRACE_STATUS.addr[0] - (self.pm4.PACKET3_SET_UCONFIG_REG_START if self.target[0] == 9 else 0) if self.target[0] != 9: self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0) self.sqtt_config(tracing=False) self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0) self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH)) # Copy WPTR to memory (src_sel = perf, dst_sel = tc_l2, wr_confirm = True) self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, wptrs + se * 4) self.set_grbm() if self.target[0] != 9: self.spi_config(tracing=False) self.memory_barrier() ### exec def kernargs(self, call:UOp, prg:UOp, data:AMDProgramData) -> UOp: if not data.kernargs_segment_size and not data.enable_dispatch_ptr: return UOp.const(0, dtypes.uint64) # no kernargs, nothing to point at args = get_call_kernel_args(call, prg, self.devs) words = pack_args(layout_args(args), data.kernargs_segment_size) + (dispatch_packet(data, prg.arg) if data.enable_dispatch_ptr else []) return UOp(Ops.LINEAR, src=tuple(words), arg="kernargs").getaddr(self.devs) def exec(self, call:UOp, prg:UOp): data, lib = amd_build_program(self.dev, prg, self.devs) prog_addr = lib.getaddr(self.devs) + data.entry_point_offset args_addr = self.kernargs(call, prg, data) scratch_addr = UOp.alloc((data.private_segment_size,), dtypes.uint8, 0, device=self.devs[0], name="scratch").getaddr(self.devs) user_regs:list = [] if data.enable_private_segment_sgpr: user_regs = [scratch_addr | (1 << 63), 0xffffffff, 0x20c14000] if data.enable_dispatch_ptr: user_regs += [args_addr + data.kernargs_segment_size] user_regs += [args_addr] dispatch_init = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode( **({'cs_w32_en': int(data.wave32)} if self.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1) if data.private_segment_size: self.exec_barrier() # kernels share the scratch ring, each with its own TMPRING wave size self.acquire_mem(gli=0, gl2=0) if self.dev.prof_enabled: slot = self.prof_start(prg.arg, lib) self.wreg(self.gc.regCOMPUTE_PGM_LO, prog_addr >> 8) self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2) self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3) self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size)) for xcc_id in range(self.dev.xccs): # architected flat scratch: each xcc gets its part with self.pred_exec(xcc_mask=1 << xcc_id): self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // self.dev.xccs * xcc_id) >> 8) self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0) self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs) self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, self.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH"))) self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *prg.arg.local_size, 0, 0) self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *prg.arg.global_size, dispatch_init) if self.dev.sqtt_enabled: self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_MARKER) | self.pm4.EVENT_INDEX(0)) if self.dev.prof_enabled: self.prof_stop(slot) def exec_barrier(self): self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH)) def wait(self, signal:UOp, value:UOp): self.wait_reg_mem(value.cast(dtypes.uint32), mem=signal.getaddr(self.devs)) def timestamp(self, signal:UOp): with self.pred_exec(xcc_mask=0b1): self.release_mem(signal.getaddr(self.devs) + UOp.const(8, dtypes.uint64), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter, self.pm4.int_sel__mec_release_mem__none) def signal(self, signal:UOp, value:UOp): with self.pred_exec(xcc_mask=0b1): self.release_mem(signal.getaddr(self.devs), value, self.pm4.data_sel__mec_release_mem__send_32_bit_low, self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True) def submit(self, cmdbuf:UOp) -> UOp: # the ring gets an indirect buffer packet: 4 dwords, put stays aligned so it never wraps mid packet base, off = unwrap_view(cmdbuf) blob = struct.pack("IIII", self.pm4.PACKET3(self.pm4.PACKET3_INDIRECT_BUFFER, 2), 0, 0, cmdbuf.max_numel() // 4 | self.pm4.INDIRECT_BUFFER_VALID) ib = UOp.alloc((16,), dtypes.uint8, device=self.dev.host, name=to_name("ib", self.queue)) return amd_push(self.prof_bump(cmdbuf), patch(ib, [(4, base.getaddr(self.devs) + off)], blob), *_queue_args(self, self.dev.compute_queue)) class AMDComputeAQLQueue(AMDComputeQueue): # the ring holds 64 byte aql packets: a dispatch per kernel, the pm4 between them wrapped as an ib def __init__(self, submit): super().__init__(submit) self.cmd_addr = UOp.variable("cmdbuf", 0, 2**48, dtypes.uint64) # the packets point into the cmdbuf, its address binds at submit self.pkts:list[UOp] = [] self.run_start = 0 def close_run(self, end:int): if end > self.run_start: hdr = AQL_HDR | (hsa.HSA_PACKET_TYPE_VENDOR_SPECIFIC << hsa.HSA_PACKET_HEADER_TYPE) | (1 << 16) ib = [self.pm4.PACKET3(self.pm4.PACKET3_INDIRECT_BUFFER, 2), self.cmd_addr + self.run_start, (end - self.run_start) // 4 | self.pm4.INDIRECT_BUFFER_VALID] self.pkts += [UOp.const(w, dtypes.uint32) if isinstance(w, int) else w for w in [hdr, *ib, 10, *[0] * 10]] self.run_start = end def signal(self, signal:UOp, value:UOp): self.close_run(len(self.blob)) # close pm4, so we can sync xccs. without that wait-signal can race super().signal(signal, value) def exec_barrier(self): pass def exec(self, call:UOp, prg:UOp): data, lib = amd_build_program(self.dev, prg, self.devs) self.dev.scratch_buffer(data.private_segment_size) # the queue descriptor holds the scratch if self.dev.prof_enabled: slot = self.prof_start(prg.arg, lib) self.close_run(len(self.blob)) self.pkts += [UOp.const(w, dtypes.uint32) if isinstance(w, int) else w for w in dispatch_packet(data, prg.arg, lib.getaddr(self.devs) + data.desc_offset, self.kernargs(call, prg, data))] self.run_start = len(self.blob) if self.dev.prof_enabled: self.prof_stop(slot) def submit(self, cmdbuf:UOp) -> UOp: # the doorbell is the last packet's index self.close_run(cmdbuf.max_numel()) base, off = unwrap_view(cmdbuf) self.blob, self.patches = bytearray(), [] # q again, for the aql stream self.q(*UOp.sink(*self.pkts).substitute({self.cmd_addr: base.getaddr(self.devs) + off}).src) aql = encode_cmdbuf(self, name="aql", device=self.dev.host) return amd_push(self.prof_bump(cmdbuf), aql, *_queue_args(self, self.dev.compute_queue), unit=64, doorbell_lag=1) # ***************** # SDMA class AMDSDMAQueue(HWQueue): dev:AMDDevice def __init__(self, submit): super().__init__(submit) self.sdma, self.target, self.max_copy_size = self.dev.sdma, self.dev.target, self.dev.max_copy_size def copy(self, dst:UOp, src:UOp, sz:int): for off in range(0, sz, self.max_copy_size): self.q(self.sdma.SDMA_OP_COPY|self.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_COPY_LINEAR), min(sz-off, self.max_copy_size)-1, 0, *(a + UOp.const(off, dtypes.uint64) if off else a for a in (src.getaddr(self.devs), dst.getaddr(self.devs)))) def memory_barrier(self): (req, done), op = self.dev.hdp_flush, self.sdma.SDMA_OP_POLL_REGMEM | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_HDP_FLUSH(1) self.q(op | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_EQ), done.addr[0] << 2, req.addr[0] << 2, m:=req.fields_mask('sdma0'), m, self.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(10) | self.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)) def wait(self, signal:UOp, value:UOp, eq:bool=False): func = WAIT_REG_MEM_FUNCTION_EQ if eq else WAIT_REG_MEM_FUNCTION_GEQ op = self.sdma.SDMA_OP_POLL_REGMEM | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(func) | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1) self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), (1 << 8 * min(value.dtype.itemsize, 4)) - 1, self.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | self.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)) def timestamp(self, signal:UOp): self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL), signal.getaddr(self.devs) + UOp.const(8, dtypes.uint64)) def write(self, dst:UOp, *words): self.q(self.sdma.SDMA_OP_WRITE | self.sdma.SDMA_PKT_WRITE_UNTILED_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_WRITE_LINEAR), dst.getaddr(self.devs), _dw(words) - 1, *words) def signal(self, signal:UOp, value:UOp): op = self.sdma.SDMA_OP_FENCE | (self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.target[0] != 9 else 0) self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), self.sdma.SDMA_OP_TRAP, 0) def submit(self, cmdbuf:UOp) -> UOp: return amd_sdma_submit(cmdbuf, *_queue_args(self, unwrap(self.dev.sdma_queue(int(self.queue.split(":")[1]))))) def encode(self) -> UOp: return self.submit(encode_cmdbuf(self, self.lin, device=self.dev.host)) @dataclass(frozen=True) class AMDProgramData: desc_offset:int; entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool # noqa: E702 private_segment_size:int; group_segment_size:int; kernargs_segment_size:int # noqa: E702 enable_dispatch_ptr:int; enable_private_segment_sgpr:int # noqa: E702 def amd_build_program(dev, prg:UOp, devs:tuple[str, ...]) -> tuple[AMDProgramData, UOp]: data, image = _amd_program_image(dev, prg.src[3].arg) return data, patch(make_program(prg, len(image), devs[0]), [], image) def _amd_program_image(dev, lib:bytes) -> tuple[AMDProgramData, bytes]: image, sections, relocs = elf_loader(lib) rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata") for off, sym, typ, addent in relocs: assert typ == 5, f"unknown AMD reloc {typ}" # R_AMDGPU_REL64 image[off:off+8] = struct.pack(' (dev.iface.props['lds_size_in_kb']*1024)//512: raise RuntimeError("Too many resources requested: group_segment_size") edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR data = AMDProgramData(desc_offset=rodata, entry_point_offset=rodata + desc.kernel_code_entry_byte_offset, rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0), # priv=1 on gfx11 for cwsr rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3, wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, group_segment_size=desc.group_segment_fixed_size, kernargs_segment_size=desc.kernarg_size, enable_dispatch_ptr=edp, enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER) return data, bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords class AMDAllocator(Allocator['AMDDevice']): def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue) def _free(self, storage:BufferStorage, options:BufferSpec): self.dev.synchronize() self.dev.iface.free(storage) def _map(self, buf:Buffer) -> BufferStorage: return self.dev.iface.map(buf) def _unmap(self, mapping:BufferStorage): self.dev.iface.unmap(mapping) def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset @dataclass class AMDQueueDesc: ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer # noqa: E702 eop_buffer: Buffer|None = None; cwsr_buffer: Buffer|None = None; params: tuple|None = None # noqa: E702 class KFDIface: kfd:FileIOInterface|None = None event_page:Buffer gpus:list[FileIOInterface] = [] count:int = 0 def _is_usable_gpu(self, gpu_id): with contextlib.suppress(OSError): return int(gpu_id.read()) != 0 return False def __init__(self, dev, device_id): self.dev = dev kfd_topo_path = "/sys/devices/virtual/kfd/kfd/topology/nodes" # Initialize KFD interface during first run if KFDIface.kfd is None: KFDIface.kfd = FileIOInterface("/dev/kfd", os.O_RDWR) gpus = [g for g in FileIOInterface(kfd_topo_path).listdir() if self._is_usable_gpu(FileIOInterface(f"{kfd_topo_path}/{g}/gpu_id"))] KFDIface.gpus = filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])), "AMD") KFDIface.count = len(KFDIface.gpus) if device_id >= len(KFDIface.gpus): raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?") self.gpu_id = int(FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/gpu_id").read()) self.props = {(p:=l.split())[0]: int(p[1]) for l in FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/properties").read().splitlines()} self.dev_sysfs_path = f"/sys/class/drm/renderD{self.props['drm_render_minor']}/device" ip_base = f"{self.dev_sysfs_path}/ip_discovery/die/0" id2ip = {am.GC_HWID: am.GC_HWIP, am.SDMA0_HWID: am.SDMA0_HWIP, am.NBIF_HWID: am.NBIF_HWIP} ip_hw = [(id2ip[int(hwid)], int(hwid)) for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip] self.ip_versions = {ip:tuple(int(FileIOInterface(f'{ip_base}/{hw}/0/{part}').read()) for part in ['major','minor','revision']) for ip,hw in ip_hw} self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR) self.kfd_ver = ((ver_st:=kfd.AMDKFD_IOC_GET_VERSION(KFDIface.kfd)).major_version, ver_st.minor_version) kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id) if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0) def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> BufferStorage: flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT else: flags |= (kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR if host else kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) # Make mapped cpu address to be uncachable if cpu_addr is not None: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR: buf = addr = cpu_addr or FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0) else: buf, addr = 0, FileIOInterface.anon_mmap(0, size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE, 0) try: mem = kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(self.kfd, va_addr=addr, size=size, gpu_id=self.gpu_id, flags=flags, mmap_offset=buf) except OSError as e: if e.errno == errno.EINVAL and (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) and cpu_access: raise MemoryError("Cannot allocate host-visible VRAM. Ensure the resizable BAR option is enabled on your system.") from e if e.errno == errno.ENOMEM: raise MemoryError(f"Cannot allocate {size} bytes: no memory is available.") from e raise if not (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR): buf = self.drm_fd.mmap(mem.va_addr, mem.size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_FIXED, mem.mmap_offset) assert addr == buf == mem.va_addr self._map_handle(mem.handle) return BufferStorage(mem.va_addr, mem, MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None) def free(self, storage:BufferStorage): self._unmap_handle(storage.meta.handle) if storage.buf: FileIOInterface.munmap(storage.buf, storage.meta.size) kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=storage.meta.handle) def unmap(self, mapping:BufferStorage): handle, owned = mapping.meta self._unmap_handle(handle) if owned: kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=handle) def map(self, buf:Buffer) -> BufferStorage: if buf.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}: if buf._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address") return replace(mem:=self.alloc(buf.nbytes, host=True, cpu_addr=buf._buf), meta=(mem.meta.handle, True)) self._map_handle(buf.meta.handle) return BufferStorage(buf._buf, (buf.meta.handle, False)) def _map_handle(self, handle): gpus = (ctypes.c_int32 * 1)(self.gpu_id) stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1) assert stm.n_success == 1 def _unmap_handle(self, handle): gpus = (ctypes.c_int32 * 1)(self.gpu_id) stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1) assert stm.n_success == 1 def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0, idx=0): if not hasattr(self, 'queue_event_arr'): if not hasattr(KFDIface, 'event_page'): KFDIface.event_page = Buffer(self.dev.device, 0x8000, options=BufferSpec(uncached=True), preallocate=True) kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle) KFDIface.event_page.get_buf(self.dev.device) self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(*[kfd.struct_kfd_event_data(event_id=kfd.AMDKFD_IOC_CREATE_EVENT( KFDIface.kfd, event_type=t, auto_reset=int(t == kfd.KFD_IOC_EVENT_SIGNAL)).event_id) for t in (kfd.KFD_IOC_EVENT_SIGNAL, kfd.KFD_IOC_EVENT_MEMORY, kfd.KFD_IOC_EVENT_HW_EXCEPTION)]) queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf, ring_size=ring.nbytes, gpu_id=self.gpu_id, queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7), eop_buffer_address=eop_buffer._buf if eop_buffer else 0, eop_buffer_size=eop_buffer.nbytes if eop_buffer else 0, ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size, write_pointer_address=gart._buf+wptr, read_pointer_address=gart._buf+rptr+8*xcc_id) if not hasattr(self, 'doorbells'): self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages self.doorbells = cast(FileIOInterface, KFDIface.kfd).mmap(0, 0x2000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, self.doorbells_base) (put_value := Buffer("CPU", 8, preallocate=True)).host.view(fmt='Q')[0] = 0 doorbell = Buffer("CPU", 8, options=BufferSpec(external_ptr=self.doorbells + queue.doorbell_offset - self.doorbells_base), preallocate=True) return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(8, rptr+8*xcc_id).ensure_allocated(), write_ptr=gart.view(8, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer) def sleep(self, tm:int): kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=ctypes.addressof(self.queue_event_arr), num_events=3, wait_for_all=0, timeout=tm) if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: self.on_device_hang() def on_device_hang(self): def _str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_) # try to collect fault info if not already set from sleep(). if not self.queue_event_arr[1].memory_exception_data.gpu_id and not self.queue_event_arr[2].hw_exception_data.gpu_id: with contextlib.suppress(RuntimeError): self.sleep(tm=1) report = [] if self.queue_event_arr[1].memory_exception_data.gpu_id: report += [f"MMU fault: 0x{self.queue_event_arr[1].memory_exception_data.va:X} | {_str(self.queue_event_arr[1].memory_exception_data.failure)}"] if self.queue_event_arr[2].hw_exception_data.gpu_id: report += [f"HW fault: {_str(self.queue_event_arr[2].hw_exception_data)}"] raise RuntimeError("\n".join(report)) def require_profile_mode(self, can_set_mode=True): if self.dev.target[0] == 9: return fn = f'{self.dev_sysfs_path}/power_dpm_force_performance_level' if (perflevel:=FileIOInterface(fn).read().strip()) != 'profile_standard': if can_set_mode: atexit.register(lambda: os.system(f"echo '{perflevel}' | sudo tee {fn} > /dev/null")) os.system(f"echo 'profile_standard' | sudo tee {fn} > /dev/null") self.require_profile_mode(can_set_mode=False) else: raise RuntimeError("PMC/SQTT requires stable power state: run `amd-smi set -l stable_std` for KFD iface") @functools.cached_property def drm_dev_info(self) -> amdgpu_drm.struct_drm_amdgpu_info_device: amdgpu_drm.DRM_IOCTL_AMDGPU_INFO(self.drm_fd, query=amdgpu_drm.AMDGPU_INFO_DEV_INFO, return_pointer=ctypes.addressof(inf:=amdgpu_drm.struct_drm_amdgpu_info_device()), return_size=ctypes.sizeof(inf)) return inf def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return ((self.drm_dev_info.cu_bitmap[se % 4][sa + (se // 4) * 2] >> (2 * wgp)) & 0x3) == 0x3 class PCIIface(PCIIfaceBase): def __init__(self, dev, dev_id): pci_ids = {"gfx1100": (0X744C,), "gfx1102": (0x7480,), "gfx1201": (0x7550, 0x7551), "gfx1200": (0x7590,), "gfx942": (0x74A1, 0x74B5), "gfx950": (0x75A0, 0x75A8, 0x75B0, 0x75B3)} devs = ((0xffff, pci_ids[arch.split(',')[0]] if (arch:=DEV.target('AMD').arch) else tuple(i for v in pci_ids.values() for i in v)),) super().__init__(dev, dev_id, vendor=0x1002, devices=devs, vram_bar=0, va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev) self._compute_props() def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]: return ([(self.dev_impl.paddr2xgmi(p), sz) for p, sz in paddrs], AddrSpace.PEER) if self.dev_impl.is_hive() else super().p2p_paddrs(paddrs) def require_profile_mode(self): return True def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics. def _compute_props(self): self.ip_versions = self.dev_impl.ip_ver gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}") if self.dev_impl.gc_info.header.version_major == 2: cu_per_sa = self.dev_impl.gc_info.gc_num_cu_per_sh max_sh_per_se = self.dev_impl.gc_info.gc_num_sh_per_se else: cu_per_sa = 2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa) max_sh_per_se = self.dev_impl.gc_info.gc_num_sa_per_se array_count = max_sh_per_se * self.dev_impl.gc_info.gc_num_se * self.dev_impl.gfx.xccs self.props = {'cu_per_simd_array': cu_per_sa, 'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count, 'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd, 'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs, 'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)} def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0, idx=0): assert cwsr_buffer is None, "no cwsr buffer for am" rcvr_params: tuple if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA: doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr, gart._buf+wptr, idx))) else: doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr, gart._buf+wptr, eop_buffer._buf, eop_buffer.nbytes, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql))) put_value = Buffer(host:=self.dev.host, 8, initial_value=bytes(8)) doorbell = Buffer(host, 8, options=BufferSpec(external_ptr=self.dev_impl.doorbell64.addr + doorbell_index*8), preallocate=True) return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(8, rptr).ensure_allocated(), write_ptr=gart.view(8, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, params=rcvr_params) def _collect_interrupts(self, reset=False, drain_only=False): d = self.dev if drain_only: d.iface.dev_impl.ih.drain() else: d.iface.dev_impl.ih.interrupt_handler() if reset and d.iface.dev_impl.recover(force=True): for q, ip in [(d.compute_queue, d.iface.dev_impl.gfx), *[(q, d.iface.dev_impl.sdma) for q in d.sdma_queues.values()]]: for b in (q.put_value, q.read_ptr, q.write_ptr): b.host.view(fmt='Q')[0] = 0 ip.setup_ring(*q.params) for tl in (x.timeline.host.view(fmt='Q') for x in (d, Device[d.host])): tl[0] = tl[1] # the host waits on this device too d.error_state.host.view(fmt='q')[0] = 0 if d.is_usb: usb_reset(d) def sleep(self, timeout): if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))): self.pci_dev.irq_fd.read(8 * events_cnt) self._collect_interrupts() if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state") def on_device_hang(self): self._collect_interrupts(reset=self.dev.can_recover) raise RuntimeError("Device hang detected") def device_fini(self): self.dev_impl.fini() class USBAllocator(AMDAllocator): # the host program reads another device's memory in place: its bytes are the mapping def map(self, buf:Buffer) -> BufferStorage: return BufferStorage(buf.host.addr, buf.host.mv) def _unmap(self, mapping:BufferStorage): pass class USBIface(PCIIface): def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called if dev_id >= len(visible:=filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")): raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)") self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible) self.dev_impl = AMDev(self.pci_dev) self._compute_props() @functools.cached_property def ctrl(self) -> Buffer: # the controller's memory the queue and the host share, one range (usb.py slices it): the sys page at 0, the cq page at 0x1000, the sram at # 0x5000. the host's view starts at the sys page's controller address 0xa000, which puts the sram on its scsi window 0xf000 vaddr, pieces = self.dev_impl.mm.alloc_vaddr(size=0x85000), [(0x0, 0x820000, 0x1000), (0x1000, 0x822000, 0x1000), (0x5000, 0x200000, 0x80000)] for off, paddr, n in pieces: self.dev_impl.mm.map_range(vaddr + off, n, [(paddr, n)], aspace=AddrSpace.SYS, uncached=True) view = self.pci_dev.dma_view(0xa000, 0x85000) for off, n in ((0x800, 4), (0x5000, 0x80000)): view.view(off, n)[:] = bytes(n) # no stale fence or sentinel return Buffer(self.dev.device, 0x85000, options=BufferSpec(external_ptr=vaddr), opaque=BufferStorage(vaddr, host=view)) def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> BufferStorage: # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs) def sleep(self, timeout): pass def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {}) class AMDDevice(Compiled): timestamp_divider = 100.0 # AMD GPU clock: ticks/us sleep_timeout_ms = 200 max_scratch_psize = 0 ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)] def device_props(self): return self.iface.props def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,)) def __init__(self, device:str=""): self.iface = self._select_iface(device) self.is_usb = isinstance(self.iface, USBIface) self.is_vf = self.is_am() and self.iface.dev_impl.is_vf self.can_recover = self.is_am() and not self.is_vf self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100) self.arch = "gfx%d%x%x" % self.target assert (self.target in ((9,4,2),(9,5,0))) or self.target[0] in (11, 12), f"Unsupported arch: {self.arch}" if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}") self.xccs = self.iface.props.get('num_xcc', 1) self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] // self.xccs self.cu_cnt = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.xccs self.waves_per_cu = self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu'] self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target[0] != 9 else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512) self.ip_off = importlib.import_module(f"tinygrad.runtime.autogen.am.{'vega' if self.target[0] == 9 else 'navi'}_offsets") self.soc = import_soc(self.target) self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'soc15' if self.target[0] == 9 else 'nv'}") self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0))) self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP], bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)}) self.nbio = AMDIP('nbio' if self.target[0] < 12 else 'nbif', self.iface.ip_versions[am.NBIF_HWIP], bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)}) pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1' self.hdp_flush = getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ'), getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE') self.is_aql = getenv("AMD_AQL", int(self.xccs > 1)) self.max_copy_size = 0x40000000 if (4, 4, 2) <= (v:=self.iface.ip_versions[am.SDMA0_HWIP]) < (5, 0, 0) or v >= (5, 2, 0) else 0x400000 self.sdma_queues:dict = {} allocator = USBAllocator(self) if self.is_usb else AMDAllocator(self) super().__init__(device, allocator, [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, arch=self.arch) # Scratch setup self.max_private_segment_size = 0 # the ring and its pointers of every queue it can make, made on first use queues = {"COMPUTE:0": lambda: self.compute_queue} | {f"COPY:{i}": functools.partial(self.sdma_queue, i) for i in range(8)} names = ("ring", "write_ptr", "doorbell", "put_value") Compiled.link_bufs |= {to_name(self.device, n, q): lambda b, n=n, f=f: getattr(f(), n) for q, f in queues.items() for n in names} Compiled.link_bufs |= {to_name(self.device, "scratch"): lambda b: self.scratch_buffer(b.max_numel())} if self.is_usb: setup_usb_rules(self) # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them) self.pmc_enabled, self.sqtt_enabled, self.prof_enabled = PROFILE > 0 and PMC > 0, PROFILE > 0 and SQTT > 0, PROFILE > 0 and (PMC > 0 or SQTT > 0) if self.prof_enabled: self.iface.require_profile_mode() self.prof_slots, self.prof_read = getenv("PROF_SLOTS", 32), 0 self.pmc_sched:list[PMCSample] = [] self.sqtt_ses, self.sqtt_win = self.se_cnt * self.xccs, (getenv("SQTT_BUFFER_SIZE", 256) << 20) // self.prof_slots # mb, per shader engine Compiled.link_bufs |= {to_name(self.device, n): lambda b, n=n: getattr(self, n) for n in ("prof_log", "pmc_buf", "sqtt_buf", "sqtt_wptrs")} if self.pmc_enabled: self.pmc_counters = import_pmc(self.target) # validate counters: SQ for SIMD busy/instruction counts, LDS stats, GRBM for GPU cycles, L2 cache hits/misses l2, lds = ("TCC", "SQ") if self.target[0] == 9 else ("GL2C", "SQC") pmc_default = f"SQ_BUSY_CYCLES,SQ_INSTS_VALU,SQ_INSTS_SALU,{lds}_LDS_IDX_ACTIVE,{lds}_LDS_BANK_CONFLICT,GRBM_GUI_ACTIVE,{l2}_HIT,{l2}_MISS" self.pmc_names = getenv("PMC_COUNTERS", pmc_default).split(",") for k in self.pmc_names: if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}") # the pf resets a vf that holds its init access for 20s: set up every queue now and hand it back if self.is_vf: self.compute_queue # the property builds it for idx in range(min(self.iface.count, 8)): self.sdma_queue(idx) self.iface.dev_impl.release_vf_access() def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0): spec = BufferSpec(host=True, uncached=True, cpu_access=True) ring = Buffer(self.device, ring_size, options=spec, preallocate=True, allocator=self.allocator) gart = Buffer(self.device, 0x100, options=spec, initial_value=bytes(0x100), allocator=self.allocator) if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL: self.aql_gart = gart self.aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING, read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset, max_cu_id=(self.cu_cnt * self.xccs) - 1, max_wave_id=self.waves_per_cu - 1) if hasattr(self, 'scratch'): self.aql_scratch() else: self.aql_gart.host.view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc) cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.xccs, mmap.PAGESIZE) cwsr_buffer = Buffer(self.device, cwsr_buffer_size, preallocate=True, allocator=self.allocator) if ctx_save_restore_size else None eop_buffer = Buffer(self.device, eop_buffer_size, preallocate=True, allocator=self.allocator) if eop_buffer_size else None return self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset, wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer, ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx) @functools.cached_property def compute_queue(self) -> AMDQueueDesc: # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391 sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000 lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000 vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000 wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE) ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE) return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE, (1 << 20) if self.is_usb else (16 << 20), eop_buffer_size=0x1000, ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, debug_memory_size=round_up(self.wave_cnt * 32, 64)) @functools.cached_property def has_copy_queue(self) -> bool: self.has_copy_queue = False # queue setup can allocate buffers that check this property return self.sdma_queue(0) is not None def sdma_queue(self, idx:int): if getenv("AMD_DISABLE_SDMA"): return None if idx in self.sdma_queues: return self.sdma_queues[idx] with contextlib.suppress(OSError): self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 16 << 20, idx=idx) return self.sdma_queues.get(idx, None) def tmpring_size(self, private_segment_size): private_segment_size = max(private_segment_size, 128) lanes_per_wave = 64 # wave64 mem_alignment_size = 256 if self.target[0] != 9 else 1024 size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave) size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt # NOTE: xcc logic is correct only for GFX9. max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size) num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1) tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields') return int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little') def scratch_buffer(self, private_segment_size): AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize) if self.max_private_segment_size < private_segment_size: lanes_per_wave = 64 # wave64 mem_alignment_size = 256 if self.target[0] != 9 else 1024 size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave) size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt self.scratch = Buffer(self.device, size_per_xcc * self.xccs, options=BufferSpec(nolru=True), preallocate=True) self.max_private_segment_size = private_segment_size if hasattr(self, 'aql_desc'): self.aql_scratch() return self.scratch def aql_scratch(self): gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3} rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1, 'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})} rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields') rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields') base = self.scratch._buf self.aql_desc.scratch_backing_memory_location = base self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size self.aql_desc.scratch_resource_descriptor[:] = [lo32(base), int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(base), SWIZZLE_ENABLE=1), 'little'), lo32(self.scratch.nbytes // self.xccs), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')] self.aql_desc.compute_tmpring_size = self.tmpring_size(self.max_private_segment_size) self.aql_gart.host.view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc) def _prof_buffer(self, size:int, dtype, host:bool=True) -> Buffer: buf = Buffer(self.device, size * dtype.itemsize, options=BufferSpec(host=host, nolru=True, uncached=host, cpu_access=True), preallocate=True) buf.host.view(fmt='B')[:buf.nbytes] = bytes(buf.nbytes) return buf @functools.cached_property def prof_log(self) -> Buffer: return self._prof_buffer(1 + self.prof_slots, dtypes.uint64) @property def pmc_size(self) -> int: return self.pmc_sched[-1].off + self.pmc_sched[-1].size @functools.cached_property def pmc_buf(self) -> Buffer: return self._prof_buffer(self.pmc_size * self.prof_slots, dtypes.uint8) @functools.cached_property def sqtt_buf(self) -> Buffer: return self._prof_buffer(self.sqtt_win * self.prof_slots * self.sqtt_ses, dtypes.uint8, host=False) @functools.cached_property def sqtt_wptrs(self) -> Buffer: return self._prof_buffer(self.prof_slots * self.sqtt_ses, dtypes.uint32) def sqtt_trace(self, slot:int, se:int) -> bytes: off = (se * self.prof_slots + slot) * self.sqtt_win wptr = (self.sqtt_wptrs.host.view(fmt='I')[slot * self.sqtt_ses + se] & 0x1FFFFFFF) * 32 if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf + off) // 32) & 0x1FFFFFFF) * 32 assert 0 <= wptr <= self.sqtt_win, f"{wptr} > {self.sqtt_win}, should never happen" if wptr >= self.sqtt_win - 32: # the wptr stops at the last dword when the window overflows print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {se})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow")) blob = bytes(self.sqtt_buf.host.view(fmt='B')[off:off + wptr]) return (struct.pack(' 0: print(colored(f"{self.device}: Warning: {lost} kernel profiles were overwritten: synchronize more often or raise PROF_SLOTS", "yellow")) for k in range(max(self.prof_read, log[0] - self.prof_slots), log[0]): slot, tag = k % self.prof_slots, log[1 + k % self.prof_slots] if self.pmc_enabled: blob = bytes(self.pmc_buf.host.view(fmt='B')[slot * self.pmc_size:(slot + 1) * self.pmc_size]) Compiled.profile_events.append(ProfilePMCEvent(self.device, tag, self.pmc_sched, blob, k)) for se in range(self.sqtt_ses if self.sqtt_enabled else 0): itrace = bool((SQTT_ITRACE_SE_MASK.value >> se) & 1) Compiled.profile_events.append(ProfileSQTTEvent(self.device, tag, se, self.sqtt_trace(slot, se), itrace, k)) self.prof_read = log[0] super().collect_prof() def on_device_hang(self): self.iface.on_device_hang() Compiled.pm_encode += PatternMatcher([ (UPat.sink(arg=KernelInfo("amd_compute_submit")).call(UPat(name="q")).after(allow_any_len=True, name="s"), lambda s, q: (AMDComputeAQLQueue if cast(AMDDevice, Device[q.device]).is_aql else AMDComputeQueue)(s).encode()), (UPat.sink(arg=KernelInfo("amd_copy_submit")).call(UPat()).after(allow_any_len=True, name="s"), lambda s: AMDSDMAQueue(s).encode()), ])