|
|
|
@@ -65,6 +65,7 @@ class UOpMetaClass(type):
|
|
|
|
|
assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}"
|
|
|
|
|
buffers[created] = _buffer
|
|
|
|
|
if SPEC > 1:
|
|
|
|
|
if SPEC > 2: test_pyrender(created)
|
|
|
|
|
from tinygrad.uop.spec import full_spec
|
|
|
|
|
with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created)
|
|
|
|
|
if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}")
|
|
|
|
@@ -360,7 +361,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
|
|
|
|
def end(self, *src:UOp):
|
|
|
|
|
if len(src) == 0: return self
|
|
|
|
|
return UOp(Ops.END, src=(self,)+src)
|
|
|
|
|
def after(self, *src:UOp): return UOp(Ops.AFTER, self.dtype, (self,)+src)
|
|
|
|
|
def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
|
|
|
|
|
def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
|
|
|
|
|
def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
|
|
|
|
|
def alu(self, op, *src:UOp, **kwargs):
|
|
|
|
@@ -378,10 +379,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
|
|
|
|
|
if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
|
|
|
|
|
return ret
|
|
|
|
|
@staticmethod
|
|
|
|
|
def range(end:sint, *arg, dtype=dtypes.index, **kwargs):
|
|
|
|
|
def range(end:sint, *arg, dtype=dtypes.index, src=(), **kwargs):
|
|
|
|
|
if len(arg) == 0: raise RuntimeError("range needs an arg")
|
|
|
|
|
if len(arg) == 1: arg = arg+(AxisType.LOOP,)
|
|
|
|
|
return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=arg, **kwargs)
|
|
|
|
|
return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=arg, **kwargs)
|
|
|
|
|
@staticmethod
|
|
|
|
|
def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
|
|
|
|
|
def r(self, op:Ops, axis:tuple[int, ...]):
|
|
|
|
@@ -1234,44 +1235,118 @@ renderer_infer = PatternMatcher([
|
|
|
|
|
*renderer.patterns
|
|
|
|
|
])
|
|
|
|
|
|
|
|
|
|
sugar = { Ops.SINK: "sink", Ops.STORE: "store", Ops.LOAD: "load", Ops.SQRT: "sqrt", Ops.INDEX: "index", Ops.REDUCE: "reduce",
|
|
|
|
|
Ops.WHERE: "where", Ops.RECIPROCAL: "reciprocal", Ops.EXP2: "exp2", Ops.LOG2: "log2", Ops.SIN: "sin"}
|
|
|
|
|
pm_pyrender = PatternMatcher([
|
|
|
|
|
(UPat(Ops.CONST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg}, src={x.src[0].arg})")),
|
|
|
|
|
(UPat(Ops.CONST, name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg})")),
|
|
|
|
|
(UPat(Ops.END, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.end({', '.join([y.arg for y in x.src[1:]])})")),
|
|
|
|
|
(UPat(Ops.CAST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.cast({x.dtype})")),
|
|
|
|
|
(UPat(Ops.BITCAST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.bitcast({x.dtype})")),
|
|
|
|
|
(UPat({Ops.MAX, Ops.THREEFRY, Ops.CMPLT, Ops.CMPNE, Ops.POW}, src=UPat(Ops.NOOP), name="x"),
|
|
|
|
|
lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.alu({x.op}, {x.src[1].arg})")),
|
|
|
|
|
(UPat(Ops.RANGE, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=
|
|
|
|
|
f"UOp.range({x.src[0].arg}, {str(x.arg[0])}, {str(x.arg[1])}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})")),
|
|
|
|
|
(UPat(Ops.SPECIAL, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg= f"UOp.special({x.src[0].arg}, \"{x.arg}\", dtype={x.dtype})")),
|
|
|
|
|
(UPat(Ops.DEFINE_VAR, name="x"), lambda x: UOp(Ops.NOOP, arg=
|
|
|
|
|
f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})")),
|
|
|
|
|
(UPat(set(sugar.keys()), src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP,
|
|
|
|
|
arg=f"{x.src[0].arg}.{sugar[x.op]}({', '.join([y.arg for y in x.src[1:]] + ([f'arg={str(x.arg)}'] if x.arg is not None else []))})")),
|
|
|
|
|
(UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.NOOP),), name="x"),
|
|
|
|
|
lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, arg=({', '.join([str(y) for y in x.arg])}))")),
|
|
|
|
|
# *** pyrender ***
|
|
|
|
|
|
|
|
|
|
def render_marg(ctx,x:UOp):
|
|
|
|
|
if x.op in {Ops.PERMUTE, Ops.FLIP}: return str(x.marg)
|
|
|
|
|
pieces = []
|
|
|
|
|
if x.op in {Ops.RESHAPE, Ops.EXPAND}:
|
|
|
|
|
pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
|
|
|
|
|
if x.op in {Ops.PAD, Ops.SHRINK}:
|
|
|
|
|
pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
|
|
|
|
|
return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
|
|
|
|
|
|
|
|
|
|
# TODO: use this more in pyrender
|
|
|
|
|
def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
|
|
|
|
|
|
|
|
|
|
sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER,
|
|
|
|
|
Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER}
|
|
|
|
|
pm_pyrender_extra = PatternMatcher([
|
|
|
|
|
(UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
|
|
|
|
|
(UPat(Ops.CONST, name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
|
|
|
|
|
(UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
|
|
|
|
|
f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})"),
|
|
|
|
|
(UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
|
|
|
|
|
(UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
|
|
|
|
|
(UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
|
|
|
|
|
f"UOp.new_buffer({repr(d.arg)}, {x.size}, {x.dtype}, {u.arg})"),
|
|
|
|
|
(UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
|
|
|
|
|
(UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}.r({r.arg[0]}, {r.arg[1]})"),
|
|
|
|
|
# NOTE: range has srcs sometimes after control flow
|
|
|
|
|
(UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
|
|
|
|
|
"UOp.range("+', '.join([str(c.arg)] + [str(y) for y in x.arg])+
|
|
|
|
|
(f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+\
|
|
|
|
|
(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+\
|
|
|
|
|
(', tag='+str(x.tag) if x.tag is not None else '')+")"),
|
|
|
|
|
# TODO: index shouldn't mismatch dtype
|
|
|
|
|
(UPat(Ops.INDEX, src=(UPat(), UPat()), name="x"), lambda ctx,x:
|
|
|
|
|
f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, dtype={x.dtype})" if x.src[0].dtype != x.dtype else None),
|
|
|
|
|
# TODO: fix forced_reshape
|
|
|
|
|
(UPat(Ops.RESHAPE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.forced_reshape({render_marg(ctx,x)})" if x.src[0].shape == x.shape else None),
|
|
|
|
|
(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
|
|
|
|
|
# NOTE: CMPNE doesn't work cause there's no __rne__
|
|
|
|
|
(UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, name="y"), UPat(name="z")), name="x"),
|
|
|
|
|
lambda ctx,x,y,z: f"({y.arg}{syms[x.op]}{ctx[z]})"),
|
|
|
|
|
# NOTE: sub doesn't work cause it's written as add/mul
|
|
|
|
|
(UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, name="z")), name="x"), lambda ctx,x,y,z: f"({ctx[y]}{syms[x.op]}{z.arg})"),
|
|
|
|
|
(UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}{syms[x.op]}{ctx[x.src[1]]})"),
|
|
|
|
|
(UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join( \
|
|
|
|
|
([f'arg={repr(x.arg)}'] if x.arg is not None else []) + ([f'tag={repr(x.tag)}'] if x.tag is not None else []))+")"),
|
|
|
|
|
(UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
|
|
|
|
|
([f'arg={repr(x.arg)}'] if x.arg is not None else []) + ([f'tag={repr(x.tag)}'] if x.tag is not None else []))+")"),
|
|
|
|
|
])
|
|
|
|
|
|
|
|
|
|
# NOTE: you can remove pm_pyrender_extra and it'll still be correct
|
|
|
|
|
pm_pyrender = pm_pyrender_extra+PatternMatcher([
|
|
|
|
|
(UPat(Ops.KERNEL, name="u"), lambda ctx,u: "UOp(Ops.KERNEL, src="+', '.join( \
|
|
|
|
|
([f"({ctx[u.src[0]]},)"] if len(u.src) == 1 else ([f"({', '.join([ctx[x] for x in u.src])})"] if len(u.src) > 1 else []))) + \
|
|
|
|
|
f", arg=Kernel({ctx[u.arg.ast]}(), {u.arg.metadata})"+(f", tag={repr(u.tag)}" if u.tag is not None else "")+")"),
|
|
|
|
|
(UPat(GroupOp.All, name="u"), lambda ctx,u: "UOp("+', '.join([str(u.op), str(u.dtype)] + \
|
|
|
|
|
([f"({ctx[u.src[0]]},)"] if len(u.src) == 1 else ([f"({', '.join([ctx[x] for x in u.src])})"] if len(u.src) > 1 else [])) + \
|
|
|
|
|
([f"arg={repr(u.arg)}"] if u.arg is not None else []) + ([f"tag={repr(u.tag)}"] if u.tag is not None else []))+")"),
|
|
|
|
|
])
|
|
|
|
|
|
|
|
|
|
@Context(SPEC=0)
|
|
|
|
|
def pyrender(ast:UOp) -> str:
|
|
|
|
|
cmap = ast.get_consumer_map()
|
|
|
|
|
to_render = set()
|
|
|
|
|
for u in ast.toposort():
|
|
|
|
|
if u.op is Ops.STORE: to_render.add(u.src[1])
|
|
|
|
|
if len(cmap[u]) == 1 and u.op not in {Ops.DEFINE_GLOBAL, Ops.LOAD} or u.op in {Ops.CONST}: continue
|
|
|
|
|
uops = list(ast.toposort())
|
|
|
|
|
ret: dict[str, str] = {}
|
|
|
|
|
r: dict[UOp, str] = {}
|
|
|
|
|
|
|
|
|
|
not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
|
|
|
|
|
always_rendered = {Ops.DEFINE_GLOBAL, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.KERNEL, Ops.WHERE}
|
|
|
|
|
to_render: set[UOp] = {ast}
|
|
|
|
|
for u in uops:
|
|
|
|
|
if u.op in {Ops.SINK}:
|
|
|
|
|
for s in u.src: to_render.add(s)
|
|
|
|
|
if u.op is Ops.STORE: to_render.add(u.src[1])
|
|
|
|
|
if u.op in {Ops.REDUCE, Ops.REDUCE_AXIS}: to_render.add(u.src[0])
|
|
|
|
|
if u.op in not_rendered: continue
|
|
|
|
|
# checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
|
|
|
|
|
if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
|
|
|
|
|
to_render.add(u)
|
|
|
|
|
ret: list[str] = []
|
|
|
|
|
rep: dict[UOp, UOp] = {}
|
|
|
|
|
for u in ast.toposort():
|
|
|
|
|
if u not in to_render: continue
|
|
|
|
|
ret.append(f"c{len(ret)} = {u.substitute(rep).render(simplify=False, pm=pm_pyrender+renderer)}")
|
|
|
|
|
rep[u] = UOp(Ops.NOOP, arg=f"c{len(ret)-1}")
|
|
|
|
|
return "\n".join(ret[0:-1] + ["ast ="+ret[-1].split("=", 1)[1]])
|
|
|
|
|
|
|
|
|
|
kernels: dict[UOp, tuple[str, str]] = {}
|
|
|
|
|
for i,u in enumerate(uops):
|
|
|
|
|
if u.op is Ops.KERNEL:
|
|
|
|
|
if u.arg.ast not in kernels:
|
|
|
|
|
kernels[u.arg.ast] = (f"k{len(kernels)}", f"def k{len(kernels)}():\n " + pyrender(u.arg.ast).replace('\n', '\n ') + "\n return ast\n\n")
|
|
|
|
|
r[u.arg.ast] = kernels[u.arg.ast][0]
|
|
|
|
|
ren = cast(str, pm_pyrender.rewrite(u, ctx=r))
|
|
|
|
|
assert isinstance(ren, str)
|
|
|
|
|
#if u.tag is not None: ren += f".rtag({u.tag})"
|
|
|
|
|
if u not in to_render: r[u] = ren
|
|
|
|
|
else:
|
|
|
|
|
r[u] = f"c{i}" if u is not uops[-1] else "ast"
|
|
|
|
|
ret[r[u]] = ren
|
|
|
|
|
return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {v}" for k,v in ret.items()])
|
|
|
|
|
|
|
|
|
|
def eval_pyrender(code:str) -> UOp:
|
|
|
|
|
from tinygrad.dtype import AddrSpace
|
|
|
|
|
from tinygrad.codegen.opt import Opt, OptOps
|
|
|
|
|
from tinygrad.schedule.rangeify import BufferizeOpts, Kernel
|
|
|
|
|
lcls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel,
|
|
|
|
|
"Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace}
|
|
|
|
|
exec(code, None, lcls)
|
|
|
|
|
return lcls['ast']
|
|
|
|
|
|
|
|
|
|
def test_pyrender(test_ast:UOp, check_parents=True):
|
|
|
|
|
code = pyrender(test_ast)
|
|
|
|
|
ast:UOp = eval_pyrender(code)
|
|
|
|
|
if ast is not test_ast:
|
|
|
|
|
if check_parents:
|
|
|
|
|
for u in test_ast.toposort(): test_pyrender(u, check_parents=False)
|
|
|
|
|
raise RuntimeError(f"PYRENDER ISSUE:\nSTR MATCH: {str(test_ast) == str(ast)}\nUOP:\n{test_ast}\nPRODUCED:\n{ast}\nCODE:\n{code}")
|
|
|
|
|
return code
|
|
|
|
|
|
|
|
|
|
# *** what was symbolic.py ***
|
|
|
|
|
|
|
|
|
|
geohot
