diff --git a/test/external/fuzz_fast_idiv.py b/test/external/fuzz_fast_idiv.py index 85f0e0b0ee..ebcd5bbe66 100644 --- a/test/external/fuzz_fast_idiv.py +++ b/test/external/fuzz_fast_idiv.py @@ -3,7 +3,7 @@ import z3 from tinygrad import dtypes from tinygrad.uop.spec import z3_renderer, z3_cdiv from tinygrad.uop.ops import UOp, graph_rewrite -from tinygrad.uop.transcendental import fast_idiv +from tinygrad.uop.decompositions import fast_idiv random.seed(42) powers_of_two = [2**i for i in range(64)] diff --git a/test/unit/test_transcendental_helpers.py b/test/unit/test_transcendental_helpers.py index 5a14b381de..f40d593776 100644 --- a/test/unit/test_transcendental_helpers.py +++ b/test/unit/test_transcendental_helpers.py @@ -2,8 +2,8 @@ import unittest, math import numpy as np from tinygrad import dtypes from tinygrad.uop.ops import UOp, Ops -from tinygrad.uop.transcendental import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction -from tinygrad.uop.transcendental import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if +from tinygrad.uop.decompositions import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction +from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if from test.helpers import eval_uop class TestTranscendentalFunctions(unittest.TestCase): diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index c876b48309..69e17c3d8a 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -11,7 +11,7 @@ from tinygrad.codegen.lowerer import pm_lowerer, get_index from tinygrad.codegen.quantize import pm_quant from tinygrad.codegen.gpudims import pm_add_gpudims from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing -from tinygrad.uop.optional import get_late_rewrite_patterns +from tinygrad.uop.decompositions import get_late_rewrite_patterns from tinygrad.codegen.expander import migrate_indexing, expander from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \ ReduceContext, correct_load_store, pm_render @@ -82,13 +82,13 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC supported_ops = tuple(opts.code_for_op.keys()) extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([]) + # optional pre matcher + if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher")) + # decompositions pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2) ret.append(RewriteStep(pm_decomp, name="decompositions")) - # optional pre matcher - if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher")) - # final rules for the renderer (without sym) pm_final_rewrite = pm_decomp+pm_render+extra_matcher ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite")) diff --git a/tinygrad/uop/transcendental.py b/tinygrad/uop/decompositions.py similarity index 79% rename from tinygrad/uop/transcendental.py rename to tinygrad/uop/decompositions.py index d7fd060a16..d65a09770c 100644 --- a/tinygrad/uop/transcendental.py +++ b/tinygrad/uop/decompositions.py @@ -1,8 +1,9 @@ +from typing import Callable import math, functools from tinygrad.dtype import dtypes, DType, promo_lattice from tinygrad.device import is_dtype_supported -from tinygrad.helpers import polyN -from tinygrad.uop.ops import UOp +from tinygrad.helpers import polyN, getenv +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher TRANSCENDENTAL_SUPPORTED_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64) @@ -292,3 +293,59 @@ def fast_idiv(device: str, x: UOp, d: int) -> UOp|None: if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype): return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0) return None + +# ***** threefry ***** + +def threefry2x32(x: UOp, key: UOp): + # split x and key from uint64 to two uint32 + x0, x1 = (x & 0xffffffff).cast_vec(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast_vec(dtypes.uint32) + key0, key1 = (key & 0xffffffff).cast_vec(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast_vec(dtypes.uint32) + + rotations = [[13, 15, 26, 6], [17, 29, 16, 24]] + ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0] + xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]] + for i in range(5): + for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r))) + xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)] + + return xr[1].cast_vec(dtypes.uint64) * 2**32 | xr[0].cast_vec(dtypes.uint64) + +# ***** decomposition patterns ***** + +powers_of_two = {2**i:i for i in range(64)} +@functools.cache +def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False): + pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \ + ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental] + # no real hardware supports THREEFRY + pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32)) + # rewrite SQRT to xpow 0.5 + if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5)))) + # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1) + if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)] + # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y) + if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)] + if Ops.SHR in ops: + # no reason to check x<0 for uints + pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] + pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where( + c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] # (x+(x<0).where(c-1, 0)) >> v + if not getenv("DISABLE_FAST_IDIV"): + pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))] + pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)] + if Ops.NEG in ops: + pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))] + if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))] + if Ops.CMPLT in ops: + # These are late rewrites because simplex expects equalities to be a certain format + pat += [ + ((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1 x==c + ] + if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))] + if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))] + return PatternMatcher(pat) diff --git a/tinygrad/uop/optional.py b/tinygrad/uop/optional.py deleted file mode 100644 index 9cc3c952b6..0000000000 --- a/tinygrad/uop/optional.py +++ /dev/null @@ -1,44 +0,0 @@ -from typing import Callable -import functools -from tinygrad.dtype import dtypes -from tinygrad.uop.ops import Ops, UPat, PatternMatcher -from tinygrad.helpers import getenv -from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES, fast_idiv - -# ***** optional patterns ***** - -powers_of_two = {2**i:i for i in range(64)} -@functools.cache -def get_late_rewrite_patterns(ops, force_transcendental=False): - pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \ - ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental] - # rewrite SQRT to xpow 0.5 - if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5)))) - # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1) - if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)] - # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y) - if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)] - if Ops.SHR in ops: - # no reason to check x<0 for uints - pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] - pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where( - c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] # (x+(x<0).where(c-1, 0)) >> v - if not getenv("DISABLE_FAST_IDIV"): - pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))] - pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)] - if Ops.NEG in ops: - pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))] - if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))] - if Ops.CMPLT in ops: - # These are late rewrites because simplex expects equalities to be a certain format - pat += [ - ((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1 x==c - ] - if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))] - if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))] - return PatternMatcher(pat) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index f169da56c5..d405bced9c 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -225,4 +225,4 @@ def type_verify(uops:list[UOp], extra_spec:PatternMatcher|None=None): with Context(TRACK_MATCH_STATS=0): ret = check_spec.rewrite(u) if cast(bool|None, ret) is not True: if DEBUG >= 3: print_uops(uops) - raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[x.op for x in u.src]} {u.arg}") + 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}") diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 5ed4cef39b..244bf13c55 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -5,7 +5,7 @@ from collections import defaultdict from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING -from tinygrad.uop.transcendental import xpow +from tinygrad.uop.decompositions import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** @@ -71,6 +71,19 @@ symbolic_simple = PatternMatcher([ (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow), # positive const ** x (UPat.cvar("c", vec=False).alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None), + # rules for threefry + ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast_vec(dtypes.uint32)&0xFFFFFFFF), # TODO: why is the and needed? + (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), + (((UPat.var('x', dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x), + # hacks for threefry long removal when padded (TODO: genericize) + (UPat.var('x', dtypes.uint32).cast(dtypes.uint64) * UPat.var('y').where(UPat.const(dtypes.uint64, 1<<32), UPat.const(dtypes.uint64, 0)), + lambda x,y: y.where(x, 0).cast_vec(dtypes.uint64) * (1<<32)), + ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32), + lambda x,y: y.where(x.cast_vec(dtypes.uint32), 0)), + # new decomp rules for threefry + (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), + (((UPat.var('x', dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x), + (UPat.var('b').where(UPat.var('x', dtypes.uint32).cast(dtypes.uint64), UPat.const(dtypes.uint64, 0)).cast(dtypes.uint32), lambda b,x: b.where(x,0)) ]) # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ******** @@ -378,22 +391,6 @@ def simplify_valid(valid:UOp) -> UOp|None: if ret[-1] is not stmt: something_changed = True return functools.reduce(operator.and_, ret) if something_changed else None -# ***** threefry ***** - -def threefry2x32(x: UOp, key: UOp): - # split x and key from uint64 to two uint32 - x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32) - key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32) - - rotations = [[13, 15, 26, 6], [17, 29, 16, 24]] - ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0] - xr = [x0 + ks[-1], x1 + ks[0]] - for i in range(5): - for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r))) - xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)] - - return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64) - # ******** phase 3 is the complete symbolic, and deals with very complex things like loop rewriting and threefry transform ******** def reduce_mul_chain(r:UOp): @@ -428,16 +425,6 @@ sym = symbolic_flat+PatternMatcher([ # tensor core with a 0 input is acc (UPat(Ops.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc), (UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc), - # threefry + remove longs - (UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32), - ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)), # cast does truncation - (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), - (((UPat.var('x', dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x), - # hacks for threefry long removal when padded (TODO: genericize) - (UPat.var('x', dtypes.uint32).cast(dtypes.uint64) * UPat.var('y').where(UPat.const(dtypes.uint64, 1<<32), UPat.const(dtypes.uint64, 0)), - lambda x,y: y.where(x, UOp.const(dtypes.uint32, 0)).cast(dtypes.uint64) * (1<<32)), - ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32), - lambda x,y: y.where(x.cast(dtypes.uint32), UOp.const(dtypes.uint32, 0))), # ** self folding ** # x!=0 -> (bool)x (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),