From 4ef6b46b34fea79a5eb47d0a6e272ff5e0eee7ae Mon Sep 17 00:00:00 2001 From: Alisher Zhubanyshev Date: Sun, 13 Jul 2025 21:25:50 +0500 Subject: [PATCH] hcq: reduce launch overhead (#11193) * nv: improve mmio creation speed * add memoryview test * fix indents * move mv bench to `test_helpers`, remove comparison --- test/unit/test_helpers.py | 28 ++++++++++++++++++++++++++-- tinygrad/helpers.py | 2 +- tinygrad/runtime/support/hcq.py | 3 ++- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/test/unit/test_helpers.py b/test/unit/test_helpers.py index 70c965862e..01a2b952d5 100644 --- a/test/unit/test_helpers.py +++ b/test/unit/test_helpers.py @@ -1,6 +1,6 @@ -import ctypes, gzip, unittest +import ctypes, gzip, unittest, timeit from tinygrad import Variable -from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray +from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits from tinygrad.tensor import Tensor, get_shape from tinygrad.shape.view import get_contraction, get_contraction_with_reduce @@ -186,6 +186,30 @@ class TestMemoryview(unittest.TestCase): mv[0] = 2 assert base[0] == 2 + @unittest.skipIf(CI, "dangerous for CI, it allocates tons of memory") + def test_to_mv(self): + sizes = [ + (16, "16 B"), + (64, "64 B"), + (256, "256 B"), + (1024, "1 KB"), + (4 * 1024, "4 KB"), + (16 * 1024, "16 KB"), + (64 * 1024, "64 KB"), + (256 * 1024, "256 KB"), + (1 * 1024 * 1024, "1 MB"), + (10 * 1024 * 1024, "10 MB"), + (200 * 1024 * 1024, "200 MB"), + ] + + for sz, label in sizes: + buf = np.random.randint(0, 256, sz, dtype=np.uint8) + ptr = buf.ctypes.data + + iters = 100_000 + t_us = timeit.timeit(lambda: to_mv(ptr, sz), number=iters) * 1e6 / iters + print(f"Size {label:>9} | Time: {t_us:8.3f} µs") + class TestGetContraction(unittest.TestCase): def test_contraction_with_reduce(self): r = get_contraction((16, 1, 1, 1), (16, 1, 1)) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index d0ffd3f803..ff2f5b7621 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -327,7 +327,7 @@ def wait_cond(cb, value=True, timeout_ms=10000, msg="") -> bool: # TODO: make this work with read only memoryviews (if possible) def from_mv(mv:memoryview, to_type:type[ctypes._SimpleCData]=ctypes.c_char) -> ctypes.Array: return ctypes.cast(ctypes.addressof(to_type.from_buffer(mv)), ctypes.POINTER(to_type * len(mv))).contents -def to_mv(ptr:int, sz:int) -> memoryview: return memoryview(ctypes.cast(ptr, ctypes.POINTER(ctypes.c_uint8 * sz)).contents).cast("B") +def to_mv(ptr:int, sz:int) -> memoryview: return memoryview((ctypes.c_uint8 * sz).from_address(ptr)).cast("B") def mv_address(mv): return ctypes.addressof(ctypes.c_char.from_buffer(mv)) def to_char_p_p(options: list[bytes], to_type=ctypes.c_char): return (ctypes.POINTER(to_type) * len(options))(*[ctypes.cast(ctypes.create_string_buffer(o), ctypes.POINTER(to_type)) for o in options]) diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 872f3511aa..15e0acf3bd 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -320,7 +320,8 @@ class HCQProgram(Generic[HCQDeviceType]): Returns: Arguments state with the given buffers and values set for the program. """ - argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size)) + argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size), + size=self.kernargs_alloc_size) return self.args_state_t(argsbuf, self, bufs, vals=vals) def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),