models: tg sync test

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -258,7 +258,7 @@ 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=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"]),
+  info = 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=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)

openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py

@@ -39,6 +39,7 @@ def capture_frame(jit, kwargs: dict, output_container: str | None = None):
 
 
 def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict, memoryview]:
+  from tinygrad.dtype import AddrSpace
   from tinygrad.uop.ops import Ops
   sink = tempfile.TemporaryFile(dir=".")
   blob, bases, input_specs = Blob(sink), {}, {}
@@ -60,10 +61,12 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
   for call in kernel_linear.src:
     ast, arguments = call.src[0], [describe(argument) for argument in call.src[1:]]
     if ast.op is Ops.PROGRAM:
-      elf, info = ast.to_elf(), ast.arg
+      elf, info, compact = ast.to_elf(), ast.arg, {slot: i for i, slot in enumerate(ast.arg.globals)}
+      if any(p.addrspace is AddrSpace.ALU for p in ast.kernel_params):
+        raise ValueError(f"{elf.name} takes a variable, which plain artifacts don't support")
       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": [[p.arg.name, compact[p.arg.slot], repr(p.dtype), [int(x) for x in p._shape]] for p in ast.kernel_params],
                     "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