From 8a1624ab6ebb7bbbf0e8a6eb2e1b462ec39eb905 Mon Sep 17 00:00:00 2001 From: discountchubbs Date: Thu, 5 Nov 2026 10:39:49 -0800 Subject: [PATCH] modeld_v2: artifact parity with upstream tg --- .github/workflows/test_models.yml | 17 ++++++++ .../modeld_v2/convert_upstream_pkl.py | 42 ++++++++++++++----- openpilot/sunnypilot/modeld_v2/helpers.py | 3 +- 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/.github/workflows/test_models.yml b/.github/workflows/test_models.yml index bf22f14a11..7210598574 100644 --- a/.github/workflows/test_models.yml +++ b/.github/workflows/test_models.yml @@ -46,6 +46,23 @@ jobs: CHECK_OOB: 0 GMMU: 0 run: tools/op.sh test -v openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py + - name: Resolve commaai's tinygrad pin, which their precompiled chestnut model is built and converted with + id: upstream-tg + env: + GH_TOKEN: ${{ github.token }} + run: | + read -r sha url < <(gh api repos/commaai/openpilot/contents/tinygrad_repo --jq '.sha + " " + .submodule_git_url') + echo "commaai/openpilot master pins $url@$sha, ours is $(git rev-parse HEAD:tinygrad_repo)" + dir="$RUNNER_TEMP/upstream_tinygrad" + git init -q "$dir" && git -C "$dir" fetch -q --depth 1 "$url" "$sha" && git -C "$dir" checkout -q FETCH_HEAD + echo "path=$dir" >> "$GITHUB_OUTPUT" + - name: Convert a pickled USB+AMD jit under commaai's tinygrad, as the chestnut model build does + env: + DEV: MOCKUSB+AMD + CHECK_OOB: 0 + GMMU: 0 + PYTHONPATH: ${{ steps.upstream-tg.outputs.path }}:${{ github.workspace }} + run: tools/op.sh test -v openpilot/sunnypilot/modeld_v2/tests/test_convert_upstream_pkl.py - name: Where tinygrad moved, for the helper or failure point named above if: failure() && github.event_name == 'pull_request' run: | diff --git a/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py b/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py index 9ae62f0f5d..89fe65066d 100755 --- a/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py +++ b/openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py @@ -66,38 +66,58 @@ def _patched(buf, topo) -> tuple[bytearray, dict]: return blob, rows -def _evaluate(u, env): +def _packed_args(program) -> dict[int, int]: # since tinygrad af8961676 a program reads its params from one packed "args": byte offset -> slot + from tinygrad.device import TinyELF + params = getattr(program.arg, "params", None) or () + return {off: p.arg.slot for (off, _), p in zip(TinyELF.iter_sig(tuple((None, p.addrspace, p.dtype, ()) for p in params)), params, strict=True)} + + +def _outer_slot(u, packed: dict[int, int]) -> int | None: # the submit argument u reads: a param of its own, a REFPARAM, or a word of "args" + from tinygrad.uop.ops import Ops + if u.op is Ops.PARAM: + return None if packed else u.arg.slot + if u.op.name == "REFPARAM": + return u.arg.slot + if packed and u.op is Ops.LOAD and u.src[0].op is Ops.INDEX and _root(u.src[0].src[0]).op is Ops.PARAM: + index = _evaluate(u.src[0].src[1], {}, packed) + return None if index is None else packed.get(index * u.dtype.itemsize) + return None + + +def _evaluate(u, env, packed: dict[int, int]): from tinygrad.uop.ops import GroupOp, Ops, exec_alu if u.op in (Ops.AFTER, Ops.CAST, Ops.BITCAST): - return _evaluate(u.src[0], env) + return _evaluate(u.src[0], env, packed) if u.op is Ops.CONST: return u.arg - if u.op in (Ops.PARAM, Ops.RANGE): - return env.get(u.arg.slot if u.op is Ops.PARAM else u) + if u.op is Ops.RANGE: + return env.get(u) + if (slot := _outer_slot(u, packed)) is not None: + return env.get(slot) if u.op is Ops.LOAD: - table, index = _root(u.src[0].src[0]), _evaluate(u.src[0].src[1], env) + table, index = _root(u.src[0].src[0]), _evaluate(u.src[0].src[1], env, packed) 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]): + if u.op in GroupOp.ALU and None not in (values := [_evaluate(x, env, packed) for x in u.src]): return exec_alu(u.op, u.dtype, values, *([False] if "truncate_output" in inspect.signature(exec_alu).parameters else [])) return None def _input_pokes(submit, kernargs) -> dict[int, tuple[int, int]]: from tinygrad.uop.ops import Ops - outer = dict(enumerate(submit.src[1:])) + outer, packed = dict(enumerate(submit.src[1:])), _packed_args(submit.src[0]) kernargs_params = {i for i, a in outer.items() if a.op is Ops.GETADDR and _root(a.src[0]) is kernargs} inputs = {i: _root(a.src[0]).arg.slot for i, a in outer.items() if a.op is Ops.GETADDR and _root(a.src[0]).op is Ops.PARAM} pokes = {} for call in (u for u in submit.src[0].src[0].toposort() if u.op is Ops.CALL and len(u.src) > 3): _, address, value, *_ = call.src[1:] - used = list({x.arg.slot for x in value.toposort() if x.op is Ops.PARAM} & inputs.keys()) - if len(used) != 1 or not kernargs_params & {x.arg.slot for x in address.toposort() if x.op is Ops.PARAM}: + used = list({_outer_slot(x, packed) for x in value.toposort()} & inputs.keys()) + if len(used) != 1 or not kernargs_params & {_outer_slot(x, packed) for x in address.toposort()}: continue ranges = list(dict.fromkeys(x for x in (*address.toposort(), *value.toposort()) if x.op is Ops.RANGE)) for trip in itertools.product(*[range(int(r.vmax) + 1) for r in ranges]): env = dict.fromkeys(outer, 0) | dict(zip(ranges, trip, strict=True)) - if (offset := _evaluate(address, env)) % 8 == 0: - pokes[offset] = (inputs[used[0]], _evaluate(value, env)) + if (offset := _evaluate(address, env, packed)) % 8 == 0: + pokes[offset] = (inputs[used[0]], _evaluate(value, env, packed)) return pokes diff --git a/openpilot/sunnypilot/modeld_v2/helpers.py b/openpilot/sunnypilot/modeld_v2/helpers.py index 7b79c1736d..4601708cfd 100644 --- a/openpilot/sunnypilot/modeld_v2/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/helpers.py @@ -264,7 +264,8 @@ def _program_uop(call: dict, lib: bytes): 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=_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"]), + # params: since tinygrad af8961676 the launch lays out only the args listed here + info = _fields(ProgramInfo, global_size=tuple(call["global_size"]), local_size=tuple(call["local_size"]), params=tuple(parameters), 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)