From 7191f88551da3765a989d9958afc74fd0be142c0 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sun, 13 Apr 2025 16:50:18 +0800 Subject: [PATCH] add asserts for KERNEL op ast [pr] (#9868) --- tinygrad/engine/grouper.py | 3 ++- tinygrad/spec.py | 7 ++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tinygrad/engine/grouper.py b/tinygrad/engine/grouper.py index 758d9212e4..50c86fd3f0 100644 --- a/tinygrad/engine/grouper.py +++ b/tinygrad/engine/grouper.py @@ -226,7 +226,8 @@ class Kernel: ast: UOp metadata: tuple[Metadata, ...] = () def __repr__(self): - return f"" + ast_rep = f"SINK{tuple(s.op for s in self.ast.src)}" if self.ast.op is Ops.SINK else repr(self.ast.op) + return f"" @dataclass(frozen=True) class KernelContext: diff --git a/tinygrad/spec.py b/tinygrad/spec.py index d3b03e8b13..205bedd2a3 100644 --- a/tinygrad/spec.py +++ b/tinygrad/spec.py @@ -141,8 +141,13 @@ spec = PatternMatcher([ # *** this is the spec of a Kernel in UOp *** +def validate_kernel(k:UOp): + assert k.arg.ast.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.SINK}, f"must end with SINK/COPY/BUFFER_VIEW {k.arg}" + if k.arg.ast.op is Ops.SINK: assert all(s.op is Ops.STORE for s in k.arg.ast.src), f"SINK must end with STORE {k.arg.ast}" + return True + kernel_spec = buffer_spec+PatternMatcher([ - (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN))), lambda: True), + (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN)), name="k"), validate_kernel), # assign has a buffer and kernel source, it can optionally depend on other assigns (UPat(Ops.ASSIGN, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.KERNEL, Ops.ASSIGN))), lambda: True), (UPat(GroupOp.All-{Ops.SINK}), lambda: False),