From c631c72f22bb65f69af2aa62002566d5c0386f98 Mon Sep 17 00:00:00 2001 From: uuuvn <83587632+uuuvn@users.noreply.github.com> Date: Sun, 23 Mar 2025 16:30:38 +0500 Subject: [PATCH] HCQ: Increment timeline signal before submitting (#9550) `AMDComputeQueue.__del__` frees `hw_page` which is safe because `AMDAllocator._free` does `self.dev.synchronize()` which is supposed to wait for execution of IB to finish, however that doesn't happen if AMDComputeQueue is dropped right after submit before timeline signal is incremented, which it is in most places leading to a race if .bind() is also used (required for multi-xcc because bug in mec fw treats all PACKET3_PRED_EXECs outside IBs as if they had EXEC_COUNT of zero). --- docs/developer/hcq.md | 3 +-- tinygrad/runtime/graph/hcq.py | 3 +-- tinygrad/runtime/ops_amd.py | 6 ++---- tinygrad/runtime/ops_nv.py | 3 +-- tinygrad/runtime/support/hcq.py | 37 ++++++++++++++------------------- 5 files changed, 21 insertions(+), 31 deletions(-) diff --git a/docs/developer/hcq.md b/docs/developer/hcq.md index a2c725499a..33655ab787 100644 --- a/docs/developer/hcq.md +++ b/docs/developer/hcq.md @@ -115,9 +115,8 @@ HCQ-compatible devices use a global timeline signal for synchronizing all operat ```python HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \ .exec(...) - .signal(your_device.timeline_signal, your_device.timeline_value) \ + .signal(your_device.timeline_signal, your_device.next_timeline()) \ .submit(your_device) -your_device.timeline_value += 1 # Optionally wait for execution your_device.timeline_signal.wait(your_device.timeline_value - 1) diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index 7f879c1474..45d07bd7a9 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -184,8 +184,7 @@ class HCQGraph(MultiGraphRunner): self.comp_queues[dev].submit(dev, hcq_var_vals) if (copy_queue:=self.copy_queues.get(dev, None)) is not None: copy_queue.submit(dev, hcq_var_vals) - self.last_timeline[dev] = (dev.timeline_signal, dev.timeline_value) - dev.timeline_value += 1 + self.last_timeline[dev] = (dev.timeline_signal, dev.next_timeline()) if wait: st = time.perf_counter() diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 2f04450b59..0bcc519955 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -751,8 +751,7 @@ class AMDDevice(HCQCompiled): self.max_private_segment_size = required def invalidate_caches(self): - AMDComputeQueue().memory_barrier().signal(self.timeline_signal, self.timeline_value).submit(self) - self.timeline_value += 1 + AMDComputeQueue().memory_barrier().signal(self.timeline_signal, self.next_timeline()).submit(self) self.synchronize() def on_device_hang(self): self.dev_iface.on_device_hang() @@ -761,8 +760,7 @@ class AMDDevice(HCQCompiled): if self.sqtt_enabled: wptrs_buf = self.allocator.alloc(round_up(len(self.sqtt_buffers), 0x1000), BufferSpec(cpu_access=True, nolru=True)) wptrs = to_mv(wptrs_buf.va_addr, wptrs_buf.size) - AMDComputeQueue().stop_trace(len(self.sqtt_buffers), wptrs_buf).signal(self.timeline_signal, self.timeline_value).submit(self) - self.timeline_value += 1 + AMDComputeQueue().stop_trace(len(self.sqtt_buffers), wptrs_buf).signal(self.timeline_signal, self.next_timeline()).submit(self) self.synchronize() if DEBUG>=2: print('Saving SQTT in profile...') for i,buf0 in enumerate(self.sqtt_buffers): diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index eaf7b58ae2..6e6e19b46a 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -520,8 +520,7 @@ class NVDevice(HCQCompiled[NVSignal]): cast(NVComputeQueue, NVComputeQueue().wait(self.timeline_signal, self.timeline_value - 1)) \ .setup(local_mem=self.shader_local_mem.va_addr, local_mem_tpc_bytes=bytes_per_tpc) \ - .signal(self.timeline_signal, self.timeline_value).submit(self) - self.timeline_value += 1 + .signal(self.timeline_signal, self.next_timeline()).submit(self) def invalidate_caches(self): rmctrl.fb_flush_gpu_cache(self.fd_ctl, self.root, self.subdevice, diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 7af43d4590..2df29711d5 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -261,16 +261,14 @@ def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Type[HWQueue]|None=No if enabled and queue is not None: queue.timestamp(st) elif enabled: assert queue_type is not None - queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(st).signal(dev.timeline_signal, dev.timeline_value).submit(dev) - dev.timeline_value += 1 + queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(st).signal(dev.timeline_signal, dev.next_timeline()).submit(dev) try: yield (st, en) finally: if enabled and queue is not None: queue.timestamp(en) elif enabled: assert queue_type is not None - queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.timeline_value).submit(dev) - dev.timeline_value += 1 + queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev) if enabled and PROFILE: dev.sig_prof_records.append((cast(HCQSignal, st), cast(HCQSignal, en), desc, queue_type is dev.hw_copy_queue_t)) @@ -329,8 +327,7 @@ class HCQProgram(Generic[DeviceType]): with hcq_profile(self.dev, queue=q, desc=self.name, enabled=wait or PROFILE) as (sig_st, sig_en): q.exec(self, kernargs, global_size, local_size) - q.signal(self.dev.timeline_signal, self.dev.timeline_value).submit(self.dev) - self.dev.timeline_value += 1 + q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev) if wait: self.dev.synchronize() return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None @@ -374,6 +371,10 @@ class HCQCompiled(Compiled, Generic[SignalType]): Compiled.profile_events += [ProfileRangeEvent(self.device, name, st.timestamp, en.timestamp, cp) for st,en,name,cp in self.sig_prof_records] self.sig_prof_records = [] + def next_timeline(self): + self.timeline_value += 1 + return self.timeline_value - 1 + @classmethod def _alloc_signal_addr(cls) -> int: if not cls.signal_pool: @@ -384,8 +385,7 @@ class HCQCompiled(Compiled, Generic[SignalType]): def _at_profile_finalize(self): def _sync(d:HCQCompiled, q_t:Type[HWQueue]): - q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.timeline_value).submit(d) - d.timeline_value += 1 + q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.next_timeline()).submit(d) st = time.perf_counter_ns() d.timeline_signal.wait(d.timeline_value - 1) # average of the two et = time.perf_counter_ns() @@ -439,9 +439,8 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]): ctypes.memmove(self.b[self.b_next].va_addr, from_mv(src[i:]), lsize:=min(self.b[self.b_next].size, src.nbytes-i)) 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.timeline_value).submit(self.dev) - self.b_timeline[self.b_next] = self.dev.timeline_value - self.dev.timeline_value += 1 + .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev) + self.b_timeline[self.b_next] = self.dev.timeline_value - 1 def copy_from_disk(self, dest:HCQBuffer, src, size): def _get_temp_buf(): @@ -456,9 +455,8 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]): for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size): self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \ .copy(dest.va_addr + dst_off, batch_info[0] + src_off, copy_size) \ - .signal(self.dev.timeline_signal, self.dev.timeline_value).submit(self.dev) - self.b_timeline[batch_info[1]] = self.dev.timeline_value - self.dev.timeline_value += 1 + .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev) + self.b_timeline[batch_info[1]] = self.dev.timeline_value - 1 def _copyout(self, dest:memoryview, src:HCQBuffer): self.dev.synchronize() @@ -468,9 +466,8 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]): for i in range(0, dest.nbytes, self.b[0].size): self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \ .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.timeline_value).submit(self.dev) - self.dev.timeline_signal.wait(self.dev.timeline_value) - self.dev.timeline_value += 1 + .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) @@ -482,11 +479,9 @@ class HCQAllocator(HCQAllocatorBase, Generic[DeviceType]): src_dev.hw_copy_queue_t().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \ .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \ .copy(dest.va_addr, src.va_addr, sz) \ - .signal(src_dev.timeline_signal, src_dev.timeline_value).submit(src_dev) - src_dev.timeline_value += 1 + .signal(src_dev.timeline_signal, src_dev.next_timeline()).submit(src_dev) if src_dev != dest_dev: dest_dev.hw_compute_queue_t().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \ .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \ - .signal(dest_dev.timeline_signal, dest_dev.timeline_value).submit(dest_dev) - dest_dev.timeline_value += 1 + .signal(dest_dev.timeline_signal, dest_dev.next_timeline()).submit(dest_dev)