modeld_v2: artifact parity with upstream tg

.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: |

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
 
 

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)