diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 45dfa56491..7bd3063e41 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -527,6 +527,8 @@ jobs: TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0 - name: Run HCQ2 JIT tests run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py + - name: Run HCQ2 unit tests + run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py testmockam: name: Linux (am) diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index 0c8e99d693..fbdef9f444 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -288,10 +288,10 @@ def amd_build_program(prg:UOp) -> UOp: class AMDAllocator(HCQAllocator['AMDDevice']): def __init__(self, dev:AMDDevice): - super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb()) + super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb()) def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer: - return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue) + 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 _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque) @@ -581,7 +581,7 @@ class AMDDevice(HCQ2Compiled): self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000 self.sdma_queues:dict = {} - self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this + self.has_copy_queue = not getenv("AMD_DISABLE_SDMA") super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch) diff --git a/test/device/test_hcq2.py b/test/device/test_hcq2.py new file mode 100644 index 0000000000..19d72d7bb8 --- /dev/null +++ b/test/device/test_hcq2.py @@ -0,0 +1,14 @@ +import unittest, numpy as np +from unittest.mock import patch +from tinygrad import Device, Tensor +from tinygrad.helpers import getenv +from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in + +@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required") +class TestHCQ2(unittest.TestCase): + def test_copy_without_copy_queue(self): + with patch.object(Device[Device.DEFAULT], "has_copy_queue", False): + np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61)) + +if __name__ == "__main__": + unittest.main() diff --git a/tinygrad/device.py b/tinygrad/device.py index 570cd79813..6f796e07a6 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -336,6 +336,8 @@ class Compiled: pm_lower:Any = None pm_bufferize:Any = None + has_copy_queue:bool = True + def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None): from tinygrad.renderer import Renderer self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer] diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 9e2429e304..9b0372d035 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -96,12 +96,24 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_ def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS) +def hcq_call_devs(call:UOp) -> Any|None: return next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None) + def stage_copy(dst:UOp, src:UOp) -> UOp|None: if not (_need_staging(src, dst) or _need_staging(dst, src)): return None stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8) return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage))) -pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)]) + +def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None: + if (devs:=hcq_call_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None + d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2)) + ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1) + return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src)) + +pm_insert_copy_staging = PatternMatcher([ + (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy), + (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy) +]) # ***************** # 2. deps @@ -217,7 +229,7 @@ def sched_hcq_batches(l:UOp, profile:bool) -> UOp: srcs:list[UOp] = [] batch:list[tuple[UOp, tuple[str, ...]]] = [] for call in l.src: - if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs))) + if (devs:=hcq_call_devs(call)) is not None: batch.append((call, to_tuple(devs))) else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], [] return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))