mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
firestar5683
729 lines
30 KiB
Python
729 lines
30 KiB
Python
"""Compile StarPilot unified driving artifacts against the vendored tinygrad runtime."""
|
|
import argparse
|
|
import base64
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import pickle
|
|
import tempfile
|
|
import time
|
|
from collections import namedtuple
|
|
from functools import partial
|
|
from typing import TYPE_CHECKING, cast
|
|
|
|
if __name__ == '__main__':
|
|
from openpilot.starpilot.models.runtime_bridge.runtime import prepare_compiler
|
|
prepare_compiler()
|
|
|
|
import numpy as np
|
|
from tinygrad import Device, TinyJit, Context, Tensor
|
|
from tinygrad.dtype import _to_np_dtype, dtypes
|
|
from tinygrad.device import Buffer, BufferSpec
|
|
from tinygrad.uop.ops import UOp
|
|
from tinygrad_repo.examples.openpilot.helpers import dump_pickle, load_pickle
|
|
if TYPE_CHECKING:
|
|
from openpilot.starpilot.models.runtime_bridge.legacy_types import LegacyBufferFactory, LegacyUOpFactory
|
|
|
|
COMPILER_REVISION = "9d0446a4ba8a532c8b674fb6ad795af015cd9dcf"
|
|
ARTIFACT_ABI = "tinygrad_single_v1_arena"
|
|
|
|
|
|
def make_metadata_dict(path, runner):
|
|
from tinygrad.nn.onnx import OnnxPBParser
|
|
|
|
parser = OnnxPBParser(path, load_external_data=False)
|
|
properties, output_shapes = {}, {}
|
|
for field, wire_type in parser._parse_message(parser.reader.len):
|
|
if field == 14:
|
|
entry = parser._parse_StringStringEntryProto()
|
|
properties[entry['key']] = entry['value']
|
|
elif field == 7:
|
|
for field, wire_type in parser._parse_message(parser._decode_end_pos()):
|
|
if field == 12:
|
|
value = parser._parse_ValueInfoProto()
|
|
output_shapes[value['name']] = value['parsed_type'].shape if value['parsed_type'] is not None else ()
|
|
else:
|
|
parser.reader.skip_field(wire_type)
|
|
else:
|
|
parser.reader.skip_field(wire_type)
|
|
if "output_slices" not in properties:
|
|
raise ValueError(f"Missing output_slices metadata: {path}")
|
|
return {
|
|
"model_checkpoint": properties.get("model_checkpoint"),
|
|
"output_slices": pickle.loads(base64.b64decode(properties["output_slices"])),
|
|
"input_shapes": {name: tuple(spec.shape) for name, spec in runner.graph_inputs.items()},
|
|
"output_shapes": output_shapes,
|
|
}
|
|
|
|
|
|
ARTIFACT_FORMAT_VERSION = 1
|
|
MODEL_TYPES = ("vision_policy", "vision_multi_policy", "supercombo")
|
|
NV12Frame = namedtuple("NV12Frame", ["width", "height", "stride", "y_height", "uv_height", "size"])
|
|
IMAGE_HISTORY_IN_WARP = "warp"
|
|
IMAGE_HISTORY_IN_POLICY = "policy"
|
|
IMAGE_HISTORY_PIPELINES = (IMAGE_HISTORY_IN_WARP, IMAGE_HISTORY_IN_POLICY)
|
|
LEGACY_WARP_INPUTS = ("img_q", "big_img_q", "tfm", "big_tfm")
|
|
FAST_WARP_INPUTS = ("tfm", "big_tfm")
|
|
BASE_POLICY_INPUTS = ("feat_q", "desire_q", "packed_npy_inputs")
|
|
FAST_POLICY_INPUTS = ("img_q", "big_img_q", *BASE_POLICY_INPUTS)
|
|
WARP_INPUTS = LEGACY_WARP_INPUTS
|
|
SPLIT_POLICY_INPUTS = BASE_POLICY_INPUTS
|
|
SUPERCOMBO_POLICY_INPUTS = BASE_POLICY_INPUTS
|
|
WARP_DEV = os.getenv("WARP_DEV")
|
|
|
|
|
|
|
|
def _detect_desire_key(input_shapes):
|
|
return next((key for key in input_shapes if key.startswith("desire")), None)
|
|
|
|
|
|
def _detect_vision_keys(input_shapes):
|
|
image_keys = sorted(key for key in input_shapes if "img" in key)
|
|
road_key = next((key for key in image_keys if "big" not in key), None)
|
|
wide_key = next((key for key in image_keys if "big" in key), None)
|
|
if road_key is None or wide_key is None:
|
|
raise ValueError(f"Cannot determine road/wide image keys from {list(input_shapes)}")
|
|
return road_key, wide_key
|
|
|
|
|
|
def derive_frame_skip(input_shapes):
|
|
features_shape = input_shapes.get("features_buffer")
|
|
if features_shape is None:
|
|
return 1
|
|
return 1 if features_shape[1] >= 99 else 4
|
|
|
|
|
|
def make_random_images(keys, shape, device=None, *, rng=None):
|
|
return {key: Tensor.randint(shape, low=0, high=256, dtype="uint8", device=device).realize() for key in keys}
|
|
|
|
|
|
def make_random_blob_images(keys, size, device=None):
|
|
keepalive: list[np.ndarray] = []
|
|
|
|
def make_inputs(*, rng):
|
|
nonlocal keepalive
|
|
keepalive = []
|
|
tensors = {}
|
|
for key in keys:
|
|
frame = (32 * rng.standard_normal(size).astype(np.float32) + 128).clip(0, 255).astype(np.uint8)
|
|
keepalive.append(frame)
|
|
tensors[key] = Tensor.from_blob(frame.ctypes.data, (size,), dtype="uint8", device=device).realize()
|
|
return tensors
|
|
|
|
return make_inputs
|
|
|
|
|
|
def warp_perspective_tinygrad(src_flat, matrix_inverse, dst_shape, src_shape, stride_pad, border_fill_val=None):
|
|
width_dst, height_dst = dst_shape
|
|
height_src, width_src = src_shape
|
|
|
|
x = Tensor.arange(width_dst).reshape(1, width_dst).expand(height_dst, width_dst).reshape(-1)
|
|
y = Tensor.arange(height_dst).reshape(height_dst, 1).expand(height_dst, width_dst).reshape(-1)
|
|
|
|
src_x = matrix_inverse[0, 0] * x + matrix_inverse[0, 1] * y + matrix_inverse[0, 2]
|
|
src_y = matrix_inverse[1, 0] * x + matrix_inverse[1, 1] * y + matrix_inverse[1, 2]
|
|
src_w = matrix_inverse[2, 0] * x + matrix_inverse[2, 1] * y + matrix_inverse[2, 2]
|
|
src_x = src_x / src_w
|
|
src_y = src_y / src_w
|
|
|
|
x_round = Tensor.round(src_x)
|
|
y_round = Tensor.round(src_y)
|
|
x_nn_clipped = x_round.clip(0, width_src - 1).cast("int")
|
|
y_nn_clipped = y_round.clip(0, height_src - 1).cast("int")
|
|
sampled = src_flat[y_nn_clipped * (width_src + stride_pad) + x_nn_clipped]
|
|
|
|
if border_fill_val is None:
|
|
return sampled
|
|
|
|
in_bounds = ((x_round >= 0) & (x_round <= width_src - 1) &
|
|
(y_round >= 0) & (y_round <= height_src - 1)).cast(sampled.dtype)
|
|
return sampled * in_bounds + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - in_bounds)
|
|
|
|
|
|
def frames_to_tensor(frames):
|
|
height = (frames.shape[0] * 2) // 3
|
|
width = frames.shape[1]
|
|
return Tensor.cat(
|
|
frames[0:height:2, 0::2],
|
|
frames[1:height:2, 0::2],
|
|
frames[0:height:2, 1::2],
|
|
frames[1:height:2, 1::2],
|
|
frames[height:height + height // 4].reshape((height // 2, width // 2)),
|
|
frames[height + height // 4:height + height // 2].reshape((height // 2, width // 2)),
|
|
dim=0,
|
|
).reshape((6, height // 2, width // 2))
|
|
|
|
|
|
def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
|
|
cam_w, cam_h, stride, y_height, uv_height, _ = nv12
|
|
uv_offset = stride * y_height
|
|
stride_pad = stride - cam_w
|
|
|
|
def frame_prepare(input_frame, matrix_inverse):
|
|
matrix_inverse_uv = matrix_inverse * Tensor(
|
|
[[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]],
|
|
device=WARP_DEV,
|
|
)
|
|
uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
|
|
with Context(SPLIT_REDUCEOP=0):
|
|
y = warp_perspective_tinygrad(
|
|
input_frame[:cam_h * stride], matrix_inverse, (model_w, model_h), (cam_h, cam_w), stride_pad,
|
|
).realize()
|
|
u = warp_perspective_tinygrad(
|
|
uv[:cam_h // 2, :cam_w:2].flatten(), matrix_inverse_uv,
|
|
(model_w // 2, model_h // 2), (cam_h // 2, cam_w // 2), 0,
|
|
).realize()
|
|
v = warp_perspective_tinygrad(
|
|
uv[:cam_h // 2, 1:cam_w:2].flatten(), matrix_inverse_uv,
|
|
(model_w // 2, model_h // 2), (cam_h // 2, cam_w // 2), 0,
|
|
).realize()
|
|
return frames_to_tensor(y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)))
|
|
|
|
return frame_prepare
|
|
|
|
|
|
def make_warp_input_queues(vision_input_shapes, frame_skip, device):
|
|
road_key, _ = _detect_vision_keys(vision_input_shapes)
|
|
image_shape = vision_input_shapes[road_key]
|
|
frame_count = image_shape[1] // 6
|
|
image_buffer_shape = (frame_skip * (frame_count - 1) + 1, 6, image_shape[2], image_shape[3])
|
|
|
|
npy = {
|
|
"tfm": np.zeros((3, 3), dtype=np.float32),
|
|
"big_tfm": np.zeros((3, 3), dtype=np.float32),
|
|
}
|
|
queues = {
|
|
"img_q": Tensor(np.zeros(image_buffer_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
|
"big_img_q": Tensor(np.zeros(image_buffer_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
|
**{key: Tensor(value, device="NPY").realize() for key, value in npy.items()},
|
|
}
|
|
return queues, npy
|
|
|
|
|
|
def _packed_policy_shapes(input_shapes, include_prev_feature=False):
|
|
desire_key = _detect_desire_key(input_shapes)
|
|
if desire_key is None:
|
|
raise ValueError(f"No desire input found in {list(input_shapes)}")
|
|
|
|
shapes = {"desire": (input_shapes[desire_key][2],)}
|
|
for key, shape in input_shapes.items():
|
|
if key in ("features_buffer", desire_key) or "img" in key:
|
|
continue
|
|
shapes[key] = tuple(shape)
|
|
if include_prev_feature:
|
|
features_shape = input_shapes["features_buffer"]
|
|
shapes["prev_feat"] = (features_shape[0], math.prod(features_shape[2:]))
|
|
return shapes, [math.prod(shape) for shape in shapes.values()]
|
|
|
|
|
|
def make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
|
|
queues, npy = make_warp_input_queues(vision_input_shapes, frame_skip, device)
|
|
features_shape = policy_input_shapes["features_buffer"]
|
|
feature_dim = math.prod(features_shape[2:])
|
|
desire_key = _detect_desire_key(policy_input_shapes)
|
|
desire_shape = policy_input_shapes[desire_key]
|
|
packed_shapes, packed_sizes = _packed_policy_shapes(policy_input_shapes)
|
|
|
|
packed_inputs = np.zeros(sum(packed_sizes), dtype=np.float32)
|
|
npy.update({
|
|
key: value.reshape(shape)
|
|
for (key, shape), value in zip(
|
|
packed_shapes.items(), np.split(packed_inputs, np.cumsum(packed_sizes[:-1])), strict=True,
|
|
)
|
|
})
|
|
queues.update({
|
|
"feat_q": Tensor(
|
|
np.zeros((frame_skip * (features_shape[1] - 1) + 1, features_shape[0], feature_dim), dtype=np.float32),
|
|
device=device,
|
|
).contiguous().realize(),
|
|
"desire_q": Tensor(
|
|
np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32),
|
|
device=device,
|
|
).contiguous().realize(),
|
|
"packed_npy_inputs": Tensor(packed_inputs, device="NPY").realize(),
|
|
})
|
|
return queues, npy
|
|
|
|
|
|
def make_supercombo_input_queues(input_shapes, frame_skip, device):
|
|
queues, npy = make_warp_input_queues(input_shapes, frame_skip, device)
|
|
features_shape = input_shapes["features_buffer"]
|
|
feature_dim = math.prod(features_shape[2:])
|
|
desire_key = _detect_desire_key(input_shapes)
|
|
desire_shape = input_shapes[desire_key]
|
|
packed_shapes, packed_sizes = _packed_policy_shapes(input_shapes, include_prev_feature=True)
|
|
|
|
packed_inputs = np.zeros(sum(packed_sizes), dtype=np.float32)
|
|
npy.update({
|
|
key: value.reshape(shape)
|
|
for (key, shape), value in zip(
|
|
packed_shapes.items(), np.split(packed_inputs, np.cumsum(packed_sizes[:-1])), strict=True,
|
|
)
|
|
})
|
|
queues.update({
|
|
"feat_q": Tensor(
|
|
np.zeros((frame_skip * features_shape[1], features_shape[0], feature_dim), dtype=np.float32),
|
|
device=device,
|
|
).contiguous().realize(),
|
|
"desire_q": Tensor(
|
|
np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32),
|
|
device=device,
|
|
).contiguous().realize(),
|
|
"packed_npy_inputs": Tensor(packed_inputs, device="NPY").realize(),
|
|
})
|
|
return queues, npy
|
|
|
|
|
|
def shift_and_sample(buffer, new_value, sample_fn):
|
|
buffer.assign(buffer[1:].cat(new_value, dim=0).contiguous())
|
|
return sample_fn(buffer)
|
|
|
|
|
|
def sample_skip(buffer, frame_skip):
|
|
return buffer[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0)
|
|
|
|
|
|
def sample_desire(buffer, frame_skip):
|
|
return buffer.reshape(-1, frame_skip, *buffer.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
|
|
|
|
|
|
def make_warp(nv12, model_w, model_h, frame_skip, image_history_pipeline=IMAGE_HISTORY_IN_POLICY):
|
|
frame_prepare = make_frame_prepare(nv12, model_w, model_h)
|
|
|
|
if image_history_pipeline == IMAGE_HISTORY_IN_POLICY:
|
|
def warp(tfm, big_tfm, frame, big_frame):
|
|
tfm = tfm.to(WARP_DEV)
|
|
big_tfm = big_tfm.to(WARP_DEV)
|
|
Tensor.realize(tfm, big_tfm)
|
|
|
|
return Tensor.cat(
|
|
frame_prepare(frame, tfm).unsqueeze(0),
|
|
frame_prepare(big_frame, big_tfm).unsqueeze(0),
|
|
)
|
|
|
|
return warp
|
|
|
|
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
|
|
|
def warp_enqueue(img_q, big_img_q, tfm, big_tfm, frame, big_frame):
|
|
tfm = tfm.to(WARP_DEV)
|
|
big_tfm = big_tfm.to(WARP_DEV)
|
|
Tensor.realize(tfm, big_tfm)
|
|
|
|
warped = Tensor.cat(
|
|
frame_prepare(frame, tfm).unsqueeze(0),
|
|
frame_prepare(big_frame, big_tfm).unsqueeze(0),
|
|
).to(Device.DEFAULT)
|
|
img = shift_and_sample(img_q, warped[0:1], sample_skip_fn)
|
|
big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn)
|
|
return img, big_img
|
|
|
|
return warp_enqueue
|
|
|
|
|
|
def make_run_split_policy(vision_runner, policy_runners, metadata, policy_order, frame_skip,
|
|
image_history_pipeline=IMAGE_HISTORY_IN_POLICY):
|
|
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
|
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
|
vision_metadata = metadata["vision"]
|
|
policy_metadata = metadata[policy_order[0]]
|
|
vision_features_slice = vision_metadata["output_slices"]["hidden_state"]
|
|
desire_key = _detect_desire_key(policy_metadata["input_shapes"])
|
|
packed_shapes, packed_sizes = _packed_policy_shapes(policy_metadata["input_shapes"])
|
|
road_key, wide_key = _detect_vision_keys(vision_metadata["input_shapes"])
|
|
|
|
def run_model(img, big_img, feat_q, desire_q, packed_npy_inputs):
|
|
unpacked = {
|
|
key: tensor.reshape(shape)
|
|
for (key, shape), tensor in zip(
|
|
packed_shapes.items(), packed_npy_inputs.split(packed_sizes), strict=True,
|
|
)
|
|
}
|
|
desire_buffer = shift_and_sample(
|
|
desire_q, unpacked.pop("desire").reshape(1, 1, -1), sample_desire_fn,
|
|
)
|
|
vision_output = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast("float32")
|
|
new_feature = vision_output[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
|
|
features_buffer = shift_and_sample(feat_q, new_feature, sample_skip_fn)
|
|
features_buffer = features_buffer.reshape(policy_metadata["input_shapes"]["features_buffer"])
|
|
|
|
policy_inputs = {
|
|
"features_buffer": features_buffer,
|
|
desire_key: desire_buffer,
|
|
**unpacked,
|
|
}
|
|
policy_outputs = [
|
|
next(iter(policy_runners[key](policy_inputs).values())).cast("float32")
|
|
for key in policy_order
|
|
]
|
|
return (vision_output, *policy_outputs)
|
|
|
|
if image_history_pipeline == IMAGE_HISTORY_IN_POLICY:
|
|
def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
|
|
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT)
|
|
warped = warped.to(Device.DEFAULT)
|
|
Tensor.realize(packed_npy_inputs, warped)
|
|
img = shift_and_sample(img_q, warped[0:1], sample_skip_fn)
|
|
big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn)
|
|
return run_model(img, big_img, feat_q, desire_q, packed_npy_inputs)
|
|
|
|
return run_policy
|
|
|
|
def run_policy(img, big_img, feat_q, desire_q, packed_npy_inputs):
|
|
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT).realize()
|
|
return run_model(img, big_img, feat_q, desire_q, packed_npy_inputs)
|
|
|
|
return run_policy
|
|
|
|
|
|
def make_run_supercombo(model_runner, metadata, frame_skip, image_history_pipeline=IMAGE_HISTORY_IN_POLICY):
|
|
input_shapes = metadata["model"]["input_shapes"]
|
|
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
|
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
|
desire_key = _detect_desire_key(input_shapes)
|
|
packed_shapes, packed_sizes = _packed_policy_shapes(input_shapes, include_prev_feature=True)
|
|
road_key, wide_key = _detect_vision_keys(input_shapes)
|
|
|
|
def run_model(img, big_img, feat_q, desire_q, packed_npy_inputs):
|
|
unpacked = {
|
|
key: tensor.reshape(shape)
|
|
for (key, shape), tensor in zip(
|
|
packed_shapes.items(), packed_npy_inputs.split(packed_sizes), strict=True,
|
|
)
|
|
}
|
|
desire_buffer = shift_and_sample(
|
|
desire_q, unpacked.pop("desire").reshape(1, 1, -1), sample_desire_fn,
|
|
)
|
|
previous_feature = unpacked.pop("prev_feat")
|
|
features_buffer = shift_and_sample(
|
|
feat_q, previous_feature.reshape(1, 1, -1), sample_skip_fn,
|
|
)
|
|
features_buffer = features_buffer.reshape(input_shapes["features_buffer"])
|
|
model_inputs = {
|
|
road_key: img,
|
|
wide_key: big_img,
|
|
"features_buffer": features_buffer,
|
|
desire_key: desire_buffer,
|
|
**unpacked,
|
|
}
|
|
model_output = next(iter(model_runner(model_inputs).values())).cast("float32")
|
|
return model_output,
|
|
|
|
if image_history_pipeline == IMAGE_HISTORY_IN_POLICY:
|
|
def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
|
|
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT)
|
|
warped = warped.to(Device.DEFAULT)
|
|
Tensor.realize(packed_npy_inputs, warped)
|
|
img = shift_and_sample(img_q, warped[0:1], sample_skip_fn)
|
|
big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn)
|
|
return run_model(img, big_img, feat_q, desire_q, packed_npy_inputs)
|
|
|
|
return run_policy
|
|
|
|
def run_policy(img, big_img, feat_q, desire_q, packed_npy_inputs):
|
|
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT).realize()
|
|
return run_model(img, big_img, feat_q, desire_q, packed_npy_inputs)
|
|
|
|
return run_policy
|
|
|
|
|
|
def stateful_image_shapes(metadata):
|
|
shape = tuple(metadata['input_shapes']['new_img'])
|
|
if len(shape) != 4 or shape[:2] != (2, 6):
|
|
raise ValueError(f"Unsupported stateful image shape: {shape}")
|
|
return dict.fromkeys(('img', 'big_img'), (1, *shape[1:]))
|
|
|
|
|
|
def stateful_host_shapes(metadata):
|
|
return {name: shape for name, shape in metadata['input_shapes'].items()
|
|
if name != 'new_img' and name not in metadata['state_pairs']}
|
|
|
|
|
|
def make_stateful_input_queues(metadata, device):
|
|
queues, npy = make_warp_input_queues(stateful_image_shapes(metadata), 1, device)
|
|
shapes = stateful_host_shapes(metadata)
|
|
sizes = [math.prod(shape) for shape in shapes.values()]
|
|
packed = np.zeros(sum(sizes), dtype=np.float32)
|
|
npy.update({name: value.reshape(shape) for (name, shape), value in
|
|
zip(shapes.items(), np.split(packed, np.cumsum(sizes[:-1])), strict=True)})
|
|
queues['packed_npy_inputs'] = Tensor(packed, device='NPY').realize()
|
|
if metadata.get('mapped_image_input', False):
|
|
image_shape = tuple(metadata['input_shapes']['new_img'])
|
|
from openpilot.starpilot.models.runtime_bridge.runtime import selected_legacy
|
|
if selected_legacy():
|
|
buffer = cast("LegacyBufferFactory", Buffer)(device, math.prod(image_shape), dtypes.uint8,
|
|
options=BufferSpec(cpu_access=True), preallocate=True)
|
|
view = cast("LegacyUOpFactory", UOp).from_buffer(buffer)
|
|
else:
|
|
buffer = Buffer(device, math.prod(image_shape) * dtypes.uint8.itemsize,
|
|
options=BufferSpec(cpu_access=True), preallocate=True)
|
|
view = UOp.from_buffer(buffer, dtypes.uint8)
|
|
queues['mapped_gpu_image'] = Tensor(view).reshape(image_shape)
|
|
for name in metadata['state_pairs']:
|
|
queues[name] = Tensor(np.zeros(metadata['input_shapes'][name], dtype=metadata['input_dtypes'][name]),
|
|
device=device).contiguous().realize()
|
|
return queues, npy
|
|
|
|
|
|
def make_run_stateful_supercombo(model_runner, metadata):
|
|
shapes = stateful_host_shapes(metadata)
|
|
sizes = [math.prod(shape) for shape in shapes.values()]
|
|
|
|
def run_policy(warped, packed_npy_inputs, mapped_gpu_image=None, **state):
|
|
packed = packed_npy_inputs.to(Device.DEFAULT).realize()
|
|
inputs = {name: value.reshape(shape).cast(model_runner.graph_inputs[name].dtype)
|
|
for (name, shape), value in zip(shapes.items(), packed.split(sizes), strict=True)}
|
|
if metadata.get('mapped_image_input', False):
|
|
if mapped_gpu_image is None:
|
|
raise ValueError('Mapped stateful image queue is missing')
|
|
image = warped.to('CPU').realize().to(Device.DEFAULT)
|
|
image = mapped_gpu_image.assign(image).realize()
|
|
else:
|
|
image = warped.to(Device.DEFAULT)
|
|
inputs['new_img'] = image.cast(model_runner.graph_inputs['new_img'].dtype)
|
|
outputs = {name: value.contiguous() for name, value in model_runner(inputs | state).items()}
|
|
for name, next_name in metadata['state_pairs'].items():
|
|
if outputs[next_name].dtype != state[name].dtype:
|
|
raise ValueError(f'State dtype mismatch: {name} -> {next_name}')
|
|
Tensor.realize(*outputs.values())
|
|
Tensor.realize(*(state[name].assign(outputs[next_name]) for name, next_name in metadata['state_pairs'].items()))
|
|
return outputs['outputs'].cast('float32'),
|
|
|
|
return run_policy
|
|
|
|
|
|
def compile_jit(jit, make_random_inputs, input_keys, make_queues, validation_runs=1):
|
|
seed = 42
|
|
|
|
def random_inputs_run(fn, current_seed, test_values=None, test_buffers=None, expect_match=True):
|
|
input_queues, npy = make_queues(Device.DEFAULT)
|
|
rng = np.random.default_rng(current_seed)
|
|
Tensor.manual_seed(current_seed)
|
|
testing = test_values is not None or test_buffers is not None
|
|
run_count = validation_runs if testing else max(3, validation_runs)
|
|
values, buffers = [], []
|
|
|
|
for index in range(run_count):
|
|
for value in npy.values():
|
|
value[:] = rng.standard_normal(value.shape).astype(value.dtype)
|
|
Device.default.synchronize()
|
|
random_inputs = make_random_inputs(rng=rng)
|
|
start = time.perf_counter()
|
|
outputs = fn(**{key: input_queues[key] for key in input_keys}, **random_inputs)
|
|
mid = time.perf_counter()
|
|
Device.default.synchronize()
|
|
end = time.perf_counter()
|
|
print(f" [{index + 1}/{run_count}] enqueue {(mid - start) * 1e3:6.2f} ms -- total {(end - start) * 1e3:6.2f} ms")
|
|
|
|
if index < validation_runs:
|
|
values.extend(np.copy(value.numpy()) for value in outputs)
|
|
buffers.extend(np.copy(value.numpy()) for value in input_queues.values())
|
|
if not all(np.isfinite(value).all() for value in values):
|
|
raise ValueError("Compiled JIT produced non-finite outputs")
|
|
|
|
if test_values is not None:
|
|
match = all(np.array_equal(lhs, rhs) for lhs, rhs in zip(values, test_values, strict=True))
|
|
assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={current_seed})"
|
|
if test_buffers is not None:
|
|
match = all(np.array_equal(lhs, rhs) for lhs, rhs in zip(buffers, test_buffers, strict=True))
|
|
assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={current_seed})"
|
|
return values, buffers
|
|
|
|
print("capture + replay")
|
|
test_values, test_buffers = random_inputs_run(jit, seed)
|
|
print("pickle round trip")
|
|
with tempfile.TemporaryDirectory(prefix="starpilot-jit-") as temporary:
|
|
path = Path(temporary) / "captured.pkl"
|
|
dump_pickle(jit, path)
|
|
loaded = load_pickle(path, out_of_band=True)
|
|
random_inputs_run(loaded, seed, test_values, test_buffers, expect_match=True)
|
|
random_inputs_run(loaded, seed + 1, test_values, test_buffers, expect_match=False)
|
|
return jit
|
|
|
|
|
|
def _parse_size(value):
|
|
width, height = value.lower().split("x")
|
|
return int(width), int(height)
|
|
|
|
|
|
def validate_metadata(metadata):
|
|
output_shapes = metadata.get("output_shapes", {})
|
|
output_shape = output_shapes.get("outputs")
|
|
if not output_shape or len(output_shape) < 2:
|
|
raise ValueError(f"Invalid model output shape metadata: {output_shapes}")
|
|
output_size = output_shape[-1]
|
|
for name, output_slice in metadata.get("output_slices", {}).items():
|
|
start, stop, step = output_slice.indices(output_size)
|
|
if step != 1 or start < 0 or stop < start or stop > output_size:
|
|
raise ValueError(f"Invalid output slice {name}={output_slice} for output size {output_size}")
|
|
|
|
|
|
def main():
|
|
from openpilot.starpilot.models.runtime_bridge.runtime import require_compiler
|
|
require_compiler()
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model-type", choices=MODEL_TYPES, required=True)
|
|
parser.add_argument("--model-size", type=_parse_size, required=True)
|
|
parser.add_argument("--camera-resolutions", type=_parse_size, nargs="+", required=True)
|
|
parser.add_argument("--frame-skip", type=int)
|
|
parser.add_argument("--behavior-version")
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--out-of-band", action="store_true", help="Compatibility flag; all artifacts use the current buffer arena format.")
|
|
parser.add_argument("--vision-onnx")
|
|
parser.add_argument("--policy-onnx")
|
|
parser.add_argument("--off-policy-onnx")
|
|
parser.add_argument("--on-policy-onnx")
|
|
parser.add_argument("--supercombo-onnx")
|
|
parser.add_argument(
|
|
"--image-history-pipeline",
|
|
choices=IMAGE_HISTORY_PIPELINES,
|
|
default=IMAGE_HISTORY_IN_POLICY,
|
|
help="Where img/big_img history queues are updated. 'policy' is the newer faster ABI; 'warp' reproduces legacy v22 artifacts.",
|
|
)
|
|
args = parser.parse_args()
|
|
if args.frame_skip is not None and args.frame_skip < 1:
|
|
parser.error("--frame-skip must be positive")
|
|
|
|
from tinygrad.nn.onnx import OnnxRunner
|
|
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
|
|
|
output = {
|
|
"format_version": ARTIFACT_FORMAT_VERSION,
|
|
"artifact_abi": ARTIFACT_ABI,
|
|
"compiler_revision": COMPILER_REVISION,
|
|
"execution_device": Device.DEFAULT,
|
|
"warp_device": WARP_DEV or Device.DEFAULT,
|
|
"model_type": args.model_type,
|
|
"metadata": {},
|
|
"image_history_pipeline": args.image_history_pipeline,
|
|
}
|
|
if args.behavior_version:
|
|
output["behavior_version"] = args.behavior_version
|
|
|
|
if args.model_type == "supercombo":
|
|
if not args.supercombo_onnx:
|
|
parser.error("--supercombo-onnx is required for supercombo")
|
|
model_path = Path(args.supercombo_onnx)
|
|
model_runner = OnnxRunner(model_path)
|
|
output["metadata"]["model"] = make_metadata_dict(model_path, model_runner)
|
|
validate_metadata(output["metadata"]["model"])
|
|
policy_shapes = output["metadata"]["model"]["input_shapes"]
|
|
if 'new_img' in policy_shapes:
|
|
if args.image_history_pipeline != IMAGE_HISTORY_IN_POLICY:
|
|
parser.error('ONNX-managed history requires --image-history-pipeline policy')
|
|
metadata = output['metadata']['model']
|
|
metadata['state_pairs'] = {name: f'next_{name}' for name in policy_shapes
|
|
if f'next_{name}' in metadata['output_shapes']}
|
|
if not metadata['state_pairs']:
|
|
raise ValueError('Stateful supercombo is missing next-state outputs')
|
|
metadata['input_dtypes'] = {name: np.dtype(_to_np_dtype(spec.dtype)).name for name, spec in model_runner.graph_inputs.items()}
|
|
for name, next_name in metadata['state_pairs'].items():
|
|
if policy_shapes[name] != metadata['output_shapes'][next_name]:
|
|
raise ValueError(f'State shape mismatch: {name} -> {next_name}')
|
|
metadata['mapped_image_input'] = True
|
|
frame_skip = 1
|
|
make_policy_queues = partial(make_stateful_input_queues, metadata)
|
|
run_policy = make_run_stateful_supercombo(model_runner, metadata)
|
|
image_shapes = stateful_image_shapes(metadata)
|
|
policy_input_keys = ('packed_npy_inputs', 'mapped_gpu_image', *metadata['state_pairs'])
|
|
else:
|
|
frame_skip = args.frame_skip or derive_frame_skip(policy_shapes)
|
|
make_policy_queues = partial(make_supercombo_input_queues, policy_shapes, frame_skip)
|
|
run_policy = make_run_supercombo(
|
|
model_runner, output["metadata"], frame_skip, args.image_history_pipeline,
|
|
)
|
|
image_shapes = policy_shapes
|
|
policy_input_keys = FAST_POLICY_INPUTS if args.image_history_pipeline == IMAGE_HISTORY_IN_POLICY else SUPERCOMBO_POLICY_INPUTS
|
|
else:
|
|
if not args.vision_onnx:
|
|
parser.error("--vision-onnx is required for split models")
|
|
policy_paths = {}
|
|
if args.policy_onnx:
|
|
policy_paths["policy"] = args.policy_onnx
|
|
if args.off_policy_onnx:
|
|
policy_paths["off_policy"] = args.off_policy_onnx
|
|
if args.on_policy_onnx:
|
|
policy_paths["on_policy"] = args.on_policy_onnx
|
|
if args.model_type == "vision_policy" and set(policy_paths) != {"policy"}:
|
|
parser.error("vision_policy requires --policy-onnx")
|
|
if args.model_type == "vision_multi_policy" and not policy_paths:
|
|
parser.error("vision_multi_policy requires at least one policy ONNX")
|
|
|
|
vision_path = Path(args.vision_onnx)
|
|
resolved_policy_paths = {key: Path(path) for key, path in policy_paths.items()}
|
|
vision_runner = OnnxRunner(vision_path)
|
|
policy_runners = {key: OnnxRunner(path) for key, path in resolved_policy_paths.items()}
|
|
output["metadata"]["vision"] = make_metadata_dict(vision_path, vision_runner)
|
|
validate_metadata(output["metadata"]["vision"])
|
|
for key, path in resolved_policy_paths.items():
|
|
output["metadata"][key] = make_metadata_dict(path, policy_runners[key])
|
|
validate_metadata(output["metadata"][key])
|
|
|
|
policy_order = [key for key in ("on_policy", "off_policy", "policy") if key in policy_runners]
|
|
output["policy_order"] = policy_order
|
|
first_policy_shapes = output["metadata"][policy_order[0]]["input_shapes"]
|
|
for key in policy_order[1:]:
|
|
if output["metadata"][key]["input_shapes"] != first_policy_shapes:
|
|
raise ValueError(f"Policy input shapes differ between {policy_order[0]} and {key}")
|
|
frame_skip = args.frame_skip or derive_frame_skip(first_policy_shapes)
|
|
make_policy_queues = partial(
|
|
make_split_input_queues,
|
|
output["metadata"]["vision"]["input_shapes"],
|
|
first_policy_shapes,
|
|
frame_skip,
|
|
)
|
|
run_policy = make_run_split_policy(
|
|
vision_runner, policy_runners, output["metadata"], policy_order, frame_skip,
|
|
args.image_history_pipeline,
|
|
)
|
|
image_shapes = output["metadata"]["vision"]["input_shapes"]
|
|
policy_input_keys = FAST_POLICY_INPUTS if args.image_history_pipeline == IMAGE_HISTORY_IN_POLICY else SPLIT_POLICY_INPUTS
|
|
|
|
output["frame_skip"] = frame_skip
|
|
output["policy_input_keys"] = policy_input_keys
|
|
warp_input_keys = FAST_WARP_INPUTS if args.image_history_pipeline == IMAGE_HISTORY_IN_POLICY else LEGACY_WARP_INPUTS
|
|
output["warp_input_keys"] = warp_input_keys
|
|
run_policy_jit = TinyJit(run_policy, prune=True)
|
|
road_key, wide_key = _detect_vision_keys(image_shapes)
|
|
if args.image_history_pipeline == IMAGE_HISTORY_IN_POLICY:
|
|
make_random_model_inputs = partial(
|
|
make_random_images,
|
|
keys=["warped"],
|
|
shape=(2, 6, *image_shapes[road_key][2:]),
|
|
device=WARP_DEV,
|
|
)
|
|
else:
|
|
make_random_model_inputs = partial(
|
|
make_random_images,
|
|
keys=[road_key, wide_key],
|
|
shape=image_shapes[road_key],
|
|
)
|
|
output["run_policy"] = compile_jit(
|
|
run_policy_jit, make_random_model_inputs, policy_input_keys, make_policy_queues,
|
|
validation_runs=5 if output['metadata'].get('model', {}).get('state_pairs') else 1,
|
|
)
|
|
|
|
model_w, model_h = args.model_size
|
|
for cam_w, cam_h in args.camera_resolutions:
|
|
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
|
warp_enqueue = TinyJit(
|
|
make_warp(nv12, model_w, model_h, frame_skip, args.image_history_pipeline),
|
|
prune=True,
|
|
)
|
|
make_random_warp_inputs = make_random_blob_images(
|
|
keys=["frame", "big_frame"], size=nv12.size, device=WARP_DEV,
|
|
)
|
|
make_warp_queues = partial(make_warp_input_queues, image_shapes, frame_skip)
|
|
output[(cam_w, cam_h)] = compile_jit(
|
|
warp_enqueue, make_random_warp_inputs, warp_input_keys, make_warp_queues,
|
|
)
|
|
|
|
dump_pickle(output, args.output)
|
|
print(f"Saved JITs to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|