- (UPat(Ops.STORE, dtype=dtypes.void, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
- (UPat(Ops.STORE, dtype=dtypes.void, src=(index_pat, UPat(name="val")), allow_any_len=True), validate_store),
+ (UPat(Ops.STORE, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
+ (UPat(Ops.STORE, src=(index_pat, UPat(name="val")), allow_any_len=True), validate_store),
# most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
(UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),
- (UPat((Ops.CMPLT, Ops.CMPNE), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
+ (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
# and SHL/SHR, the shift distance can be an int
(UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"), lambda a,x,y: a.dtype == x.dtype and y.dtype in (x.dtype, dtypes.uint)),
(UPat((Ops.IDIV, Ops.MOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
@@ -207,16 +224,7 @@ spec = PatternMatcher([
# *** this is the UOp AST spec ***
-def verify_sink_dims(sink:UOp):
- if not all_same([s.shape for s in sink.src]): return False
- for dims in zip(*[x.shape for x in sink.toposort() if x.op is Ops.VIEW]):
- if len(n_dims:={s for s in dims if resolve(s!=1)}) > 1:
- print(f"# INVALID KERNEL DIMS: can only have 1 or n in each dimension: {n_dims}")
- return False
-
ast_spec = PatternMatcher([
- # shapes must have either 1 or n in each dimension
- (UPat(Ops.SINK, src=UPat(Ops.STORE), name="sink"), verify_sink_dims),
# VIEW can only exist in the edges
(UPat(Ops.VIEW, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL),))), lambda: True),
(UPat(Ops.VIEW, name="view"), lambda view: len(view.src) == 0),
@@ -232,4 +240,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_repo/tinygrad/uop/symbolic.py b/tinygrad_repo/tinygrad/uop/symbolic.py
index 867b1b61..f20f6f6e 100644
--- a/tinygrad_repo/tinygrad/uop/symbolic.py
+++ b/tinygrad_repo/tinygrad/uop/symbolic.py
@@ -1,11 +1,11 @@
# all of symbolic lives here now
-from typing import Any, Literal, cast
+from typing import Any, cast
import math, operator, struct, functools
from collections import defaultdict
from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace
+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 ********
@@ -41,6 +41,8 @@ symbolic_simple = PatternMatcher([
(UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
(UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x),
(UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x),
+ (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()),
+ (UPat.var("x", dtype=dtypes.ints+(dtypes.bool,)).trunc(), lambda x: x),
# ** zero folding **
(UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
(UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
@@ -65,10 +67,25 @@ symbolic_simple = PatternMatcher([
(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)),
(UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
(UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
+ # b.cast(a).cast(b) -> b if a preserves all values in b
+ (UPat.var('x').cast().named('a').cast().named('b'), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
# ** pow **
(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(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(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), 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 ********
@@ -123,65 +140,91 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
ret.append(u)
return functools.reduce(operator.add, ret) if changed else None
-def div_and_mod_folding(x: UOp, y: UOp, which: Literal[Ops.MOD, Ops.IDIV], split_rem: bool=False) -> UOp|None:
- # simplify x // y or x % y, None means no change
- # simple cancel div/mod case
+def cancel_divmod(d: UOp, x: UOp, y: UOp) -> UOp|None:
+ # simple cancel div/mod case when the range of the numerator lies within a single denominator interval
x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
-
+ if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
- return x - q*y if which is Ops.MOD else x.const_like(q)
+ return x - q*y if d.op is Ops.MOD else d.const_like(q)
+ return None
- if (y.op is not Ops.CONST) or ((c := y.arg) < 0) or (x.dtype.count > 1): return None
- if y.arg == 0: raise ZeroDivisionError(f"{'Division' if which is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(which, y)}")
-
- svars, factors, quotients, remainders, gcd, div, const, something_changed = [], [], [], [], c, 1, 0, False
+def remove_nested_mod(m: UOp, x: UOp, y: UOp) -> UOp|None:
+ # remove nested mod in case the inner mod is a multiple of the outer mod
+ # example: (a%4 + b)%2 -> (a+b)%2
+ if ((c := y.arg) < 0) or x.vmin<0: return None
+ new_xs = []
+ something_changed = False
for u in split_uop(x, Ops.ADD):
- if u.op is Ops.MOD and which is Ops.MOD and u.src[1].op is Ops.CONST and u.src[1].arg%c == 0:
- u = u.src[0]
- something_changed = True
- v: UOp = u.divides(f:=u.const_factor())
- q, r = divmod(f, c)
- if r==0 or ((which is Ops.MOD or split_rem or u.op is Ops.CONST) and r!=f): something_changed = True
- if u.op is Ops.CONST: const += f
- else: # div is the smallest common divisor of all terms
- if f > 1 and c % f == 0 and (div == 1 or div > f): div = f
- gcd = math.gcd(r, gcd)
- factors.append(f); svars.append(v); quotients.append(q); remainders.append(r) # noqa: E702
+ if u.op is Ops.MOD:
+ if u.src[1].divides(c) is not None:
+ something_changed = True
+ u = u.src[0]
+ new_xs.append(u)
+ new_x: UOp = functools.reduce(operator.add, new_xs)
+ if something_changed and new_x.vmin>=0: return new_x % y
+ return None
+def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None:
# we can fold if the expression has only one non-constant term and this term can only take on two values
- if len(svars)==1 and (v:=svars[0]).vmax-v.vmin == 1:
- y1 = cmod(factors[0]*v.vmin+const, c) if which is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
- y2 = cmod(factors[0]*v.vmax+const, c) if which is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
+ if ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+ x,const = x.pop_const()
+ terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
+ if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
+ y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c) # type: ignore
+ y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c) # type: ignore
return (y2-y1)*(v-v.vmin) + y1
+ return None
- if not CORRECT_DIVMOD_FOLDING or x_min>=0:
- # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c
- # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c
- rems = [min(r, r-c, key=abs) for r in remainders]
- if (rem:=sum(r*v for r,v in zip(rems,svars))+const%c).vmin//c==rem.vmax//c and all(f > 0 for f in factors):
- if which is Ops.MOD: return rem - rem.vmin//c*c
- return sum((f-r)//c * v for f,r,v in zip(factors,rems,svars)) + (const-const%c+rem.vmin//c*c)//c
+def fold_divmod_congruence(d: UOp, x: UOp, y: UOp) -> UOp|None:
+ # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c
+ if (x.vmin<0 and CORRECT_DIVMOD_FOLDING) or ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+ x,const = x.pop_const()
+ terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
+ # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c
+ rems = [min((r:=f%c), r-c, key=abs) for f in factors]
+ if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c!=rem.vmax//c: return None
+ if d.op is Ops.MOD: return rem - rem.vmin//c*c
+ return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
- if (g:=math.gcd(gcd, const))!=1:
- ret = UOp(which, x.dtype, src=(sum(f//g * v for f,v in zip(factors, svars)) + const//g, x.const_like(c//g)))
- return ret*g if which is Ops.MOD else ret
+def divide_by_gcd(d: UOp, x: UOp, y: UOp) -> UOp|None:
+ # x//y -> (x//gcd)//(y//gcd) or x%y -> gcd*(x//gcd)%(y//gcd)
+ terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
+ if (gcd := math.gcd(y.arg, *factors)) == 1: return None
+ ret = sum(f//gcd * v for f,v in zip(factors, terms)).alu(d.op, y.const_like(y.arg//gcd))
+ return ret*gcd if d.op is Ops.MOD else ret
+
+def nest_div_by_smallest_factor(d: UOp, x: UOp, y: UOp) -> UOp|None:
+ # we try and nest the div and see if it allows the numerator to be simplified
+ if ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+ factors = [u.const_factor() for u in split_uop(x.pop_const()[0], Ops.ADD)]
+ # div is the smallest factor of the denominator (greater than 1) out of all "factors"
+ # TODO: there are better ways to pick `div`, this sometimes adds extra divisions
+ # TODO: add same optimization for mod
+ div = min([y.arg]+[abs(f) for f in factors if abs(f) > 1 and (c%f)==0])
+ if (1 < div < c) and (newxs:=(newx:=(x//div)).simplify()) is not newx and x.vmin>=0 and newx.vmin>=0: return newxs//(c//div)
+ return None
+
+def simplify_remainder(d: UOp, x: UOp, y: UOp) -> UOp|None:
+ # we try and take out the quotient and see if it allows the numerator to be simplified
+ if ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+ x_no_const,const = x.pop_const()
+ terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x_no_const, Ops.ADD)])
+ quotients, remainders = zip(*[divmod(f, c) for f in factors])
+ gcd = math.gcd(c, *remainders) # gcd without const!
+ if const%c==const and gcd==1 and not any(r==0 or (r!=f and d.op is Ops.MOD) for r,f in zip(remainders, factors)): return None
- if gcd != 1: something_changed = True
- if not something_changed:
- if which is Ops.IDIV and (1 < div < c) and (newx:=div_and_mod_folding(x, x.const_like(div), Ops.IDIV)) is not None: return newx//(c//div)
- return None
quo, rem = x.const_like(const//c), x.const_like((const%c)//gcd)
- for q,r,f,v in zip(quotients, remainders, factors, svars):
- if which is Ops.IDIV and (not split_rem) and r!=0:
+ for q,r,f,v in zip(quotients, remainders, factors, terms):
+ if d.op is Ops.IDIV and r!=0:
rem += f//gcd * v
else:
rem += r//gcd * v
quo += q * v
# if numerator before/after is negative, and it has remainder, don't simplify because C divmod is different from python divmod.
- if (x_min < 0 or rem.vmin < 0) and remainders: return None
- if which is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd
+ if (x.vmin < 0 or rem.vmin < 0) and remainders: return None
+ if d.op is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd
return rem//(c//gcd)+quo
def gep_through_wmma(gep:UOp, wmma:UOp):
@@ -199,18 +242,18 @@ def gep_through_wmma(gep:UOp, wmma:UOp):
gep_pushing = PatternMatcher([
# GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST
- (UPat(Ops.GEP, src=(UPat(Ops.GEP, name='g2'),), name='g1'),
+ (UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'),
lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))),
- (UPat(Ops.GEP, src=(UPat(Ops.VECTORIZE, name="vec"),), name="gep"),
+ (UPat(Ops.VECTORIZE, name='vec').f(Ops.GEP, name='gep'),
lambda gep, vec: UOp(Ops.VECTORIZE, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
- (UPat(Ops.GEP, src=(UPat.cvar("c", vec=False),), name="gep"), lambda gep, c: gep.const_like(c.arg)),
- (UPat(Ops.GEP, src=(UPat(Ops.VCONST, name="c"),), name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
+ (UPat.cvar("c", vec=False).f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
+ (UPat(Ops.VCONST, name="c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
# GEP on void is skipped
(UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
# GEP in order is removed
(UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
# push all GEPs through ALUs (fix arange stuff)
- (UPat(Ops.GEP, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu'),), name='gep'),
+ (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, name='gep'),
lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \
if not isinstance(gep.dtype, PtrDType) else None),
# CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later)
@@ -219,7 +262,7 @@ gep_pushing = PatternMatcher([
# VECTORIZE on same GEP
(UPat(Ops.VECTORIZE, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))),
# push some GEPs through WMMAs
- (UPat(Ops.GEP, src=(UPat(Ops.WMMA, name="wmma"),), name="gep"), gep_through_wmma),
+ (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
])
commutative = PatternMatcher([
@@ -236,6 +279,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
((UPat.var("y") + UPat.var("x") * UPat.cvar("c0")) + UPat.var("x") * UPat.cvar("c1"), lambda x,y,c0,c1: y+x*(c0+c1)),
(UPat.var("x") + UPat.var("x") * UPat.cvar("c"), lambda x,c: x*(c+1)), # (x+x*c)-> x*(c+1)
((UPat.var("y") + UPat.var("x")) + UPat.var("x") * UPat.cvar("c"), lambda x,y,c: y+x*(c+1)),
+ ((UPat.var("y") + UPat.var("x") * UPat.cvar("c")) + UPat.var("x"), lambda x,y,c: y+x*(c+1)),
(UPat.var("x") + UPat.var("x"), lambda x: x*2), # (x+x)-> x*2
((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
@@ -286,15 +330,21 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
# div folding
((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
if c.vmin>0 and d.vmin>0 and ((x.vmin>=0 and a.vmin>=0) or (x.vmax<=0 and a.vmax<=0)) else None), # (x//c+a)//d -> (x+a*c)//(c*d)
- (UPat.var("x", dtypes.sints) // UPat.var("y"), lambda x,y: div_and_mod_folding(x,y,Ops.IDIV)),
+ (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.var("y"))), cancel_divmod),
+ (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_binary_numerator),
+ (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_divmod_congruence),
+ (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), divide_by_gcd),
+ (UPat(Ops.MOD, dtypes.sints, name="m", src=(UPat.var("x"), UPat.cvar("y", vec=False))), remove_nested_mod),
+ (UPat((Ops.IDIV), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), nest_div_by_smallest_factor),
+ (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), simplify_remainder),
(UPat.var("x") // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
(UPat.var("x") // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <=0 else None),
((UPat.var("x", dtypes.sints)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
# ** mod **
# mod folding
- (UPat.var("x") % UPat.var("y"), lambda x,y: div_and_mod_folding(x,y,Ops.MOD)),
- (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <=0 else None),
+ (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
+ (UPat.var("x") % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None),
])+gep_pushing
symbolic_flat = symbolic+PatternMatcher([
@@ -375,22 +425,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):
@@ -425,17 +459,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.uint32).cast(dtypes.uint64).cast(dtypes.uint32), lambda x: x), # cast there and back is noop (TODO: genericize)
- ((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))),
@@ -466,6 +489,7 @@ sym = symbolic_flat+PatternMatcher([
if any(x.op in REMOVE_FROM_SINK for x in root.src) else None),
((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()), # 1/(x^c) -> (1/x)^c
((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
+ ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)
(UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")), lambda x,d: 1-d), # x*/(1+x) -> 1-1/(1+x)
(UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)),
(UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y),
diff --git a/tinygrad_repo/tinygrad/viz/index.html b/tinygrad_repo/tinygrad/viz/index.html
index f25e65bc..dd1368d9 100644
--- a/tinygrad_repo/tinygrad/viz/index.html
+++ b/tinygrad_repo/tinygrad/viz/index.html
@@ -55,6 +55,11 @@
opacity: 0.4;
pointer-events: none;
}
+ label {
+ display: inline-flex;
+ align-items: center;
+ gap: 4px;
+ }
.graph svg {
width: 100%;
height: 100%;
@@ -68,7 +73,6 @@
user-select: auto;
}
g.tag circle {
- r: 5;
fill: #FFD700;
stroke: #B8860B;
stroke-width: 0.8;
@@ -76,10 +80,9 @@
g.tag text {
text-anchor: middle;
font-size: 6px;
- fill: black;
+ fill: #08090e;
}
.label :is(text, p) {
- color: #08090e;
font-weight: 350;
}
.edgePath {
@@ -173,7 +176,7 @@
background-color: #1a1b26;
border: 1px solid #4a4b56;
color: #f0f0f5;
- border-radius: 8px;
+ border-radius: 4px;
padding: 6px;
cursor: pointer;
height: 32px;
@@ -184,7 +187,6 @@
}
.btn:hover {
background-color: #2a2b36;
- border-color: #5a5b66;
}
.collapsed .container {
display: none;
@@ -203,10 +205,9 @@
pre code.hljs {
overflow-y: auto;
max-height: 30vh;
- border-radius: 8px;
padding: 8px;
}
- .progress-message {
+ #progress-message {
position: absolute;
z-index: 2;
left: 50%;
@@ -223,10 +224,11 @@
pointer-events: none;
display: none;
font-size: 10px;
+ white-space: pre;
}
#device-list > div {
min-height: 32px;
- max-width: 100px;
+ max-width: 132px;
overflow-x: auto;
overflow-y: hidden;
white-space: nowrap;
@@ -239,6 +241,8 @@
padding: 0 8px;
width: 100%;
height: 100%;
+ max-height: 100vh;
+ overflow-x: auto;
}
.raw-text code {
max-height: none !important;
@@ -247,14 +251,12 @@
width: 100%;
border-collapse: separate;
border-spacing: 0;
- border-radius: 8px;
- overflow: hidden;
background-color: #1a1b26;
color: #f0f0f5;
font-size: 0.95em;
}
table td {
- border-bottom: 1px solid #2c2f40;
+ border-bottom: 1px solid #4a4b56;
vertical-align: top;
}
table tr:last-child > td {
@@ -269,6 +271,40 @@
tr.main-row > td, tr.sub-row > td {
padding: 8px 12px;
}
+ tr.code-row > td:first-child {
+ font-family: monospace;
+ }
+ td.pct-row > div {
+ height: 12px;
+ width: 100%;
+ display: flex;
+ }
+ td.pct-row > div > div {
+ height: 100%;
+ }
+ thead {
+ position: sticky;
+ top: 0;
+ z-index: 10;
+ background-color: #20222e;
+ }
+ thead th {
+ text-align: left;
+ padding: 10px 12px;
+ font-weight: 600;
+ border-bottom: 1px solid #4a4b56;
+ font-size: 0.95em;
+ letter-spacing: 0.03em;
+ }
+ .legend {
+ display: flex;
+ align-items: center;
+ }
+ .legend > div {
+ width: 0.95em;
+ height: 0.95em;
+ margin-right: 4px;
+ }
@@ -283,7 +319,7 @@
-
+
diff --git a/tinygrad_repo/tinygrad/viz/js/index.js b/tinygrad_repo/tinygrad/viz/js/index.js
index 1a8a1a16..a6f9eec2 100644
--- a/tinygrad_repo/tinygrad/viz/js/index.js
+++ b/tinygrad_repo/tinygrad/viz/js/index.js
@@ -10,6 +10,16 @@ const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/
const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect();
+let timeout = null;
+const updateProgress = ({ show=true }) => {
+ clearTimeout(timeout);
+ const msg = document.getElementById("progress-message");
+ if (show) {
+ msg.innerText = "Rendering new graph...";
+ timeout = setTimeout(() => { msg.style.display = "block"; }, 2000);
+ } else msg.style.display = "none";
+}
+
// ** UOp graph
function intersectRect(r1, r2) {
@@ -22,13 +32,15 @@ function intersectRect(r1, r2) {
return {x:r1.x+dx*scale, y:r1.y+dy*scale};
}
-let [workerUrl, worker, timeout] = [null, null, null];
+function addTags(root) {
+ root.selectAll("circle").data(d => [d]).join("circle").attr("r", 5);
+ root.selectAll("text").data(d => [d]).join("text").text(d => d).attr("dy", "0.35em");
+}
+
+let [workerUrl, worker] = [null, null];
async function renderDag(graph, additions, recenter=false) {
// start calculating the new layout (non-blocking)
- const progressMessage = document.querySelector(".progress-message");
- progressMessage.innerText = "Rendering new graph...";
- if (timeout != null) clearTimeout(timeout);
- timeout = setTimeout(() => {progressMessage.style.display = "block"}, 2000);
+ updateProgress({ show:true });
if (worker == null) {
const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/js/worker.js"].map(u => fetch(u)));
workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
@@ -40,8 +52,7 @@ async function renderDag(graph, additions, recenter=false) {
worker.postMessage({graph, additions, ctxs});
worker.onmessage = (e) => {
displayGraph("graph");
- progressMessage.style.display = "none";
- clearTimeout(timeout);
+ updateProgress({ show:false });
const g = dagre.graphlib.json.read(e.data);
// draw nodes
const STROKE_WIDTH = 1.4;
@@ -65,10 +76,8 @@ async function renderDag(graph, additions, recenter=false) {
return [ret];
}).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
.attr("fill", d => d.color).text(d => d.st).attr("xml:space", "preserve");
- const tags = nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
- .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`);
- tags.selectAll("circle").data(d => [d]).join("circle");
- tags.selectAll("text").data(d => [d.tag]).join("text").text(d => d).attr("dy", "0.35em");
+ addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
+ .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
// draw edges
const line = d3.line().x(d => d.x).y(d => d.y).curve(d3.curveBasis);
d3.select("#edges").selectAll("path.edgePath").data(g.edges()).join("path").attr("class", "edgePath").attr("d", (e) => {
@@ -78,7 +87,7 @@ async function renderDag(graph, additions, recenter=false) {
points.push(intersectRect(g.node(e.w), points[points.length-1]));
return line(points);
}).attr("marker-end", "url(#arrowhead)");
- const edgeLabels = d3.select("#edge-labels").selectAll("g").data(g.edges().filter(e => g.edge(e).label != null)).join("g").attr("transform", (e) => {
+ addTags(d3.select("#edge-labels").selectAll("g").data(g.edges().filter(e => g.edge(e).label != null)).join("g").attr("transform", (e) => {
// get a point near the end
const [p1, p2] = g.edge(e).points.slice(-2);
const dx = p2.x-p1.x;
@@ -92,9 +101,7 @@ async function renderDag(graph, additions, recenter=false) {
const x = p2.x - ux * offset;
const y = p2.y - uy * offset;
return `translate(${x}, ${y})`
- }).attr("class", "tag");
- edgeLabels.selectAll("circle").data(e => [g.edge(e).label]).join("circle");
- edgeLabels.selectAll("text").data(e => [g.edge(e).label]).join("text").text(d => d).attr("dy", "0.35em");
+ }).attr("class", "tag").datum(e => g.edge(e).label));
if (recenter) document.getElementById("zoom-to-fit-btn").click();
};
@@ -109,86 +116,157 @@ function formatTime(ts, dur=ts) {
}
const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
-const devColors = {"TINY":["rgb(27 87 69)", "rgb(53 79 82)", "rgb(53 79 82)", "rgb(70 172 194)", "rgb(29, 46, 98)"],
- "DEFAULT":["rgb(29,31,42)","rgb(42,45,61)","rgb(55,59,79)","rgb(68,72,98)","rgb(18,19,26)","rgb(47,50,68)","rgb(59,63,84)","rgb(74,78,101)","rgb(24,26,35)","rgb(35,37,50)","rgb(49,53,72)","rgb(64,68,89)"],}
-const bufColors = ["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"];
+const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#1d2e62", "#63b0cd"],
+ DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
+ BUFFER:["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"],
+ CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
+const cycleColors = (lst, i) => lst[i%lst.length];
-const lighten = (rgb, depth, step=0.08) => rgb.replace(/\d+/g, n => Math.round(parseInt(n)+(255-parseInt(n)) * Math.min(1, depth*step)));
+const rescaleTrack = (source, tid, k) => {
+ for (const e of source.shapes) {
+ for (let i=0; i {
+ ctx.beginPath();
+ ctx.moveTo(x[0], y[0]);
+ ctx.lineTo(x[1], y[1]);
+ ctx.fillStyle = ctx.strokeStyle = "#f0f0f5";
+ ctx.stroke();
+}
+
+var data, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
async function renderProfiler() {
displayGraph("profiler");
d3.select(".metadata").html("");
+ // layout once!
+ if (data != null) return;
const profiler = d3.select(".profiler").html("");
- const deviceList = profiler.append("div").attr("id", "device-list").node();
+ const buf = await (await fetch("/get_profile")).arrayBuffer();
+ const view = new DataView(buf);
+ let offset = 0;
+ const u8 = () => { const ret = view.getUint8(offset); offset += 1; return ret; }
+ const u32 = () => { const ret = view.getUint32(offset, true); offset += 4; return ret; }
+ const u64 = () => { const ret = new Number(view.getBigUint64(offset, true)); offset += 8; return ret; }
+ const f32 = () => { const ret = view.getFloat32(offset, true); offset += 4; return ret; }
+ const optional = (i) => i === 0 ? null : i-1;
+ const dur = u32(), peak = u64(), indexLen = u32(), layoutsLen = u32();
+ const textDecoder = new TextDecoder("utf-8");
+ const { strings, dtypeSize } = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
+ // place devices on the y axis and set vertical positions
+ const [tickSize, padding] = [10, 8];
+ const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+"px");
const canvas = profiler.append("canvas").attr("id", "timeline").node();
// NOTE: scrolling via mouse can only zoom the graph
canvas.addEventListener("wheel", e => (e.stopPropagation(), e.preventDefault()), { passive:false });
- if (profileRet == null) profileRet = await (await fetch("/get_profile")).json()
- const { layout, st, et } = profileRet;
- // place devices on the y axis and set vertical positions
- const [tickSize, padding] = [10, 8];
- deviceList.style.paddingTop = `${tickSize+padding}px`;
const ctx = canvas.getContext("2d");
const canvasTop = rect(canvas).top;
// color by key (name/category/device)
const colorMap = new Map();
- const data = {shapes:[], axes:{}};
- const areaScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.mem.peak), 0)]).range([4,maxArea=100]);
- for (const [k, { timeline, mem }] of Object.entries(layout)) {
- if (timeline.shapes.length === 0 && mem.shapes.length == 0) continue;
- const div = deviceList.appendChild(document.createElement("div"));
- div.innerText = k;
- div.style.padding = `${padding}px`;
- div.onclick = () => { // TODO: make this feature more visible
- focusedDevice = k === focusedDevice ? null : k;
- const prevScroll = profiler.node().scrollTop;
- renderProfiler();
- if (prevScroll) profiler.node().scrollTop = prevScroll;
- }
- const { y:baseY, height:baseHeight } = rect(div);
- const levelHeight = baseHeight-padding;
+ data = {tracks:new Map(), axes:{}};
+ const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]);
+ for (let i=0; i e.st >= levelEt);
+ const et = e.st+Math.trunc(e.dur);
+ if (depth === -1) {
+ depth = levels.length;
+ levels.push(et);
+ } else levels[depth] = et;
+ if (depth === 0) colorKey = e.cat ?? e.name;
+ if (!colorMap.has(colorKey)) colorMap.set(colorKey, cycleColors(colorScheme[k] ?? colorScheme.DEFAULT, colorMap.size));
+ const fillColor = d3.color(colorMap.get(colorKey)).brighter(depth).toString();
+ const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
+ if (e.ref != null) ref = {ctx:e.ref, step:0};
+ else if (ref != null) {
+ const start = ref.step>0 ? ref.step+1 : 0;
+ const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
+ ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
+ }
+ const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
+ // offset y by depth
+ shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor });
}
- const fillColor = lighten(colorMap.get(colorKey), e.depth);
- const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
- if (e.ref != null) ref = {ctx:e.ref, step:0};
- else if (ref != null) {
- const start = ref.step>0 ? ref.step+1 : 0;
- const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
- ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
+ div.style("height", levelHeight*levels.length+padding+"px").style("pointerEvents", "none");
+ } else {
+ const peak = u64();
+ const height = heightScale(peak);
+ const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
+ let x = 0, y = 0;
+ const buf_shapes = new Map(), temp = new Map();
+ const timestamps = [];
+ for (let j=0; j timestamps[s]);
+ const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}`};
+ shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
+ }
+ data.tracks.set(k, { shapes, offsetY, height, peak, scaleFactor:maxheight*4/height });
+ div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
+ const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
+ let offset = 0;
+ for (const [tid, track] of data.tracks) {
+ track.offsetY += offset;
+ if (tid === newFocus) offset += rescaleTrack(track, tid, track.scaleFactor);
+ else if (tid === focusedDevice) offset += rescaleTrack(track, tid, 1/track.scaleFactor);
+ }
+ data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.height, t.offsetY], fmt:"B" } : null;
+ focusedDevice = newFocus;
+ return resize();
+ });
}
- // position shapes on the canvas and scale to fit fixed area
- const startY = offsetY+(levelHeight*timeline.maxDepth)+padding/2;
- let area = mem.shapes.length === 0 ? 0 : areaScale(mem.peak);
- if (area === 0) div.style.pointerEvents = "none";
- else div.style.cursor = "pointer";
- if (k === focusedDevice) {
- // expand memory graph for the focused device
- area = maxArea*4;
- data.axes.y = { domain:[0, mem.peak], range:[startY+area, startY], fmt:"B" };
- }
- const yscale = d3.scaleLinear().domain([0, mem.peak]).range([startY+area, startY]);
- for (const [i,e] of mem.shapes.entries()) {
- const x = e.x.map((i,_) => (mem.timestamps[i] ?? et)-st);
- const y0 = e.y.map(yscale);
- const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
- const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
- data.shapes.push({ x, y0, y1, arg, fillColor:bufColors[i%bufColors.length] });
- }
- // lastly, adjust device rect by number of levels
- div.style.height = `${Math.max(levelHeight*timeline.maxDepth, baseHeight)+area+padding}px`;
}
+ updateProgress({ "show":false });
// draw events on a timeline
const dpr = window.devicePixelRatio || 1;
const ellipsisWidth = ctx.measureText("...").width;
@@ -199,7 +277,7 @@ async function renderProfiler() {
ctx.save();
ctx.clearRect(0, 0, canvas.clientWidth, canvas.clientHeight);
// rescale to match current zoom
- const xscale = d3.scaleLinear().domain([0, et-st]).range([0, canvas.clientWidth]);
+ const xscale = d3.scaleLinear().domain([0, dur]).range([0, canvas.clientWidth]);
xscale.domain(xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale));
const zoomDomain = transform != null ? xscale.domain() : null;
let yscale = null;
@@ -207,77 +285,69 @@ async function renderProfiler() {
yscale = d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range);
}
// draw shapes
- for (const e of data.shapes) {
- const [start, end] = e.width != null ? [e.x, e.x+e.width] : [e.x[0], e.x[e.x.length-1]];
- if (zoomDomain != null && (start>zoomDomain[1]|| end=0; i--) ctx.lineTo(x[i], e.y1[i]);
- ctx.closePath();
- ctx.fill();
- // NOTE: y coordinates are in reverse order
- for (let i = 0; i < x.length - 1; i++) rectLst.push({ x0:x[i], x1:x[i+1], y0:e.y1[i], y1:e.y0[i], arg:e.arg });
- continue;
- }
- // contiguous rect
- const x = xscale(start);
- const width = xscale(end)-x;
- ctx.fillRect(x, e.y, width, e.height);
- rectLst.push({ y0:e.y, y1:e.y+e.height, x0:x, x1:x+width, arg:e.arg });
- // add label
- if (e.label == null) continue;
- ctx.textAlign = "left";
- ctx.textBaseline = "middle";
- let [labelX, labelWidth] = [x+2, 0];
- const labelY = e.y+e.height/2;
- for (const [i,l] of e.label.entries()) {
- if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
- if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
- break;
+ for (const [_, { offsetY, shapes }] of data.tracks) {
+ for (const e of shapes) {
+ const [start, end] = e.width != null ? [e.x, e.x+e.width] : [e.x[0], e.x[e.x.length-1]];
+ if (zoomDomain != null && (start>zoomDomain[1]|| end=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]);
+ ctx.closePath();
+ ctx.fill();
+ // NOTE: y coordinates are in reverse order
+ for (let i = 0; i < x.length - 1; i++) {
+ let tooltipText = e.arg.tooltipText;
+ if (yscale != null && ((yaxisVal=yscale.invert(offsetY+e.y1[i]))>0)) {
+ tooltipText += `\nTotal: ${formatUnit(yaxisVal, data.axes.y.fmt)}`;
+ }
+ rectLst.push({ x0:x[i], x1:x[i+1], y0:offsetY+e.y1[i], y1:offsetY+e.y0[i], arg:{...e.arg, tooltipText} });
+ }
+ continue;
+ }
+ // contiguous rect
+ const x = xscale(start);
+ const width = xscale(end)-x;
+ ctx.fillRect(x, offsetY+e.y, width, e.height);
+ rectLst.push({ y0:offsetY+e.y, y1:offsetY+e.y+e.height, x0:x, x1:x+width, arg:e.arg });
+ // add label
+ if (e.label == null) continue;
+ ctx.textAlign = "left";
+ ctx.textBaseline = "middle";
+ let [labelX, labelWidth] = [x+2, 0];
+ const labelY = offsetY+e.y+e.height/2;
+ for (const [i,l] of e.label.entries()) {
+ if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
+ if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
+ break;
+ }
+ ctx.fillStyle = l.color;
+ ctx.fillText(l.st, labelX, labelY);
+ labelWidth += l.width;
+ labelX += l.width;
}
- ctx.fillStyle = l.color;
- ctx.fillText(l.st, labelX, labelY);
- labelWidth += l.width;
- labelX += l.width;
}
}
// draw axes
- ctx.beginPath();
- ctx.moveTo(0, 0);
- ctx.lineTo(canvas.clientWidth, 0);
- ctx.fillStyle = ctx.strokeStyle = "#f0f0f5";
- ctx.lineWidth = 1;
- ctx.stroke();
- const ticks = xscale.ticks();
- for (const [i, tick] of ticks.entries()) {
- ctx.beginPath();
+ drawLine(ctx, xscale.range(), [0, 0]);
+ for (const tick of xscale.ticks()) {
// tick line
const x = xscale(tick);
- ctx.moveTo(x, 0);
- ctx.lineTo(x, tickSize);
- ctx.stroke();
+ drawLine(ctx, [x, x], [0, tickSize])
// tick label
ctx.textBaseline = "top";
- ctx.textAlign = i === ticks.length-1 ? "right" : "left";
- const padding = i === ticks.length-1 ? -1 : 1;
- ctx.fillText(formatTime(tick, et-st), x+(ctx.lineWidth+2)*padding, tickSize);
+ ctx.textAlign = "left";
+ ctx.fillText(formatTime(tick, dur), x+ctx.lineWidth+2, tickSize);
}
if (yscale != null) {
- ctx.beginPath();
- ctx.moveTo(0, yscale.range()[1]);
- ctx.lineTo(0, yscale.range()[0]);
- ctx.stroke();
+ drawLine(ctx, [0, 0], yscale.range());
for (const tick of yscale.ticks()) {
const y = yscale(tick);
- ctx.beginPath();
- ctx.moveTo(0, y);
- ctx.lineTo(tickSize, y);
- ctx.stroke();
+ drawLine(ctx, [0, tickSize], [y, y]);
ctx.textAlign = "left";
ctx.textBaseline = "middle";
ctx.fillText(formatUnit(tick, data.axes.y.fmt), tickSize+2, y);
@@ -377,11 +447,16 @@ function codeBlock(st, language, { loc, wrap }={}) {
return ret;
}
-function appendRow(table, name, value, unit, cls) {
+function appendTd(tr, value, unit=null) {
+ const fmt = (typeof value === "number" && !Number.isInteger(value)) ? value.toFixed(2) : value;
+ tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? "");
+}
+
+function appendRow(table, name, value, unit=null, cls="main-row") {
const tr = table.appendChild(document.createElement("tr"));
tr.className = cls;
tr.appendChild(document.createElement("td")).innerText = name;
- tr.appendChild(document.createElement("td")).innerText = unit === "us" ? formatTime(value) : value.toFixed(2)+(unit != null ? " "+unit : "%");
+ appendTd(tr, value, unit);
return tr;
}
@@ -495,10 +570,47 @@ async function main() {
if (ckey.startsWith("/disasm")) {
if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
displayGraph("profiler");
- document.querySelector(".metadata").innerHTML = "";
const root = document.createElement("div");
root.className = "raw-text";
- root.appendChild(codeBlock(ret.src, "x86asm"));
+ const metadata = document.querySelector(".metadata");
+ metadata.innerHTML = "";
+ // detailed assembly view
+ if (ret.cols != null) {
+ const asm = root.appendChild(document.createElement("table"));
+ const thead = asm.appendChild(document.createElement("thead"));
+ for (const c of ret.cols) thead.appendChild(document.createElement("th")).innerText = c.title ?? c;
+ for (const r of ret.rows) {
+ const tr = asm.appendChild(document.createElement("tr"));
+ tr.className = "main-row code-row";
+ for (const [i,value] of r.entries()) {
+ // string format scalar values
+ if (!Array.isArray(value)) appendTd(tr, value);
+ // display arrays in a bar graph
+ else {
+ const segmentsTd = tr.appendChild(document.createElement("td"));
+ segmentsTd.className = "pct-row";
+ const usageBar = segmentsTd.appendChild(document.createElement("div"));
+ for (const [k, v, width] of value) {
+ const seg = usageBar.appendChild(document.createElement("div"));
+ seg.style.width = width+"%";
+ seg.title = `${ret.cols[i].labels[k]} ${v}`;
+ seg.style.background = cycleColors(colorScheme.CATEGORICAL, parseInt(k));
+ }
+ }
+ }
+ }
+ const summary = metadata.appendChild(document.createElement("table"));
+ for (const s of ret.summary) {
+ const tr = summary.appendChild(document.createElement("tr"));
+ tr.className = "main-row";
+ const td = tr.appendChild(document.createElement("td"));
+ const div = td.appendChild(document.createElement("div"));
+ div.className = "legend";
+ div.appendChild(document.createElement("div")).style.background = cycleColors(colorScheme.CATEGORICAL, s.idx);
+ div.appendChild(document.createElement("p")).textContent = s.label;
+ appendTd(tr, s.value);
+ }
+ } else root.appendChild(codeBlock(ret.src, "x86asm"));
return document.querySelector(".profiler").replaceChildren(root);
}
// ** UOp view (default)
diff --git a/tinygrad_repo/tinygrad/viz/serve.py b/tinygrad_repo/tinygrad/viz/serve.py
index e0c60f1d..48e6efb1 100755
--- a/tinygrad_repo/tinygrad/viz/serve.py
+++ b/tinygrad_repo/tinygrad/viz/serve.py
@@ -1,13 +1,14 @@
#!/usr/bin/env python3
-import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io
+import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io, struct
+import subprocess, ctypes
from contextlib import redirect_stdout
from decimal import Decimal
from http.server import BaseHTTPRequestHandler
from urllib.parse import parse_qs, urlparse
from typing import Any, TypedDict, Generator
-from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey
-from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint
-from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfilePointEvent, Device
+from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent
+from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp, srender, sint, sym_infer
+from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
from tinygrad.renderer import ProgramSpec
from tinygrad.dtype import dtypes
@@ -17,7 +18,8 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
**{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
- Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D"}
+ Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
+ Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e"}
# VIZ API
@@ -29,10 +31,13 @@ def get_metadata(keys:list[TracingKey], contexts:list[list[TrackedGraphRewrite]]
for i,(k,v) in enumerate(zip(keys, contexts)):
steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc),
"query":f"/ctxs?ctx={i}&idx={j}"} for j,s in enumerate(v)]
- if isinstance(k.ret, ProgramSpec): steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
- ret.append(r:={"name":k.display_name, "fmt":k.fmt, "steps":steps})
+ ret.append(r:={"name":k.display_name, "steps":steps})
# use the first key to get runtime profiling data about this context
if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0])
+ # program spec metadata
+ if isinstance(k.ret, ProgramSpec):
+ steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
+ r["fmt"] = k.ret.src
for key in k.keys: ref_map[key] = i
return ret
@@ -54,7 +59,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
excluded: set[UOp] = set()
for u in (toposort:=x.toposort()):
# always exclude DEVICE/CONST/UNIQUE
- if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE}: excluded.add(u)
+ if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u)
# only exclude CONST VIEW source if it has no other children in the graph
if u.op is Ops.CONST and len(u.src) != 0 and all(cr.op is Ops.CONST for c in u.src[0].children if (cr:=c()) is not None and cr in toposort):
excluded.update(u.src)
@@ -68,13 +73,15 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
if u.dtype != dtypes.void: label += f"\n{u.dtype}"
for idx,x in enumerate(u.src):
if x in excluded:
- if x.op is Ops.CONST and dtypes.is_float(u.dtype): label += f"\nCONST{idx} {x.arg:g}"
- else: label += f"\n{x.op.name}{idx} {x.arg}"
+ arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(u.dtype) else f"{x.arg}"
+ label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
try:
if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None:
label += f"\n{shape_to_str(u.shape)}"
+ elif len(rngs:=u.ranges):
+ label += f"\n{str(sorted([x.arg[0] for x in rngs]))}"
except Exception:
- label += "\n"
+ label += "\n"
if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
# NOTE: kernel already has metadata in arg
if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+repr(u.metadata)
@@ -84,9 +91,9 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
@functools.cache
def _reconstruct(a:int):
- op, dtype, src, arg, tag = contexts[2][a]
+ op, dtype, src, arg, *rest = contexts[2][a]
arg = type(arg)(_reconstruct(arg.ast), arg.metadata) if op is Ops.KERNEL else arg
- return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, tag)
+ return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)
def get_details(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}
@@ -99,6 +106,15 @@ def get_details(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None,
"diff":list(difflib.unified_diff(str(u0).splitlines(), str(u1).splitlines())), "upat":(upat_loc, printable(upat_loc))}
if not ctx.bottom_up: next_sink = new_sink
+# encoder helpers
+
+def enum_str(s, cache:dict[str, int]) -> int:
+ if (cret:=cache.get(s)) is not None: return cret
+ cache[s] = ret = len(cache)
+ return ret
+
+def option(s:int|None) -> int: return 0 if s is None else s+1
+
# Profiler API
device_ts_diffs:dict[str, tuple[Decimal, Decimal]] = {}
@@ -108,75 +124,78 @@ DevEvent = ProfileRangeEvent|ProfileGraphEntry|ProfilePointEvent
def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decimal, DevEvent], None, None]:
for e in profile:
if isinstance(e, ProfileRangeEvent): yield (e.st+(diff:=cpu_ts_diff(e.device, e.is_copy)), (e.en if e.en is not None else e.st)+diff, e)
- elif isinstance(e, ProfilePointEvent): yield (e.st, e.st, e)
+ elif isinstance(e, ProfilePointEvent): yield (e.ts, e.ts, e)
elif isinstance(e, ProfileGraphEvent):
cpu_ts = []
for ent in e.ents: cpu_ts += [e.sigs[ent.st_id]+(diff:=cpu_ts_diff(ent.device, ent.is_copy)), e.sigs[ent.en_id]+diff]
yield (st:=min(cpu_ts)), (et:=max(cpu_ts)), ProfileRangeEvent(f"{e.ents[0].device.split(':')[0]} Graph", f"batched {len(e.ents)}", st, et)
for i,ent in enumerate(e.ents): yield (cpu_ts[i*2], cpu_ts[i*2+1], ent)
-# timeline layout stacks events in a contiguous block. When a late starter finishes late, there is whitespace in the higher levels.
-def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
- shapes:list[dict] = []
- levels:list[int] = []
- for st,et,dur,e in events:
+# normalize event timestamps and attach kernel metadata
+def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, scache:dict[str, int]) -> bytes|None:
+ events:list[bytes] = []
+ exec_points:dict[str, dict] = {}
+ category_enum:dict[str, int] = {}
+ for st,et,dur,e in dev_events:
+ if isinstance(e, ProfilePointEvent) and e.name == "exec": exec_points[e.key] = e.arg
if dur == 0: continue
- # find a free level to put the event
- depth = next((i for i,level_et in enumerate(levels) if st>=level_et), len(levels))
- if depth < len(levels): levels[depth] = et
- else: levels.append(et)
- name, cat = e.name, None
- if (ref:=ref_map.get(name)) is not None: name = ctxs[ref]["name"]
+ name, cat, info = e.name, None, None
+ if (ref:=ref_map.get(name)) is not None:
+ name = ctxs[ref]["name"]
+ if isinstance(p:=contexts[0][ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
+ info = f"{sym_infer(p.estimates.ops, ei['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei['var_vals'])/t:4.1f}"+ \
+ f"|{sym_infer(p.estimates.lds,ei['var_vals'])/t:.1f} GB/s\n{ei['metadata']}"
elif isinstance(e.name, TracingKey):
name, cat = e.name.display_name, e.name.cat
ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
- shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat})
- return {"shapes":shapes, "maxDepth":len(levels)}
+ events.append(struct.pack(" dict:
- step, peak, mem = 0, 0, 0
- shps:dict[int, dict] = {}
- temp:dict[int, dict] = {}
- timestamps:list[int] = []
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
+ scache:dict[str, int]) -> bytes|None:
+ peak, mem = 0, 0
+ temp:dict[int, int] = {}
+ bufs:list[bytes] = []
for st,_,_,e in events:
if not isinstance(e, ProfilePointEvent): continue
if e.name == "alloc":
- shps[e.ref] = temp[e.ref] = {"x":[step], "y":[mem], "arg":e.arg}
- timestamps.append(int(e.st))
- step += 1
- mem += e.arg["nbytes"]
+ bufs.append(struct.pack(" peak: peak = mem
if e.name == "free":
- timestamps.append(int(e.st))
- step += 1
- mem -= (removed:=temp.pop(e.ref))["arg"]["nbytes"]
- removed["x"].append(step)
- removed["y"].append(removed["y"][-1])
- for k,v in temp.items():
- if k > e.ref:
- v["x"] += [step, step]
- v["y"] += [v["y"][-1], v["y"][-1]-removed["arg"]["nbytes"]]
- for v in temp.values():
- v["x"].append(step)
- v["y"].append(v["y"][-1])
- return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
+ bufs.append(struct.pack(" bytes|None:
# start by getting the time diffs
for ev in profile:
if isinstance(ev,ProfileDeviceEvent): device_ts_diffs[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
# map events per device
dev_events:dict[str, list[tuple[int, int, float, DevEvent]]] = {}
- min_ts:int|None = None
- max_ts:int|None = None
+ start_ts:int|None = None
+ end_ts:int|None = None
for ts,en,e in flatten_events(profile):
dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e))
- if min_ts is None or st < min_ts: min_ts = st
- if max_ts is None or et > max_ts: max_ts = et
+ if start_ts is None or st < start_ts: start_ts = st
+ if end_ts is None or et > end_ts: end_ts = et
+ if start_ts is None: return None
# return layout of per device events
- for events in dev_events.values(): events.sort(key=lambda v:v[0])
- dev_layout = {k:{"timeline":timeline_layout(v), "mem":mem_layout(v)} for k,v in dev_events.items()}
- return json.dumps({"layout":dev_layout, "st":min_ts, "et":max_ts}).encode("utf-8")
+ layout:dict[str, bytes|None] = {}
+ scache:dict[str, int] = {}
+ peaks:list[int] = []
+ dtype_size:dict[str, int] = {}
+ for k,v in dev_events.items():
+ v.sort(key=lambda e:e[0])
+ layout[k] = timeline_layout(v, start_ts, scache)
+ layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)
+ ret = [b"".join([struct.pack(" list[dict]:
ret:list[dict] = []
@@ -185,11 +204,40 @@ def get_runtime_stats(key) -> list[dict]:
ret.append({"device":e.device, "data":[{"name":"Duration", "value":float(e.en-e.st), "unit":"us"}]})
return ret
+# ** Assembly analyzers
+
+def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
+ target_args = [f"-mtriple={mtriple}", f"-mcpu={mcpu}"]
+ # disassembly output can include headers / metadata, skip if llvm-mca can't parse those lines
+ data = json.loads(subprocess.check_output(["llvm-mca","-skip-unsupported-instructions=parse-failure","--json","-"]+target_args, input=asm.encode()))
+ cr = data["CodeRegions"][0]
+ resource_labels = data["TargetInfo"]["Resources"]
+ rows:list = [[instr] for instr in cr["Instructions"]]
+ # add scheduler estimates
+ for info in cr["InstructionInfoView"]["InstructionList"]: rows[info["Instruction"]].append(info["Latency"])
+ # map per instruction resource usage
+ instr_usage:dict[int, dict[int, int]] = {}
+ for d in cr["ResourcePressureView"]["ResourcePressureInfo"]:
+ instr_usage.setdefault(i:=d["InstructionIndex"], {}).setdefault(r:=d["ResourceIndex"], 0)
+ instr_usage[i][r] += d["ResourceUsage"]
+ # last row is the usage summary
+ summary = [{"idx":k, "label":resource_labels[k], "value":v} for k,v in instr_usage.pop(len(rows), {}).items()]
+ max_usage = max([sum(v.values()) for i,v in instr_usage.items() if i