diff --git a/examples/hlb_cifar10.py b/examples/hlb_cifar10.py index 9f393bd592..e1e1021522 100644 --- a/examples/hlb_cifar10.py +++ b/examples/hlb_cifar10.py @@ -3,7 +3,7 @@ # tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py # https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/ # https://siboehm.com/articles/22/CUDA-MMM -import random, time +import random, time, math import numpy as np from typing import Optional from extra.lr_scheduler import OneCycleLR @@ -220,13 +220,10 @@ def train_cifar(): return X_cutmix, Y_cutmix def random_permutation(rows:int, cols:int) -> Tensor: - # Alternating modular additions are bijections on the rows x cols domain. - idx = Tensor.arange(rows * cols) - row, col = idx // cols, idx % cols - for _ in range(4): - row = (row + col * col + random.randrange(rows) * col + random.randrange(rows)) % rows - col = (col + row * row + random.randrange(cols) * row + random.randrange(cols)) % cols - return row * cols + col + size = rows * cols + # An affine map is a permutation when its stride is coprime to the domain size. + while math.gcd(stride:=random.randrange(1, size), size) != 1: pass + return (Tensor.arange(size) * stride + random.randrange(size)) % size def shuffled_augmentations(X:Tensor, Y:Tensor): perms = random_permutation(X.shape[0] // BS, BS) @@ -398,10 +395,9 @@ def train_cifar(): while i <= STEPS: if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"): # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True - corrects = [] - corrects_ema = [] - losses = [] - losses_ema = [] + correct_sum = loss_sum = None + correct_sum_ema = loss_sum_ema = None + correct_len = correct_len_ema = eval_batches = 0 for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False): if len(GPUS) > 1: Xt.shard_(GPUS, axis=0) @@ -411,23 +407,31 @@ def train_cifar(): out = eval_forward_jitted(model, Xt_contiguous).clone().realize() out_flipped = eval_forward_jitted(model, Xt_contiguous[..., ::-1].contiguous().realize()) correct, loss = eval_step_jitted(out, out_flipped, Yt) - losses.append(loss.numpy().tolist()) - corrects.extend(correct.numpy().tolist()) + batch_correct, batch_loss = correct.sum().realize(), loss.clone().realize() + correct_sum = batch_correct if correct_sum is None else correct_sum + batch_correct + loss_sum = batch_loss if loss_sum is None else loss_sum + batch_loss + correct_len += correct.numel() + eval_batches += 1 if model_ema: out_ema = eval_forward_ema_jitted(model_ema.net_ema, Xt_contiguous).clone().realize() out_flipped_ema = eval_forward_ema_jitted(model_ema.net_ema, Xt_contiguous[..., ::-1].contiguous().realize()) correct_ema, loss_ema = eval_step_ema_jitted(out_ema, out_flipped_ema, Yt) - losses_ema.append(loss_ema.numpy().tolist()) - corrects_ema.extend(correct_ema.numpy().tolist()) + batch_correct_ema, batch_loss_ema = correct_ema.sum().realize(), loss_ema.clone().realize() + correct_sum_ema = batch_correct_ema if correct_sum_ema is None else correct_sum_ema + batch_correct_ema + loss_sum_ema = batch_loss_ema if loss_sum_ema is None else loss_sum_ema + batch_loss_ema + correct_len_ema += correct_ema.numel() # collect accuracy across ranks - correct_sum, correct_len = sum(corrects), len(corrects) - if model_ema: correct_sum_ema, correct_len_ema = sum(corrects_ema), len(corrects_ema) + assert correct_sum is not None and loss_sum is not None + correct_count, eval_loss = correct_sum.item(), (loss_sum / eval_batches).item() + if model_ema: + assert correct_sum_ema is not None and loss_sum_ema is not None + correct_count_ema, eval_loss_ema = correct_sum_ema.item(), (loss_sum_ema / eval_batches).item() - eval_acc_pct = correct_sum/correct_len*100.0 - if model_ema: acc_ema = correct_sum_ema/correct_len_ema*100.0 - print(f"eval {correct_sum}/{correct_len} {eval_acc_pct:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)") - if model_ema: print(f"eval ema {correct_sum_ema}/{correct_len_ema} {acc_ema:.2f}%, {(sum(losses_ema)/len(losses_ema)):7.2f} val_loss STEP={i}") + eval_acc_pct = correct_count/correct_len*100.0 + if model_ema: acc_ema = correct_count_ema/correct_len_ema*100.0 + print(f"eval {correct_count}/{correct_len} {eval_acc_pct:.2f}%, {eval_loss:7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)") + if model_ema: print(f"eval ema {correct_count_ema}/{correct_len_ema} {acc_ema:.2f}%, {eval_loss_ema:7.2f} val_loss STEP={i}") if STEPS == 0 or i == STEPS: break diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 721d2993ba..dc516e5025 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -8,11 +8,11 @@ from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext -from tinygrad.dtype import dtypes, AddrSpace +from tinygrad.dtype import dtypes, AddrSpace, Invalid # import all pattern matchers here from tinygrad.codegen.gpudims import pm_add_gpudims -from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid +from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_simplify_valid, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid from tinygrad.uop.movement import mop_cleanup from tinygrad.codegen.decomp.dtype import pm_dtype_decomps from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns @@ -22,6 +22,7 @@ from tinygrad.codegen.opt.postrange import apply_opts from tinygrad.codegen.late.gater import pm_move_gates_from_index from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse from tinygrad.schedule.rangeify import pm_mops +from tinygrad.schedule.indexing import apply_movement_op from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite from tinygrad.codegen.late.coalese import memory_coalesing, pm_simplify_add_image @@ -29,6 +30,8 @@ from tinygrad.helpers import all_same, flatten, argsort, partition from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element from tinygrad.schedule.rangeify import BufferizeOpts +empty_matcher = PatternMatcher([]) + def do_number_param(ctx:list[int], x:UOp): if x.arg.slot != -1: return None ctx[0] += 1 @@ -133,16 +136,40 @@ unbroadcast = pm_wmma_add+PatternMatcher([ ]) def do_devectorize(ctx, b:UOp): - ren = ctx[-1] if isinstance(ctx, tuple) else ctx + ren = ctx.ren if isinstance(ctx, DevectorizeContext) else (ctx[-1] if isinstance(ctx, tuple) else ctx) if b.op in GroupOp.Elementwise and b.dtype in dtypes.floats and ren.supports_float4: return None - if (shape:=b.shape) == (): return None + if (shape:=b._shape) is None or shape == (): return None # broadcasting needs to be already unpacked - if not all_same([x.shape for x in b.src]): return None + if any(x._shape != shape for x in b.src): return None src = [] - for idx in itertools.product(*[range(x) for x in b.shape]): - idx_c = [UOp.const(dtypes.index, i) for i in idx] - src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src]))) - return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src) + for idx_c in shape_indexes(shape): + new_src = tuple(index_lane(ctx, x, idx_c) if isinstance(ctx, DevectorizeContext) else + (x.src[idx_c[0].arg] if len(idx_c) == 1 and x.op is Ops.STACK else UOp(Ops.INDEX, x.dtype, (x,)+idx_c)) for x in b.src) + src.append(UOp(b.op, b.dtype, new_src, b.arg, b.tag)) + return UOp(Ops.STACK, b.dtype, tuple(src)).reshape(shape) if b.op is not Ops.STORE else UOp.group(*src) + +@dataclass +class DevectorizeContext: + ren: Renderer + lanes: dict[tuple[UOp, tuple[UOp, ...]], UOp] + @property + def rewrite_cache_key(self): return (type(self.ren), self.ren.target) + +devectorize_caches: list[dict[tuple[UOp, tuple[UOp, ...]], UOp]] = [] +class scoped_devectorize_cache: + def __enter__(self): devectorize_caches.append({}) + def __exit__(self, *args): devectorize_caches.pop() + +def index_lane(ctx:DevectorizeContext, x:UOp, idxs:tuple[UOp, ...]) -> UOp: + key = (x, idxs) + if (ret:=ctx.lanes.get(key)) is not None: return ret + if x.op is Ops.STACK and idxs and idxs[0].op is Ops.CONST: + ret = index_lane(ctx, x.src[idxs[0].arg], idxs[1:]) if len(idxs) > 1 else x.src[idxs[0].arg] + elif x.op in GroupOp.Movement and len(idxs) == len(x.shape): + ret = index_lane(ctx, x.src[0], apply_movement_op(x.op, x.src[0].shape, x.marg, idxs)) + else: ret = UOp(Ops.INDEX, x.dtype, (x,)+idxs) + ctx.lanes[key] = ret + return ret def index_elementwise(x:UOp, idx:UOp): indexes = idx.src[1:] @@ -317,14 +344,18 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, symbolic_simple+unbroadcast+pm_add_loads, name="*** unbroadcast / add loads") # devectorize - sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2") + sink = graph_rewrite(sink, symbolic_simple+devectorizer2, + ctx=DevectorizeContext(ren, devectorize_caches[-1] if devectorize_caches else {}), name="devectorize2") # some coalesing misses without this sink = graph_rewrite(sink, sym, name="early symbolic") # do memory coalesing (late) sink = memory_coalesing(sink, ren) - sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True) + if IMAGE: sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True) + + has_invalid = (sink.op is Ops.CONST and sink.arg is Invalid) or any(u.op is Ops.CONST and u.arg is Invalid for u in sink.backward_slice) + if has_invalid: sink = graph_rewrite(sink, pm_simplify_valid, name="simplify valid after coalescing") # lower index dtype # NOTE: we need indexing_simplify to remove the cast to long using the Invalid @@ -343,24 +374,24 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # late decomps + move gates from unrenderable INVALID where candidate_dtypes = {*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong} emulated_dtypes = set(EMULATED_DTYPES.tolist(dtypes)) | (candidate_dtypes - ren.supported_dtypes()) - if any(u.dtype in emulated_dtypes for u in sink.toposort()): + if sink.dtype in emulated_dtypes or any(u.dtype in emulated_dtypes for u in sink.backward_slice): sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") pm_decomp = pm_decomp+\ get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\ get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="late decompositions") - sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index") + if has_invalid: sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index") # final rules for the renderer (without sym) - extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) - pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends+pm_no_index + extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else empty_matcher + pm_final_rewrite = symbolic_simple+extra_matcher+pm_split_ends+pm_no_index sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite") # this was the linearizer sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True) # put unnumbered variable PARAMs in slots - num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1]) + num_params = sum(x.op is Ops.PARAM and x.arg.slot != -1 for x in sink.backward_slice_with_self) sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True) if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST") @@ -434,7 +465,7 @@ pm_to_program = PatternMatcher([ @track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True) @Context(ALLOW_DEVICE_USAGE=0) -def do_to_program(ast:UOp, renderer:Renderer) -> UOp: +def do_to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp: """ Transform an AST into a compiled PROGRAM. May trigger BEAM search. @@ -467,14 +498,14 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: if prg.src[0].arg.estimates is None and (estimated:=do_estimates(prg, prg.src[0], prg.src[1])) is not None: prg = estimated if len(prg.src) == 2: prg = do_assemble(renderer, prg, prg.src[1]) if isinstance(renderer, ISARenderer) else do_render(renderer, prg, prg.src[1]) - if len(prg.src) == 3 and (compiled:=do_compile(renderer, prg, prg.src[2])) is not None: prg = compiled + if compile_binary and len(prg.src) == 3 and (compiled:=do_compile(renderer, prg, prg.src[2])) is not None: prg = compiled return prg to_program_cache: dict[tuple, UOp] = {} -def to_program(ast:UOp, renderer:Renderer) -> UOp: +def to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp: config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32) # UOps are structurally interned, so identity is already a collision-free structural # cache key within this process and avoids recursively hashing every kernel graph. - key = (ast, type(renderer), renderer.target, *[x.value for x in config]) - if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer) + key = (ast, type(renderer), renderer.target, compile_binary, *[x.value for x in config]) + if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer, compile_binary=compile_binary) return prg diff --git a/tinygrad/device.py b/tinygrad/device.py index 637089cdd2..70f38d775c 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -284,6 +284,8 @@ class Compiler: lib = self.compile(src) if self.cachekey is not None: diskcache_put(self.cachekey, src, lib) return lib + def compile_cached_batch(self, srcs:list[tuple[str, str]]) -> list[tuple[str, str, bytes]]: + return [(name, src, self.compile_cached(src)) for name,src in srcs] def disassemble(self, lib:bytes): pass class Compiled: diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 0a096ed2d9..4f279c8324 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -4,10 +4,11 @@ import time, random, itertools, math, contextlib, weakref, array from dataclasses import dataclass, replace, field from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events -from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo -from tinygrad.device import Device, Buffer, MultiBuffer +from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo, scoped_rewrite_cache +from tinygrad.device import Device, Buffer, MultiBuffer, Compiler from tinygrad.renderer import Estimates -from tinygrad.codegen import to_program +from tinygrad.codegen import to_program, scoped_devectorize_cache +from tinygrad.uop.spec import scoped_type_verify_cache from tinygrad.codegen.opt.postrange import bufs_from_ast # **************** Helpers **************** @@ -240,11 +241,33 @@ pm_beam = PatternMatcher([ lambda ctx,call,sink: call.replace(src=(sink.replace(arg=replace(sink.arg, beam=ctx)), *call.src[1:])) if sink.arg.beam == 0 else None), ]) +def compile_call(ctx, call:UOp, ast:UOp): + dev = call.device if isinstance(call.device, str) else call.device[0] + return call.replace(src=(to_program(ast, Device[dev].renderer, compile_binary=not ctx), *call.src[1:])) + pm_compile = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), lambda call,ast: - call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))), + (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), compile_call), ]) +def batch_compile(linear:UOp) -> UOp: + groups:dict[Compiler, list[UOp]] = {} + replacements = {} + for call in linear.toposort(): + if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or len(call.src[0].src) != 3: continue + dev = call.device if isinstance(call.device, str) else call.device[0] + compiler = Device[dev].compiler + cache = call.src[0].__dict__.setdefault('_compiled_programs', {}) + if (cached:=cache.get(compiler)) is not None: replacements[call.src[0]] = cached + else: groups.setdefault(compiler, []).append(call.src[0]) + for compiler,prgs in groups.items(): + prgs = list(dict.fromkeys(prgs)) + compiled = compiler.compile_cached_batch([(p.arg.function_name, p.src[2].arg) for p in prgs]) + for prg,(name,src,lib) in zip(prgs, compiled): + replacements[prg] = compiled_prg = prg.replace(src=prg.src[:2]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)), + arg=replace(prg.arg, name=name)) + prg.__dict__['_compiled_programs'][compiler] = compiled_prg + return linear.substitute(replacements, walk=True, enter_calls=True) if replacements else linear + pm_optimize_local_size = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size), ]) @@ -262,7 +285,9 @@ pm_exec = PatternMatcher([ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp: if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True) if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True) - linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True) + with scoped_rewrite_cache(), scoped_devectorize_cache(), scoped_type_verify_cache(): + linear = graph_rewrite(linear, pm_compile, ctx=True, name="render kernels", walk=True) + linear = batch_compile(linear) if getenv("HCQ2"): from extra.hcq2.hcq2 import hcq_compile linear = hcq_compile(linear, input_uops, jit=jit) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 92b336279a..72d09fa263 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -413,6 +413,13 @@ def diskcache_get(table:str, key:dict|str|int) -> Any: if (val:=res.fetchone()) is not None: return pickle.loads(val[0]) return None +def diskcache_get_batch(table:str, keys:list[str]) -> dict[str, Any]: + if CACHELEVEL < 1 or not keys: return {} + cur = db_connection().cursor() + try: rows = cur.execute(f"SELECT key, val FROM '{table}_{VERSION}' WHERE key IN ({','.join('?' for _ in keys)})", keys).fetchall() + except sqlite3.OperationalError: return {} + return {key:pickle.loads(val) for key,val in rows} + _db_tables = set() def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False): if CACHELEVEL < 1: return val @@ -430,6 +437,16 @@ def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False): cur.close() return val +def diskcache_put_batch(table:str, items:list[tuple[str, Any]]) -> None: + if CACHELEVEL < 1 or not items: return + conn, cur = db_connection(), db_connection().cursor() + if table not in _db_tables: + cur.execute(f"CREATE TABLE IF NOT EXISTS '{table}_{VERSION}' (key text, val blob, PRIMARY KEY (key))") + _db_tables.add(table) + cur.executemany(f"REPLACE INTO '{table}_{VERSION}' (key, val) VALUES (?, ?)", [(key, pickle.dumps(val)) for key,val in items]) + conn.commit() + cur.close() + def diskcache(func:Callable[..., T]): def wrapper(*args, **kwargs) -> T: table, key = f"cache_{func.__name__}", hashlib.sha256(pickle.dumps((args, kwargs))).hexdigest() diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 99f06cd52a..900d540708 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -75,6 +75,8 @@ class Renderer: compiler: Compiler = Compiler() def __init__(self, target:Target): self.target = target + @property + def rewrite_cache_key(self): return (type(self), self.target) def __reduce__(self): return self.__class__, (self.target,) def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer") def asm(self, prg:UOp, lin:UOp) -> bytes: raise NotImplementedError("needs an assembler") diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index ab997443e0..236ff022a8 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -13,7 +13,7 @@ from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm from tinygrad.runtime.autogen.am import am -from tinygrad.runtime.support.elf import elf_loader +from tinygrad.runtime.support.elf import elf_loader, elf_symbols from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE @@ -559,24 +559,37 @@ class AMDCopyQueue(HWQueue): sdma_queue.signal_doorbell(dev) +class AMDProgramModule: + __slots__ = ("lib_gpu", "image", "symbols", "rodata_entry", "__weakref__") + def __init__(self, dev:AMDDevice, lib_gpu:HCQBuffer, image:bytes, spec:BufferSpec, symbols:dict[str, int], rodata_entry:int): + self.lib_gpu, self.image, self.symbols, self.rodata_entry = lib_gpu, image, symbols, rodata_entry + weakref.finalize(self, HCQProgram._fini, dev, lib_gpu, spec) + class AMDProgram(HCQProgram): def __init__(self, dev:AMDDevice, name:str, lib:bytes, **kwargs): # TODO; this API needs the type signature of the function and global_size/local_size self.dev, self.name, self.lib = dev, name, lib - image, sections, relocs = elf_loader(self.lib) + module_key = hashlib.sha256(self.lib).digest() + if (module:=self.dev.program_modules.get(module_key)) is None: + image, sections, relocs = elf_loader(self.lib) + symbols = elf_symbols(self.lib) - rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1) + for apply_image_offset, rel_sym_offset, typ, addent in relocs: + if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('= 0, ".rodata section not found" - for apply_image_offset, rel_sym_offset, typ, addent in relocs: - if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack(' bytes: +def compile_hip(prg:str, arch="gfx1100", asm=False, use_device_libs=True, backend_opt=3) -> bytes: check(comgr.amd_comgr_create_action_info(ctypes.byref(action_info := comgr.amd_comgr_action_info_t()))) check(comgr.amd_comgr_action_info_set_language(action_info, comgr.AMD_COMGR_LANGUAGE_HIP)) check(comgr.amd_comgr_action_info_set_isa_name(action_info, b"amdgcn-amd-amdhsa--" + arch.encode())) @@ -80,7 +80,8 @@ def compile_hip(prg:str, arch="gfx1100", asm=False, use_device_libs=True) -> byt if status != 0: print(_get_comgr_data(data_set_bc, comgr.AMD_COMGR_DATA_KIND_LOG).decode()) raise RuntimeError("compile failed") - check(set_options(action_info, b"-O3 -mllvm -amdgpu-internalize-symbols")) + check(set_options(action_info, f"-O{backend_opt} -mllvm -vectorize-loops=false " + "-mllvm -vectorize-slp=false -mllvm -unroll-threshold=0".encode())) check(comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_CODEGEN_BC_TO_RELOCATABLE, action_info, data_set_bc, data_set_reloc)) check(set_options(action_info, b"")) @@ -99,6 +100,58 @@ class HIPCompiler(Compiler): def compile(self, src:str) -> bytes: try: return compile_hip(src, self.arch, src.split('\n', 1)[0].strip() == '.text', use_device_libs="__ocml_" in src) except RuntimeError as e: raise CompileError(e) from e + def compile_cached_batch(self, srcs:list[tuple[str, str]]) -> list[tuple[str, str, bytes]]: + batch_size = getenv("AMD_COMPILE_BATCH_SIZE", 256) + if batch_size <= 1 or any(src.split('\n', 1)[0].strip() == '.text' for _,src in srcs): return super().compile_cached_batch(srcs) + + renamed = [] + for name,src in srcs: + new_name = f"{name}_{hashlib.sha256(src.encode()).hexdigest()[:8]}" + renamed.append((new_name, src.replace(f" {name}(", f" {new_name}(", 1))) + + ret:list[tuple[str, str, bytes]|None] = [None] * len(srcs) + missing:dict[tuple[int, bool], list[int]] = {} + cache_writes:list[tuple[str, tuple[str, str]]] = [] + module_writes:list[tuple[str, bytes]] = [] + cache_srcs = [f"batchO1\n{src}" for _,src in renamed] + cached_sources = diskcache_get_batch(self.cachekey, cache_srcs) if self.cachekey is not None else {} + refs = {v[1] for v in cached_sources.values() if isinstance(v, tuple) and v[0] == "batch"} + module_keys = {key:f"__batch__{key}" for key in refs} + same_table_modules = diskcache_get_batch(self.cachekey, list(module_keys.values())) if self.cachekey is not None else {} + cached_modules = {key:same_table_modules[module_keys[key]] for key in refs if module_keys[key] in same_table_modules} + if self.cachekey is not None and (old_refs:=refs-cached_modules.keys()): + cached_modules.update(diskcache_get_batch(f"{self.cachekey}_batch", list(old_refs))) + for i,(name,src) in enumerate(renamed): + opt = 1 + cached = cached_sources.get(cache_srcs[i]) + if isinstance(cached, tuple) and cached[0] == "batch": cached = cached_modules.get(cached[1]) + if cached is not None: ret[i] = (name, src, cached) + else: missing.setdefault((opt, "__ocml_" in src), []).append(i) + for (opt,use_device_libs),indices in missing.items(): + for start in range(0, len(indices), batch_size): + batch = indices[start:start+batch_size] + preamble = dict.fromkeys(line for i in batch for line in + renamed[i][1].split('extern "C" __attribute__((global))', 1)[0].splitlines()) + kernels = [] + for i in batch: + _, kernel = renamed[i][1].split('extern "C" __attribute__((global))', 1) + kernels.append('extern "C" __attribute__((global))'+kernel) + combined = '\n'.join((*preamble, *kernels)) + module_key = hashlib.sha256(f"O{opt}\n{combined}".encode()).hexdigest() + lib = diskcache_get(self.cachekey, f"__batch__{module_key}") if self.cachekey is not None else None + if lib is None and self.cachekey is not None: lib = diskcache_get(f"{self.cachekey}_batch", module_key) + if lib is None: + try: lib = compile_hip(combined, self.arch, use_device_libs=use_device_libs, backend_opt=opt) + except RuntimeError as e: raise CompileError(e) from e + if self.cachekey is not None: module_writes.append((f"__batch__{module_key}", lib)) + for i in batch: + name,src = renamed[i] + if self.cachekey is not None: cache_writes.append((f"batchO{opt}\n{src}", ("batch", module_key))) + ret[i] = (name, src, lib) + if self.cachekey is not None: + diskcache_put_batch(self.cachekey, module_writes+cache_writes) + assert all(x is not None for x in ret) + return [x for x in ret if x is not None] def disassemble(self, lib:bytes): amdgpu_disassemble(lib) class HIPCCCompiler(Compiler): diff --git a/tinygrad/runtime/support/elf.py b/tinygrad/runtime/support/elf.py index 41917b2cb9..40ad1e0674 100644 --- a/tinygrad/runtime/support/elf.py +++ b/tinygrad/runtime/support/elf.py @@ -6,6 +6,22 @@ from tinygrad.runtime.autogen import libc @dataclass(frozen=True) class ElfSection: name:str; header:libc.Elf64_Shdr|libc.Elf32_Shdr; content:bytes # noqa: E702 +def elf_symbols(blob:bytes) -> dict[str, int]: + ecls = {libc.ELFCLASS32: "Elf32", libc.ELFCLASS64: "Elf64"}[blob[libc.EI_CLASS]] + header = getattr(libc, f"{ecls}_Ehdr").from_buffer_copy(blob) + section_headers = (getattr(libc, f"{ecls}_Shdr") * header.e_shnum).from_buffer_copy(blob[header.e_shoff:]) + sym_t = getattr(libc, f"{ecls}_Sym") + symbols = {} + for sh in section_headers: + if sh.sh_type not in (libc.SHT_SYMTAB, libc.SHT_DYNSYM): continue + str_sh = section_headers[sh.sh_link] + strtab = blob[str_sh.sh_offset:str_sh.sh_offset+str_sh.sh_size] + for sym in (sym_t * (sh.sh_size // sh.sh_entsize)).from_buffer_copy(blob[sh.sh_offset:]): + symbols[strtab[sym.st_name:strtab.find(b'\0', sym.st_name)].decode()] = sym.st_value + return symbols + +def elf_symbol_address(blob:bytes, name:str) -> int|None: return elf_symbols(blob).get(name) + def link_sym(sym:str, libs:list[str]) -> int: for lib in libs: try: return unwrap(ctypes.cast(getattr(ctypes.CDLL(ctypes.util.find_library(lib)), sym), ctypes.c_void_p).value) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index c8e328cec7..1561700874 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses -from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC +from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC, prod ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.SLICE, Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM, @@ -167,6 +167,18 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO rngs = tuple(r if (sz == sh and off == 0) else (r-off).valid(graph_rewrite((r >= off) & (r < (sh+off)), symbolic+pm_simplify_valid, name="pad")) for r,sh,(off,sz) in zip(rngs, in_shape, arg)) case Ops.RESHAPE: + if all(isinstance(s, int) for s in in_shape+arg) and tuple(s for s in in_shape if s != 1) == tuple(s for s in arg if s != 1): + out_rngs = iter(r for s,r in zip(arg, rngs) if s != 1) + return tuple(UOp.const(dtypes.index, 0) if s == 1 else next(out_rngs) for s in in_shape) + if all(isinstance(s, int) and s > 0 for s in in_shape+arg) and all(r.op is Ops.CONST and isinstance(r.arg, int) for r in rngs): + flat_idx = sum(int(r.arg)*prod(arg[i+1:]) for i,r in enumerate(rngs)) + ret = [] + for s in cast(tuple[int, ...], in_shape)[::-1]: + ret.append(UOp.const(dtypes.index, flat_idx % s)) + flat_idx //= s + return tuple(ret[::-1]) + if len(in_shape) == 1: + return (UOp.const(dtypes.index, 0).usum([r*prod(arg[i+1:]) for i,r in enumerate(rngs)]),) sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER, dtype=r.dtype) for i,r in enumerate(sink.ranges)} rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 18687803f0..5a2b3e6b6e 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1589,15 +1589,21 @@ if TRACK_MATCH_STATS or PROFILE: # A pure Python sentinel, but *typed* as UOp so it fits all the dict annotations SENTINEL: Final[UOp] = cast(UOp, object()) class BottomUpGate(Exception): pass +rewrite_caches: list[dict[Any, dict[UOp, UOp]]] = [] +class scoped_rewrite_cache: + def __enter__(self): rewrite_caches.append({}) + def __exit__(self, *args): rewrite_caches.pop() + class RewriteContext: - __slots__ = ("pm", "bpm", "bpm_cache", "ctx", "replace", "enter_calls") - def __init__(self, pm, bpm, ctx=None, enter_calls=False): + __slots__ = ("pm", "bpm", "bpm_cache", "ctx", "replace", "enter_calls", "shared_cache") + def __init__(self, pm, bpm, ctx=None, enter_calls=False, shared_cache=None): self.pm: PatternMatcher|None = pm self.bpm: PatternMatcher|None = bpm self.bpm_cache: dict[UOp, UOp|None] = {} self.ctx = ctx self.replace: dict[UOp, UOp] = {} self.enter_calls = enter_calls + self.shared_cache: dict[UOp, UOp]|None = shared_cache # no cache needed: pm_rewrite is called at most once per UOp due to the replace dict check in unified_rewrite def pm_rewrite(self, x:UOp) -> UOp|None: return unwrap(self.pm).rewrite(x, self.ctx) @@ -1636,7 +1642,7 @@ class RewriteContext: stack: list[tuple[UOp, int, UOp]] = [(root, 0, root)] on_stack = {root} # all UOps either on the stack or in self.replace, i.e. dont have to be placed again waitlist: dict[UOp, list[tuple[UOp, int, UOp]]] = {} # UOps waiting on a dependency to be in self.replace - replace, ctx, stack_limit, enter_calls = self.replace, self.ctx, REWRITE_STACK_LIMIT.value, self.enter_calls + replace, ctx, stack_limit, enter_calls, shared_cache = self.replace, self.ctx, REWRITE_STACK_LIMIT.value, self.enter_calls, self.shared_cache bpm, cached_bpm_rewrite = self.bpm, self.cached_bpm_rewrite pm_rewrite = self.pm.rewrite if self.pm is not None else None while stack: @@ -1644,6 +1650,10 @@ class RewriteContext: n, stage, new_n = stack.pop() if n in replace: continue # skip any nodes we have seen if stage == 0: + if shared_cache is not None and (cached:=shared_cache.get(n, SENTINEL)) is not SENTINEL: + replace[n] = cached + if n in waitlist: stack.extend(waitlist.pop(n)) + continue # if bottom up, we rewrite this node early. in both cases, we add its srcs to the stack if bpm is not None: # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match @@ -1682,6 +1692,7 @@ class RewriteContext: # if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict if pm_rewrite is None or (new_src_n:=pm_rewrite(new_n, ctx)) is None: replace[n] = new_n + if shared_cache is not None: shared_cache[n] = new_n if n in waitlist: stack.extend(waitlist.pop(n)) continue else: @@ -1698,12 +1709,17 @@ class RewriteContext: else: # otherwise we are done replace[n] = replaced_new_n + if shared_cache is not None: shared_cache[n] = replaced_new_n if n in waitlist: stack.extend(waitlist.pop(n)) return replace[root] @profile_matches def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp: - rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls) + cache_ctx = None if ctx is None else getattr(ctx, 'rewrite_cache_key', SENTINEL) + cache_key = pm if ctx is None else (pm, cache_ctx) + shared_cache = rewrite_caches[-1].setdefault(cache_key, {}) if rewrite_caches and cache_ctx is not SENTINEL and \ + not bottom_up and bpm is None and not walk else None + rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls, shared_cache) return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink) def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index ba543938cb..eb7574ed68 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -33,15 +33,23 @@ def validate_index(uidx:UOp, gate:UOp|None=None): return validate_index_with_z3(sz, idx, gate) def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): - lst = list(ast.toposort()) if isinstance(ast, UOp) else ast + lst = [*ast.backward_slice, ast] if isinstance(ast, UOp) else ast if SPEC > 1: test_pyrender(lst[-1]) # assume this is the sink + verified = type_verify_caches[-1].setdefault(check_spec, set()) if type_verify_caches else None with Context(TRACK_MATCH_STATS=0): for i,u in enumerate(lst): + if verified is not None and u in verified: continue ret: bool|None = check_spec.rewrite(u) if ret is not True: if DEBUG >= 3: print_uops(lst) raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}") + if verified is not None: verified.add(u) + +type_verify_caches: list[dict[PatternMatcher, set[UOp]]] = [] +class scoped_type_verify_cache: + def __enter__(self): type_verify_caches.append({}) + def __exit__(self, *args): type_verify_caches.pop() # ***** new specs ***** diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index ff47523c79..2a2536d126 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -51,6 +51,7 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: # q == b//div -> b*mul (full recombine) # q == (b//div)%d -> (b%(div*d))*mul (partial recombine into a wider mod, needs d>0) terms = list(x.split_uop(Ops.ADD)) + if not any(u.op is Ops.FLOORMOD or (u.op is Ops.MUL and any(s.op is Ops.FLOORMOD for s in u.src)) for u in terms): return None for i,u in enumerate(terms): mod, mul = u.pop_const(Ops.MUL) if mod.op is not Ops.FLOORMOD or mod.src[1].op is not Ops.CONST: continue @@ -315,6 +316,10 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=False) -> UOp: expr, is_upper, c = res bounds[expr][int(is_upper)] = c if not bounds: return uop + if not try_simplex: + uop_nodes = uop.backward_slice_with_self + bounds = defaultdict(lambda: [None, None], ((expr, v) for expr,v in bounds.items() if expr in uop_nodes)) + if not bounds: return uop # simplify uop given that valid is True all_candidates = [] @@ -347,7 +352,7 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=False) -> UOp: def _valid_priority(v: UOp, valids:list[UOp]) -> int: # we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified - return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids) + return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.backward_slice_with_self else 0 for other in valids) def simplify_valid(valid:UOp) -> UOp|None: if valid.op_in_backward_slice_with_self(Ops.INDEX): return None # this should only be for indexing, skip if there's a INDEX @@ -402,7 +407,7 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None: if x.dtype is not dtypes.index: return None # Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD): return None - return cond.where(uop_given_valid(cond, x, try_simplex=False), i) + return cond.where(new_x, i) if (new_x:=uop_given_valid(cond, x, try_simplex=False)) is not x else None # TODO: this is O(number of WHERE * number of node) # def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None: