From 3b3de8df61343b65483af43efa70e9de285cdf43 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Thu, 24 Jul 2025 17:36:19 +0300 Subject: [PATCH] hcq: graphed copies (#11302) * fast copies p2 * upd and fix * graph supports * fixes * fixes * fixes * fix * fix * fix mockgpu * fix alignment * smaller in ci --- test/test_copy_speed.py | 33 ++++++++++++++++++++++++++---- tinygrad/engine/jit.py | 1 + tinygrad/runtime/graph/hcq.py | 33 +++++++++++++++++++++++------- tinygrad/runtime/support/system.py | 2 +- 4 files changed, 57 insertions(+), 12 deletions(-) diff --git a/test/test_copy_speed.py b/test/test_copy_speed.py index 5d0cc69c26..ed2080922d 100644 --- a/test/test_copy_speed.py +++ b/test/test_copy_speed.py @@ -1,10 +1,9 @@ -import unittest -from tinygrad import Tensor -from tinygrad import Device +import unittest, numpy as np +from tinygrad import Tensor, Device, TinyJit from tinygrad.helpers import Timing, CI, OSX import multiprocessing.shared_memory as shared_memory -N = 4096 +N = 256 if CI else 4096 class TestCopySpeed(unittest.TestCase): @classmethod def setUpClass(cls): Device[Device.DEFAULT].synchronize() @@ -49,6 +48,32 @@ class TestCopySpeed(unittest.TestCase): with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"): t.to('CPU').realize() + def testCopyDefaulttoCPUJit(self): + if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op") + + @TinyJit + def _do_copy(t): return t.to('CPU').realize() + + t = Tensor.randn(N, N, 4).contiguous().realize() + for _ in range(5): + with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"): + x = _do_copy(t) + Device[Device.DEFAULT].synchronize() + np.testing.assert_equal(t.numpy(), x.numpy()) + + def testCopytoCPUtoDefaultJit(self): + if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op") + + @TinyJit + def _do_copy(x): return t.to(Device.DEFAULT).realize() + + for _ in range(5): + t = Tensor.randn(N, N, 4, device="CPU").contiguous().realize() + with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"): + x = _do_copy(t) + Device[Device.DEFAULT].synchronize() + np.testing.assert_equal(t.numpy(), x.numpy()) + @unittest.skipIf(CI, "CI doesn't have 6 GPUs") @unittest.skipIf(Device.DEFAULT != "GPU", "only test this on GPU") def testCopyCPUto6GPUs(self): diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index e9ab08bf6e..7166b2c578 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -44,6 +44,7 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer] match ji.prg: case CompiledRunner(): ji_graph_dev = ji.prg.dev case BufferXfer(): ji_graph_dev = Device[unwrap(ji.bufs[0]).device] + case BufferCopy(): ji_graph_dev = next((Device[unwrap(b).device] for b in ji.bufs if unwrap(b).device not in {"CPU", "LLVM"}), None) case ViewOp(): continue # ViewOps are just ignored case _: ji_graph_dev = None # Everything else is not graphed and flushes existing graph if it's being constructed diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index 614b1174c3..7f1c3ba5b4 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -1,11 +1,11 @@ import collections, time from typing import Any, cast -from tinygrad.helpers import round_up, PROFILE, merge_dicts +from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Variable -from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner +from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner, BufferCopy from tinygrad.engine.jit import MultiGraphRunner class HCQGraph(MultiGraphRunner): @@ -13,6 +13,9 @@ class HCQGraph(MultiGraphRunner): super().__init__(jit_cache, input_rawbuffers, var_vals) self.devices = list(set(cast(HCQCompiled, d) for ji in jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs])) + # CPU Device is always last + self.devices = sorted(self.devices, key=lambda x: 1 if x._is_cpu() else 0) + # Replace input buffers with variables. self.hcq_bufs = [[cast(Buffer, x)._buf for x in ji.bufs] for ji in jit_cache] self.input_replace_to_var: dict[tuple[int, int], Variable] = {} @@ -48,7 +51,8 @@ class HCQGraph(MultiGraphRunner): self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices} self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation - self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices}, **{"KICK": self.devices[0].new_signal(value=0)}} + self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if dev.device != "CPU"}, + **{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev.device == "CPU"}} self.kickoff_value: int = 0 self.kickoff_var = UOp.variable("kickoff_var", 0, 0xffffffff, dtype=dtypes.uint32) @@ -64,10 +68,15 @@ class HCQGraph(MultiGraphRunner): for dev, queue in self.comp_queues.items(): dev_access[queue].add(dev) + self.input_replace_map: dict[HCQCompiled, set[int]] = collections.defaultdict(set) self.fixedvars: dict[HCQCompiled, dict[Variable, int]] = {} for j,ji in enumerate(jit_cache): - enqueue_dev: HCQCompiled = ji.prg.dev if (is_exec_prg:=isinstance(ji.prg, CompiledRunner)) else Device[ji.bufs[1].device] #type:ignore + if is_exec_prg:=isinstance(ji.prg, CompiledRunner): enqueue_dev: HCQCompiled = ji.prg.dev + else: + # For copy ops prioritize enqeueuing on the dest device, so reverse the buffers. + for b in cast(list[Buffer], ji.bufs[::-1]): + if (enqueue_dev:=cast(HCQCompiled, Device[b.device])).hw_copy_queue_t is not None: break # set any fixedvars on the device self.fixedvars[enqueue_dev] = merge_dicts([self.fixedvars.get(enqueue_dev, {}), ji.fixedvars]) @@ -148,10 +157,11 @@ class HCQGraph(MultiGraphRunner): # Encode main commands based on ji type. if isinstance(ji.prg, CompiledRunner): enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(ji.prg.p.global_size or (1,1,1)), tuple(ji.prg.p.local_size or (1,1,1))) - elif isinstance(ji.prg, BufferXfer): + elif isinstance(ji.prg, (BufferXfer, BufferCopy)): dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]] - cast(HCQAllocator, Device[src.device].allocator).map(dest._buf) - + for bufid, src in enumerate(cast(list[Buffer], ji.bufs)): + if (inprep_idx:=self.input_replace.get((j, bufid))) is not None: self.input_replace_map[enqueue_dev].add(inprep_idx) + else: cast(HCQAllocator, enqueue_dev.allocator).map(self.hcq_bufs[j][bufid]) enqueue_queue.copy(self.hcq_bufs[j][0].va_addr, self.hcq_bufs[j][1].va_addr, dest.nbytes) self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device])) @@ -177,6 +187,9 @@ class HCQGraph(MultiGraphRunner): for sig in self.queue_signals_to_reset: sig.value = 0 self.signals['KICK'].value = self.kickoff_value + for dev in self.devices: + for idx_to_map in self.input_replace_map[dev]: cast(HCQAllocator, dev.allocator).map(input_rawbuffers[idx_to_map]._buf) + if PROFILE and self.kickoff_value > 1: self.collect_timestamps() hcq_var_vals = {self.kickoff_var: self.kickoff_value, **var_vals, @@ -210,3 +223,9 @@ class HCQGraph(MultiGraphRunner): if PROFILE and self.kickoff_value >= 1: self.collect_timestamps() for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True)) + + @staticmethod + def supports_exec_item(dev, ei:ExecItem) -> bool: + # MOCKGPU is not supported, since it can't execute commands in parallel + copy = (isinstance(ei.prg, BufferCopy) and cast(HCQCompiled, dev).hw_copy_queue_t is not None) and not getenv("MOCKGPU") + return all(issubclass(type(Device[b.device]), HCQCompiled) for b in ei.bufs if b) and (isinstance(ei.prg, (CompiledRunner, BufferXfer)) or copy) diff --git a/tinygrad/runtime/support/system.py b/tinygrad/runtime/support/system.py index 438a054d26..845c18ab61 100644 --- a/tinygrad/runtime/support/system.py +++ b/tinygrad/runtime/support/system.py @@ -167,4 +167,4 @@ class PCIIfaceBase: snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}") - self.dev_impl.mm.map_range(cast(int, b.va_addr), b.size, paddrs, system=True, snooped=snooped, uncached=uncached) + self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, system=True, snooped=snooped, uncached=uncached)