sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

scheduler: don't trade complexity for speed (#8370)

Browse Source 
* scheduler: don't trade complexity for speed

* don't need is_scheduled

* make those tests real world

* graph_rewrite dedup
This commit is contained in:
qazal
2024-12-22 03:30:51 +02:00
committed by GitHub
parent 991b91d4d6
commit 88bc51385c
2 changed files with 21 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_schedule.py
+11 -21
View File
@@ -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)
tinygrad/engine/schedule.py
+10 -11
View File
@@ -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: