# 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_). 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