diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 239a04c3db..5d6b750cc9 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -1,7 +1,7 @@ from typing import cast import itertools from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC -from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, GroupOp +from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat from tinygrad.uop.spec import type_verify, program_spec, kernel_spec from tinygrad.renderer import Renderer from tinygrad.dtype import dtypes @@ -19,19 +19,13 @@ from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_s from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize -pm_preprocess = PatternMatcher([ - (UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True), - lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)), - (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])), -]) - def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp: if ren is None: ren = Renderer() if SPEC: type_verify(sink, kernel_spec) # preprocess - sink = graph_rewrite(sink, pm_preprocess+pm_mops, name="early movement ops") + sink = graph_rewrite(sink, pm_mops, name="early movement ops") # first we optimize if optimize: diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 376689294d..b85e4a143a 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -18,6 +18,10 @@ sys.setrecursionlimit(10000) pm_mops = PatternMatcher([ (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)), # type: ignore + # move movement ops after AFTER + (UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True), + lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:], tag=None),)+r.src[1:], r.arg, tag=a.tag)), + (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])), ]) # ***************** diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 4b1b4bf2de..9aec0a9912 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -783,6 +783,11 @@ class KernelInfo: @property def function_name(self): return to_function_name(self.name) +@dataclass(frozen=True) +class Kernel: + ast: UOp + metadata: tuple[Metadata, ...] = () + # ******** ops in python ******** def safe_exp2(x): diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index d9dbd32c53..b1b4b62013 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -52,9 +52,12 @@ movement_ops = PatternMatcher([ # inputs to movement ops (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True), (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True), + + # AFTER on Movement Op + (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement),), allow_any_len=True), lambda: True), ]) -tensor_spec = movement_ops+PatternMatcher([ +tensor_spec = PatternMatcher([ # buffer spec (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True), (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d: @@ -106,7 +109,7 @@ tensor_spec = movement_ops+PatternMatcher([ # AFTER if things were kernelized (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True), -])+shared_spec +])+movement_ops+shared_spec # ***** UOp spec in codegen shared between kernel and program ***** @@ -154,10 +157,7 @@ shared_codegen_spec = PatternMatcher([ # ***** UOp spec in kernel graph ***** -kernel_spec = movement_ops+PatternMatcher([ - # AFTER on Movement Op - (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement),), allow_any_len=True), lambda: True), - +kernel_spec = PatternMatcher([ # index is allowed here (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True), @@ -169,7 +169,7 @@ kernel_spec = movement_ops+PatternMatcher([ # reduce must be on ranges (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])), -])+shared_codegen_spec+shared_spec +])+movement_ops+shared_codegen_spec+shared_spec # ***** UOp spec in linearized programs ***** diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 5ca508a759..99bd7f0fd2 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -514,9 +514,8 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([ lambda x,y,alu: UOp(Ops.VECTORIZE, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)), # VECTORIZE of a single element is just that element (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x), - # VECTORIZE void is SINK - (UPat(Ops.VECTORIZE, dtype=dtypes.void, src=UPat(Ops.BARRIER, name='b')), lambda b: b), - (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp(Ops.SINK, dtypes.void, x.src)), + # VECTORIZE void is GROUP + (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp.group(*x.src)), # 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),