From 88bc51385ca3946c33cfc35fe48dd8fb78eb4758 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sun, 22 Dec 2024 03:30:51 +0200 Subject: [PATCH] scheduler: don't trade complexity for speed (#8370) * scheduler: don't trade complexity for speed * don't need is_scheduled * make those tests real world * graph_rewrite dedup --- test/test_schedule.py | 32 +++++++++++--------------------- tinygrad/engine/schedule.py | 21 ++++++++++----------- 2 files changed, 21 insertions(+), 32 deletions(-) diff --git a/test/test_schedule.py b/test/test_schedule.py index 5bde4bdfdb..dadca5b2c1 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -16,7 +16,7 @@ from tinygrad.shape.view import View from tinygrad.ops import PatternMatcher, UOp, Ops, UPat, graph_rewrite, track_rewrites, view_supported_devices from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, GlobalCounters, flatten, getenv, SPLIT_REDUCEOP, unwrap, prod, Context from tinygrad.codegen.kernel import Kernel, verify_ast -from tinygrad.engine.schedule import BUF_LIMIT, ScheduleContext, ScheduleItem, create_schedule, view_right, view_left, do_realize, remove_movement_ops +from tinygrad.engine.schedule import BUF_LIMIT, ScheduleItem, create_schedule, view_right, view_left, remove_movement_ops from tinygrad.engine.realize import CompiledRunner, get_runner, run_schedule from extra.models.llama import precompute_freqs_cis @@ -247,6 +247,11 @@ class TestSchedule(unittest.TestCase): run_schedule(sched) self.assertIsNot(a.lazydata.realized, b.lazydata.realized) + def test_dedup_outputs(self): + a = Tensor.full((4, 4), 1.).contiguous().realize() + b = Tensor.full((4, 4), 1.).contiguous().realize() + check_schedule([a+b, a+b], 1) + def test_fold_double_unary(self): y = Tensor.empty(2) out = y.sum(keepdim=True).sqrt().__neg__() @@ -1972,29 +1977,14 @@ class TestView(unittest.TestCase): run_schedule(sched) np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:]) -@track_rewrites(named=True) -def big_graph_rewrite(big_graph:UOp, ctx) -> UOp: return graph_rewrite(big_graph, do_realize, ctx) class TestBigGraph(unittest.TestCase): def test_sink_childless_const(self): - x = UOp.const(dtypes.int, 0) - big_graph = big_graph_rewrite(x.sink(), ctx:=ScheduleContext()) - self.assertIs(big_graph, UOp(Ops.NOOP)) - self.assertEqual(len(ctx.realizes), 0) - - def test_sink_childless_const_alt(self): - x = UOp.const(dtypes.int, 0) - y = UOp(Ops.VIEW, dtypes.int, (UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp.const(dtypes.int, 0)), ShapeTracker.from_shape(())) - big_graph = big_graph_rewrite(UOp.sink(x, y), ctx:=ScheduleContext()) - self.assertIs(big_graph, UOp(Ops.NOOP)) - self.assertEqual(len(ctx.realizes), 0) + x = Tensor(0) + check_schedule(x, 0) def test_sink_childless_const_alt_expanded(self): - # this is a real STORE of CONST (post expand) - y = UOp(Ops.VIEW, dtypes.int, (UOp.new_buffer(Device.DEFAULT, 1, dtypes.int), UOp.const(dtypes.int, 0)), ShapeTracker.from_shape(())) - out = UOp(Ops.VIEW, dtypes.int, (UOp.new_buffer(Device.DEFAULT, 2, dtypes.int), y.reshape((1,)).expand((2,)).contiguous(),), ShapeTracker.from_shape((2,))) - big_graph = big_graph_rewrite(out.sink(), ctx:=ScheduleContext()) - self.assertIs(big_graph, out.sink()) - self.assertEqual(len(ctx.realizes), 1) + x = Tensor.zeros(4, 4).contiguous() + check_schedule(x, 1) tensor_const_pm = PatternMatcher([ (UPat(Ops.VIEW, src=(UPat(Ops.DEVICE), UPat(Ops.CONST, src=()))), lambda: True), @@ -2091,7 +2081,7 @@ class TestConst(unittest.TestCase): # ** part 3: Tensor variable bindings - @unittest.expectedFailure # TODO: should schedule assert if you try to realize a Variable? + #@unittest.expectedFailure # TODO: should schedule assert if you try to realize a Variable? def test_var_schedule(self): vv = UOp.variable("a", 0, 10).bind(1) a = Tensor(vv) diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index dd1d1ed004..7d2b23a913 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -134,6 +134,7 @@ def is_constant(u:UOp): return u.op is Ops.VIEW and len(u.src) == 2 and u.src[1] def to_uop(buf:UOp, ctx:ScheduleContext, cache:dict[UOp, UOp]) -> UOp: if (r:=cache.get(buf)) is not None: return r + if buf.op is Ops.SINK: return UOp.sink(*[to_uop(x, ctx, cache) for x in buf.src]) # shapeless op is passthrough # realized is passthrough # constants are passthrough @@ -525,8 +526,9 @@ def fold_img_cast(ctx:ScheduleContext, xb:UOp, view:UOp, b:UOp, to_cast:UOp, **k del ctx.realizes[b] return to_cast.view(unwrap(view.st)) -def init_big_graph(sink:UOp) -> UOp|None: - new_src = tuple(x.base for x in sink.src if is_scheduled(x.base)) +def init_big_graph(ctx:ScheduleContext, sink:UOp) -> UOp|None: + new_src = tuple(x.base for x in sink.src if x.base.realized is None and not is_constant(x.base)) + for x in new_src: realize(ctx, x.buf_uop, x) return None if new_src == sink.src else UOp(Ops.NOOP) if len(new_src) == 0 else UOp.sink(*new_src) do_realize = PatternMatcher([ @@ -588,19 +590,16 @@ remove_movement_ops = PatternMatcher([(UPat(GroupOp.Movement, name="x"), lambda @track_rewrites(named=True) def create_schedule_with_vars(outs:list[UOp], skip_check:bool=not __debug__) -> tuple[list[ScheduleItem], dict[Variable, int]]: if not skip_check: type_verify(list(UOp.sink(*outs).toposort), extra_spec=tensor_uop_spec) - if len(outs:=dedup(x.base for x in outs if x.base.realized is None and not x.base.is_unrealized_const())) == 0: return [], {} - # create the big graph - ctx = ScheduleContext() - cache: dict[UOp, UOp] = {} # to_uop is removing (many) of the movement ops - for u in (big_graph:=UOp.sink(*(to_uop(x, ctx, cache) for x in outs))).src: ctx.realizes[u.buf_uop] = u - big_graph = graph_rewrite(big_graph, remove_movement_ops+ops_folding+do_realize, ctx) - big_graph = graph_rewrite(big_graph, merge_bufs, ctx) + sink = to_uop(UOp.sink(*outs), ctx:=ScheduleContext(), cache={}) + # const folding and fusion + sink = graph_rewrite(sink, remove_movement_ops+ops_folding+do_realize, ctx) + sink = graph_rewrite(sink, merge_bufs, ctx) # create the scheduler context - graph_rewrite(big_graph, create_ctx, ctx) + graph_rewrite(sink, create_ctx, ctx) # group realizes into kernels store_groups = group_realizes(ctx) - graph_rewrite(big_graph, break_sched, ctx) + graph_rewrite(sink, break_sched, ctx) # preschedule realize groups prescheduled: list[ScheduleItem] = [] for store_uops in store_groups: