From 359ef4722822cdea2b4c57e567cc5abbe0c7267e Mon Sep 17 00:00:00 2001 From: James Vecellio-Grant <159560811+Discountchubbs@users.noreply.github.com> Date: Thu, 22 Oct 2026 10:43:42 -0700 Subject: [PATCH] models: self healing test for tg (#2072) * models: self healing test for tg * rm container * ew * set env * set var (cherry picked from commit 9d029d683dea3fd034659d6bc291e3eb46669d7e) --- .../modeld_v2/convert_upstream_pkl.py | 17 +- openpilot/sunnypilot/modeld_v2/helpers.py | 12 +- .../modeld_v2/plain_artifact_compiler.py | 26 ++- .../tests/test_convert_upstream_pkl.py | 121 +++++++++++++ .../modeld_v2/tests/test_plain_artifact.py | 162 ++++++++++++++++++ 5 files changed, 327 insertions(+), 11 deletions(-) create mode 100644 openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py create mode 100644 openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py diff --git a/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py b/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py index 74d3b860d..9ae62f0f5 100755 --- a/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py +++ b/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py @@ -7,6 +7,7 @@ See the LICENSE.md file in the root directory for more details. """ import argparse import dataclasses +import inspect import itertools import mmap import re @@ -45,6 +46,11 @@ def _address(v): return _root(v.src[0]), add + start, size +def _buffer(u): + from tinygrad.uop.ops import Ops + return u.buffer if u.op is Ops.BUFFER else getattr(u.arg, "buffer", None) + + def _patched(buf, topo) -> tuple[bytearray, dict]: from tinygrad.uop.ops import Ops blob, rows = bytearray(buf.max_numel() * buf.dtype.itemsize), {} @@ -72,7 +78,7 @@ def _evaluate(u, env): table, index = _root(u.src[0].src[0]), _evaluate(u.src[0].src[1], env) return struct.unpack_from(u.dtype.fmt, table.arg, index * u.dtype.itemsize)[0] if table.op is Ops.BINARY and index is not None else None if u.op in GroupOp.ALU and None not in (values := [_evaluate(x, env) for x in u.src]): - return exec_alu(u.op, u.dtype, values, False) + return exec_alu(u.op, u.dtype, values, *([False] if "truncate_output" in inspect.signature(exec_alu).parameters else [])) return None @@ -112,7 +118,7 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]: from tinygrad.uop.ops import Ops submit = jit.captured._linear.src[0].without_after kernels, topo = submit.arg.aux.kernels, list(jit.captured._linear.toposort(enter_calls=True)) - allocs = {str(u.tag): u for u in topo if u.op is Ops.ALLOC and u.tag is not None} + allocs = {str(name): u for u in topo if u.op is Ops.ALLOC and (name := getattr(u, "tag", None) or u.arg.name) is not None} (commands, patched), kernargs = _patched(allocs["cmdbuf_compute_0"], topo), allocs["kernargs_compute_0"] argument_words, pokes = _patched(kernargs, topo)[1], _input_pokes(submit, kernargs) sh = {name: getattr(amd_gpu, f"reg{name}") + amd_gpu.GC_BASE__INST0_SEG0 - pm4.PACKET3_SET_SH_REG_START @@ -158,7 +164,7 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]: args.append({"input": pokes[slot][0]}) elif slot in argument_words: base, at, nbytes = _address(argument_words[slot]) - args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.arg.buffer.nbytes, "dtype": u8}) + args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.max_numel() * base.dtype.itemsize, "dtype": u8}) outs, ins = kernel[6] calls.append({"kind": "program", "name": re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]), "lib": programs[image_alloc], "target": dataclasses.asdict(target), "signature": [[f"data{i}", i, u8, [a.get("nbytes", 1)]] for i, a in enumerate(args)], "global_size": grid, "local_size": local, @@ -166,12 +172,13 @@ def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]: needs_init |= {args[i]["buffer"] for i in ins if "buffer" in args[i] and args[i]["buffer"] not in written} written |= {args[i]["buffer"] for i in outs if "buffer" in args[i]} + assert all(_buffer(b) is not None for b, i in buffers.items() if i in needs_init), "a kernel reads a scratch before any kernel writes it" names = list(jit.captured.expected_names) inputs = [{"slot": slot, "name": names[slot] if slot < len(names) else None, "dtype": repr(u.dtype), "numel": int(u.max_numel()), "device": u.device} for slot, u in enumerate(_prepare_jit_inputs((), kwargs)[0])] header = {"input_names": names, "calls": calls, "returns": {"kind": "NoneType", "tensors": []}, "inputs": inputs, - "buffers": [{"nbytes": b.arg.buffer.nbytes, "device": b.arg.buffer.device, "dtype": u8} - | ({"init": blob.add(b.arg.buffer.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()], + "buffers": [{"nbytes": bb.nbytes if (bb := _buffer(b)) else b.max_numel() * b.dtype.itemsize, "device": bb.device if bb else b.device, "dtype": u8} + | ({"init": blob.add(bb.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()], "check": {"inputs": [[path, store(array), device] for path, array, device in frame["inputs"]], "outputs": [store(array) for array in frame["outputs"]], "output_container": frame["output_container"]}} with sink: diff --git a/openpilot/sunnypilot/modeld_v2/helpers.py b/openpilot/sunnypilot/modeld_v2/helpers.py index f45584a61..7b79c1736 100644 --- a/openpilot/sunnypilot/modeld_v2/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/helpers.py @@ -250,6 +250,12 @@ def _allocate_buffers(header: dict, blob: memoryview, device_map: dict | None, s return buffers +def _fields(what, /, **values): # a tinygrad dataclass built or replaced with the fields it has today: a dropped field is fine, a new required one is not + known = {field.name for field in dataclasses.fields(what)} + values = {name: value for name, value in values.items() if name in known} + return what(**values) if isinstance(what, type) else dataclasses.replace(what, **values) + + def _program_uop(call: dict, lib: bytes): from tinygrad.helpers import Target from tinygrad.uop.ops import KernelInfo, Ops, ProgramInfo, UOp @@ -257,9 +263,9 @@ def _program_uop(call: dict, lib: bytes): for name, compact_slot, dtype, shape in call["signature"]: parameter = UOp.param(call["globals"][compact_slot], _dtype(dtype), math.prod(shape)) image = (shape[0], shape[1]) if len(shape) == 3 and shape[2] == 4 else None - parameters.append(parameter.replace(arg=dataclasses.replace(parameter.arg, name=name, image=image))) - info = ProgramInfo(global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), vars=(), globals=tuple(call["globals"]), - outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=Target(**call["target"])) + parameters.append(parameter.replace(arg=_fields(parameter.arg, name=name, image=image))) + info = _fields(ProgramInfo, global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), globals=tuple(call["globals"]), + outs=tuple(call["outs"]), ins=tuple(call["ins"]), target=_fields(Target, **call["target"])) return UOp(Ops.PROGRAM, src=(UOp(Ops.SINK, arg=KernelInfo(name=call["name"])), UOp(Ops.LINEAR, src=tuple(parameters)), UOp(Ops.SOURCE, arg=""), UOp(Ops.BINARY, arg=lib)), arg=info) diff --git a/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py b/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py index 01dd292e2..db761a054 100644 --- a/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py +++ b/openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py @@ -23,7 +23,7 @@ def record_kernel_linear(): recorded, hcq_compile, compile_linear = [], realize.hcq_compile, jit.compile_linear def recording_compile_linear(*args, **kwargs): - with mock.patch.object(realize, "hcq_compile", lambda linear, *a, **k: (recorded.append(linear), hcq_compile(linear, *a, **k))[1]): + with mock.patch.object(realize, "hcq_compile", lambda linear, *a, **k: (linear.src and recorded.append(linear), hcq_compile(linear, *a, **k))[1]): return compile_linear(*args, **kwargs) with mock.patch.object(jit, "compile_linear", recording_compile_linear): @@ -38,6 +38,18 @@ def capture_frame(jit, kwargs: dict, output_container: str | None = None): return {"inputs": inputs, "outputs": [tensor.numpy().copy() for tensor in outputs], "output_container": output_container}, result +def kernel_signature(ast, elf) -> list: + from tinygrad.dtype import AddrSpace + if hasattr(ast, "kernel_params"): + if any(p.arg.addrspace is AddrSpace.ALU for p in ast.kernel_params): + raise ValueError(f"{elf.name} takes a variable, which plain artifacts don't support") + compact = {slot: i for i, slot in enumerate(ast.arg.globals)} + return [[p.arg.name, compact[p.arg.slot], repr(p.dtype), [int(x) for x in p._shape]] for p in ast.kernel_params] + if not all(isinstance(slot, int) for _, slot, _, _ in elf.signature): + raise ValueError("TinyELF.signature lost its slots and this tinygrad has no UOp.kernel_params: update kernel_signature") + return [[name, slot, repr(dtype), [int(x) for x in shape]] for name, slot, dtype, shape in elf.signature] + + def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, memoryview]: from tinygrad.uop.ops import Ops sink = tempfile.TemporaryFile(dir=".") @@ -51,6 +63,12 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, if uop.op is Ops.PARAM: input_specs.setdefault(uop.arg.slot, {"dtype": repr(uop.dtype), "numel": int(uop.max_numel()), "device": uop.device}) return {"input": uop.arg.slot} + root, offset = uop, 0 + while root.op in (Ops.AFTER, Ops.BITCAST, Ops.SHRINK): + offset, root = offset + (root.src[1].arg * root.src[0].dtype.itemsize if root.op is Ops.SHRINK else 0), root.src[0] + if root.op is Ops.ALLOC: + base = bases.setdefault(id(root), (len(bases), root, root.dtype)) + return {"buffer": base[0], "offset": offset, "nbytes": int(uop.max_numel()) * uop.dtype.itemsize, "dtype": repr(uop.dtype)} buffer = uop.buffer root = buffer._base if buffer._base is not None else buffer base = bases.setdefault(id(root), (len(bases), root, uop.base.dtype if buffer._base is None else uop.dtype)) @@ -63,7 +81,7 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, elf, info = ast.to_elf(), ast.arg global_size, local_size = info.launch_dims({}) calls.append({"kind": "program", "name": elf.name, "lib": blob.add(bytes(elf.lib)), "target": dataclasses.asdict(elf.target), - "signature": [[name, slot, repr(dtype), [int(x) for x in shape]] for name, slot, dtype, shape in elf.signature], + "signature": kernel_signature(ast, elf), "global_size": [int(x) for x in global_size], "local_size": [int(x) for x in local_size], "globals": list(info.globals), "outs": list(info.outs), "ins": list(info.ins), "args": arguments}) outs, ins = info.outs, info.ins @@ -77,7 +95,9 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, ret = jit.captured.ret returns = {"kind": type(ret).__name__, "tensors": [describe(t.uop.base) | {"shape": [int(x) for x in t.shape]} for t in flatten(ret)]} - buffers = [{"nbytes": base.nbytes, "device": base.device, "dtype": repr(dtype)} + from tinygrad.device import Buffer + assert not [base for index, base, _ in bases.values() if index in needs_init and not isinstance(base, Buffer)], "a kernel reads a scratch before writing it" + buffers = [{"nbytes": base.nbytes if isinstance(base, Buffer) else int(base.max_numel()) * dtype.itemsize, "device": base.device, "dtype": repr(dtype)} | ({"init": blob.add(base.as_memoryview())} if index in needs_init else {}) for index, base, dtype in bases.values()] names = list(jit.captured.expected_names) if check is not None: diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py b/openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py new file mode 100644 index 000000000..3deffcfc7 --- /dev/null +++ b/openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py @@ -0,0 +1,121 @@ +""" +Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors. + +This file is part of sunnypilot and is licensed under the MIT License. +See the LICENSE.md file in the root directory for more details. +""" +import os +import re +import sys +import tempfile +import unittest +from unittest import mock +import numpy as np + +os.environ.setdefault("CHECK_OOB", "0") +os.environ.setdefault("GMMU", "0") +if sys.platform == "darwin" and os.environ.get("DEV", "").startswith("MOCKUSB"): # the mock's VRAM is a memfd, which macOS lacks: an unlinked file is the same + from tinygrad.runtime.autogen import libc + + def _memfd_create(name, flags): + fd, path = tempfile.mkstemp() + os.unlink(path) + return fd + libc.memfd_create = _memfd_create + +from openpilot.common.test import OpenpilotTestCase +from openpilot.sunnypilot.modeld_v2 import convert_upstream_pkl +from openpilot.sunnypilot.modeld_v2.convert_upstream_pkl import convert +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, load_model +from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame + +IMG, DESIRE, STATE, OUT = (6, 8, 8), (8,), (4, 16), (1, 32) + + +class TestConvertUpstreamPkl(OpenpilotTestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + if not os.environ.get("TINYGRAD_TESTS"): + raise unittest.SkipTest("TINYGRAD_TESTS not set. skipping converter test") + from tinygrad import Device, Tensor + assert Device.DEFAULT == "AMD" and getattr(Device["AMD"], "is_usb", False), f"{Device.DEFAULT} is not an AMD device behind the USB bridge" + assert (Tensor([1, 2, 3]) * 2).tolist() == [2, 4, 6] + cls.tmp = tempfile.TemporaryDirectory() + cls.rng = np.random.default_rng(0) + cls.weights = Tensor(cls.rng.standard_normal((16, 32)).astype(np.float32)).realize() + cls.bias = Tensor(cls.rng.standard_normal((1, 32)).astype(np.float32)).realize() + cls.input_specs = {"input_imgs": (IMG, "uint8", Device.DEFAULT), "desire": (DESIRE, "float32", Device.DEFAULT), + "features": (STATE, "float32", Device.DEFAULT)} + cls.output_specs = {"outputs": (OUT, "float32", Device.DEFAULT), "next_features": (STATE, "float32", Device.DEFAULT)} + + @classmethod + def tearDownClass(cls): + cls.tmp.cleanup() + super().tearDownClass() + + def model(self, input_imgs, desire, features, output_buffers): + from tinygrad import Tensor + feat = features.mean(axis=0).reshape(1, 16) + input_imgs.float().mean() + desire.sum() + Tensor.realize(output_buffers["outputs"].assign((feat @ self.weights + self.bias).relu()), + output_buffers["next_features"].assign(features[1:].cat(feat * 2))) + + def inputs(self, seed: int) -> dict: + from examples.openpilot.helpers import allocate_inputs + rng = np.random.default_rng(seed) + + def initialize(arrays: dict): + for value in arrays.values(): + value[...] = rng.standard_normal(value.shape) if value.dtype.kind == "f" else rng.integers(0, 256, value.shape) + return allocate_inputs(self.input_specs, initialize) | {"output_buffers": allocate_inputs(self.output_specs)} + + def jit(self): + from tinygrad import TinyJit + run = TinyJit(self.model, prune=True) + kwargs = self.inputs(1) + run(**kwargs) + run(**kwargs) + return run + + def assert_matches_eager(self, plain, seed: int): + fresh, expected = self.inputs(seed), self.inputs(seed) + plain(**fresh) + self.model(**expected) + for name in self.output_specs: + assert np.array_equal(fresh["output_buffers"][name].numpy(), expected["output_buffers"][name].numpy()), name + + def convert(self, jit) -> tuple[dict, memoryview]: + from tinygrad import Device + kwargs = self.inputs(2) + frame, _ = capture_frame(jit, kwargs, "output_buffers") + header, blob = convert(jit, kwargs, frame, Device[Device.DEFAULT].renderer.target) + kernels = jit.captured._linear.src[0].without_after.arg.aux.kernels + assert [call["name"] for call in header["calls"]] == [re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]) for kernel in kernels] + assert all(re.fullmatch(r"\w+", call["name"]) for call in header["calls"]), "kernel names: failure point 1" + assert [(tuple(call["outs"]), tuple(call["ins"])) for call in header["calls"]] == [tuple(kernel[6]) for kernel in kernels], "outs/ins: failure point 1" + bound = {argument["input"] for call in header["calls"] for argument in call["args"] if "input" in argument} + assert bound == {spec["slot"] for spec in header["inputs"]}, f"input slots bound by pokes {sorted(bound)}: failure point 4" + assert [spec["name"] for spec in header["inputs"]][:len(self.input_specs)] == sorted(self.input_specs) + return header, blob + + def test_converted_jit_matches_eager(self): # the submit decoded into plain calls, replayed against its own frame + jit = self.jit() + plain = PlainJit(*self.convert(jit)) + self.assert_matches_eager(plain, 3) + self.assert_matches_eager(plain, 4) + + def test_pickle_round_trip_through_main(self): + from examples.openpilot.helpers import dump_pickle, load_pickle + from tinygrad.engine.realize import lower_and_compile + pkl, out = os.path.join(self.tmp.name, "tiny_tinygrad.pkl"), os.path.join(self.tmp.name, "tiny_plain.pkl") + dump_pickle({"run": self.jit(), "input_specs": self.input_specs, "output_specs": self.output_specs, "metadata": {"name": "tiny"}}, pkl) + + artifact = load_pickle(pkl, out_of_band=True) + artifact["run"].captured._linear = lower_and_compile(artifact["run"].captured._linear) + self.assert_matches_eager(PlainJit(*self.convert(artifact["run"])), 5) + + with mock.patch.object(sys, "argv", ["convert_upstream_pkl", pkl, out]): + convert_upstream_pkl.main() + bundle = load_model(out, cache=False, verify=True) + assert bundle["metadata"] == {"name": "tiny"} and bundle["input_specs"] == self.input_specs and bundle["output_specs"] == self.output_specs + self.assert_matches_eager(bundle["run"], 6) diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py b/openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py new file mode 100644 index 000000000..80bebc119 --- /dev/null +++ b/openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py @@ -0,0 +1,162 @@ +""" +Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors. + +This file is part of sunnypilot and is licensed under the MIT License. +See the LICENSE.md file in the root directory for more details. +""" +import dataclasses +import functools +import inspect +import os +import sys +import tempfile +from unittest import mock +import numpy as np +import unittest + +import openpilot.sunnypilot.modeld_v2.helpers as helpers +from openpilot.common.test import OpenpilotTestCase +from openpilot.sunnypilot.modeld_v2 import compile_warp +from openpilot.sunnypilot.modeld_v2.frame_fusion import frame_part, pack_layout +from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, cache_is_fresh, load_model, read_parts, write_bundle +from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit + +CAM_W, CAM_H, MODEL_W, MODEL_H = 32, 16, 16, 8 +FRAME = f"{CAM_W},{CAM_H},{CAM_W},{CAM_H},{CAM_H // 2},{CAM_W * CAM_H * 3 // 2}" +FRAME_SIZE = CAM_W * CAM_H * 3 // 2 +HOST = {"desire": (8,), "traffic_convention": (2,)} +STATE, OUT = (4, 16), (1, 32) + + +def compile_jit(fn, inputs: dict, output_container: str | None = None) -> tuple[dict, memoryview]: + from tinygrad import TinyJit + run = TinyJit(fn, prune=True) + with record_kernel_linear() as recorded: + run(**inputs) + run(**inputs) + check, _ = capture_frame(run, inputs, output_container) + assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}" + return serialize_jit(run, recorded[0], check) + + +class TestPlainArtifact(OpenpilotTestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + if not os.environ.get("TINYGRAD_TESTS"): + raise unittest.SkipTest("TINYGRAD_TESTS env var not set, skipping artifact test.") + from tinygrad import Tensor + cls.tmp = tempfile.TemporaryDirectory() + cls.rng = np.random.default_rng(0) + cls.weights = Tensor(cls.rng.standard_normal((16, 32)).astype(np.float32)).realize() + cls.bias = Tensor(cls.rng.standard_normal((1, 32)).astype(np.float32)).realize() + + @classmethod + def tearDownClass(cls): + cls.tmp.cleanup() + super().tearDownClass() + + def model(self, x, h): # a two input, one weight model: the shape of a policy without the size + return ((x @ self.weights + self.bias).relu() + h.mean()).contiguous() + + def tensors(self, seed: int): + from tinygrad import Tensor + rng = np.random.default_rng(seed) + return {"x": Tensor(rng.standard_normal((1, 16)).astype(np.float32)).realize(), "h": Tensor(rng.standard_normal((4, 16)).astype(np.float32)).realize()} + + def artifact(self, name: str = "model") -> str: + path = os.path.join(self.tmp.name, f"{name}_tinygrad.pkl") + if not os.path.exists(path): + write_bundle(path, {"run": compile_jit(self.model, self.tensors(1))}, {"metadata": {"name": name}}) + return path + + def test_round_trip_matches_jit(self): # compiler -> bundle -> PlainJit, replay verified, then fresh inputs bit exact with eager tinygrad + bundle = load_model(self.artifact(), cache=False, verify=True) + assert bundle["metadata"] == {"name": "model"} and isinstance(bundle["run"], PlainJit) + fresh = self.tensors(2) + assert np.array_equal(bundle["run"](**fresh).numpy(), self.model(**fresh).numpy()) + + def test_replay_check_catches_a_changed_kernel(self): + parts, extras = read_parts(self.artifact()) + header, blob = parts["run"] + poisoned = dict(header, check=dict(header["check"], outputs=[header["check"]["inputs"][0][1]])) + with self.assertRaises(helpers.ArtifactMismatch): + PlainJit(poisoned, blob) + + def test_cache_round_trip_holds_no_weights_and_is_rebuilt_when_stale(self): + path = self.artifact("cached") + expected = load_model(path, cache=False, verify=False)["run"](**self.tensors(3)).numpy() + load_model(path) + assert cache_is_fresh(path) + with open(f"{path}.cache", "rb") as file: + cache = file.read() + assert self.weights.numpy().tobytes() not in cache and len(cache) < os.path.getsize(path) + with mock.patch.object(helpers, "PlainJit", side_effect=AssertionError("a fresh cache must not rebuild")): + cached = load_model(path) + assert np.array_equal(cached["run"](**self.tensors(3)).numpy(), expected) + with mock.patch.object(helpers, "tinygrad_identity", return_value="another tinygrad"): + assert not cache_is_fresh(path) + assert isinstance(load_model(path)["run"], PlainJit) + assert cache_is_fresh(path) + assert not cache_is_fresh(path) # keyed on the tinygrad that wrote it, so the next boot after a bump rebuilds once + os.utime(path) + with mock.patch.object(helpers, "tinygrad_identity", return_value="another tinygrad"): + assert not cache_is_fresh(path) + + def test_fused_with_warp_matches_warp_then_model(self): + from tinygrad import Tensor + warp_path = os.path.join(self.tmp.name, "warp_tinygrad.pkl") + with mock.patch.object(sys, "argv", ["compile_warp", "--frame", FRAME, "--warp-to", f"{MODEL_W}x{MODEL_H}", "--output", warp_path]): + compile_warp.main() + warp = load_model(warp_path, cache=False)["run"] + + def native(new_img, desire, traffic_convention, state_feat_q, output_buffers): + feat = state_feat_q.mean(axis=0).reshape(1, 16) + new_img.float().mean() + desire.sum() + traffic_convention[0] + Tensor.realize(output_buffers["outputs"].assign((feat @ self.weights + self.bias).relu()), + output_buffers["next_state_feat_q"].assign(state_feat_q[1:].cat(feat * 2))) + + def model_inputs(new_img): + return {"new_img": new_img} | {name: Tensor(np.ones(shape, np.float32) * 0.5).realize() for name, shape in HOST.items()} | \ + {"state_feat_q": Tensor(self.rng.standard_normal(STATE).astype(np.float32)).realize(), + "output_buffers": {"outputs": Tensor.zeros(OUT).contiguous().realize(), "next_state_feat_q": Tensor.zeros(STATE).contiguous().realize()}} + frames = Tensor(self.rng.integers(0, 256, (2, FRAME_SIZE), dtype=np.uint8)).realize() + tfm = Tensor(np.tile(np.eye(3, dtype=np.float32), (2, 1, 1))).realize() + path = os.path.join(self.tmp.name, "native_tinygrad.pkl") + part = compile_jit(native, model_inputs(warp(input_frame=frames, M_inv=tfm).realize()), "output_buffers") + extras = {"input_specs": {"new_img": ((2, 6, MODEL_H // 2, MODEL_W // 2), "|u1", "CPU"), "state_feat_q": (STATE, "= {"global_size", "local_size", "globals", "outs", "ins", "target"}, "ProgramInfo changed: update helpers._program_uop" + assert fields[ParamArg] >= {"slot", "dtype", "name", "image"}, "ParamArg changed: update helpers._program_uop" + assert fields[Target] >= {"device", "renderer", "arch"}, "Target changed: update the 'target' key in plain_artifact_compiler.serialize_jit" + assert "signature" in getattr(TinyELF, "__annotations__", {}) or hasattr(UOp, "kernel_params"), \ + "kernel signatures moved: update plain_artifact_compiler.serialize_jit" + assert callable(getattr(UOp, "param", None)) and callable(getattr(UOp, "from_buffer", None)), "UOp.param/from_buffer: update helpers.PlainJit" + assert callable(getattr(UOp, "store_call", None)) and callable(getattr(UOp, "call", None)), "UOp.call/store_call: update helpers.PlainJit" + assert callable(getattr(Buffer, "view", None)) and callable(getattr(Buffer, "ensure_allocated", None)), "Buffer.view: update helpers._upload" + assert callable(getattr(jit, "compile_linear", None)) and callable(getattr(jit, "_prepare_jit_inputs", None)), "engine.jit: update helpers._Jit" + assert "expected_names" in inspect.signature(jit.CapturedJit).parameters, "CapturedJit: update helpers.PlainJit and plain_artifact_compiler.serialize_jit"