diff --git a/test/test_scan.py b/test/test_fold.py similarity index 87% rename from test/test_scan.py rename to test/test_fold.py index 1b91b1aba9..29922b316f 100644 --- a/test/test_scan.py +++ b/test/test_fold.py @@ -8,10 +8,10 @@ class TestScan(unittest.TestCase): a_red = a.sum(axis=1) np.testing.assert_allclose(a_red.numpy(), a.numpy().sum(axis=1), atol=1e-6) - def test_scan_add(self): + def test_fold_add(self): a = Tensor.randn(10, 10).realize() init = Tensor.zeros(10, 1).contiguous() - a_red = (a+init).scan(init).reshape(10) + a_red = (a+init).fold(init).reshape(10) np.testing.assert_allclose(a_red.numpy(), a.numpy().sum(axis=1), atol=1e-6) if __name__ == '__main__': diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 5b904efee3..3d20c20c38 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -308,7 +308,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp): if len(reduce_range) == 0: return ret return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)).index(UOp.const(dtypes.int, 0)) -def scan_to_store(x:UOp): +def fold_to_store(x:UOp): _, acc, ranges = x.src[0], x.src[1], x.src[2:] assert acc.op is Ops.INDEX buf = acc.src[0] @@ -319,8 +319,8 @@ def scan_to_store(x:UOp): pm_reduce = PatternMatcher([ # REDUCE -> DEFINE_ACC+STORE (UPat(Ops.REDUCE, name="red"), reduce_to_acc), - # SCAN -> STORE - (UPat(Ops.SCAN, name="x"), scan_to_store), + # FOLD -> STORE + (UPat(Ops.FOLD, name="x"), fold_to_store), # tensor core built in accumulate (UPat(Ops.WMMA, name="wmma") + UPat.var("add"), lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)), diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index f23445578a..91cafbf48d 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -2,7 +2,8 @@ from __future__ import annotations import math, itertools from collections import defaultdict from typing import cast, Final -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp +from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos from tinygrad.device import Buffer from tinygrad.dtype import dtypes, ImageDType from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten @@ -12,10 +13,6 @@ from tinygrad.renderer import Renderer remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) -# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters -axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3, - AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5} - class Scheduler: def __init__(self, ast:UOp, ren:Renderer): self.ast, self.ren = ast, ren diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index fcafc7544c..65c03d9977 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -67,7 +67,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp): new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None) if x in ctx.range_map: # for scan we use the output ranges on the 2nd arg - new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][int(x.op is Ops.SCAN and i == 1)]) if i in realized_ranges]) + new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][int(x.op is Ops.FOLD and i == 1)]) if i in realized_ranges]) new_srcs.append(new_src) # NOTE: do we need this? return x.replace(src=tns) if x.src != (tns:=tuple(new_srcs)) else None @@ -107,7 +107,7 @@ pm_apply_rangeify = PatternMatcher([ # REDUCE_AXIS -> REDUCE (UPat(Ops.REDUCE_AXIS, name="x"), convert_reduce_axis_to_reduce_with_ranges), # SCAN -> SCAN (with new ranges) - (UPat(Ops.SCAN, name="x"), add_ranges_to_scan), + (UPat(Ops.FOLD, name="x"), add_ranges_to_scan), # PAD -> WHERE (UPat(Ops.PAD, name="x"), convert_pad_to_where_to_keep_behavior_local), # add third op to assign @@ -255,8 +255,8 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # REDUCE_AXIS creates ranges for the axes it is reducing if x.op is Ops.REDUCE_AXIS: rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape))) - if x.op is Ops.SCAN: - rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if resolve(x.src[1].shape[i] == 1) else r \ + if x.op is Ops.FOLD: + rngs = tuple(rctx.new_range(s, axistype=AxisType.FOLD) if resolve(x.src[1].shape[i] == 1) else r \ for i,(r,s) in enumerate(zip(rngs, x.src[0].shape))) if debug: diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 743937bce6..61dd2555d9 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1509,8 +1509,8 @@ class Tensor(OpMixin): # ***** reduce ops ***** - def scan(self, init:Tensor) -> Tensor: - return self._apply_uop(UOp.scan, init) + def fold(self, init:Tensor) -> Tensor: + return self._apply_uop(UOp.fold, init) def _reduce(self, op:Ops, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor: axis = tuple(self._resolve_dim(x) for x in (range(self.ndim) if axis is None else make_tuple(axis, 1))) diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index fcc39692d4..3d708d1999 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -93,7 +93,7 @@ class Ops(FastEnum): REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto() # scan - SCAN = auto() + FOLD = auto() # errors/placeholders REWRITE_ERROR = auto(); SENTINEL = auto() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 0fd1ede7a2..4c602784ee 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -13,14 +13,19 @@ if TYPE_CHECKING: class AxisType(Enum): def __repr__(self): return str(self) - GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702 + GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); FOLD = auto() # noqa: E702 + UPCAST = auto(); UNROLL = auto() # noqa: E702 THREAD = auto() axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u", - AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"} + AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.FOLD: "F", AxisType.UNROLL: "r"} axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE", - AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"} + AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.FOLD: "red", AxisType.UNROLL: "magenta"} -range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.SCAN: 2} +# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters +axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3, + AxisType.GROUP_REDUCE: 2, AxisType.FOLD: 4, AxisType.REDUCE: 5, AxisType.UNROLL: 6} + +range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.FOLD: 2} # https://en.wikipedia.org/wiki/Identity_element def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt) @@ -214,7 +219,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,) # shape of init - case Ops.SCAN: + case Ops.FOLD: return self.src[1]._shape # passthrough ops @@ -442,7 +447,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid) def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs) - def scan(self, init:UOp): return UOp(Ops.SCAN, self.dtype, src=(self, init)) + def fold(self, init:UOp): return UOp(Ops.FOLD, self.dtype, src=(self, init)) def is_contiguous(self): # TODO: this is is_realized diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 2c4086fc15..3d2ca6348d 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -107,6 +107,9 @@ _tensor_spec = PatternMatcher([ # REDUCE with an outerworld range (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])), + # FOLD in Tensor graph is between two shaped objects + (UPat(Ops.FOLD, src=(UPat(), UPat())), lambda: True), + # AFTER if things were kernelized (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True), ])+movement_ops+shared_spec @@ -172,8 +175,9 @@ kernel_spec = PatternMatcher([ # bufferize can be on anything (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True), - # reduce must be on ranges + # reduce/fold 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:])), + (UPat(Ops.FOLD, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[2:])), ])+movement_ops+shared_codegen_spec+shared_spec # ***** UOp spec in linearized programs ***** diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 00a6095128..9b24f7eb6d 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -17,7 +17,7 @@ from tinygrad.dtype import dtypes uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", Ops.DEFINE_GLOBAL:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B", Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", - Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.SCAN: "#FF7B7B", + Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.FOLD: "#FF7B7B", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",