tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-30 23:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

nv: fix copies larger than 4g (#10117)

Browse Source
This commit is contained in:
nimlgen
2025-04-30 18:43:17 +03:00
committed by GitHub
parent d81acbeef6
commit 2ec3b722e2
2 changed files with 24 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_hcq.py
+20 -1
View File
@@ -1,6 +1,6 @@
import unittest, ctypes, struct, os
from tinygrad import Device, Tensor, dtypes
from tinygrad.helpers import getenv
from tinygrad.helpers import getenv, CI
from tinygrad.device import Buffer, BufferSpec
from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
from tinygrad.engine.realize import get_runner, CompiledRunner
@@ -218,6 +218,25 @@ class TestHCQ(unittest.TestCase):
mv_buf1 = buf1.as_buffer().cast('Q')
for i in range(sz//8): assert mv_buf1[i] == 0x0101010101010101, f"offset {i*8} differs, not all copied, got {hex(mv_buf1[i])}"
@unittest.skipIf(CI, "skip in CI")
def test_copy_6g(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
sz = 6 << 30
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 1, sz)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf.va_addr, buf2._buf.va_addr, sz) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
mv_buf1 = buf1.as_buffer().cast('Q')
for i in range(0, sz//8, (2 << 20)): assert mv_buf1[i] == 0x0101010101010101, f"offset {i*8} differs, not all copied, got {hex(mv_buf1[i])}"
def test_update_copy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
tinygrad/runtime/ops_nv.py
+4 -3
View File
@@ -176,9 +176,10 @@ class NVComputeQueue(NVCommandQueue):
class NVCopyQueue(NVCommandQueue):
def copy(self, dest:sint, src:sint, copy_size:int):
self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src), *data64(dest))
self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, copy_size)
self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, 0x182) # TRANSFER_TYPE_NON_PIPELINED | DST_MEMORY_LAYOUT_PITCH | SRC_MEMORY_LAYOUT_PITCH
for off in range(0, copy_size, 0xffffffff):
self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src+off), *data64(dest+off))
self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, min(copy_size-off, 0xffffffff))
self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, 0x182) # TRANSFER_TYPE_NON_PIPELINED | DST_MEMORY_LAYOUT_PITCH | SRC_MEMORY_LAYOUT_PITCH
return self
def signal(self, signal:NVSignal, value:sint=0):