| |
| """Python API: StreamPETR (Autoware camera_streampetr) on one Tenstorrent Blackhole p150. |
| |
| from tt_streampetr import StreamPETR, load_sample |
| |
| with StreamPETR.from_pretrained(device_id=0) as model: # weights -> HF cache, device open, traces captured |
| out = model(**load_sample("code/tt_streampetr/samples/synthetic_5cam.json")) # 5 cameras + calibration |
| # a sequence: one dict per frame with the same stream id (the temporal memory carries over), e.g. |
| # model(images=..., calibration=..., stream={"id": "rig", "timestamp_s": t, "T_world_from_ego": T}) |
| print(out.to_dict()) # the same JSON as POST /predict |
| |
| The contract shared by every bundle of the Autoware collection is the vendored ``ttaw.api_base.ModelBase`` |
| (BUNDLE_CONVENTIONS.md section 8): ``from_pretrained`` resolves the pinned weights before it claims the chip, opens it |
| (ETH dispatch, 12x10), builds the graph and captures every trace variant in ``warmup_variants``, so the first call is |
| as fast as the later ones; calls are serialised by a lock (one chip, batch 1) and fill ``timing_ms``; ``close()`` is |
| idempotent, also runs at interpreter exit, and closes the chip only if the model opened it. The HTTP server |
| (``tt_streampetr.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit. |
| |
| Inputs per call (``_prepare``; host code in ``tt_streampetr.host``, exact Autoware semantics): |
| |
| - ``images``: the five cameras ``CAM_FRONT``, ``CAM_FRONT_LEFT``, ``CAM_BACK_LEFT``, ``CAM_FRONT_RIGHT``, |
| ``CAM_BACK_RIGHT`` (any order, any size; resized / cropped / normalised like the node); |
| - ``calibration``: per camera ``intrinsics`` (``CameraInfo.P`` or K of the image as sent; optional full 3x4 ``P``) and |
| ``T_ref_from_camera`` (camera optical frame -> base_link), or ``{"preset": name}``; |
| - ``stream``: ``{"id", "reset", "timestamp_s", "T_world_from_ego"}`` drives the temporal object memory: a new id or |
| ``reset`` starts like a fresh node (zero memory), ``timestamp_s`` (the CAM_FRONT stamp) gives the memory ages and |
| ``T_world_from_ego`` (map <- base_link) the ego-motion compensation (a stream frame without it is skipped, as the |
| node skips a cycle whose TF lookup fails: ``meta.skipped = "no_ego_pose"``). Without ``stream``: one frame from zero |
| memory. Per-camera ``timestamp_s`` feed the node's 0.15 s camera-sync check (a failure resets the stream and |
| returns no objects, ``meta.skipped``). |
| |
| Device graph (``tt_streampetr.tt``, PORT_LOG.md 1.1): ONE trace ``frame`` = VoVNet-99-eSE + CPFPN, the head (tokens, |
| temporal alignment, 6 decoder layers, shared branches), the top-256 selection and the memory update, with the |
| memory (embedding bf16, geometry fp32 in the previous ego frame) kept in persistent device states. The host runs |
| the node's pre-processing (C14), the position embedding and the spatial-alignment coefficients once per |
| calibration, the stream clock / ego poses / ages, and the post-processing (all-layer decode, IoU-BEV NMS). |
| Load-time knobs: :data:`KNOBS` (``STREAMPETR_*``). Importing this module has no side effects (torch / ttnn are |
| imported in the hooks). |
| """ |
| from __future__ import annotations |
|
|
| from collections import OrderedDict |
| from typing import Any, Dict, Mapping, Optional |
|
|
| import numpy as np |
|
|
| from . import io as tio |
| from .device import DEVICE_DEFAULTS |
| from .reference.config import CAMERA_ORDER, CLASS_NAMES, NORMALIZATION_PRESETS |
| from .ttaw.api_base import ModelBase |
| from .ttaw.knobs import Knob, Knobs |
| from .ttaw.outputs import Detections3D |
|
|
| __all__ = ["StreamPETR", "Output", "KNOBS"] |
|
|
| |
| Output = Detections3D |
|
|
| |
| KNOBS = Knobs("STREAMPETR", [ |
| Knob("NORMALIZATION", "autoware_main", "input normalisation preset (PLAN D7; SPEC 3.4): autoware_main = universe " |
| "main RGB planes + RGB stats; autoware_0.52 = an rgb8 camera on 0.52.1; " |
| "awml_training = BGR planes + BGR stats", choices=NORMALIZATION_PRESETS), |
| Knob("USE_TEMPORAL", True, "carry the object memory between frames (the node's use_temporal; False wipes it)"), |
| Knob("MAX_STREAMS", 1, "temporal streams kept (PLAN D16; another stream id evicts the least recent one)"), |
| Knob("AUTOWARE_COMPAT", True, "reproduce the node's age quirks (default, PORT_LOG section 6 Q1): a sync-failure " |
| "wipe keeps the stale ages, the 600 s clock-origin reset shifts them (SPEC 6 / risk " |
| "11); False zeroes / rebases the ages instead (the SPEC's clean variant)"), |
| Knob("PREPROCESS_FMA", False, "model nvcc's fmaf contraction in the resize kernel (ttaw.image_triangle; <= 3e-5 " |
| "difference)", experiment=True), |
| Knob("PREPROCESS_WORKERS", 4, "host threads of the five-camera resize / normalise"), |
| Knob("HOST_PRE", "numba", "host pre-processing (OPT_PLAN item 1): numba = host.fast.preprocess_rows_bf16 (the C14 " |
| "triangle resize, crop, normalisation and the bf16 cast of the upload rows fused in one " |
| "pass per camera, written into a persistent ttnn host buffer; bit-identical upload); " |
| "numpy = ttaw.image_triangle float32 NCHW + img_rows + bf16 conversion (the release " |
| "path; also the automatic fallback without numba)", choices=("numba", "numpy")), |
| Knob("HOST_POST", "numba", "host post-processing (OPT_PLAN item 4): numba = the compiled IoU-BEV NMS " |
| "(host.fast.iou_bev_nms, the same keep decisions); numpy = ttaw.nms.iou_bev_nms (also " |
| "the fallback without numba)", choices=("numba", "numpy")), |
| Knob("STEM_IN", "kin", "device input layout (OPT_PLAN item 2): kin = the image as 3,840-byte bf16 rows " |
| "[5*480, 640*3] and per camera a K-IN generic_op (tt/kin.py) writing stem conv 0's L1 " |
| "shards (conv on the L1 path); rows = [5*480*640, 3] 6-byte pixel pages and the stem's " |
| "DRAM-sliced conv (the release path)", choices=("kin", "rows")), |
| Knob("CONV_L1", "s3r45n", "3x3 conv placement (OPT_PLAN item 3, tt/placement.py): the OSA 3x3 convs of the listed " |
| "stages (n = the CPFPN 3x3) on ttnn's L1 path (the op moves its DRAM input into L1 shards; " |
| "no TILE reshape pair), DRAM slicing as the first-call fallback. s3r45n (default) = stages " |
| "4-5 + neck + stage 3 except each block's first conv (whose 256 / 512-channel input fits " |
| "L1 only with a changed K split: s345* fails the frozen sample_stream.centre_err gate, " |
| "OPT_REPORT R2.1); s3r45n / s45n are bit-identical to none = every 3x3 conv DRAM-sliced " |
| "(the release path)", choices=("none", "s5", "s45", "s45n", "s3r45n", "s345", "s345n", "s2345n")), |
| Knob("CONV_L1_ABH", "256,128,64,32", "with CONV_L1: activation block heights (comma list, multiples of 32) retried " |
| "in order when a conv's L1 path does not fit with ttnn's own blocks"), |
| Knob("ESE", "kernel", "eSE gate application (OPT_PLAN items 5 / 5b): kernel = one ESG generic_op per OSA block " |
| "(tt/ese.py: the gate tiles built in L1 per camera run, the stock FPU multiply and the " |
| "identity add fused; bit-identical); bcast = one row-broadcast multiply of the " |
| "[1, N, HW, C] view by the per-camera gate rows in the stages whose H*W is a multiple of 32 " |
| "(2, 3; exact: the same products); matmul = the gate broadcast to every row by the " |
| "transposed camera-indicator matmul, then a multiply (the release path)", |
| choices=("matmul", "bcast", "kernel")), |
| Knob("ESE_MEAN", "kernel", "eSE per-camera sums (OPT_PLAN item 5b): kernel = ESM partial sums on the whole grid " |
| "(tt/ese.py) + a fixed-order fp32 reduce (ESR; the fp32 sums are regrouped, " |
| "closer to the exact sums than the matmul's; gated, the 8-frame stream outputs stayed " |
| "bit-identical); matmul = the camera-indicator matmul (the release path)", |
| choices=("matmul", "kernel")), |
| Knob("OSA_CAT", "kmm", "OSA concatenation [x, y1..y5] (OPT_PLAN item 6b, tt/kcat.py): kernel = one KCAT " |
| "data-movement generic_op (output row tiles over the whole grid; a copy: bit-identical); " |
| "kmm = (OPT round 9, tt/kmm.py) stages 2-3 read the parts straight from the concat " |
| "1x1 matmul's in0 reader (the stock 1D matmul program ported to a generic_op, same " |
| "compute: bit-identical; no concat tensor), KCAT in stages 4-5; " |
| "concat = ttnn.concat (the release path)", choices=("concat", "kernel", "kmm")), |
| Knob("CONV_CFG", "adb,wdb", "data-movement Conv2dConfig flags of every backbone 3x3 conv (OPT round 4, " |
| "tt/backbone.conv_flags): comma list of adb (activation double buffer, always " |
| "applied), wdb (weights double buffer), split (forced split reader), reuse " |
| "(activation reuse; NOT exact: it changes the stride-1 DRAM-sliced convs' outputs); " |
| "wdb / split / reuse are tried on the first call and dropped where a conv rejects " |
| "them (never by shrinking activation blocks). adb,wdb is bit-identical; none = " |
| "ttnn's defaults (the release path)"), |
| Knob("CONV_CHAIN", "kcat45", "L1 conv chaining in the OSA blocks (OPT_PLAN item 19, OPT round 7, " |
| "tt/backbone.OSABlock._convs; kcat45 default, bit-identical): " |
| "chain = a conv on ttnn's L1 path whose successor is on the L1 path too keeps its " |
| "sharded output as the successor's input (no InterleavedToSharded; the concat part is " |
| "its DRAM copy); kcat = chain, and every L1-path conv output stays sharded as its " |
| "own concat part (the KCAT-SH reader reads it from L1: no ShardedToInterleaved " |
| "either; needs OSA_CAT=kernel, else DRAM copies); none = every L1 conv returns DRAM " |
| "interleaved (the round-6 path); kcat45 = kcat in stages 4-5, chain in stage 3 (whose " |
| "five outputs do not fit L1 together with the conv buffers)", |
| choices=("none", "chain", "kcat45", "kcat")), |
| Knob("RESHAPE", "kernel", "input re-tiling of the DRAM-sliced 3x3 convs (OPT_PLAN item 20, OPT round 8, " |
| "tt/krt.py): kernel = the flat [1, 1, N*H*W, C] map unflattened to [N, H, W, C] by one " |
| "KRT data-movement generic_op (16-row half-tile copies, whole grid) so ttnn.conv2d's own " |
| "ReshapeView is a no-op (stage-3 block-first convs, W = 80; a copy: bit-identical); " |
| "stock = ttnn.conv2d's ReshapeView (the round-7 path)", choices=("stock", "kernel")), |
| Knob("KMM_BUF", "auto", "KMM circular-buffer depth (OPT round 9, R9.2, tt/kmm.buffers; OSA_CAT=kmm only): auto = " |
| "deeper in0 / in1 buffers (stage 2 12 / 12 blocks, stage 3 4 / 4; data movement only: " |
| "bit-identical); 2 = the stock factory's double buffers (R9.1)", choices=("2", "auto")), |
| Knob("OSA_MM", "rowblock", "OSA concat 1x1 matmul config (OPT_PLAN item 6a, tt/layers.row_block_config): rowblock = " |
| "ttnn's 1D config with the output in row blocks (stage 2: 25 x 8 -> 5 x 8 fp32 L1 partials; " |
| "the same K order: bit-identical); auto = ttnn's own config (the release path)", |
| choices=("auto", "rowblock")), |
| Knob("KV_EPI", True, "split-key K / V operands (OPT_PLAN item 7, tt/kvx.py): True = single-width K / V " |
| "projections and one KVX generic_op per K / V pair writing the per-layer [hi | lo] / [v | v] " |
| "head tensors (replaces the 5-program key split over the doubled width, the head split and " |
| "the per-layer slices; bit-identical); False = the doubled projections + stock programs " |
| "(the release path)"), |
| Knob("SELF_KV", "buffer", "self-attention keys (OPT_PLAN item 9, with KV_EPI): buffer = KVX writes the memory and " |
| "query K / V rows straight into per-layer [memory; queries] buffers (bit-identical); " |
| "concat = untilize, ROW_MAJOR concat and tilize per layer (the release path)", |
| choices=("concat", "buffer")), |
| Knob("S26_KERNEL", True, "exact-product slicing of the temporal chain's sliced maps (OPT_PLAN item 8, tt/ksl.py): " |
| "True = the x slices and their pair-block concatenation as one KSL generic_op per map " |
| "input (the three maps on the sines share one; bit-identical); False = 5 S - 1 stock " |
| "elementwise programs + ttnn.concat per map (the release path)"), |
| Knob("SEL", "mk", "top-256 selection + memory update (OPT rounds 5 / 10): rank = (tt/sel.py) exact ranks from fp32 " |
| "comparisons on the whole grid, a one-hot matrix and one-hot matmul gathers (the fp32 rows " |
| "as three bf16 terms; the same set and order as the stable sort: bit-identical); sort = " |
| "ttnn.sort on one core + ttnn.embedding / ttnn.gather (the release path); mk = MK-SEL " |
| "(OPT round 10, tt/mksel.py): the selection + the whole memory update (gathers, ego identity, " |
| "ring shift, state writes) as one generic_op on 72 cores (bit-identical; the taps path keeps rank)", |
| choices=("sort", "rank", "mk")), |
| Knob("PRECISION", "HiFi4+fp32+l1acc:w=fp32", "precision of every device conv / linear (ttaw.precision spec; " |
| "PORT_LOG S10: fp32 weights and biases, HiFi4, fp32 accumulation)"), |
| Knob("ATTN_EXP", "exact", "SDPA softmax exponential (OPT round 5): approx = the SFPU's approximate exp " |
| "(exp_approx_mode + math_approx_mode) in every SDPA; exp = the approximate exp only; " |
| "cross / self = the approximate exp in those SDPAs only. The exact exp makes the SDPAs " |
| "SFPU-bound (cross 643 -> 207 us approximate) but the numerics move: gated (approx " |
| "fails sample_stream); exact = the release path (PORT_LOG r3)", |
| choices=("exact", "approx", "exp", "cross", "self")), |
| Knob("XATTN", "x8", "SDPA implementation (OPT round 6, item 11b, tt/spx.py): sdpa = ttnn's SDPA; spx = the SPX " |
| "generic_op (the stock SDPA program factory ported, stock reader / writer, a copy of the " |
| "compute kernel) with the stock exponential (bit-identical); x8 = SPX whose softmax " |
| "exponential runs the stock exp_fp32_accurate(x * scale) operations as one hand-scheduled " |
| "TTI pass (26 SFPU instructions per row instead of a multiply pass + ~37; bit-identical). " |
| "Needs ATTN_EXP=exact and ATTN_FP32_ACC (else sdpa)", |
| choices=("sdpa", "spx", "x8")), |
| Knob("ATTN_FP32_ACC", True, "SDPA with fp32 accumulation (probe P7: half the max error)"), |
| Knob("HEAD_ACT", "float32", "dtype of the head's token / query / residual activations (PORT_LOG S25; SDPA " |
| "operands stay bf16)", choices=("bfloat16", "float32")), |
| Knob("EXACT_MAP", "gather", "fp32 maps of the temporal alignment (PORT_LOG S18): column gather (exact bf16 " |
| "embedding of split terms) + fp32 elementwise, or the bf16-split matmul", |
| choices=("gather", "matmul")), |
| Knob("KEY_SPLIT", True, "SDPA keys as hi + lo bf16 terms in a doubled head dim (PORT_LOG S25: the bf16 rounding " |
| "of the keys dominates the head's error otherwise)"), |
| Knob("TEMPORAL_EXACT", True, "the temporal positional chain (query embedding MLP, MLN, time embedding, their " |
| "LayerNorms) as exact-product sliced maps (PORT_LOG S26: its matrix-engine error " |
| "dominates the self-attention error otherwise)"), |
| ]) |
|
|
|
|
| class StreamPETR(ModelBase): |
| """StreamPETR (Autoware camera_streampetr) on one Blackhole p150. Create it with :meth:`from_pretrained`.""" |
|
|
| MODEL_NAME = "streampetr-p150" |
| ENV_PREFIX = "STREAMPETR" |
| DEFAULT_REPO = "AutowareFoundation/camera_streampetr" |
| DEFAULT_TAG = "v1.0" |
| |
| DEFAULT_REVISION = "90f37d64186686ebc458509d69d3d4ca16d185c7" |
| ALLOW_PATTERNS = ["simplify_extract_img_feat.onnx", "simplify_position_embedding.onnx", |
| "simplify_pts_head_memory.onnx", "ml_package_camera_streampetr.param.yaml"] |
| VARIANTS = ["default"] |
| DEFAULT_VARIANT = "default" |
| INPUT_KIND = "multicam" |
| CAMERA_ORDER = CAMERA_ORDER |
| POINT_FIELDS = tio.DEFAULT_POINT_FIELDS |
| |
| LABELS = list(CLASS_NAMES) |
| |
| RUNTIME_PARAMS = { |
| "score_threshold": (float, 0.0, 1.0, None), |
| "iou_nms_threshold": (float, 0.0, 1.0, 0.5), |
| "iou_nms_search_distance": (float, 0.0, 1000.0, 10.0), |
| "circle_nms_distance": (float, 0.0, 100.0, 0.0), |
| "decode_layers": (int, 1, 6, 6), |
| "awml_decode": (bool, None, None, False), |
| "max_detections": (int, 1, 5400, 500), |
| } |
| EXTRA_INPUTS = () |
| DEVICE_DEFAULTS = DEVICE_DEFAULTS |
|
|
| |
| def _build(self) -> None: |
| """Weights (``reference.weights.StreamPETRWeights``: the three ONNX files by consuming node), the host state |
| and the device graph (``tt.model.TtStreamPETR``, the ``frame`` variant: backbone + head + in-trace memory |
| update). The position embedding and the spatial-alignment coefficients run on the host once per |
| calibration (PORT_LOG S5). No capture here (``_warm_one``).""" |
| from .reference.model import Head, PositionEmbedding |
| from .reference.weights import StreamPETRWeights, find_weights_dir |
| from .tt.model import TtStreamPETR |
|
|
| self._build_host() |
| wd = find_weights_dir(self.weights_path) |
| if wd is None: |
| raise FileNotFoundError(f"{self.weights_path}: the three camera_streampetr v1.0 ONNX files are needed") |
| self.net_weights = StreamPETRWeights(wd) |
| self.ref_head = Head(self.net_weights, self.cfg) |
| self.ref_pe = PositionEmbedding(self.net_weights, self.cfg) |
| self.tt = TtStreamPETR(self.device, self.net_weights, self.cfg, variants=("frame",), |
| precision=str(self.knobs.PRECISION), attn_fp32_acc=bool(self.knobs.ATTN_FP32_ACC), |
| head_act=str(self.knobs.HEAD_ACT), key_split=bool(self.knobs.KEY_SPLIT), |
| exact_map=str(self.knobs.EXACT_MAP), |
| temporal_exact=bool(self.knobs.TEMPORAL_EXACT), |
| state_banks=int(self.knobs.MAX_STREAMS), |
| head=self.ref_head, stem_in=str(self.knobs.STEM_IN), |
| conv_l1=str(self.knobs.CONV_L1), conv_l1_abh=str(self.knobs.CONV_L1_ABH), |
| ese=str(self.knobs.ESE), osa_mm=str(self.knobs.OSA_MM), |
| kv_epi=bool(self.knobs.KV_EPI), self_kv=str(self.knobs.SELF_KV), |
| s26_kernel=bool(self.knobs.S26_KERNEL), ese_mean=str(self.knobs.ESE_MEAN), |
| osa_cat=str(self.knobs.OSA_CAT), conv_cfg=str(self.knobs.CONV_CFG), |
| sel=str(self.knobs.SEL), attn_exp=str(self.knobs.ATTN_EXP), |
| xattn=str(self.knobs.XATTN), conv_chain=str(self.knobs.CONV_CHAIN), |
| reshape=str(self.knobs.RESHAPE), kmm_buf=str(self.knobs.KMM_BUF)) |
| self.runner = self.tt.runner |
| self._img_stage = None |
| if self._fast_pre(): |
| from .host import fast as hfast |
| from .ttaw.tensors import HostStaging |
|
|
| stage = HostStaging(self.tt.img_shape, "bfloat16") |
| self._img_stage = stage if stage.zero_copy else None |
| hfast.warm() |
| self.calib_fn = self._calibration_constants |
| self._active_calib: Optional[str] = None |
| self._device_owner: Optional[str] = None |
| self._banks: "OrderedDict[str, int]" = OrderedDict() |
|
|
| def _warm_one(self, variant: Any) -> None: |
| """Capture the ``frame`` trace (warm-up runs of the eager graph first; ``TraceRunner.capture``).""" |
| self.tt.capture() |
|
|
| def _calibration_constants(self, rig: Any) -> Dict[str, Any]: |
| """Host fp32 PE graph (``reference.model.PositionEmbedding``) and the spatial-alignment MLN coefficients of |
| its ``cone`` (``Head.spatial_alignment_coeffs``) -> the device params of one calibration.""" |
| import torch |
|
|
| with torch.no_grad(): |
| pos, cone = self.ref_pe(rig.img_metas_pad, rig.intrinsics, rig.img2lidar) |
| gamma, beta = self.ref_head.spatial_alignment_coeffs(cone) |
| params = self.tt.calibration_params(pos.numpy(), gamma.numpy(), beta.numpy()) |
| return {"key": rig.key(), "params": params} |
|
|
| def _build_host(self) -> None: |
| """The host half of the model: configuration, load-time knobs, the per-calibration cache and the per-stream |
| temporal states (no device, no weights; the host tests call it directly).""" |
| from .reference.config import default_config |
|
|
| params = {k.upper(): v for k, v in getattr(self, "compile_params", {}).items() if k.upper() in KNOBS.knobs} |
| self.knobs = KNOBS.read(**params) |
| self.cfg = default_config(normalization=str(self.knobs.NORMALIZATION), |
| use_temporal=bool(self.knobs.USE_TEMPORAL)) |
| self._streams: "OrderedDict[str, Any]" = OrderedDict() |
| self._calib: "OrderedDict[str, Dict[str, Any]]" = OrderedDict() |
| self.calib_fn = None |
| self._img_bits: Optional[np.ndarray] = None |
| self._img_stage = getattr(self, "_img_stage", None) |
|
|
| def _fast_pre(self) -> bool: |
| """HOST_PRE=numba applies (numba importable; else the numpy path).""" |
| from .host import fast as hfast |
|
|
| return str(self.knobs.HOST_PRE) == "numba" and hfast.available() |
|
|
| def _img_target(self) -> np.ndarray: |
| """Where HOST_PRE=numba writes the bf16 upload rows: the persistent ttnn host buffer (zero copy) when the |
| device model exists, else a reused numpy buffer. Calls are serialized by the model lock, and each frame's |
| upload completes inside its ``_forward``, so one buffer suffices.""" |
| if self._img_stage is not None: |
| return self._img_stage.view |
| n, (h, w) = self.cfg.num_cams, tuple(self.cfg.input_hw) |
| if self._img_bits is None: |
| self._img_bits = np.empty((n * h * w, 3), np.uint16) |
| return self._img_bits |
|
|
| def _img_input(self, prepared: Dict[str, Any]) -> Any: |
| """The value of the ``img`` trace input: the staging host tensor, a ttnn host tensor of the bf16 bits, or |
| (HOST_PRE=numpy) the float32 NHWC rows.""" |
| if "img_bits" not in prepared: |
| from .tt.model import img_rows |
|
|
| return self.tt.img_input(img_rows(prepared["img"])) |
| if self._img_stage is not None and prepared["img_bits"] is self._img_stage.view: |
| return self._img_stage.tensor |
| import torch |
| import ttnn |
|
|
| bits = np.ascontiguousarray(prepared["img_bits"]).reshape(self.tt.img_shape) |
| t = torch.from_numpy(bits.view(np.int16)).view(torch.bfloat16) |
| return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT) |
|
|
| def _stream_state(self, sid: str, reset: bool) -> Any: |
| """The host temporal state of ``sid`` (``host.temporal.StreamState``); a new id or ``reset`` starts fresh; |
| beyond ``MAX_STREAMS`` the least recently used stream is dropped (PLAN D16, first publish: one stream).""" |
| from .host.temporal import StreamState |
|
|
| if reset or sid not in self._streams: |
| self._streams[sid] = StreamState(self.cfg, autoware_compat=bool(self.knobs.AUTOWARE_COMPAT), |
| use_temporal=bool(self.knobs.USE_TEMPORAL)) |
| self._streams.move_to_end(sid) |
| while len(self._streams) > int(self.knobs.MAX_STREAMS): |
| self._streams.popitem(last=False) |
| return self._streams[sid] |
|
|
| def _calibration(self, rig: Any) -> Optional[Dict[str, Any]]: |
| """Once-per-calibration constants (PE, cone, spatial-alignment gamma / beta), cached by the rig's hash.""" |
| if self.calib_fn is None: |
| return None |
| key = rig.key() |
| if key not in self._calib: |
| self._calib[key] = self.calib_fn(rig) |
| while len(self._calib) > 8: |
| self._calib.popitem(last=False) |
| self._calib.move_to_end(key) |
| return self._calib[key] |
|
|
| def _prepare(self, points: Any = None, *, images: Any = None, calibration: Optional[Mapping[str, Any]] = None, |
| stream: Optional[Mapping[str, Any]] = None, **_: Any) -> Dict[str, Any]: |
| """Five images + calibration + stream -> the host values of one frame: the network input ``img`` |
| [1, 5, 3, 480, 640] float32 (preset ``NORMALIZATION``), the rig calibration (+ its cached PE constants once |
| the device model is built) and the stream's frame plan (ts, E / E^-1, the ages, fresh / skipped).""" |
| from .host.inputs import prepare_frame |
| from .host.preprocess import preprocess_images |
| from .host.temporal import StreamState |
|
|
| if points is not None: |
| raise tio.InputError("this model takes five camera images, not a point cloud") |
| frame = prepare_frame(images, calibration, stream, self.cfg) |
| st = frame.stream |
| if st.stateless: |
| sid, state = None, StreamState(self.cfg, autoware_compat=bool(self.knobs.AUTOWARE_COMPAT), |
| use_temporal=bool(self.knobs.USE_TEMPORAL)) |
| else: |
| sid, state = st.id, self._stream_state(st.id, st.reset) |
| anchor = frame.anchor_stamp if frame.anchor_stamp is not None else 0.0 |
| T = np.eye(4) if st.stateless else st.T_world_from_ego |
| plan = state.begin(anchor, T, frame.camera_stamps) |
| out: Dict[str, Any] = {"frame": frame, "plan": plan, "state": state, "stream_id": sid} |
| if plan.skipped is None: |
| if self._fast_pre(): |
| from .host import fast as hfast |
|
|
| out["img_bits"] = hfast.preprocess_rows_bf16( |
| frame.images, preset=self.cfg.normalization, dst_hw=tuple(self.cfg.input_hw), |
| fma=bool(self.knobs.PREPROCESS_FMA), threads=int(self.knobs.PREPROCESS_WORKERS), |
| out=self._img_target()) |
| else: |
| out["img"] = preprocess_images(frame.images, self.cfg, fma=bool(self.knobs.PREPROCESS_FMA), |
| workers=int(self.knobs.PREPROCESS_WORKERS)) |
| out["calib"] = self._calibration(frame.rig) |
| return out |
|
|
| def _select_device_memory(self, prepared: Dict[str, Any]) -> None: |
| """Bring the right temporal memory into the device states before the frame (PLAN D16, PORT_LOG S11): |
| the memory of the stream whose frame runs (saved to / loaded from its DRAM bank when another stream or a |
| stateless frame ran in between), zeros for a fresh stream or a stateless frame, and the node's per-frame |
| wipe when ``USE_TEMPORAL`` is off.""" |
| r, sid, plan, state = self.runner, prepared["stream_id"], prepared["plan"], prepared["state"] |
| owner = self._device_owner |
| if owner is not None and owner != sid and owner in self._streams and owner in self._banks: |
| for name in ("emb", "geo"): |
| r.save_state(name, self._banks[owner]) |
| if sid is not None and sid not in self._banks: |
| used = set(self._banks.values()) |
| free_banks = [b for b in range(int(self.knobs.MAX_STREAMS)) if b not in used] |
| if not free_banks: |
| old, bank = self._banks.popitem(last=False) |
| free_banks = [bank] |
| self._banks[sid] = free_banks[0] |
| if sid is not None: |
| self._banks.move_to_end(sid) |
| for gone in [k for k in self._banks if k not in self._streams]: |
| del self._banks[gone] |
| if sid is None or plan.fresh: |
| self.tt.set_memory() |
| elif not state.use_temporal: |
| ref = state.wiped_device_memory() |
| n = self.cfg.pre_memory_length |
| self.tt.set_memory(None, np.tile(ref, (n, 1)), np.zeros((n, 4, 4), np.float32), |
| np.zeros((n, 2), np.float32)) |
| elif owner != sid: |
| for name in ("emb", "geo"): |
| r.load_state(name, self._banks[sid]) |
| self._device_owner = sid |
|
|
| def _forward(self, prepared: Dict[str, Any]) -> Any: |
| """The traced device graph: upload ``img`` and the per-frame ages / relative motion (and the calibration |
| constants when the calibration changed), replay ``frame`` (backbone + head + memory update), read the packed |
| cls / bbox. A skipped frame (camera sync) runs nothing.""" |
| import time |
|
|
| from .tt.model import TtStreamPETR |
| from .tt.params import age_input |
|
|
| plan = prepared["plan"] |
| if plan.skipped is not None: |
| return None |
| t0 = time.perf_counter() |
| self._select_device_memory(prepared) |
| params = self.tt.geometry_params(plan.M) |
| calib = prepared["calib"] |
| if calib["key"] != self._active_calib: |
| params.update(calib["params"]) |
| self._active_calib = calib["key"] |
| inputs = {"img": self._img_input(prepared), |
| "ages": age_input(plan.pre["pre_memory_timestamp"], device_shape=True)} |
| t1 = time.perf_counter() |
| packed = self.runner("frame", inputs=inputs, params=params) |
| t2 = time.perf_counter() |
| prepared["state"].end_device(plan) |
| out = TtStreamPETR.unpack(packed["cls"], packed["reg"]) |
| out["timing_ms"] = {"host_in": (t1 - t0) * 1e3, "trace": (t2 - t1) * 1e3} |
| return out |
|
|
| def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output: |
| """Head outputs (``raw``: ``all_cls_scores`` [1, 5, 5400], ``all_bbox_preds`` [1, 10, 5400]) -> the node's |
| decode of all 6 layers, thresholds, NMS -> ``Detections3D``; a skipped frame (camera sync) -> no objects.""" |
| from .host import postprocess as hpost |
| from .reference.pipeline import to_detections |
|
|
| plan = prepared["plan"] |
| meta = {"stream": prepared["stream_id"], "fresh": bool(plan.fresh), "skipped": plan.skipped, |
| "ts": float(plan.ts), "frame_index": int(plan.frame_index)} |
| if plan.skipped is not None or raw is None: |
| empty = {"objects": hpost.decode(np.zeros((5, 0), np.float32), np.zeros((10, 0), np.float32), self.cfg), |
| "num_pre_nms": 0} |
| return to_detections(empty, self.cfg, model=self.MODEL_NAME, meta=meta) |
| res = hpost.autoware_postprocess(raw["all_cls_scores"], raw["all_bbox_preds"], self.cfg, |
| hpost.params_from(params, self.cfg), |
| fast_nms=str(self.knobs.HOST_POST) == "numba") |
| out = to_detections(res, self.cfg, model=self.MODEL_NAME, meta=meta) |
| out.timing_ms.update({k: round(v, 3) for k, v in (raw.get("timing_ms") or {}).items()}) |
| return out |
|
|
| def _release(self) -> None: |
| """Release the traces and persistent device tensors; also called when ``from_pretrained`` fails half-way.""" |
| runner = getattr(self, "runner", None) |
| if runner is not None: |
| runner.release() |
|
|
| def extra_info(self) -> Dict[str, Any]: |
| """Additions to ``model.info`` and ``/info``: the trace variants, CQs and persistent tensors of the runner, |
| the load-time knobs.""" |
| info: Dict[str, Any] = {} |
| runner = getattr(self, "runner", None) |
| if runner is not None: |
| info["trace"] = runner.describe() |
| knobs = getattr(self, "knobs", None) |
| if knobs is not None: |
| info["knobs"] = knobs.as_dict() |
| return info |
|
|
