geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

cleanup ast rewriter [pr] (#9012)

Browse Source
This commit is contained in:
qazal
2025-02-10 19:07:59 +01:00
committed by GitHub
parent 586e48d696
commit bf0bcb2d5a
2 changed files with 11 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_uops.py
+2 -2
View File
@@ -10,7 +10,7 @@ from tinygrad.device import Buffer, Device
from tinygrad.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
from tinygrad.spec import spec
from tinygrad.renderer import ProgramSpec
from tinygrad.engine.schedule import to_si
from tinygrad.engine.schedule import fix_kernel_ops
from tinygrad.engine.realize import CompiledRunner, lower_schedule_item, get_kernel
from tinygrad.codegen.linearize import linearize_uop
from tinygrad.codegen.rewriter import full_graph_rewrite, sym
@@ -487,7 +487,7 @@ class TestIndexingOrdering(unittest.TestCase):
class TestUPatHelpers(unittest.TestCase):
def test_location(self):
self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "rewriter.py")
self.assertEqual(to_si.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "schedule.py")
self.assertEqual(fix_kernel_ops.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "schedule.py")
self.assertEqual(spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "ops.py")
with self.assertRaises(AssertionError): # TODO: location UPat files created in test/*?
test_upat = UPat(Ops.CONST, dtypes.bool)
tinygrad/engine/schedule.py
+9 -10
View File
@@ -345,19 +345,18 @@ def check_load_st(glbl:UOp, view:UOp):
raise RuntimeError("self operand of augmented assign must be contiguous.\nhelp: consider using .contiguous():\n"
+colored(" - a += a.T\n", "red")+colored(" + a += a.T.contiguous()", "green"))
to_si = PatternMatcher([
# BUFFER -> DEFINE_GLOBAL
fix_kernel_ops = PatternMatcher([
# BUFFER becomes DEFINE_GLOBAL
(UPat(Ops.BUFFER, name="x"), _append_buf),
# simplify and unbind the final VIEWs
# BIND in shapetracker becomes DEFINE_VAR
(UPat(Ops.VIEW, name="x"), _append_st_vars),
# don't need SINK on COPY or BUFFER_VIEW
# remove SINK from COPY and BUFFER_VIEW
(UPat(Ops.SINK, src=(UPat.store(UPat.var("b"), UPat(), UPat((Ops.COPY, Ops.BUFFER_VIEW), name="x")),)), lambda b,x: x.replace(src=(b, *x.src))),
# don't need contiguous or assign anymore
# remove CONTIGUOUS/ASSIGN/DEVICE
(UPat(Ops.CONTIGUOUS, src=(UPat.var("x"),)), lambda x: x),
(UPat(Ops.ASSIGN, src=(UPat(), UPat.var("x"),)), lambda x: x),
# don't need DEVICE anymore
(UPat(Ops.VIEW, name="view", src=(UPat(Ops.DEVICE),)), lambda view: view.replace(src=())),
# once images are loaded they become the base dtype
# no ImageDType after load
(UPat(GroupOp.All-{Ops.DEFINE_GLOBAL}, name="x"), lambda x: x.replace(dtype=x.dtype.base) if isinstance(x.dtype, ImageDType) else None),
# if this kernel also assigns to the loaded buffer, ensure we can index it correctly
(UPat(Ops.LOAD, src=(UPat.var("glbl"), UPat.var("view"))), check_load_st),
@@ -371,11 +370,11 @@ unbind_vars = PatternMatcher([(UPat(Ops.BIND, name="bind", src=(UPat.var("var"),
def schedule_uop(pre:UOp, ctx:ScheduleContext) -> ScheduleItem:
# unbind_vars + push views to edges
sink = graph_rewrite(graph_rewrite(pre, unbind_vars+view_left, ctx=ctx.var_vals), view_right)
# remove extra uops from SINK + substitue BUFFER with DEFINE_GLOBAL
ast = graph_rewrite(sink, to_si, si_ctx:=KernelContext(ctx.var_vals))
# fix_kernel_ops
sink = graph_rewrite(sink, fix_kernel_ops, si_ctx:=KernelContext(ctx.var_vals))
# NOTE: we only add the metadata for fused tensors
metadata = tuple(dedup(m for x in pre.toposort if x.op is not Ops.BUFFER and (m:=ctx.ops_metadata.get(x)) is not None))
return ScheduleItem(ast, tuple(u.buffer for u in si_ctx.bufs), metadata)
return ScheduleItem(sink, tuple(u.buffer for u in si_ctx.bufs), metadata)
PROCESS_REPLAY_CAPTURE:dict[str, bytes] = {}
if CAPTURE_PROCESS_REPLAY: