From a8e54df3630c1cbcafa7e7c5ee211a237242e8eb Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 6 Feb 2025 13:35:34 +0800 Subject: [PATCH] benchmark single kernel launch (#8921) * benchmark kernel launch * don't realize unneeded * faster * faster metal * fix mypy * without sync * no div 0 * lru cache that * no sync in the profile --- .../external_benchmark_kernel_launch.py | 38 +++++++++++++++++++ test/test_speed_v_torch.py | 8 +++- tinygrad/engine/jit.py | 2 +- tinygrad/ops.py | 4 +- tinygrad/runtime/ops_metal.py | 10 +++-- tinygrad/shape/shapetracker.py | 1 + 6 files changed, 54 insertions(+), 9 deletions(-) create mode 100644 test/external/external_benchmark_kernel_launch.py diff --git a/test/external/external_benchmark_kernel_launch.py b/test/external/external_benchmark_kernel_launch.py new file mode 100644 index 0000000000..46da20d196 --- /dev/null +++ b/test/external/external_benchmark_kernel_launch.py @@ -0,0 +1,38 @@ +import time +from tinygrad import Tensor, TinyJit, Device, Context +from tinygrad.helpers import Profiling, Timing, GlobalCounters + +# python3 test/test_speed_v_torch.py TestSpeed.test_add_a + +@TinyJit +def plus(a:Tensor, b:Tensor): return a+b + +if __name__ == "__main__": + a = Tensor([1]).realize() + b = Tensor([1]).realize() + for i in range(5): + with Timing(prefix=f"{i}:"): + c = plus(a,b) + Device[c.device].synchronize() + assert c.item() == 2 + for i in range(5): + st = time.perf_counter() + c = plus(a,b) + et = time.perf_counter() - st + print(f"nosync {i}: {et*1e6:.2f} us") + Device[c.device].synchronize() + for i in range(5): + st = time.perf_counter() + c = plus(a,b) + Device[c.device].synchronize() + et = time.perf_counter() - st + print(f"precise {i}: {et*1e6:.2f} us") + assert GlobalCounters.time_sum_s == 0 + with Context(DEBUG=2): + st = time.perf_counter() + c = plus(a,b) + Device[c.device].synchronize() + et = time.perf_counter() - st + print(f"kernel {GlobalCounters.time_sum_s*1e3:.2f} ms / full {et*1e3:.2f} ms -- {et/(GlobalCounters.time_sum_s+1e-12):.2f} x") + with Profiling(): + c = plus(a,b) diff --git a/test/test_speed_v_torch.py b/test/test_speed_v_torch.py index 1cf1a3e044..6025c97c41 100644 --- a/test/test_speed_v_torch.py +++ b/test/test_speed_v_torch.py @@ -202,8 +202,12 @@ class TestSpeed(unittest.TestCase): def f(a, b): return (a*b).sum() helper_test_generic_square('mul_sum', 4096, f, f) - def test_add(self): - for N in [1, 1024, 4096]: + def test_add_a(self): + def f(a, b): return a + b + helper_test_generic_square('add', 1, f, f) + + def test_add_big(self): + for N in [1024, 4096]: def f(a, b): return a + b helper_test_generic_square('add', N, f, f) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index afb9d726a4..5655066019 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -193,7 +193,7 @@ class CapturedJit(Generic[ReturnType]): def _prepare_jit_inputs(args, kwargs): input_tensors: list[tuple[int|str, Tensor]] = [(name,t) for name,t in list(enumerate(args))+sorted(kwargs.items()) if t.__class__ is Tensor] names, tensors = [name for name,_ in input_tensors], [t for _,t in input_tensors] - if tensors: Tensor.realize(*tensors) + if len(unrealized_tensors := [x for x in tensors if not x.lazydata.is_realized]): Tensor.realize(*unrealized_tensors) # TODO: should we be unpacking multi here? lbs: list[UOp] = flatten([t.lazydata.src if t.lazydata.op is Ops.MULTI else [t.lazydata] for t in tensors]) input_buffers: list[Buffer] = [lb.base.realized for lb in lbs if lb.base.realized is not None] diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 4fe73dce59..6fbcced63e 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -512,8 +512,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @property def base(self) -> UOp: - if self.op in GroupOp.Movement: return self.src[0].base - return self.src[0].base if self.op is Ops.VIEW and len(self.src) == 1 else self + if (self.op is Ops.VIEW and len(self.src) == 1) or self.op in GroupOp.Movement: return self.src[0].base + return self def view(self, new_st:ShapeTracker) -> UOp: return UOp(Ops.VIEW, self.dtype, (self.base,), new_st) def _mop(self, op:Ops, arg): diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index 020cafd546..c0ebf2ab15 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -42,6 +42,7 @@ def msg(ptr: objc_id, selector: str, /, *args: Any, restype: type[T] = objc_id) sender.restype = restype return sender(ptr, sel(selector), *args) +@functools.lru_cache(None) def to_ns_str(s: str): return msg(libobjc.objc_getClass(b"NSString"), "stringWithUTF8String:", s.encode(), restype=objc_instance) def from_ns_str(s): return bytes(msg(s, "UTF8String", restype=ctypes.c_char_p)).decode() @@ -146,13 +147,14 @@ class MetalProgram: self.pipeline_state = msg(self.dev.sysdevice, "newComputePipelineStateWithDescriptor:options:reflection:error:", descriptor, MTLPipelineOption.MTLPipelineOptionNone, None, ctypes.byref(error_pipeline_creation:=objc_instance()), restype=objc_instance) error_check(error_pipeline_creation) + # cache these msg calls + self.max_total_threads: int = cast(int, msg(self.pipeline_state, "maxTotalThreadsPerThreadgroup", restype=ctypes.c_ulong)) def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False): - max_total_threads = msg(self.pipeline_state, "maxTotalThreadsPerThreadgroup", restype=ctypes.c_ulong) - if prod(local_size) > cast(int, max_total_threads): + if prod(local_size) > self.max_total_threads: exec_width = msg(self.pipeline_state, "threadExecutionWidth", restype=ctypes.c_ulong) memory_length = msg(self.pipeline_state, "staticThreadgroupMemoryLength", restype=ctypes.c_ulong) - raise RuntimeError(f"local size {local_size} bigger than {max_total_threads} with exec width {exec_width} memory length {memory_length}") + raise RuntimeError(f"local size {local_size} bigger than {self.max_total_threads} with exec width {exec_width} memory length {memory_length}") command_buffer = msg(self.dev.mtl_queue, "commandBuffer", restype=objc_instance) encoder = msg(command_buffer, "computeCommandEncoder", restype=objc_instance) msg(encoder, "setComputePipelineState:", self.pipeline_state) @@ -160,7 +162,7 @@ class MetalProgram: for i,a in enumerate(vals, start=len(bufs)): msg(encoder, "setBytes:length:atIndex:", bytes(ctypes.c_int(a)), 4, i) msg(encoder, "dispatchThreadgroups:threadsPerThreadgroup:", to_struct(*global_size), to_struct(*local_size)) msg(encoder, "endEncoding") - msg(command_buffer, "setLabel:", to_ns_str(self.name)) + msg(command_buffer, "setLabel:", to_ns_str(self.name)) # TODO: is this always needed? msg(command_buffer, "commit") self.dev.mtl_buffers_in_flight.append(command_buffer) if wait: diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 2569c22968..cf347acbc1 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -109,6 +109,7 @@ class ShapeTracker: def unbind(self) -> tuple[ShapeTracker, dict[Variable, int]]: unbound_views, var_vals = zip(*[v.unbind() for v in self.views]) + if all(len(x) == 0 for x in var_vals): return self, {} return ShapeTracker(tuple(unbound_views)), merge_dicts(var_vals) def real_strides(self, ignore_valid=False) -> tuple[Optional[sint], ...]: return views_to_real_strides(self.views, ignore_valid)