code/tt_bevformer/api.py

# SPDX-License-Identifier: Apache-2.0
"""Python API: BEVFormer-small (Autoware tensorrt_bevformer) on one Tenstorrent Blackhole p150.

    from tt_bevformer import BEVFormer

    with BEVFormer.from_pretrained(device_id=0) as model:  # weights -> HF cache, device open, traces captured
        out = model(**load_sample("code/tt_bevformer/samples/synthetic_6cam.json"))  # see PYTHON.md for every input form
    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_bevformer.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit.

This module holds the identity of the port and its hooks, which delegate to ``tt_bevformer.reference`` (host pre- and
post-processing ported from Autoware) and ``tt_bevformer.tt`` (the ttnn graph, run by a ``ttaw.trace.TraceRunner``).
Importing it has no side effects: import ttnn / torch inside the hooks, never at module level.
"""
from __future__ import annotations

from typing import Any, Dict, Mapping, Optional

import numpy as np

from . import io as tio
from .device import DEVICE_DEFAULTS
from .ttaw.api_base import ModelBase
from .ttaw.knobs import Knob, Knobs
from .ttaw.outputs import Detections3D

__all__ = ["BEVFormer", "Output", "KNOBS"]

# The result class of this model (ttaw.outputs: Detections3D, Detections2D, Segmentation3D, Mask2D, Trajectory).
Output = Detections3D

# Load-time switches (read once at build; env BEVFORMER_<NAME>). Each optimization adds its knob here (A/B switch).
KNOBS = Knobs("BEVFORMER", [
    Knob("PRECISION", "HiFi2+fp32+l1acc:w=bf16", "precision of every conv / linear (ttaw.precision spec; PLAN D20 "
                                                 "reduced-precision default): HiFi2, fp32 accumulation, packer L1 "
                                                 "accumulation (fp32 K-block partials), bf16 weights and biases; bf16 "
                                                 "activations. The full-precision path (the frozen gates): "
                                                 "HiFi4+fp32+l1acc:w=fp32 (serve profile 'accurate')"),
    Knob("PRECISION_RULES", "*.conv_offset=HiFi2+fp32+l1acc:w=bf16;img_backbone.layer3.*=HiFi2+fp32+l1acc:w=bfp8",
         "per-module precision overrides (PLAN D20 reduced-precision path): "
                                "'glob=spec;glob=spec' on the module names (img_backbone.layer3.0.conv2, "
                                "img_backbone.layer3.0.conv2.conv_offset, encoder.layers.0.ffn.0, ...), first match "
                                "wins, ttaw.precision specs; modules no rule matches keep PRECISION. A fused "
                                "mmres_k block needs one weight dtype for its conv3 / downsample / next conv1. Default: bfp8 weights for "
         "layer 3 (its DCN offset convs bf16); 'none' = PRECISION everywhere"),
    Knob("DEC_ATTN", "sdpa", "decoder self-attention core: matmul (fp32 scores QK^T + mask, fp32 softmax, bf16 "
                               "probabilities @ V, concat heads) or sdpa (ttnn's flash SDPA with a bf16 pad-key mask, "
                               "heads concatenated by the op; PLAN D20 reduced-precision path)",
         choices=("matmul", "sdpa")),
    Knob("MSDA", "kernel", "multi-scale deformable attention: kernel (msda_k, OPT_PLAN #1: Q16.16 positions, fp32 "
                         "corner and softmax weights, fp32 sums), grid (the fp32-grid decomposition: grid_sample + "
                         "exact 0/1 matmuls) or fused (the stock bf16 op); per call site 'kernel,dec=grid' etc. "
                         "(sites tsa, sca, dec; tt.msda.site_modes)"),
    Knob("MSDA_PX", True, "msda_k sites in pixel units: sampling weights x 2^16 and pixel bases, so the int32 "
                           "positions are one add + typecast (no broadcast scale / offset passes over the per-head "
                           "grids)"),
    Knob("MSDA_VALUE", "tile", "msda_k value source of the TSA / decoder sites: rm (the head maps untilized into L1) or "
                              "tile (the TILE head split in L1, read as face rows: no untilize)", choices=("rm", "tile")),
    Knob("MSDA_LOC", "compact", "msda_k sampling positions of the TSA / SCA sites (with MSDA_PX): heads (the sampling "
                              "projection in 32-column head blocks, split per head, base / shift adds and the int32 "
                              "cast on the split) or compact (only the 2P real columns per head, msda_k reads the "
                              "merged map: no head split, 4x / 2x fewer columns; the same values)",
         choices=("heads", "compact")),
    Knob("MSDA_SOFTMAX", "kernel", "softmax of the TSA / SCA attention logits (with MSDA_LOC=compact): op (head-block "
                               "projection, head split, ttnn.softmax) or kernel (compact projection, msda_k rebuilds "
                               "each head's logit tile and runs ttnn.softmax's LLK sequence itself: same weights)",
         choices=("op", "kernel")),
    Knob("SCA_GATHER", "tile", "SCA rebatching gathers (query rows per camera, the scatter-add's 3 gathers): rm (untilize + "
                             "zero-row pad + ttnn.embedding + tilize) or tile (rowgather_k: rows copied straight "
                             "between TILE maps; the same bytes)", choices=("rm", "tile")),
    Knob("BEV_MM", "fast", "BEV-side linears: default (ttnn.linear's own configs) or tuned (tt.linconv.TUNED_BEV: the "
                              "same K blocking with faster M / N blocking or a 2D grid; bit-identical) or fast (PLAN D20: 1D configs with K "
                              "blocks of up to 8 tiles for the >= 4096-row linears; not bit-identical)",
         choices=("default", "tuned", "fast")),
    Knob("MSDA_PAIRS", "all", "value maps as pixel-pair maps (tt.pairmap_k: 128-byte pages of two neighbouring sticks): "
                              "msda_k reads the two upper / lower corners of a point with one request and tilizes them "
                              "(4x fewer reads; bit-identical): off, sca (the SCA value state) or all (SCA + TSA)",
         choices=("off", "sca", "all")),
    Knob("PAIRMAP", "v2", "pairmap_k (the TSA / SCA value pair maps): v1 (4 face-row reads of 32 bytes + one write per "
                          "128-byte page) or v2 (per tile of 32 pixels x one head: one tile read, a pack_untilize, one "
                          "128-byte write per page: ~1.1 NoC requests per page instead of 5; the same bytes)",
         choices=("v1", "v2")),
    Knob("MSDA_SKIP", False, "msda_k: sampling points wholly outside the map (and the SCA's padding rows) issue no "
                             "reads (weight 0), items strided over the cores"),
    Knob("DCN", "kernel", "DCNv2 sampling of the 26 ResNet conv2: kernel (dcn_k, OPT_PLAN #3: Q16.16 tap positions, "
                       "fp32 bilinear x mask weights, the modulated im2col in one program) or grid (grid_sample + "
                       "mask expand matmul + multiply)", choices=("grid", "kernel")),
    Knob("DCN_IM2COL", "l1", "DCN modulated im2col (kernel mode + tuned GEMM): dram (one [NHW, 9C] tensor in DRAM, "
                               "102 MB per layer-3 call) or l1 (row chunks of <= ~300 KB per core written by dcn_k into "
                               "an L1 height-sharded tensor the GEMM reads in place; same K blocking: bit-identical; "
                               "layer 3 only) or l1_all (layer 4 as well: D20, bf16 weights)",
         choices=("dram", "l1", "l1_all")),
    Knob("DCN_SCALE", "full", "dcn_k position scale: row (a [1, 32] row broadcast multiply) or full (the scale repeated "
                             "over the rows: no broadcast, the same fp32 products)", choices=("row", "full")),
    Knob("DCN_OFFSET_L1", True, "DCN offset convs on ttnn's L1 conv path (no DRAM-slicing unflatten / flatten "
                                 "reshape copies of the conv1 rows)"),
    Knob("BACKBONE_MM", "exact2", "backbone matmul configs (tt.linconv.TUNED, measured K blocks of 4-8 tiles): conv (stock: "
                                "ttnn.conv2d's matmul for the 1x1 convs, ttnn.linear's default for the DCN GEMM), dcn "
                                "(the tuned config for the DCN im2col GEMM only), exact (dcn + the 1x1 / stride-1 convs "
                                "whose ttnn.linear config is bit-identical to conv2d's: tt.linconv.EXACT) or linear "
                                "(dcn + every 1x1 / stride-1 conv as tuned ttnn.linear); exact2: exact plus 2D-grid configs with "
                                "the same K blocking (tt.linconv.EXACT2 / TUNED2; bit-identical)",
         choices=("conv", "dcn", "exact", "exact2", "linear")),
    Knob("BACKBONE_S2", "prod", "the 1x1 / stride-2 convs of layer3.0 / layer4.0 (conv1, downsample): conv (ttnn.conv2d on "
                                "the DRAM-slicing path) or sub (one exact even-pixel row gather shared by both, then "
                                "ttnn.linear; not bit-identical: other K blocking) or prod (PLAN D20 round 3: also "
                                "layer2.0's, and the last block of layers 1-3 writes only the even pixels' sticks "
                                "from its fused mmres_k program: no full-resolution stage output, no untilize / row "
                                "gather / conv2d DRAM slicing; layer 2 not bit-identical, layers 3-4 as sub)",
         choices=("conv", "sub", "prod")),
    Knob("MSDA_FRAC16", True, "msda_mk's readers write each tap fraction q as the fp32 bits of 1 + q / 2^16 (a shift "
                               "and an or instead of the leading-zero fp32 conversion per point) and the compute "
                               "kernel's SFPU subtracts 1 (exact: bit-identical)"),
    Knob("SCA_COPY", "kernel", "the SCA value state's DRAM -> L1 copy before each SCA call: op (ttnn.to_memory_config) or "
                           "kernel (repage_k: batched page reads / writes on both data-movement RISC-Vs; bit-identical)",
         choices=("op", "kernel")),
    Knob("DEC_VALUE", "fused", "the decoder's cross-attention values: split (value projection to DRAM, then "
                               "nlp_create_qkv_heads into L1) or fused (the projection straight to L1, read by msda_k "
                               "head-interleaved: no head split; bit-identical)", choices=("split", "fused")),
    Knob("S2_SKIP", True, "with BACKBONE_S2=prod: the producing mmres_k computes / reads / writes only the tile rows "
                           "with a pixel on an even image row (the others feed no even pixel; the same values: "
                           "bit-identical)"),
    Knob("BACKBONE_C3", "l1_2", "the layer-1 3x3 convs: slice (ttnn.conv2d's DRAM-slicing path) or l1 (the L1 path with "
                                 "8-tile activation blocks; bit-identical) or l1_2 (also the layer-2 3x3 convs on the "
                                 "L1 path, their outputs read in place by mmres_k)", choices=("slice", "l1", "l1_2")),
    Knob("RES_L1", True, "layer 3's conv3 (1x1, 256 -> 1024) writes its output in L1 shards (conv2d's K blocking) and "
                          "the residual add + ReLU reads it in place, writing the block output to DRAM (bit-identical)"),
    Knob("BACKBONE_RES", "next", "conv3 + residual add + ReLU of the layer-1 / layer-2 bottlenecks: stock (ttnn.linear "
                                 "with the EXACT2 config, then ttnn.add: the 180 / 90 MB maps go to DRAM and back) or "
                                 "fused (tt.mmres_k: one kernel replays the linear's and the add's compute steps, the "
                                 "layer1.0 downsample in the same kernel; bit-identical) or next (fused, and the same "
                                 "kernel also computes the next block's conv1 from the rows it wrote; bit-identical)",
         choices=("stock", "fused", "next")),
    Knob("DEC_ATTN_L1", True, "decoder self-attention: the [8, 928, 928] scores, masked scores and probabilities in L1 "
                                "instead of DRAM (the same ops and configs; bit-identical)"),
    Knob("KERNEL_TILIZE", "std", "tilize inside the custom kernels dcn_k / msda_k: std (tilize_block) or fast (the "
                                 "Blackhole fast tilize; a lossless bf16 reorder: bit-identical)", choices=("std", "fast")),
    Knob("DCN_GLUE", "kernel", "DCN offset glue (scale multiply, base add, int32 cast, sigmoid on the [NHW, 32] offset "
                             "rows): dram or l1 (the intermediates in L1; the same ops: bit-identical) or kernel "
                             "(dcn_k runs glue_k's steps per block on the offset conv's own L1 shards: no glue "
                             "program, no sharded-to-interleaved copy; bit-identical)",
         choices=("dram", "l1", "kernel")),
    Knob("DCN_SUM", "mul", "dcn_k's weighted sum of the 4 bilinear corners (with DCN_BATCH): mul (fp32 column-broadcast "
                           "multiplies with packer-L1 accumulation, then a bf16 copy pass) or mm (PLAN D20: the corner "
                           "weights as diagonal fp32 tiles, out = sum_c diag(w_c) @ X_c as one DEST accumulation of "
                           "matmuls, one bf16 pack)", choices=("mul", "mm")),
    Knob("DCN_SAMPLE", "bilinear", "dcn_k's sampling of each DCN tap: bilinear (4 corners; the reference semantics) or "
                                  "nearest (PLAN D20 round 3: the nearest pixel, round half up: one gathered stick "
                                  "per tap, the items alternating between the two gather RISC-Vs, mask x tile, no "
                                  "corner weights / sum; not bit-identical, judged by the task gates) or nearest_l4 "
                                  "(layer 4's DCNs only)",
         choices=("bilinear", "nearest", "nearest_l4")),
    Knob("DCN_PAIRS", False, "dcn_k reads a horizontal pixel-pair map of its input (dcnpair_k: page (n, y, i) = pixels "
                             "i - 1 and i): one 2C-wide request per sampling row and corner pair instead of two "
                             "(half the gather requests; the same values: bit-identical)"),
    Knob("DCN_OUT", "l1", "with RES3_PARTS: where the layer-3 DCN GEMM writes its row chunks for the fused block's "
                             "mmres_k: dram or l1 (interleaved L1: no 11 MB DRAM round trip per block; the same "
                             "values)", choices=("dram", "l1")),
    Knob("DCN_BATCH", True, "dcn_k: the 4 corners of an item tilized back to back, then the 4 weighted multiplies "
                              "(one tilize and one multiply init per item instead of four; the same products and "
                              "accumulation order: bit-identical)"),
    Knob("BACKBONE_C3_OUT", "shard", "with BACKBONE_RES fused / next: the layer-1 3x3 conv outputs stay height-sharded "
                                    "in L1 (shard) and mmres_k reads them in place, or go to DRAM (dram)",
         choices=("dram", "shard")),
    Knob("STEM", "plain", "stem input: plain (NHWC 3-channel rows, the 7x7 / stride-2 conv) or s2d (space-to-depth "
                          "16-channel rows uploaded as 368 x 640 maps, the equivalent 4x4 / stride-1 conv: 4x fewer "
                          "and 32-byte upload pages / DRAM slices)", choices=("plain", "s2d")),
    Knob("UPLOAD", "rows", "image upload layout: pixels (one 6-byte page per pixel: H2D 141 ms) or rows (one 7,680-byte "
                             "page per image row: H2D 4 ms, then one exact 7.8 ms ROW_MAJOR reshape on the device inside "
                             "the image trace)", choices=("pixels", "rows")),
    Knob("STEM_IN", "kpre", "stem input with UPLOAD=rows and STEM=plain: slice (the image rows reshaped to 6-byte pixel "
                             "pages, the stem conv on ttnn's DRAM-slicing path) or kin (per camera a K-IN kernel copies "
                             "the rows into the conv's L1 input shards, the conv on the L1 path, the pool per camera; "
                             "bit-identical; tt.kin) or kpre (the uint8 node images go up, one input per camera, "
                             "and tt.kpre runs the C14 pre-processing on the device straight into those shards: "
                             "bit-identical to HOST_PRE + kin; the host keeps only the node resize of other sizes)",
         choices=("slice", "kin", "kpre")),
    Knob("OUTPUTS", "tile", "the bev trace's outputs: pack (pack_outputs: one ROW_MAJOR row, one D2H copy) or tile (the "
                            "two fp32 TILE maps read as they are: no device reshape)", choices=("pack", "tile")),
    Knob("MAX_STREAMS", 4, "BEV-history streams kept on the device (PLAN D16: one prev_bev bank per stream id; a new "
                           "id beyond it evicts the least recently used one)"),
    Knob("MAX_GAP_S", 2.0, "a stream restarts its history (use_prev_bev = 0) when its timestamp jumps by more than "
                           "this many seconds or goes backwards (PLAN D13: reset on a large time gap)"),
    Knob("SHIFT_MODE", "surrogate", "TSA reference shift: DerryHub's atan2 surrogate (the deployed graph) or atan2 "
                                    "(training); they differ only for dx == 0 or dx < 0 == dy (SPEC 4.3)",
         choices=("surrogate", "atan2")),
    Knob("CAN_BUS_LAYOUT", "training", "can_bus[7:16] slot order of the ego dynamics: BEVFormer training (accel, "
                                       "rotation rate, velocity) or the Autoware node's literal KinematicState "
                                       "order (PORT_LOG.md E3)", choices=("training", "autoware")),
    Knob("HOST_PRE", "bf16", "host pre-processing (node resize + C14 preset): bf16 (host.fast.preprocess_rows_bf16: "
                             "the node resize, the preset and the bf16 cast of the upload fused in one numba pass per "
                             "camera, bit-exact), numba (host.fast: bit-exact twins, NHWC float32 rows written "
                             "directly, RGB read as decoded) or numpy (host.preprocess); the numba paths fall back to "
                             "numpy when numba is not installed (bf16 also to numba with STEM=s2d)",
         choices=("numpy", "numba", "bf16")),
    Knob("HOST_MALLOC", "keep", "host heap: default (glibc's) or keep (mallopt: large per-frame buffers reuse the "
                                   "retained heap instead of fresh mmap pages; host.fast.keep_heap)",
         choices=("default", "keep")),
    Knob("RES3_NG", 4, "layer-3 conv3 + residual add + ReLU (with RES_L1): 0 = L1-sharded linear + ttnn.add (stock), "
                       "2 / 4 = one mmres_k program with the fp32 conv3 weight split by columns over that many core "
                       "groups (bit-identical)", choices=(0, 2, 4, 8)),
    Knob("ACT8", True, "with BACKBONE_RES fused / next (PLAN D20 path): the layer-1 / layer-2 residual stream between "
                        "the fused blocks of a stage in bfp8 (mmres_k writes and reads it; the stage outputs stay "
                        "bf16): about half the DRAM traffic of these DRAM-bound programs"),
    Knob("ACT8_L3", True, "with RES3_NEXT (PLAN D20 path): the layer-3 residual stream between the fused blocks in bfp8 "
                           "(each block's output is read only by the next block's mmres_k; the last block's output "
                           "stays bf16)"),
    Knob("RES3_NEXT", True, "with RES3_NG > 0 (PLAN D20 path): the layer-3 mmres_k program also computes the next "
                             "block's conv1 (1024 -> 256, ReLU) from the rows it writes (one core group, whole rows; "
                             "the conv3 and conv1 weights resident in L1 together: needs bfp8 layer-3 weights, "
                             "PRECISION_RULES)"),
    Knob("RES3_RM", "on", "with RES3_NEXT (PLAN D20 path): the fused layer-3 program writes the next block's conv1 "
                           "output as the row-major L1 map that the DCN's offset conv and dcn_k read (untilized in "
                           "the kernel: no untilize program, no DRAM round trip of the [22080, 256] map; the same "
                           "values): off, on (L1 interleaved) or shard (height shards of the offset conv's layout: "
                           "no interleaved-to-sharded copy either: dcn_k's sharded reads are slow) or dual (the interleaved "
                           "map for dcn_k plus a height-sharded twin for the offset conv)",
         choices=("off", "on", "shard", "dual")),
    Knob("RES3_L1", False, "with RES3_NEXT and ACT8_L3 (PLAN D20 path): the layer-3 bfp8 residual stream between the "
                            "fused blocks (24 MB, read only by the next block's mmres_k) in L1 interleaved instead of "
                            "DRAM (the same values)"),
    Knob("L4_L1", False, "(PLAN D20 path) the layer-4 1x1 convs (conv1, conv3, the layer4.0 strided pair) write L1 "
                          "interleaved, so the residual adds keep the block outputs in L1 ([5536, 2048] maps: no DRAM "
                          "round trips; the same values)"),
    Knob("RES3_PARTS", True, "with RES3_NG > 0: the fused layer-3 block reads the DCN GEMM's row chunks in place "
                              "(no [22080, 256] chunk concat)"),
    Knob("UPLOAD_PIPE", True, "with HOST_PRE=bf16, UPLOAD=rows and STEM_IN=kin: one image-rows input per camera, "
                               "pre-processed in a worker thread while the calling thread uploads the cameras "
                               "already done (the H2D of camera i overlaps the pre-processing of camera i + 1; "
                               "api.preprocess then excludes the image pass, api.device includes it)"),
    Knob("BEV_USE", "variants", "TSA queue blend: blend (one bev trace: use * [rot(prev_bev); bev_queries] + (1 - "
                                "use) * [query; query], 3 binary ops per blend, 5 blends per frame) or variants (a "
                                "second bev trace, bev_prev, replayed when use_prev_bev == 1: the blends are the "
                                "identities, no ops; the bev trace for fresh / reset frames; bit-identical)",
         choices=("blend", "variants")),
    Knob("TSA_VIN_SHARED", True, "with BEV_USE=variants: the TSA value input [rot(prev_bev); bev_queries] (the same "
                                  "in every layer of bev_prev) is concatenated once per frame instead of once per "
                                  "layer (bit-identical)"),
    Knob("DEC_QK", "grid", "decoder self-attention QK^T [8, 928, 32] x [8, 32, 928]: default (ttnn's 2D config, 4 x 4 "
                           "cores) or grid (3 x 3 tiles per core on a 10 x 10 grid; K is one tile: bit-identical)",
         choices=("default", "grid")),
    Knob("GLUE_K", True, "the fp32 glue in front of msda_k (int32(cols + base [+ shift]): 2-3 ops per TSA / SCA call) "
                          "and dcn_k (int32(off * S + B), sigmoid(off): 4 ops per DCN) as one tt.glue_k program each, "
                          "with the stock LLK calls in the same order (bit-identical)"),
    Knob("TSA_VP_L1", True, "with MSDA_PAIRS=all: the TSA value projection writes L1 directly for pairmap_k (no DRAM "
                             "-> L1 copy of the [2, 22528, 256] projection; the same ttnn.linear program: "
                             "bit-identical)"),
    Knob("MSDA_MK", "all", "the D19 MSDA block module megakernel (tt.msda_mk): one program per encoder MSDA call runs the "
                           "query concat, the compact sampling + attention projections (fp32 weights resident in L1, "
                           "in0_block_w 1, packer L1 accumulation: ttnn.linear's steps), the loc glue (glue_k's steps) "
                           "and the msda_k gather / softmax / sum, instead of 5 (TSA) / 4 (SCA) programs; needs "
                           "MSDA=kernel, MSDA_PX, MSDA_LOC=compact, MSDA_SOFTMAX=kernel, MSDA_PAIRS: off, sca, tsa or "
                           "all (bit-identical; fp32 weights and HiFi4 whatever PRECISION says)", choices=("off", "sca", "tsa", "all")),
    Knob("SCA_ACTIVE", True, "with MSDA_MK sca / all: the SCA megakernel computes only the active 32-row blocks of "
                               "each camera's rebatched rows (a per-calibration runtime page, sca_nblk: the work "
                               "split follows the rig's coverage; the padding rows behind each camera's rows are "
                               "never gathered, so the outputs are bit-identical)"),
    Knob("SCA_BASE", "kernel", "with MSDA_MK sca / all: frame (the SCA's pixel-unit compact base map base * S + O is "
                              "derived from sca_base in every bev replay: multiply, add, untilize, slice, repeat, "
                              "tilize) or kernel (msda_mk reads the raw sca_base, its reader duplicates the point "
                              "columns into both column halves of a tile and the front applies S and O with the "
                              "same SFPU mul / add before the cols + base add: bit-identical)",
         choices=("frame", "kernel")),
    Knob("SCA_SCATTER", "kernel", "with SCA_GATHER=tile: the SCA tail (the 3 row gathers of the camera outputs, the 2 "
                                 "adds and the 1 / cameras multiply) as stock programs or as one scatter_k program "
                                 "(kernel: the stock ops' FPU add / SFPU multiply steps; bit-identical)",
         choices=("stock", "kernel")),
    Knob("MMRES_MCAST", True, "mmres_k (the fused conv3 + residual [+ next conv1] programs of layers 1-3): the resident "
                              "weights are read from DRAM once per grid row by the row's first core and multicast "
                              "along the row (on) instead of every core reading all of them (off; the layer-3 "
                              "program read 120 x 556 KB of weights per call); the same bytes: bit-identical"),
    Knob("NECK_L1", True, "the FPN 3x3 conv on ttnn's L1 path (its [5536, 256] input from DRAM: no DRAM-slicing "
                            "unflatten / flatten reshape copies)"),
    Knob("ENC_L1", "all", "encoder intermediates in L1 interleaved instead of DRAM (the same ttnn.linear / msda_mk "
                          "programs: the same values): off, on (the FFN hidden map [22528, 512] and the TSA "
                          "megakernel output) or all (also the TSA / SCA output projections and the FFN output; "
                          "max = also the FFN input LayerNorm and the query + position sum was measured: L1 clash)",
         choices=("off", "on", "all")),
    Knob("MSDA_MK_MCAST", True, "with MSDA_MK: the megakernel's resident projection weights (fp32) are read from DRAM "
                                "once per grid row by the row's first core and multicast along the row (on) instead "
                                "of every core reading all of them (off); the same bytes: bit-identical"),
    Knob("HOST_RESIZE", "fast", "with STEM_IN=kpre: the node's uint8 INTER_LINEAR resize of non-node-size cameras "
                                "(e.g. 1920x1080): fast (host.fast._node_images_fast: all cameras in one parallel "
                                "numba pass, the column phases unrolled with constant weights, 32-bit arithmetic; "
                                "node-size images copied into the staging buffers in parallel) or rows (one row-cached "
                                "pass per camera); bit-identical", choices=("fast", "rows")),
    Knob("HOST_THREADS", 4, "threads of the host image pre-processing (C14 preset, one camera per thread)"),
    Knob("VALUES_DIRECT", True, "with MSDA_PAIRS: the SCA value projection writes L1 directly and pairmap_k writes "
                                 "the sca_value states itself (no DRAM -> L1 copy of the projection, no copy of the "
                                 "maps into the states; bit-identical)"),
    Knob("SCA_VALUE", "l1", "with MSDA_PAIRS: msda_k's SCA value source: l1 (the DRAM state copied into L1 per layer) "
                            "or dram (read in place from the state; bit-identical)", choices=("l1", "dram")),
    Knob("HOST_STAGING", True, "with STEM_IN=kpre: the six node images are written straight into persistent ttnn "
                                "host buffers (a pool of up to 4 sets; ttaw HostStaging) by the node resize (or one copy "
                                "for node-size images) instead of a ttnn.from_torch copy per frame"),
    Knob("PIPELINE", 1, "concurrent calls (OPT_PLAN R5-3): 0 = one call at a time under the model lock; 1 = per-stage "
                        "ordering: the request parsing, node resize and image host tensors of frame k + 1 run while "
                        "frame k is on the device, the post-processing runs outside every lock (frames take the "
                        "device in the order they finished pre-processing; outputs equal serial calls); 2 = also "
                        "the 26 MB camera upload of frame k + 1 on CQ1 while frame k's bev trace runs (needs a "
                        "device opened with 2 command queues, BEVFORMER_NUM_CQS=2; otherwise 1)",
         choices=(0, 1, 2)),
])



def tt_model_kwargs(k: Any, *, banks: int = 0, stem_in: Optional[str] = None, img_cams: bool = False) -> Dict[str, Any]:
    """The ``tt.model.TtBEVFormer`` keyword arguments for the knob values ``k`` (``KNOBS.read()``): the one mapping
    shared by the API and the device PCC gates (``tests/test_pcc_device.py``), so the gates measure the shipped
    device path. ``stem_in`` overrides the STEM_IN mapping (the API resolves kpre against numba; the gates feed the
    golden network input through the K-IN rows path); ``img_cams`` is the API's UPLOAD_PIPE decision."""
    if stem_in is None:
        stem_in = "kin" if str(k.STEM_IN) == "kpre" else str(k.STEM_IN)
    return dict(precision=str(k.PRECISION), banks=int(banks), msda=str(k.MSDA), dcn=str(k.DCN),
                dcn_offset_l1=bool(k.DCN_OFFSET_L1), msda_px=bool(k.MSDA_PX), msda_tile=k.MSDA_VALUE == "tile",
                backbone_mm=str(k.BACKBONE_MM), msda_skip=bool(k.MSDA_SKIP), stem=str(k.STEM),
                upload=str(k.UPLOAD), stem_in=stem_in, msda_loc=str(k.MSDA_LOC),
                backbone_s2=str(k.BACKBONE_S2), msda_softmax=str(k.MSDA_SOFTMAX),
                dcn_im2col=str(k.DCN_IM2COL), sca_gather=str(k.SCA_GATHER), bev_mm=str(k.BEV_MM),
                dcn_scale=str(k.DCN_SCALE), backbone_c3=str(k.BACKBONE_C3), outputs=str(k.OUTPUTS),
                res_l1=bool(k.RES_L1), msda_pairs=str(k.MSDA_PAIRS), backbone_res=str(k.BACKBONE_RES),
                dec_attn_l1=bool(k.DEC_ATTN_L1), kernel_tilize=str(k.KERNEL_TILIZE), dcn_glue=str(k.DCN_GLUE),
                dcn_batch=bool(k.DCN_BATCH), c3_out=str(k.BACKBONE_C3_OUT), img_cams=bool(img_cams),
                res3_ng=int(k.RES3_NG), res3_parts=bool(k.RES3_PARTS), bev_prev=str(k.BEV_USE) == "variants",
                vin_shared=bool(k.TSA_VIN_SHARED), dec_qk=str(k.DEC_QK), glue_k=bool(k.GLUE_K),
                tsa_vp_l1=bool(k.TSA_VP_L1), values_direct=bool(k.VALUES_DIRECT), sca_value=str(k.SCA_VALUE),
                msda_mk=str(k.MSDA_MK), sca_active=bool(k.SCA_ACTIVE),
                sca_base=str(k.SCA_BASE), sca_scatter=str(k.SCA_SCATTER),
                precision_rules=str(k.PRECISION_RULES), dec_attn=str(k.DEC_ATTN), res3_next=bool(k.RES3_NEXT), act8=bool(k.ACT8),
                act8_l3=bool(k.ACT8_L3), mmres_mcast=bool(k.MMRES_MCAST),
                msda_mk_mcast=bool(k.MSDA_MK_MCAST), res3_rm={"off": False, "on": True}.get(str(k.RES3_RM), str(k.RES3_RM)),
                dcn_sum=str(k.DCN_SUM), dcn_pairs=bool(k.DCN_PAIRS), dcn_sample=str(k.DCN_SAMPLE), s2_skip=bool(k.S2_SKIP), msda_frac16=bool(k.MSDA_FRAC16), sca_copy=str(k.SCA_COPY), dec_value=str(k.DEC_VALUE),
                dcn_out=str(k.DCN_OUT), pairmap=str(k.PAIRMAP),
                res3_l1=bool(k.RES3_L1), l4_l1=bool(k.L4_L1),
                enc_l1=str(k.ENC_L1), neck_l1=bool(k.NECK_L1))


class BEVFormer(ModelBase):
    """BEVFormer-small (Autoware tensorrt_bevformer) on one Blackhole p150. Create it with :meth:`from_pretrained`."""

    MODEL_NAME = "bevformer-p150"
    ENV_PREFIX = "BEVFORMER"  # prefix of the environment knobs (SERVING.md section 3.4)
    # No AutowareFoundation repo exists for BEVFormer (Autoware's bevformer_small.onnx is not obtainable, SPEC
    # section 1): the user-decided mirror of the official bevformer_small_epoch_24.pth state_dict (PLAN.md D13).
    DEFAULT_REPO = "changh95/bevformer-small-weights"
    DEFAULT_TAG = None  # the mirror has no tags; DEFAULT_REVISION is the commit `main` held when it was validated
    DEFAULT_REVISION = "0e6d9d9c7b8aca95820b264491e19a72d10f0608"
    WEIGHTS_LICENSE = ("Apache-2.0 (changh95/bevformer-small-weights mirror card; BEVFormer authors' checkpoint, "
                       "trained on nuScenes, CC BY-NC-SA 4.0)")
    ALLOW_PATTERNS = ["bevformer_small_epoch_24.safetensors", "provenance.json"]
    VARIANTS = ["default"]  # load-time: selects weights files and trace shapes
    DEFAULT_VARIANT = "default"
    INPUT_KIND = "multicam"  # lidar | camera | multicam | lidar+multicam | planner
    # The network's camera order: nuScenes info order, baked into the learned cams_embeds[6] (SPEC section 2.1)
    CAMERA_ORDER = ("CAM_FRONT", "CAM_FRONT_RIGHT", "CAM_FRONT_LEFT", "CAM_BACK", "CAM_BACK_LEFT", "CAM_BACK_RIGHT")
    POINT_FIELDS = tio.DEFAULT_POINT_FIELDS  # columns the network consumes, in order
    # The Autoware label the node publishes per network class (AW/src/ros_utils.cpp:70-94: bus, trailer and
    # construction_vehicle -> TRUCK; barrier, traffic_cone -> UNKNOWN); the network class names are
    # reference.config.CLASS_NAMES (meta.class_names)
    LABELS = ["CAR", "TRUCK", "TRUCK", "TRUCK", "TRUCK", "UNKNOWN", "MOTORCYCLE", "BICYCLE", "PEDESTRIAN", "UNKNOWN"]
    # Per-request knobs: name -> (type, min, max, default). Only host-side post-processing belongs here; anything
    # that changes a device shape or the weights is a load-time option (BUNDLE_CONVENTIONS.md section 9).
    RUNTIME_PARAMS = {
        "score_threshold": (float, 0.0, 1.0, 0.2),   # post_process_params.score_thre (AW/config/bevformer.param.yaml)
        "max_num": (int, 1, 9000, 300),              # the node's top-k over 900 x 10 scores (get_parameter_or 300)
        "physical_boxes": (bool, None, None, False),  # False: the node's message (size w,l,h; yaw -rot+pi); True:
    }                                                 # size l,w,h and the length axis' heading -rot-pi/2 (SPEC 5.1)
    # Model-specific input keyword arguments handed to _prepare (e.g. PointPainting ("rois",)); planner models set
    # INPUT_SCHEMA = {name: (shape, dtype)} of their ONNX inputs (ttaw API.md section 10).
    # ``ego``: the frame's CAN dynamics {"accel", "rotation_rate", "velocity"} (base_link, training layout of
    # can_bus[7:16]; host.temporal); the ego pose itself comes in stream.T_world_from_ego
    EXTRA_INPUTS = ("ego",)
    DEVICE_DEFAULTS = DEVICE_DEFAULTS  # validated open parameters (device.py)

    # ---- port-specific hooks (called by ModelBase; keep host work out of _forward) ---------------------------
    def _build(self) -> None:
        """Weights (``reference.weights.BEVFormerWeights``: the mirror's safetensors, BN folded in float64) ->
        device-form parameters (``tt.params``) -> the device graph ``tt.model.TtBEVFormer`` (TraceRunner variants
        ``image`` and ``bev``, persistent inputs / RT-dev params / states allocated here; no capture: ``_warm_one``)
        and the host state (per-stream temporal state, calibration cache)."""
        self._build_host()
        from .reference.weights import WEIGHTS_FILE, BEVFormerWeights, find_weights
        from .tt.model import TtBEVFormer
        from .tt.params import BEVFormerParams

        path = find_weights(self.weights_path)
        if path is None:
            raise FileNotFoundError(f"{WEIGHTS_FILE} not found in {self.weights_path}")
        self.net_weights = BEVFormerWeights(path)
        self.params = BEVFormerParams(self.net_weights)
        streams = int(self.knobs.MAX_STREAMS)
        self.tt = TtBEVFormer(self.device, self.params,
                              **tt_model_kwargs(self.knobs, banks=streams if streams > 1 else 0,
                                                stem_in="kpre" if self._kpre() else None,
                                                img_cams=self._upload_pipe()))
        self.runner = self.tt.runner

    def _kpre(self) -> bool:
        """STEM_IN=kpre applies (the device pre-processing needs the plain stem and numba for the node resize)."""
        from .host import fast as hfast

        return str(self.knobs.STEM_IN) == "kpre" and str(self.knobs.STEM) == "plain" and hfast.available()

    def _upload_pipe(self) -> bool:
        """UPLOAD_PIPE applies (per-camera image inputs pre-processed while the cameras before them upload)."""
        from .host import fast as hfast

        k = self.knobs
        return (bool(k.UPLOAD_PIPE) and str(k.HOST_PRE) == "bf16" and hfast.available() and str(k.STEM) == "plain"
                and str(k.UPLOAD) == "rows" and str(k.STEM_IN) == "kin")

    def _build_host(self) -> None:
        """The host half of the model: knobs, per-stream temporal states and the calibration cache (no device, no
        weights; the host tests call it directly)."""
        params = {k.upper(): v for k, v in getattr(self, "compile_params", {}).items() if k.upper() in KNOBS.knobs}
        self.knobs = KNOBS.read(**params)
        if int(self.knobs.MAX_STREAMS) < 1:
            raise ValueError("BEVFORMER_MAX_STREAMS must be >= 1")
        from .tt.msda import site_modes

        site_modes(str(self.knobs.MSDA))             # an invalid BEVFORMER_MSDA spec fails here, before the device
        self._temporal: Dict[str, Any] = {}          # stream id -> host.temporal.TemporalState
        self._geometry: Dict[bytes, Any] = {}        # lidar2img bytes -> per-calibration geometry (small LRU)
        self._cal_key: Optional[bytes] = None        # the calibration whose tables the device inputs hold
        self.streams = None                          # ttaw.trace.StreamBanks, created after the captures
        self._pipe_init(int(self.knobs.PIPELINE))
        if str(self.knobs.HOST_MALLOC) == "keep":
            from .host.fast import keep_heap

            keep_heap()

    def _warm_one(self, variant: Any) -> None:
        """Capture every trace variant (warm-up + strict capture), create the per-stream state banks, then run one
        synthetic frame through the whole host + device path so the first real call is as fast as the later ones;
        the warm-up stream is forgotten."""
        from .tt.warmup import WARMUP_STREAM, synthetic_request
        from .ttaw.trace import StreamBanks

        self.runner.capture()
        if self.pipeline == 2 and self.runner.num_command_queues != 2:
            import warnings

            warnings.warn("BEVFORMER_PIPELINE=2 needs a device opened with 2 command queues (BEVFORMER_NUM_CQS=2); "
                          "running PIPELINE=1", RuntimeWarning)
            self._pipe_init(1)
        self.streams = StreamBanks(self.runner, ["prev_bev"], max_streams=int(self.knobs.MAX_STREAMS),
                                   on_full="evict", max_gap_s=float(self.knobs.MAX_GAP_S))
        with self._lock:
            prepared = self._prepare(**synthetic_request())
            raw = self._forward(prepared)
            self._postprocess(raw, prepared, self.validate_params({}))
            self.streams.forget(WARMUP_STREAM)
            self._temporal.pop(WARMUP_STREAM, None)
        if self._kpre() and str(self.knobs.HOST_RESIZE) == "fast":
            from .host import fast as hfast

            # HOST_RESIZE=fast: the 1920x1080 kernel compiles in the background (~6 s; numba's compiler lock), the
            # rows kernel serves that size until it is ready
            hfast.warm_node_resize(((1080, 1920),))

    def _geometry_for(self, lidar2img: np.ndarray) -> Dict[str, Any]:
        """Per-calibration host geometry (``host.geometry``: float64 projection, SCA rebatching tables) and the
        device inputs built from it, cached by the bytes of ``lidar2img`` (a static rig computes them once).
        A rig whose cameras see more BEV cells than the static SCA capacity is refused (HTTP 400; PORT_LOG.md open
        question 4)."""
        from .host import geometry as hgeo
        from .tt.model import calibration_inputs

        key = np.ascontiguousarray(lidar2img, np.float32).tobytes()
        geo = self._geometry.pop(key, None)
        if geo is None:
            ref_cam, valid = hgeo.point_sampling(lidar2img)
            try:
                tables = hgeo.sca_tables(valid)
            except hgeo.SCAOverflow as e:
                raise tio.InputError(f"calibration not supported: {e}") from None
            geo = {"ref_cam": ref_cam, **tables}
            geo["inputs"] = calibration_inputs(geo)
            geo["key"] = key
        self._geometry[key] = geo                    # most recent last
        while len(self._geometry) > 8:
            self._geometry.pop(next(iter(self._geometry)))
        return geo

    def _prepare(self, points: Any = None, *, images: Any = None, calibration: Optional[Mapping[str, Any]] = None,
                 stream: Optional[Mapping[str, Any]] = None, ego: Optional[Mapping[str, Any]] = None,
                 **_: Any) -> Dict[str, Any]:
        """Six images + calibration + stream (+ ``ego`` dynamics) -> the host values of one frame with the node's
        semantics (``tt_bevformer.host``): decode + node transport + C14 preset -> the network input; lidar2img ->
        the per-calibration geometry (cached); the stream's history (device banks + host can_bus state: a new id,
        ``stream.reset`` or a time gap restart it with ``use_prev_bev = 0``) -> can_bus, TSA shift, rotation map.
        Two halves: :meth:`_prepare_host` (stateless per frame) and :meth:`_prepare_stream` (the stream history;
        with ``PIPELINE`` it runs in the device stage, in frame order)."""
        host = self._prepare_host(points, images=images, calibration=calibration, stream=stream, ego=ego)
        try:
            return self._prepare_stream(host)
        except BaseException:
            self._release_staging(host)
            raise

    def _prepare_host(self, points: Any = None, *, images: Any = None,
                      calibration: Optional[Mapping[str, Any]] = None, stream: Optional[Mapping[str, Any]] = None,
                      ego: Optional[Mapping[str, Any]] = None, **_: Any) -> Dict[str, Any]:
        """The first half of :meth:`_prepare`: request parsing, the image pre-processing of the active path and the
        per-calibration geometry; the stream fields are validated here (client errors before the device stage).
        Touches no device state."""
        from .host.inputs import prepare_request
        from .host.preprocess import preprocess_images
        from .host.temporal import can_bus_from_transform

        if points is not None:
            raise tio.InputError("this model takes six camera images, not a point cloud")
        from .host import fast as hfast

        pre = str(self.knobs.HOST_PRE)
        use_fast = (pre in ("numba", "bf16") or self._kpre()) and hfast.available()
        frame = prepare_request(images, calibration, stream, ego, bgr=not use_fast)
        can_bus_abs = frame.can_bus_abs
        if str(self.knobs.CAN_BUS_LAYOUT) != "training":
            T = frame.stream.get("T_world_from_ego")
            T = np.eye(4) if T is None else tio.parse_transform(T, field="stream.T_world_from_ego")
            can_bus_abs = can_bus_from_transform(T, dict(ego or {}), layout=str(self.knobs.CAN_BUS_LAYOUT))
        stem = str(self.knobs.STEM)
        rows = None
        pipe = self._upload_pipe()
        kpre = self._kpre()
        node_u8 = None
        if kpre:          # STEM_IN=kpre: only the node resize (other sizes) on the host; the rest on the device
            stage = self._staging_set()
            node_u8 = hfast.node_images(frame.images_bgr, threads=int(self.knobs.HOST_THREADS),
                                        out=None if stage is None else [b.view for b in stage],
                                        kernel=str(self.knobs.HOST_RESIZE))
            rows = image = None
        elif pipe:        # UPLOAD_PIPE: _forward pre-processes the cameras while it uploads them
            rows = image = None
        elif use_fast and pre == "bf16" and stem != "s2d":   # the bf16 upload rows in one pass (uint16 bits)
            rows = hfast.preprocess_rows_bf16(frame.images_bgr, rgb=True, threads=int(self.knobs.HOST_THREADS))
            image = None
        elif use_fast:   # the upload rows straight from the RGB images (NHWC: image is their NCHW view, no copy)
            rows = hfast.preprocess_rows(frame.images_bgr, rgb=True, threads=int(self.knobs.HOST_THREADS),
                                         layout="s2d" if stem == "s2d" else "nhwc")
            image = None if stem == "s2d" else rows.transpose(0, 3, 1, 2)[None]
        else:
            image = preprocess_images(frame.images_bgr, workers=int(self.knobs.HOST_THREADS))
        geo = self._geometry_for(frame.lidar2img)
        st = dict(stream or {})
        sid = str(st.get("id", "default"))
        reset = st.get("reset", False)
        if not isinstance(reset, bool):
            raise tio.InputError("stream.reset must be a boolean")
        ts = st.get("timestamp_s")
        if ts is not None:
            try:
                ts = float(ts)
            except (TypeError, ValueError):
                raise tio.InputError("stream.timestamp_s must be a number") from None
        return {"image": image, "geometry": geo, "stream_id": sid, "reset": reset, "timestamp_s": ts,
                "can_bus_abs": can_bus_abs, "rows": rows, "node_u8": node_u8, "pipe": pipe, "stem": stem,
                "images_bgr": frame.images_bgr if pipe else None, "staging": stage if kpre else None}

    def _staging_set(self) -> Optional[list]:
        """HOST_STAGING: a free set of six persistent uint8 ttnn host buffers (``ttaw.tensors.HostStaging``) from a
        pool that grows to at most 4 sets; None (the ``from_torch`` path) when off, without a device (host tests),
        when all sets are in flight or when the buffers have no zero-copy view. Never blocks. A set returns to the
        pool once its frame's upload was enqueued (:meth:`_release_staging`); a set whose frame never uploads
        (e.g. a failed call) is simply dropped."""
        if not bool(self.knobs.HOST_STAGING) or getattr(self, "runner", None) is None:
            return None
        import queue

        pool = getattr(self, "_stage_pool", None)
        if pool is False:
            return None
        if pool is None:
            pool = self._stage_pool = queue.Queue()
            self._stage_sets = 0
        try:
            return pool.get_nowait()
        except queue.Empty:
            pass
        if self._stage_sets >= 4:
            return None
        from .reference import config as C
        from .tt.model import NODE_U8_SHAPE
        from .ttaw.tensors import HostStaging

        bufs = [HostStaging(NODE_U8_SHAPE, "uint8") for _ in range(C.NUM_CAMS)]
        if not all(b.zero_copy for b in bufs):
            self._stage_pool = False
            return None
        self._stage_sets += 1
        return bufs

    def _release_staging(self, host: Dict[str, Any]) -> None:
        stage = host.pop("staging", None) if isinstance(host, dict) else None
        if stage is not None:
            self._stage_pool.put(stage)

    def _prepare_stream(self, host: Dict[str, Any]) -> Dict[str, Any]:
        """The second half of :meth:`_prepare`: select the stream (device banks), step its host can_bus state and
        build the runner's frame inputs. Enqueues device copies on a stream switch / reset, so with ``PIPELINE`` it
        runs in the device stage after the previous frame's replays were enqueued."""
        from .host import geometry as hgeo
        from .host.temporal import TemporalState
        from .tt.model import frame_inputs

        image, geo, sid, stem = host["image"], host["geometry"], host["stream_id"], host["stem"]
        reset, ts, can_bus_abs, pipe = host["reset"], host["timestamp_s"], host["can_bus_abs"], host["pipe"]
        rows, node_u8 = host["rows"], host["node_u8"]
        out: Dict[str, Any] = {"image": image, "geometry": geo, "stream_id": sid, "staging": host.get("staging")}
        if self.streams is None:                     # host tests: no device state (the host values only)
            fresh = reset or sid not in self._temporal
        else:
            fresh = self.streams.select(sid, reset=reset, timestamp_s=ts)
            for old in [k for k in self._temporal if k not in self.streams.streams]:
                self._temporal.pop(old)
        state = self._temporal.setdefault(sid, TemporalState())
        can_bus, use = state.step(can_bus_abs, reset=fresh)
        shift = hgeo.tsa_shift(can_bus, str(self.knobs.SHIFT_MODE))
        rot = hgeo.rotation_index_map(float(can_bus[17]))
        up = frame_inputs(image, can_bus=can_bus, use_prev_bev=float(use[0]), shift=shift, rot_index=rot, stem=stem,
                          img_rows=rows, wide=str(self.knobs.UPLOAD) == "rows", img_cams=pipe, node_u8=node_u8)
        if host.get("staging") is not None:          # HOST_STAGING: the host tensors that hold the images already
            for i, b in enumerate(host["staging"]):
                up["inputs"][f"img{i}"] = b.tensor
        if pipe:
            out["pending_images"] = list(host["images_bgr"])
        out.update({"upload": up, "can_bus": can_bus, "use_prev_bev": use, "fresh": bool(fresh)})
        return out

    def _forward(self, prepared: Dict[str, Any]) -> Any:
        """Upload the frame (and the calibration tables when the calibration changed), replay ``image`` then
        ``bev`` back to back, read the packed last-layer outputs (one D2H). Returns ``outputs_classes`` /
        ``outputs_coords`` ``[1, 1, 900, 10]`` and the stage times."""
        import time

        up = prepared["upload"]
        inputs = dict(up["inputs"])
        geo = prepared["geometry"]
        if geo["key"] != self._cal_key:
            inputs.update(geo["inputs"])
        runner = self.runner
        t0 = time.perf_counter()
        try:
            runner.upload(inputs=inputs, params=up["params"])
        finally:
            self._release_staging(prepared)
        if prepared.get("pending_images") is not None:
            self._upload_cameras(prepared["pending_images"])
        self._cal_key = geo["key"]
        bev = self.tt.bev_variant(up["params"]["use"])
        runner.replay("image")
        runner.replay(bev)
        packed = runner.read(bev)
        t1 = time.perf_counter()
        out = self.tt.unpack(packed)
        out["timing_ms"] = {"device_run": (t1 - t0) * 1e3}
        return out

    # ---- pipelined calls (PIPELINE, OPT_PLAN R5-3) ---------------------------------------------------------
    def _pipe_init(self, mode: int) -> None:
        """Stage state of ``PIPELINE``: frames get a ticket when their host half is done (under ``_pre_lock``, so
        in pre-processing order) and take the device in ticket order (``_turn``); ``_img_turn`` is the last ticket
        whose image replay is enqueued (a CQ1 upload of the next frame's cameras waits for it). Every ttnn call runs
        under the model lock ``_lock``; only the waits for the device (``ttnn.event_synchronize``, which releases
        the GIL) run outside it."""
        import threading

        self.pipeline = int(mode)
        self.pipelined = self.pipeline > 0
        self._pre_lock = threading.Lock()
        self._cv = threading.Condition()
        self._ticket = 0          # next ticket
        self._turn = 0            # ticket whose device stage may run
        self._img_turn = -1       # last ticket whose image replay is enqueued (or that gave up its turn)
        self._waiting = 0         # threads in a device wait outside the model lock
        self._img_done = None     # CQ0 event after the last enqueued image replay (PIPELINE=2)

    def _wait_for(self, pred) -> None:
        with self._cv:
            while not pred():
                if getattr(self, "_closed", True):
                    raise RuntimeError(f"{type(self).__name__} is closed")
                self._cv.wait(0.5)

    def _host_images(self, host: Dict[str, Any]) -> Dict[str, Any]:
        """The camera inputs of a frame as ttnn host tensors (the 26 MB uint8 copy into ttnn's host storage,
        3-4 ms), built in the pre-processing stage so that the device stage only uploads them."""
        if host.get("node_u8") is None:
            return {}
        if host.get("staging") is not None:          # HOST_STAGING: already ttnn host tensors
            return {f"img{i}": b.tensor for i, b in enumerate(host["staging"])}
        import torch

        from .reference import config as C
        from .tt.model import NODE_U8_SHAPE
        from .ttaw.tensors import to_host_tensor

        return {f"img{i}": to_host_tensor(torch.from_numpy(np.ascontiguousarray(a, np.uint8).reshape(NODE_U8_SHAPE)),
                                          "uint8")
                for i, a in enumerate(host["node_u8"][: C.NUM_CAMS])}

    def _upload_images_cq1(self, imgs: Dict[str, Any]) -> Any:
        """PIPELINE=2: write the camera inputs on CQ1 once the previous frame's image replay has read them (CQ1
        waits for that replay's event on the device; the host returns as soon as the data is queued). Returns the
        CQ1 event the next image replay waits for. Called under the model lock."""
        import ttnn

        if self._img_done is not None:
            ttnn.wait_for_event(1, self._img_done)
        for name, host in imgs.items():
            ttnn.copy_host_to_device_tensor(host, self.tt.inputs[name], cq_id=1)
        return ttnn.record_event(self.device, 1)

    def _device_stage(self, host: Dict[str, Any], imgs: Dict[str, Any], img_event: Any, ticket: int) -> Any:
        """One frame's device stage with ``PIPELINE``: stream selection + frame inputs, uploads (the cameras
        already on CQ1 when ``img_event`` is given), both replays and a non-blocking read of the bev outputs into
        fresh host tensors, all enqueued on CQ0 under the model lock; the next frame gets its turns as soon as this
        frame's image replay (``_img_turn``: its CQ1 camera upload may follow) and read-back (``_turn``) are
        enqueued, so its replays queue right behind this frame's on the device. Then the wait outside the lock.
        Returns ``(prepared, raw)``."""
        import time

        import ttnn

        from .ttaw.tensors import to_numpy
        from .ttaw.trace import _rebuild

        runner = self.runner
        with self._lock:
            self._check_open()
            t0 = time.perf_counter()
            prepared = self._prepare_stream(host)
            up = prepared["upload"]
            inputs = dict(up["inputs"])
            inputs.update(imgs)                       # ttnn host tensors (to_host_tensor passes them through)
            if img_event is not None:
                for name in imgs:
                    inputs.pop(name, None)
            geo = prepared["geometry"]
            if geo["key"] != self._cal_key:
                inputs.update(geo["inputs"])
            try:
                runner.upload(inputs=inputs, params=up["params"])
            finally:
                self._release_staging(host)
            if prepared.get("pending_images") is not None:   # STEM_IN=kin + UPLOAD_PIPE
                self._upload_cameras(prepared["pending_images"])
            if img_event is not None:
                ttnn.wait_for_event(0, img_event)
            self._cal_key = geo["key"]
            runner.replay("image")
            if self.pipeline == 2:
                self._img_done = ttnn.record_event(self.device, 0)
            with self._cv:
                self._img_turn = max(self._img_turn, ticket)
                self._cv.notify_all()
            bev = self.tt.bev_variant(up["params"]["use"])
            runner.replay(bev)
            # the read-back enqueued now (CQ0, behind the replay; before the next frame's replays can overwrite
            # the outputs) into this frame's own host tensors
            trace = runner._trace_for(bev, runner._last_phase[bev])
            hosts = [ttnn.allocate_tensor_on_host(t.spec, self.device) for _, t, _ in trace.leaves]
            for (_, t, _), h in zip(trace.leaves, hosts):
                ttnn.copy_device_to_host_tensor(t, h, blocking=False, cq_id=0)
            done = ttnn.record_event(self.device, 0)
            with self._cv:
                self._waiting += 1
                self._turn = max(self._turn, ticket + 1)
                self._cv.notify_all()
        try:
            ttnn.event_synchronize(done)              # releases the GIL: the next frame's host half runs meanwhile
        finally:
            with self._cv:
                self._waiting -= 1
                self._cv.notify_all()
        t1 = time.perf_counter()
        values = {}
        for (path, _, layout), h in zip(trace.leaves, hosts):
            arr = to_numpy(h)
            values[path] = layout.unpack(arr) if layout is not None else arr
        structure = trace.outputs
        from .ttaw.trace import Packed

        if isinstance(structure, Packed) or not isinstance(structure, (Mapping, list, tuple)):
            packed = values[()]
        else:
            packed = _rebuild(structure, values)
        raw = self.tt.unpack(packed)
        raw["timing_ms"] = {"device_run": (t1 - t0) * 1e3}
        return prepared, raw

    def __call__(self, points: Any = None, *, sweeps: Any = None, images: Any = None,
                 calibration: Optional[Mapping[str, Any]] = None, stream: Optional[Mapping[str, Any]] = None,
                 inputs: Any = None, **params: Any):
        """One frame (``ModelBase.__call__``). With ``BEVFORMER_PIPELINE`` (OPT_PLAN R5-3) concurrent calls overlap:
        the host half of the pre-processing (and the image host tensors) under ``_pre_lock``, then the device stage
        in the same frame order (the stream history steps there, so a stream's frames keep their order), the
        post-processing under no lock. ``PIPELINE=2`` also uploads the cameras on CQ1 while the previous frame's
        bev trace runs. Same outputs as serial calls in the same frame order; ``timing_ms`` adds ``wait`` (time
        queued for a stage)."""
        if not getattr(self, "pipelined", False):
            return super().__call__(points, sweeps=sweeps, images=images, calibration=calibration, stream=stream,
                                    inputs=inputs, **params)
        import time

        self._check_open()
        extra = {k: params.pop(k) for k in list(params) if k in self.EXTRA_INPUTS}
        p = self.validate_params(params)
        t0 = time.perf_counter()
        with self._pre_lock:
            self._check_open()
            t1 = time.perf_counter()
            host = self._prepare_host(points, images=images, calibration=calibration, stream=stream, **extra)
            try:
                imgs = self._host_images(host)
            except BaseException:
                self._release_staging(host)
                raise
            ticket = self._ticket
            self._ticket += 1
            t2 = time.perf_counter()
        try:
            img_event = None
            if self.pipeline == 2 and imgs:
                self._wait_for(lambda: self._img_turn >= ticket - 1)
                with self._lock:
                    self._check_open()
                    img_event = self._upload_images_cq1(imgs)
                self._release_staging(host)
            t3 = time.perf_counter()
            self._wait_for(lambda: self._turn == ticket)
            t4 = time.perf_counter()
            prepared, raw = self._device_stage(host, imgs, img_event, ticket)
            t5 = time.perf_counter()
        finally:
            self._release_staging(host)               # (no-op once the upload released it)
            with self._cv:                            # a failed frame gives up its turns too
                # ... but only once the frames before it had theirs: a frame that failed before its device turn
                # (e.g. the CQ1 upload) must not advance _turn past a frame still in its device stage
                while self._turn < ticket and not getattr(self, "_closed", True):
                    self._cv.wait(0.5)
                self._img_turn = max(self._img_turn, ticket)
                self._turn = max(self._turn, ticket + 1)
                self._cv.notify_all()
        out = self._postprocess(raw, prepared, p)
        t6 = time.perf_counter()
        out.timing_ms.update({"preprocess": (t2 - t1) * 1e3, "upload_cq1": (t3 - t2) * 1e3, "device": (t5 - t4) * 1e3,
                              "postprocess": (t6 - t5) * 1e3, "wait": ((t1 - t0) + (t4 - t3)) * 1e3,
                              "total": (t6 - t0) * 1e3})
        return out

    predict = __call__

    def image_rows(self, images: Any, on_camera: Any = None) -> np.ndarray:
        """Six RGB uint8 node-order images -> the bf16 upload rows (uint16 bits ``[6, 736, 3840]``,
        ``host.fast.preprocess_rows_bf16``) in a buffer reused across frames."""
        from .host import fast as hfast
        from .reference import config as C

        buf = getattr(self, "_rows_buf", None)
        if buf is None:
            buf = self._rows_buf = np.empty((C.NUM_CAMS, C.IMG_HW[0], C.IMG_HW[1] * 3), np.uint16)
        return hfast.preprocess_rows_bf16(images, rgb=True, threads=int(self.knobs.HOST_THREADS), out=buf,
                                          on_camera=on_camera)

    def _upload_cameras(self, images: Any) -> None:
        """UPLOAD_PIPE: a worker thread pre-processes the cameras one after the other (numba, nogil) while this
        thread uploads each camera's rows as soon as they are written (the H2D of camera i overlaps camera i + 1)."""
        import queue
        import threading

        import torch

        from .reference import config as C
        from .tt.model import CAM_ROWS_SHAPE

        q: "queue.Queue" = queue.Queue()

        def work() -> None:
            try:
                self.image_rows(images, on_camera=lambda i, r: q.put((i, r)))
            except BaseException as e:  # noqa: BLE001 -- re-raised in the calling thread
                q.put((-1, e))

        th = threading.Thread(target=work, name="bevformer-pre", daemon=True)
        th.start()
        try:
            for _ in range(C.NUM_CAMS):
                i, r = q.get()
                if i < 0:
                    raise r
                t = torch.from_numpy(r.view(np.int16)).view(torch.bfloat16).reshape(CAM_ROWS_SHAPE)
                self.runner.upload(inputs={f"img{i}": t})
        finally:
            th.join()

    def _postprocess(self, raw: Any, prepared: Dict[str, Any], params: Dict[str, Any]) -> Output:
        """The node's post-processing (``host.postprocess``: last decoder layer, sigmoid, top ``max_num``, score >
        ``score_threshold``, ``post_center_range``, no NMS) -> ``Detections3D`` in LIDAR_TOP with the Autoware labels;
        ``raw`` holds ``outputs_classes`` / ``outputs_coords`` (the last layer suffices); ``meta.stream`` says which
        history the frame used (``use_prev_bev``)."""
        from .host import postprocess as hpost

        dets = hpost.decode(raw["outputs_classes"], raw["outputs_coords"], score_threshold=params["score_threshold"],
                            max_num=params["max_num"])
        out = hpost.to_output(dets, model=self.MODEL_NAME, physical=params["physical_boxes"])
        if isinstance(prepared, dict) and "stream_id" in prepared:
            out.meta["stream"] = {"id": prepared["stream_id"], "fresh": bool(prepared.get("fresh", False)),
                                  "use_prev_bev": float(np.asarray(prepared.get("use_prev_bev", [0.0])).reshape(-1)[0])}
        out.timing_ms.update({k: round(float(v), 3) for k, v in (raw.get("timing_ms") or {}).items()})
        return out

    def _release(self) -> None:
        """Release the traces, the persistent device tensors and the device weights; also called when
        ``from_pretrained`` fails half-way. With ``PIPELINE`` it first waits (up to 60 s) for the device waits of
        frames still in flight (``close()`` has marked the model closed, so they stop at their next stage)."""
        import time

        cv = getattr(self, "_cv", None)
        if cv is not None:
            end = time.monotonic() + 60.0
            with cv:
                while self._waiting > 0 and time.monotonic() < end:
                    cv.wait(0.5)
        tt = getattr(self, "tt", None)
        if tt is not None:
            tt.release()
        else:
            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 and the stream banks."""
        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()
        streams = getattr(self, "streams", None)
        if streams is not None:
            info["streams"] = streams.describe()
        return info