From f635c4d2736a63fe3b3e7359ff204c9b28c91f67 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 May 2024 19:32:56 -0700 Subject: [PATCH] fix define global (#4383) * fix define global * remove name from DEFINE_GLOBAL * fix fuzzing * fix ptx * fix python --- test/external/fuzz_schedule.py | 4 ++-- test/test_uops.py | 12 ++++++------ tinygrad/codegen/linearizer.py | 2 +- tinygrad/codegen/uops.py | 4 ++-- tinygrad/device.py | 14 ++++++++------ tinygrad/engine/realize.py | 18 +++++++++++------- tinygrad/features/search.py | 21 +++++++++++---------- tinygrad/renderer/assembly.py | 8 ++++---- tinygrad/renderer/cstyle.py | 5 ++--- tinygrad/runtime/ops_python.py | 2 +- 10 files changed, 48 insertions(+), 42 deletions(-) diff --git a/test/external/fuzz_schedule.py b/test/external/fuzz_schedule.py index 86e388399b..eb4eaaef66 100644 --- a/test/external/fuzz_schedule.py +++ b/test/external/fuzz_schedule.py @@ -2,7 +2,7 @@ import itertools import numpy as np from typing import DefaultDict, Dict, List, Set, Tuple, TypeVar from tinygrad.buffer import Buffer -from tinygrad.engine.realize import CustomOp, ExecItem, capturing, lower_schedule_item +from tinygrad.engine.realize import CustomOp, capturing, lower_schedule_item from tinygrad.helpers import DEBUG, MULTIOUTPUT, colored, getenv from tinygrad.lazy import LazyBuffer from tinygrad.engine.schedule import _graph_schedule, _LBScheduleItem @@ -67,7 +67,7 @@ def fuzz_schedule(outs: List[LazyBuffer]): raise e def _exec_si(si: ScheduleItem, seed:int): - ei = ExecItem(lower_schedule_item(si), list(si.bufs)) + ei = lower_schedule_item(si) if len(capturing): capturing[0].add(ei) if isinstance(ei.prg, CustomOp): Tensor._seed = seed ei.run() diff --git a/test/test_uops.py b/test/test_uops.py index 09c0c3724c..e178af8c6f 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -14,7 +14,7 @@ from test.helpers import is_dtype_supported def _uops_to_prg(uops): src = Device[Device.DEFAULT].compiler.render("test", uops) has_local = Device[Device.DEFAULT].compiler.compiler_opts.has_local - return CompiledRunner("test", src, Device.DEFAULT, [1] if has_local else None, [1] if has_local else None) + return CompiledRunner("test", src, Device.DEFAULT, [1] if has_local else None, [1] if has_local else None, uops=uops) def uop(uops:List[UOp], uop:UOps, dtype:Optional[DType], vin:Tuple[UOp, ...], arg:Any=None) -> UOp: uops.append(UOp(uop, dtype, tuple(vin), arg)) @@ -23,8 +23,8 @@ def uop(uops:List[UOp], uop:UOps, dtype:Optional[DType], vin:Tuple[UOp, ...], ar def _test_single_value(vals, op, dts): uops = [] output_dtype = dts[-1] if op is TernaryOps.WHERE else dtypes.bool if op is BinaryOps.CMPLT else dts[0] - buf_store = uop(uops, UOps.DEFINE_GLOBAL, PtrDType(output_dtype), (), (0, 'data0',True)) - buf_loads = [uop(uops, UOps.DEFINE_GLOBAL, PtrDType(dtype), (), (i+1, f'data{i+1}',False)) for i,dtype in enumerate(dts)] + buf_store = uop(uops, UOps.DEFINE_GLOBAL, PtrDType(output_dtype), (), (0, True)) + buf_loads = [uop(uops, UOps.DEFINE_GLOBAL, PtrDType(dtype), (), (i+1, False)) for i,dtype in enumerate(dts)] loads = (uop(uops, UOps.LOAD, dtype, [buf_loads[i], uop(uops, UOps.CONST, dtypes.int32, (), 0)]) for i,dtype in enumerate(dts)) alu = uop(uops, UOps.ALU, output_dtype, loads, op) uop(uops, UOps.STORE, None, (buf_store, uop(uops, UOps.CONST, dtypes.int32, (), 0), alu)) @@ -39,7 +39,7 @@ def _test_single_value(vals, op, dts): def _test_single_value_const(vals, op, dts): uops = [] output_dtype = dts[-1] if op is TernaryOps.WHERE else dtypes.bool if op is BinaryOps.CMPLT else dts[0] - buf_store = uop(uops, UOps.DEFINE_GLOBAL, PtrDType(output_dtype), (), (0, 'data0',True)) + buf_store = uop(uops, UOps.DEFINE_GLOBAL, PtrDType(output_dtype), (), (0, True)) loads = (uop(uops, UOps.CONST, dtype, [], a) for a,dtype in zip(vals, dts)) alu = uop(uops, UOps.ALU, output_dtype, loads, op) uop(uops, UOps.STORE, None, (buf_store, uop(uops, UOps.CONST, dtypes.int32, (), 0), alu)) @@ -52,7 +52,7 @@ def _test_single_value_const(vals, op, dts): def _test_uops_result(output_dtype, uops, res): # uops = [] - buf_store = uop(uops, UOps.DEFINE_GLOBAL, PtrDType(output_dtype), (), (0, 'data0',True)) + buf_store = uop(uops, UOps.DEFINE_GLOBAL, PtrDType(output_dtype), (), (0, True)) # res = output_fn(uops) uop(uops, UOps.STORE, None, (buf_store, uop(uops, UOps.CONST, dtypes.int32, (), 0), res)) buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate() @@ -238,7 +238,7 @@ class TestAssembly(unittest.TestCase): def test_pointer_arithmetics_caching(self): from tinygrad.renderer.assembly import ptr_ar uops = UOpGraph() - u1 = uops.add(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), tuple(), (0, 'data0', True)) + u1 = uops.add(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), tuple(), (0, True)) u2 = uops.add(UOps.SPECIAL, dtypes.int, tuple(), (0, 'gidx0', 9)) u3 = uops.add(UOps.CONST, dtypes.int, tuple(), arg=42) u4 = uops.add(UOps.ALU, dtypes.int, (u2, u3), BinaryOps.MUL) diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index ee12878c80..88e73d5d77 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -349,7 +349,7 @@ class Linearizer(Kernel): if isinstance(buf, MemBuffer): self.buf_uops[i] = self.uops.add(UOps.DEFINE_GLOBAL, buf.dtype if isinstance(buf.dtype, ImageDType) else PtrDType(buf.dtype), (), - (buf.idx, f"data{buf.idx}", any(buf.idx == x.idx for x in self.outbufs))) + (buf.idx, any(buf.idx == x.idx for x in self.outbufs))) # add var vals for i,var in enumerate(self.vars): assert var.expr is not None diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index f5e43ba59c..d3ac933cb9 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -138,6 +138,7 @@ class UOpGraph: def __iter__(self): return iter(self.uops) def vars(self) -> List[Variable]: return [x.arg for x in self.uops if x.uop is UOps.DEFINE_VAR] + def globals(self) -> List[Tuple[int, bool]]: return [x.arg for x in self.uops if x.uop is UOps.DEFINE_GLOBAL] def graph(self): from tinygrad.features.graph import graph_uops @@ -361,8 +362,7 @@ class UOpGraph: self.optimize_loops() # (recursively) remove childless uops - # TODO: remove DEFINE_GLOBAL from here - self.remove_childless(set(x for x in self.uops if x.uop in {UOps.DEFINE_GLOBAL, UOps.STORE})) + self.remove_childless(set(x for x in self.uops if x.uop is UOps.STORE)) # store float4 upcasts directly if possible self.fix_to_store_directly() diff --git a/tinygrad/device.py b/tinygrad/device.py index 9dcdd9d451..faa6a6aeef 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -144,24 +144,26 @@ class Compiler: class CompiledRunner(Runner): def __init__(self, name:str, prg:str, dname:str, global_size:Optional[List[int]]=None, local_size:Optional[List[int]]=None, - variables:Optional[List[Variable]]=None, op_estimate:sint=0, mem_estimate:sint=0, precompiled:Optional[bytes]=None, outcount:int=1): + uops:Optional[UOpGraph]=None, op_estimate:sint=0, mem_estimate:sint=0, precompiled:Optional[bytes]=None): if DEBUG >= 4: print(prg) if global_size is not None: global_size = global_size + [1]*(3-len(global_size)) if local_size is not None: local_size = local_size + [1]*(3-len(local_size)) self.name, self.prg, self.global_size, self.local_size, self.first_run = \ to_function_name(name), prg, global_size, local_size, True lib:bytes = precompiled if precompiled is not None else cast(Compiler, Device[dname].compiler).compile_cached(prg) - self.lib, self.clprg, self.outcount = lib, Device[dname].runtime(self.name, lib), outcount - self.vars: List[Variable] = [] if variables is None else variables + self.uops = uops + self.vars: List[Variable] = [] if uops is None else uops.vars() + self.globals: List[Tuple[int, bool]] = [] if uops is None else uops.globals() + self.lib, self.clprg, self.outcount = lib, Device[dname].runtime(self.name, lib), sum(x[1] for x in self.globals) super().__init__(name, dname, op_estimate, mem_estimate) def to_other_device(self, dname:str): return CompiledRunner(self.display_name, self.prg, dname, self.global_size, self.local_size, - self.vars, self.op_estimate, self.mem_estimate, self.lib, self.outcount) + self.uops, self.op_estimate, self.mem_estimate, self.lib) def __reduce__(self): return self.__class__, (self.display_name, self.prg, self.dname, self.global_size, self.local_size, - self.vars, self.op_estimate, self.mem_estimate, self.lib, self.outcount) + self.uops, self.op_estimate, self.mem_estimate, self.lib) def launch_dims(self, var_vals): global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else self.global_size @@ -195,7 +197,7 @@ class Compiled: run_count = prod((k.global_size if k.global_size else []) + (k.local_size if k.local_size else [])) # NOTE: we use min here to ignore the indexing FLOPS ret = CompiledRunner(k.name, self.compiler.render(to_function_name(k.name), k.uops), self.dname, k.global_size, k.local_size, - k.uops.vars(), min(info.flops, ops * run_count), min(info.mem_estimate, mem * run_count), outcount=len(k.outbufs)) + k.uops, min(info.flops, ops * run_count), min(info.mem_estimate, mem * run_count)) return ret def get_linearizer(self, *ast:LazyOp) -> Linearizer: diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index cc3c7a03b1..b00544a0dd 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -36,23 +36,27 @@ class EmptyOp(Runner): def __init__(self, buf:Buffer): super().__init__(colored(f"empty {buf.size:10d} {buf.dtype}", "yellow"), buf.device) def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False): pass -def lower_schedule_item(si:ScheduleItem) -> Runner: +def lower_runner(runner:Runner, bufs) -> ExecItem: + # TODO: globals isn't on the stupid diskrunner, remove the need for it + return ExecItem(runner, [bufs[x[0]] for x in runner.globals] if hasattr(runner, 'globals') else bufs) + +def lower_schedule_item(si:ScheduleItem) -> ExecItem: assert len(set(x.device for x in si.bufs)) == 1 or si.ast[0].op is LoadOps.COPY - if si.ast[0].op is BufferOps.STORE: return Device[si.outputs[0].device].get_runner(*si.ast) + if si.ast[0].op is BufferOps.STORE: return lower_runner(Device[si.outputs[0].device].get_runner(*si.ast), si.bufs) assert len(si.ast) == 1 and len(si.outputs) == 1, "only ASTRunner supports multioutput" out, ast = si.outputs[0], si.ast[0] if ast.op is LoadOps.COPY: kernel_type = BufferCopy if hasattr(Device[out.device].allocator, 'transfer') and out.device.split(":")[0] == si.inputs[0].device.split(":")[0]: - if getenv("USE_COPY_KERNEL"): return Device[out.device].get_runner(copy_ast(ast.arg)) + if getenv("USE_COPY_KERNEL"): return lower_runner(Device[out.device].get_runner(copy_ast(ast.arg)), si.bufs) kernel_type = BufferXfer - return kernel_type(ast.arg, out.device, si.inputs[0].device) - if ast.op is LoadOps.CUSTOM: return CustomOp(ast.arg) - if ast.op is LoadOps.EMPTY: return EmptyOp(out) + return ExecItem(kernel_type(ast.arg, out.device, si.inputs[0].device), list(si.bufs)) + if ast.op is LoadOps.CUSTOM: return ExecItem(CustomOp(ast.arg), list(si.bufs)) + if ast.op is LoadOps.EMPTY: return ExecItem(EmptyOp(out), list(si.bufs)) raise RuntimeError(f"don't know how to lower {ast}") def lower_schedule(schedule:List[ScheduleItem]) -> Generator[ExecItem, None, None]: - while len(schedule): yield ExecItem(lower_schedule_item(si:=schedule.pop(0)), list(si.bufs)) + while len(schedule): yield lower_schedule_item(schedule.pop(0)) capturing: List = [] # put classes with an add method in here diff --git a/tinygrad/features/search.py b/tinygrad/features/search.py index cfc057c4fb..ff24341934 100644 --- a/tinygrad/features/search.py +++ b/tinygrad/features/search.py @@ -7,8 +7,9 @@ from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, di from tinygrad.dtype import ImageDType from tinygrad.codegen.linearizer import Linearizer from tinygrad.codegen.kernel import Opt, OptOps, KernelOptError +from tinygrad.codegen.uops import UOpGraph from tinygrad.tensor import Tensor -from tinygrad.shape.symbolic import sym_infer, Variable +from tinygrad.shape.symbolic import sym_infer actions = [Opt(op=OptOps.UPCAST, axis=axis, amt=amt) for amt in [0,2,3,4,5,7] for axis in range(6)] actions += [Opt(op=OptOps.UNROLL, axis=axis, amt=amt) for amt in [0,4,7] for axis in range(4)] @@ -30,12 +31,12 @@ def _get_test_global_size(global_size, max_global_size, var_vals): break return test_global_size, factor -def _time_program(variables:List[Variable], outcount:int, rdev:Compiled, lib:bytes, global_size, local_size, var_vals, rawbufs, +def _time_program(uops, rdev:Compiled, lib:bytes, global_size, local_size, var_vals, rawbufs, early_stop=None, max_global_size=65536, clear_l2=False, cnt=3, name="test"): factor = 1 if global_size is not None and max_global_size is not None: global_size, factor = _get_test_global_size(global_size, max_global_size, var_vals) - try: car = CompiledRunner(name, "", rdev.dname, global_size, local_size, variables=variables, precompiled=lib, outcount=outcount) + try: car = CompiledRunner(name, "", rdev.dname, global_size, local_size, uops, precompiled=lib) except AssertionError: return [math.inf] * cnt tms = [] for _ in range(cnt): @@ -46,7 +47,7 @@ def _time_program(variables:List[Variable], outcount:int, rdev:Compiled, lib:byt return tms def _compile_linearizer(compiler:Compiler, lin:Linearizer, name:Optional[str]=None, enforce_max:bool=False) \ - -> Tuple[bytes, Optional[List[int]], Optional[List[int]], List[Variable], int, float, int]: + -> Tuple[bytes, Optional[List[int]], Optional[List[int]], UOpGraph, float]: lin.linearize() if enforce_max and len(lin.uops.uops) >= getenv("BEAM_UOPS_MAX", 3000) > 0: raise RuntimeError("too many uops") src = compiler.render(name if name is not None else to_function_name(lin.name), lin.uops) # NOTE: these all have the same name for deduping @@ -54,7 +55,7 @@ def _compile_linearizer(compiler:Compiler, lin:Linearizer, name:Optional[str]=No st = time.perf_counter() prog = compiler.compile(src) et = time.perf_counter() - st - return prog, lin.global_size, lin.local_size, lin.uops.vars(), len(lin.outbufs), et, len(lin.uops.uops) + return prog, lin.global_size, lin.local_size, lin.uops, et def _try_compile_linearized_w_idx(x:Tuple[int,Linearizer], compiler:Compiler): try: return x[0], _compile_linearizer(compiler, x[1], "test", enforce_max=True) @@ -128,14 +129,14 @@ def beam_search(lin:Linearizer, rawbufs:List[Buffer], amt:int, allow_test_size=T _compile_fn = functools.partial(_try_compile_linearized_w_idx, compiler=dev.compiler) for i,proc in (map(_compile_fn, enumerate(acted_lins)) if beam_pool is None else beam_pool.imap_unordered(_compile_fn, enumerate(acted_lins))): if proc is None: continue - lib, global_size, local_size, vars, outcount, compile_et, num_uops = proc + lib, global_size, local_size, uops, compile_et = proc if lib in seen_libs: continue #print(acted_lins[i].colored_shape(), acted_lins[i].applied_opts) # for debugging BEAMs that segfault seen_libs.add(lib) - try: tms = _time_program(vars, outcount, dev, lib, global_size, local_size, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0) + try: tms = _time_program(uops, dev, lib, global_size, local_size, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0) except RuntimeError: continue # for runtime issues timed_lins.append((acted_lins[i], min(tms))) - if getenv("BEAM_LOG") > 0: print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {num_uops:5d} uops {compile_et*1e6:12.2f} us compile/{timed_lins[-1][1]*1e6:12.2f} us run {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}") # noqa: E501 + if getenv("BEAM_LOG") > 0: print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(uops.uops):5d} uops {compile_et*1e6:12.2f} us compile/{timed_lins[-1][1]*1e6:12.2f} us run {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}") # noqa: E501 elif DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s: {timed_lins[-1][1]*1e6:12.2f} us {len(timed_lins):4d}/{len(acted_lins):4d} {timed_lins[-1][0].colored_shape()}\033[K", end="") # noqa: E501 # done @@ -173,8 +174,8 @@ def time_linearizer(lin:Linearizer, rawbufs:List[Buffer], allow_test_size=True, rawbufs = _ensure_buffer_alloc(rawbufs) var_vals = {k:(k.max+k.min)//2 for k in lin.ast[0].vars()} - lib, global_size, local_size, vars, outcount, _, _ = _compile_linearizer(dev.compiler, lin) - tms = _time_program(vars, outcount, dev, lib, global_size, local_size, var_vals, rawbufs, max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name)) # noqa: E501 + lib, global_size, local_size, uops, _ = _compile_linearizer(dev.compiler, lin) + tms = _time_program(uops, dev, lib, global_size, local_size, var_vals, rawbufs, max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name)) # noqa: E501 if CACHELEVEL >= 2: diskcache_put("time_linearizer", key, tms) return min(tms) diff --git a/tinygrad/renderer/assembly.py b/tinygrad/renderer/assembly.py index 7de83620da..a84d1b8dcf 100644 --- a/tinygrad/renderer/assembly.py +++ b/tinygrad/renderer/assembly.py @@ -94,7 +94,7 @@ def uops_to_asm(lang:AssemblyLanguage, function_name:str, _uops:UOpGraph) -> str matcher.rewrite_graph(uops) for pointer_op in list(filter(lambda uop: uop.uop in [UOps.LOAD, UOps.STORE], uops.uops)): ptr_ar(pointer_op, uops) - uops.remove_childless(set(x for x in uops if x.uop in {UOps.DEFINE_GLOBAL, UOps.PHI, UOps.ENDIF, UOps.ENDLOOP, UOps.STORE})) + uops.remove_childless(set(x for x in uops if x.uop in {UOps.PHI, UOps.ENDIF, UOps.ENDLOOP, UOps.STORE})) uops.optimize_loops() def kk(*s: str): kernel.append("\n".join(s)) @@ -199,11 +199,11 @@ def uops_to_asm(lang:AssemblyLanguage, function_name:str, _uops:UOpGraph) -> str r[u] = f"%{args.expr}" if lang.load_global: kk(*lang.render_load(args.expr, ssa('dat', u, lang.types[dtype]), dtype, ss=".param")) elif uop is UOps.DEFINE_GLOBAL: - bufs.append((args[1], dtype)) - r[u] = f"%{args[1]}" + bufs.append((nm:=f"data{args[0]}", dtype)) + r[u] = f"%{nm}" if lang.load_global: dt = dtypes.ulong if dtype.__class__ == PtrDType else dtype - kk(*lang.render_load(args[1], ssa('dat', u, lang.types[dt]), dt, ss=".param")) + kk(*lang.render_load(nm, ssa('dat', u, lang.types[dt]), dt, ss=".param")) elif uop is UOps.WMMA: wmma = [] for vv in vin[:2]: diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 32f434cb85..3b81da302a 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -161,9 +161,8 @@ def uops_to_cstyle(lang:CStyleLanguage, function_name:str, uops:UOpGraph) -> str bufs.append((args.expr, (dtype,False))) r[u] = args.expr elif uop is UOps.DEFINE_GLOBAL: - assert len(bufs) == args[0], f"missed a global buffer {len(bufs)} {args}" - bufs.append((args[1], (dtype,args[2]))) - r[u] = args[1] + bufs.append((nm:=f"data{args[0]}", (dtype,args[1]))) + r[u] = nm elif uop is UOps.WMMA: kk(f"{lang.render_dtype(dtype)} {ssa('wmma',u)} = __{args[0]}({r[vin[0]]}, {r[vin[1]]}, {r[vin[2]]});") elif uop is UOps.DEFINE_ACC: kk(f"{lang.render_dtype(dtype)} {ssa('acc',u)} = {lang.render_const(args, dtype)};") elif uop is UOps.CONST: r[u] = lang.render_const(args, dtype) if args >= 0 else f"({lang.render_const(args, dtype)})" diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 45b25dcb8b..2d7fe1b22d 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -71,7 +71,7 @@ class PythonProgram: dl[i] = dtype if uop is UOps.DEFINE_GLOBAL: assert dtype.fmt is not None - ul[i] = [pbufs[arg[0]].cast(dtype.fmt)] * warp_size + ul[i] = [pbufs.pop(0).cast(dtype.fmt)] * warp_size elif uop is UOps.DEFINE_LOCAL: assert dtype.fmt is not None lbuf = memoryview(bytearray(arg[1]*dtype.itemsize))