diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 23c075ca7d..50e51b1dbf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -375,6 +375,8 @@ jobs: - name: Test openpilot model kernel count and gate usage run: | ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx + - name: Test openpilot model with rangeify + run: FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot alt model correctness (float32) run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot fastvits model correctness (float32) diff --git a/test/test_schedule.py b/test/test_schedule.py index 3235d8b148..397136e147 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -358,7 +358,7 @@ class TestSchedule(unittest.TestCase): out1 = r1 + y schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4) reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}] - assert len(reduceops) == (3 if RANGEIFY else 2) + assert len(reduceops) == 2 def test_div_collapse_buffer(self): a = Tensor.full((4,), 4.0).contiguous().realize() diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 743d17d54d..f0125528d3 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -1,10 +1,11 @@ from typing import Any, cast import functools, operator from dataclasses import dataclass, field +from collections import defaultdict from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, graph_rewrite_map from tinygrad.uop.symbolic import sym, symbolic_simple -from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup +from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup, unwrap from tinygrad.schedule.multi import multi_pm from tinygrad.schedule.kernelize import Kernel @@ -96,9 +97,12 @@ remove_contig_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.r # 2. mark all children @dataclass -class ChildrenContext: children: dict[UOp, list[UOp]]|None = None +class ChildrenContext: + children: dict[UOp, list[UOp]]|None = None + realize_roots: defaultdict[UOp, list[UOp]] = field(default_factory=lambda: defaultdict(list)) def extract_children(ctx:ChildrenContext, x:UOp): if ctx.children is not None: return + children_map = x.get_children_map() ctx.children = {} for k,v in children_map.items(): @@ -107,6 +111,16 @@ def extract_children(ctx:ChildrenContext, x:UOp): # NOTE: this gate shouldn't be here if any(x.op is Ops.REDUCE_AXIS for x in k.toposort()) and any(x.op in {Ops.BUFFER, Ops.CONTIGUOUS} for x in k.toposort()): ctx.children[k] = non_sink_children + # if a node is in the toposort of multiple realizes, it will be indexed by different indices and we can bufferize early + # this prevents index_child: "children not making progress" error on big graphs + for r in [u for u in x.toposort() if u.op is Ops.REALIZE and (RANGEIFY<2 or u.arg is None)]: # ignore partial realizes + for u in r.toposort(gate=lambda x: x is not Ops.REALIZE or (RANGEIFY>1 and u.arg is not None)): + ctx.realize_roots[u].append(r) + +def bufferize_early(ctx:ChildrenContext, x:UOp): + # this will also change the sources such that mark_children wont add the children/child uops anymore + new_srcs = [s.realize() if s in unwrap(ctx.children) and len(ctx.realize_roots[s])>1 else s for s in x.src] + return x.replace(src=tuple(new_srcs)) def mark_children(ctx:ChildrenContext, x:UOp): assert ctx.children is not None @@ -116,7 +130,8 @@ def mark_children(ctx:ChildrenContext, x:UOp): pm_children = PatternMatcher([ (UPat(Ops.SINK, name="x"), extract_children), - (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN, Ops.SINK}, name="x"), mark_children), + (UPat(GroupOp.All-{Ops.CHILDREN, Ops.SINK, Ops.REALIZE}, name="x"), bufferize_early), + (UPat(GroupOp.All-{Ops.CHILDREN, Ops.SINK, Ops.REALIZE}, name="x"), mark_children), ]) # *****************