From 1ad6598963dd91db332f68ffa67ab276af500994 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 21 Oct 2025 20:52:24 +0800 Subject: [PATCH] amd: trace all instructions (#12831) --- extra/sqtt/roc.py | 10 +++++++--- tinygrad/runtime/ops_amd.py | 11 ++++++----- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/extra/sqtt/roc.py b/extra/sqtt/roc.py index 9044904e4c..221a3ecb45 100644 --- a/extra/sqtt/roc.py +++ b/extra/sqtt/roc.py @@ -27,14 +27,18 @@ class _ROCParseCtx: self.disasms[prog.base + addr] = info self.addr2prg[prog.base + addr] = prog - def next_sqtt(self): return next(self.sqtt_evs, None) + def next_sqtt(self): + x = next(self.sqtt_evs, None) + self.active_se = x.se if x is not None else None + return x + def find_program(self, addr): return self.addr2prg[addr] def on_occupancy_ev(self, ev): - if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start) + if DEBUG >= 4: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start) def on_wave_ev(self, ev): - if DEBUG >= 4: print("WAVE", ev.wave_id, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time) + if DEBUG >= 4: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time) asm = {} for j in range(ev.instructions_size): diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index ab17f425e2..4924791d0f 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -20,6 +20,7 @@ from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationM from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import +SQTT = getenv("SQTT", 0) EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h WAIT_REG_MEM_FUNCTION_EQ = 3 # == WAIT_REG_MEM_FUNCTION_NEQ = 4 # != @@ -254,9 +255,9 @@ class AMDComputeQueue(HWQueue): if (10,0,0) <= prg.dev.target < (11,0,0): self.wreg(self.gc.mmCP_COHER_START_DELAY, 0x20) self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0) - self.wreg(self.gc.regCOMPUTE_STATIC_THREAD_MGMT_SE0, 0xFFFFFFFF, 0xFFFFFFFF) - self.wreg(self.gc.regCOMPUTE_STATIC_THREAD_MGMT_SE2, 0xFFFFFFFF, 0xFFFFFFFF) - if prg.dev.target >= (11,0,0): self.wreg(self.gc.regCOMPUTE_STATIC_THREAD_MGMT_SE4, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF) + for i in range(8 if prg.dev.target >= (11,0,0) else 4): + semask = (prg.dev.sqtt_itrace_se_mask >> i) & 0b1 if prg.dev.sqtt_enabled and SQTT >= 2 else 0xffffffff + self.wreg(getattr(self.gc, f'regCOMPUTE_STATIC_THREAD_MGMT_SE{i}'), semask) self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs) self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, 0) @@ -812,7 +813,7 @@ class AMDDevice(HCQCompiled): self._ensure_has_local_memory(128) # set default scratch size to 128 bytes per thread # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them) - self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0)) + self.sqtt_enabled = PROFILE and SQTT > 0 if self.sqtt_enabled: if self.target[0] < 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}') if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000: @@ -821,7 +822,7 @@ class AMDDevice(HCQCompiled): "For more information read https://github.com/tinygrad/tinygrad/blob/master/extra/sqtt/README.md") SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(nolru=True)) for _ in range(self.se_cnt)] - self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", 2) # -1 enable all, 0 disable all, >0 bitmask for where to enable instruction tracing + self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", -1 if SQTT >= 2 else (1 << 1)) # se bitmask: -1 enable all, 0 disable all self.sqtt_next_cmd_id = itertools.count(0) cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)