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

@@ -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))

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])

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),