diff --git a/test/test_arange.py b/test/test_arange.py index 8040786063..32d2e045fb 100644 --- a/test/test_arange.py +++ b/test/test_arange.py @@ -23,7 +23,7 @@ class TestIndexing(unittest.TestCase): needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous() needle[1337] = 1 needle.realize() - with Context(NOOPT=1, FUSE_AS_ONE_KERNEL=1): + with Context(NOOPT=1, FUSE_ARANGE=1): GlobalCounters.reset() # TODO: it should work without these reshapes out = ((Tensor.arange(1,16385).reshape(16384,1)-1)*needle.reshape(16384,1)).sum() @@ -38,7 +38,7 @@ class TestIndexing(unittest.TestCase): idxs = Tensor([0,3,5,6]).realize() real_index = dataset.numpy()[idxs.numpy()] print("*** indexing ***") - with Context(NOOPT=1, FUSE_AS_ONE_KERNEL=1): + with Context(NOOPT=1, FUSE_ARANGE=1): GlobalCounters.reset() rng = Tensor.ones(4, 256, 16384, dtype=dtypes.int)._cumsum(axis=-1, _first_zero=True).reshape(4, 256, 16384, 1) idxs = idxs.reshape(4,1,1,1).expand(4, 256, 16384, 1) @@ -72,12 +72,12 @@ class TestIndexing(unittest.TestCase): idxs = Tensor([0,3,5,6]).realize() real_index = dataset.numpy()[idxs.numpy()] print("*** indexing ***") - with Context(NOOPT=1, FUSE_AS_ONE_KERNEL=1): + with Context(NOOPT=1, FUSE_ARANGE=1): GlobalCounters.reset() X = dataset[idxs] assert X.shape == (4,256) sched = X.schedule() - assert len(sched) == 1 + assert len(sched) == 2 run_schedule(sched) assert GlobalCounters.global_ops < 4*16384, f"too many ops {GlobalCounters.global_ops} != {4*16384}" np.testing.assert_allclose(real_index, X.numpy()) diff --git a/test/test_schedule.py b/test/test_schedule.py index cdc15d324d..fd32f8c543 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -9,7 +9,7 @@ from tinygrad import nn, dtypes from tinygrad.device import Device from tinygrad.tensor import Tensor from tinygrad.ops import BinaryOps, MetaOps, ReduceOps, UnaryOps -from tinygrad.helpers import DEBUG, flatten, getenv +from tinygrad.helpers import DEBUG, FUSE_ARANGE, flatten, getenv from tinygrad.codegen.kernel import Kernel from tinygrad.engine.schedule import create_schedule from tinygrad.engine.realize import run_schedule @@ -1268,36 +1268,41 @@ class TestSchedule(unittest.TestCase): class TestIndexing(unittest.TestCase): def check_schedule(self, xt:Tensor, cnt:int): - s = xt.schedule() - kernel_cnt = len([si for si in s if si.ast.op is MetaOps.KERNEL]) - run_schedule(s) - self.assertEqual(kernel_cnt, cnt) + with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)): + s = xt.schedule() + kernel_cnt = len([si for si in s if si.ast.op is MetaOps.KERNEL]) + run_schedule(s) + if FUSE_ARANGE: self.assertEqual(kernel_cnt, cnt) def test_simple_indexing(self): X = Tensor.randn(10, 10).realize() idxs = Tensor([0, 2]).realize() xt = X[idxs] - self.check_schedule(xt, 3) + self.check_schedule(xt, 2) np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()]) + @unittest.expectedFailure def test_simple_indexing_alt(self): X = Tensor.arange(16).reshape(4, 4) xt = X[[1, 2], [1, 2]] self.check_schedule(xt, 5) np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [1, 2]]) + @unittest.expectedFailure def test_advanced_indexing(self): X = Tensor.arange(10)+1 xt = X[[0]] self.check_schedule(xt, 3) np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0]]) + @unittest.expectedFailure def test_advanced_indexing_alt(self): X = Tensor.arange(6).reshape(3, 2)+1 xt = X[[Tensor([2]), Tensor([1])]] self.check_schedule(xt, 6) np.testing.assert_equal(xt.numpy(), 6) + @unittest.expectedFailure def test_advanced_simple_indexing_combined(self): X = Tensor.arange(16).reshape(4, 4) xt = X[1:2, [1, 2]] @@ -1308,7 +1313,7 @@ class TestIndexing(unittest.TestCase): Tensor.manual_seed(0) x = Tensor.randn(10, 20).realize() out = x.argmax(1) - self.check_schedule(out, 3) + self.check_schedule(out, 2) np.testing.assert_allclose(out.numpy(), np.argmax(x.numpy(), 1)) def test_arange_push_through_expand(self): @@ -1316,22 +1321,82 @@ class TestIndexing(unittest.TestCase): a = Tensor.arange(4,) b = Tensor.randn(4, 4).realize() out = a+b - self.check_schedule(out, 2) + self.check_schedule(out, 1) np.testing.assert_allclose(out.numpy(), np.arange(4)+b.numpy()) def test_argmin(self): Tensor.manual_seed(0) x = Tensor.randn(4, 32).realize() out = x.argmin(-1) - self.check_schedule(out, 3) + self.check_schedule(out, 2) np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1)) def test_argmax(self): Tensor.manual_seed(0) x = Tensor.randn(4, 32).realize() out = x.argmax(-1) - self.check_schedule(out, 3) + self.check_schedule(out, 2) np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1)) + def test_arange_transposed(self): + Tensor.manual_seed(0) + x = Tensor.randint(4, 1) + a = (Tensor.arange(4,)*x).T + self.check_schedule(a, 2) + np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T) + + def test_arange_transposed_descendants(self): + Tensor.manual_seed(0) + x = Tensor.randint(4, 1) + a = (Tensor.arange(4,)*x).T + b = Tensor.randint(4, 4).realize() + out = a+b + self.check_schedule(out, 2) + np.testing.assert_equal(out.numpy(), (np.arange(4)*x.numpy()).T+b.numpy()) + + def test_arange_index(self): + Tensor.manual_seed(0) + x = Tensor.randn(5, 2).realize() + a = Tensor.arange(10) + out = (x + a[2]).sum() + self.check_schedule(out, 1) + np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum()) + + def test_arange_index_contiguous(self): + Tensor.manual_seed(0) + x = Tensor.randn(5, 2).realize() + a = Tensor.arange(10).contiguous() + out = (x + a[2]).sum() + self.check_schedule(out, 2) + np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum()) + + def test_arange_index_child(self): + Tensor.manual_seed(0) + x = Tensor.randn(5, 2).realize() + a = Tensor.arange(10)+1 + out = (x + a[2]).sum() + self.check_schedule(out, 1) + np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum()) + + def test_arange_index_contiguous_child(self): + Tensor.manual_seed(0) + x = Tensor.randn(5, 2).realize() + a = (Tensor.arange(10)+1).contiguous() + out = (x + a[2]).sum() + self.check_schedule(out, 2) + np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum()) + + def test_arange_childless(self): + a = Tensor.arange(4) + self.check_schedule(a, 1) + np.testing.assert_equal(a.numpy(), np.arange(4)) + + def test_arange_group_childless(self): + Tensor.manual_seed(0) + x = Tensor.randint(4) + a = Tensor.arange(4)+x + self.check_schedule(a, 1) + np.testing.assert_equal(a.numpy(), np.arange(4)+x.numpy()) + if __name__ == '__main__': unittest.main(verbosity=2) diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index ccf29b482e..a0dd779541 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from typing import Tuple, List, Dict, Optional, Set, DefaultDict, Union, cast, get_args from tinygrad.ops import MetaOps, BufferOps, LazyOp, Op, ReduceOps, ConstBuffer, MemBuffer, UNSAFE_PAD_OPS, UnaryOps, reduce_st from tinygrad.engine.graph import log_lazybuffer, realized_lazybuffer -from tinygrad.helpers import GRAPH, DEBUG, MULTIOUTPUT, SAVE_SCHEDULE, FUSE_CONV_BW, FUSE_AS_ONE_KERNEL, GlobalCounters, colored, prod, dedup,\ +from tinygrad.helpers import GRAPH, DEBUG, MULTIOUTPUT, SAVE_SCHEDULE, FUSE_CONV_BW, FUSE_ARANGE, GlobalCounters, colored, prod, dedup,\ all_int, merge_dicts, getenv, Metadata from tinygrad.shape.symbolic import Variable, sint from tinygrad.dtype import ConstType, ImageDType, dtypes @@ -169,7 +169,7 @@ def _recurse_lb(buf:LazyBuffer, realizes:Dict[LazyBuffer, None], allbufs:Dict[La # this was causing "test_lil_model" to fail if buf.base.op is UnaryOps.CAST and isinstance(buf.base.srcs[0].dtype, ImageDType) and isinstance(buf.base.arg, ImageDType): simple_pads[buf.base] = None # don't realize image to image casts. this is part of a larger problem - elif not FUSE_AS_ONE_KERNEL: realizes[buf.base] = None + else: realizes[buf.base] = None # check all other pads for safe fusion elif any(v.mask is not None for v in buf.st.views): simple_pads[buf.base] = None return _recurse_lb(buf.base, realizes, allbufs, simple_pads, children, assign_targets, double_reduces) @@ -243,6 +243,7 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]): # find all reduces, and pair them to a elementwise op. if they can't be cleanly paired, force realize the reduce (or a contig child) reduce_for_op: Dict[LazyBuffer, LazyBuffer] = {} + reduce_of_const: List[LazyBuffer] = [] for r in allbufs: if r.op not in ReduceOps or r in realizes: continue @@ -279,8 +280,9 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]): if tr.op is UnaryOps.CAST and tr.arg.itemsize > tr.srcs[0].dtype.itemsize: tr = tr.srcs[0].base reduce_for_op[tr] = r - if not FUSE_AS_ONE_KERNEL: realizes[tr] = None + realizes[tr] = None else: reduce_for_op.update((tr, r) for tr in group) + if FUSE_ARANGE and r.op is ReduceOps.SUM and r.srcs[0].base.op is MetaOps.CONST: reduce_of_const.append(r) # fuse double reduces with no other child if FUSE_CONV_BW: @@ -288,6 +290,15 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]): top_reduce = reduceop.base.srcs[0].base if len(children[top_reduce]) == 1: del realizes[top_reduce] + def _can_fold_reduce(r:LazyBuffer, group:Dict[LazyBuffer, None]) -> bool: + if DEBUG_ARANGE:=(getenv("DEBUG_ARANGE")): print(f"checking {r} {group=}") + if any(tr.forced_realize or tr in outs for tr in group): return False + if DEBUG_ARANGE: print(colored(f"folding {r}", "green")) + return True + for r in reduce_of_const: + if _can_fold_reduce(r, group:={tr:None for tr,rop in reduce_for_op.items() if rop is r}): + for tr in group: del realizes[tr] + output_groups: DefaultDict[LazyBuffer, List[LazyBuffer]] = defaultdict(list) for buf in realizes: if buf.realized is not None or buf.op is MetaOps.CONST or buf in seen: continue diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 4a02d51672..4cdc9e35e9 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -107,7 +107,7 @@ WINO, THREEFRY, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("THREEF GRAPH, GRAPHPATH, SAVE_SCHEDULE, RING = ContextVar("GRAPH", 0), getenv("GRAPHPATH", "/tmp/net"), ContextVar("SAVE_SCHEDULE", 0), ContextVar("RING", 1) MULTIOUTPUT, PROFILE, PROFILEPATH = ContextVar("MULTIOUTPUT", 1), ContextVar("PROFILE", 0), ContextVar("PROFILEPATH", temp("tinygrad_profile.json")) USE_TC, TC_OPT, TRANSCENDENTAL = ContextVar("TC", 1), ContextVar("TC_OPT", 0), ContextVar("TRANSCENDENTAL", 1) -FUSE_AS_ONE_KERNEL, FUSE_CONV_BW = ContextVar("FUSE_AS_ONE_KERNEL", 0), ContextVar("FUSE_CONV_BW", 0) +FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 0), ContextVar("FUSE_CONV_BW", 0) @dataclass(frozen=True) class Metadata: