# SPDX-License-Identifier: Apache-2.0 """Python API: BEVFusion-L (Autoware autoware_bevfusion, lidar-only) on one Tenstorrent Blackhole p150. from tt_bevfusion import BEVFusion with BEVFusion.from_pretrained(device_id=0) as model: # weights -> HF cache, device open, traces captured out = model("code/tt_bevfusion/samples/test_pcd.npz") # 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_bevfusion.server.app``) calls this class, so ``/predict`` and ``model(...)`` agree bit for bit. This module holds the identity of the port and its hooks: host pre- and post-processing ported from Autoware live in ``tt_bevfusion.host`` (shared with the CPU reference ``tt_bevfusion.reference``, so device and reference see identical voxel, feature and rulebook tensors), the ttnn graph in ``tt_bevfusion.tt`` (run by a ``ttaw.trace.TraceRunner``, one trace per capacity bucket; milestones M1-M4 of PORT_LOG.md). Importing it has no side effects: ttnn / torch are imported inside the hooks. Inputs (``points``, plus optional ``sweeps`` or ``stream``; host/preprocess.py has the full contract): - a cloud in ``base_link`` with fields x, y, z (an ``intensity`` column is accepted and ignored: the v2.0 network takes (x, y, z, time_lag)); - a column named ``time_lag`` marks an already densified cloud (Autoware's 2-sweep input), used as is; - ``sweeps=[{"points", "time_lag_s", "T_current_from_sweep"}]``: past sweeps densified on the host (stateless); - ``stream={"id", "timestamp_s", "T_world_from_ego", "reset"}``: Autoware's stateful densification per stream id (``densification_num_past_frames: 1``); a missing pose refuses the frame without caching it (Autoware's TF failure). Output: ``Detections3D`` with Autoware ``DetectedObjects`` semantics: ObjectClassification label ids (``LABELS``), ``existence_probability`` as the score, the network yaw (counter-clockwise from +x, z at the box centre), (length, width, height); ``meta`` holds the voxel counts (the 256,000 clip), the sparse level counts and capacity bucket, the per-stage detection counts, the per-object orientation availability and proposal index. """ from __future__ import annotations import contextlib import queue import threading import time from typing import Any, Dict, Optional, Sequence from . import io as tio from .device import DEVICE_DEFAULTS from .host.postprocess import LABELS, RUNTIME_PARAMS, BEVFusionPostprocessor from .host.preprocess import LEVEL_NAMES, BEVFusionPreprocessor, Prepared from .reference.config import CAPACITY_BUCKETS, DEFAULT_VARIANT, VARIANTS, BEVFusionConfig from .ttaw.api_base import ModelBase from .ttaw.outputs import Detections3D __all__ = ["BEVFusion", "Output"] # The result class of this model (ttaw.outputs: Detections3D, Detections2D, Segmentation3D, Mask2D, Trajectory). Output = Detections3D class BEVFusion(ModelBase): """BEVFusion-L (Autoware autoware_bevfusion, lidar-only) on one Blackhole p150. Create it with :meth:`from_pretrained`.""" MODEL_NAME = "bevfusion-p150" ENV_PREFIX = "BEVFUSION" # prefix of the environment knobs (SERVING.md section 3.4) DEFAULT_REPO = "AutowareFoundation/bevfusion" DEFAULT_TAG = "v2.0" # the Autoware ansible artifacts pin # the commit DEFAULT_TAG points to (pinned: tags can move) DEFAULT_REVISION = "e1bf164b909d3c6c2642012f9b5ee4d827753220" ALLOW_PATTERNS = ["bevfusion_lidar.onnx", "ml_package_bevfusion_lidar.param.yaml", "detection_class_remapper.param.yaml", "deploy_metadata.yaml"] VARIANTS = list(VARIANTS) # lidar (the camera-lidar profile is planned: PLAN.md D5) DEFAULT_VARIANT = DEFAULT_VARIANT INPUT_KIND = "lidar" # lidar | camera | multicam | lidar+multicam | planner CAMERA_ORDER = () # multicam models: the fixed camera order of the network POINT_FIELDS = tio.DEFAULT_POINT_FIELDS # default column names of an unnamed (N, 4) array # output labels: autoware_perception_msgs ObjectClassification (label_id = the message value); the network's five # classes (CAR, TRUCK, BUS, BICYCLE, PEDESTRIAN) map onto them, the remapper adds TRAILER LABELS = list(LABELS) # Per-request knobs (host post-processing only; defaults are Autoware's): host/postprocess.py RUNTIME_PARAMS = RUNTIME_PARAMS EXTRA_INPUTS = () 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: """Host side (config, pre- / post-processors), then the ttnn graph of ``tt_bevfusion.tt`` (``TtBEVFusion``: C28 sparse encoder + to-dense, SECOND, SECONDFPN, dense head, local max, query head) as the variants of a ``ttaw.trace.TraceRunner``: one sparse segment per capacity bucket (its persistent inputs) writing the ``bev`` state, one dense segment (no capture here; ``_warm_one`` captures). Precision policy and knobs (``BEVFUSION_PRECISION``, ``BEVFUSION_*_TERMS``) are read once here.""" from .reference.weights import load_weights from .tt.model import TtBEVFusion, build_runner self._build_host() self.graph = TtBEVFusion(self.device, load_weights(self.weights_path, self.config)) self.runner = build_runner(self.graph, variants=("dense",), name=f"{self.MODEL_NAME}/{self.variant}") def _build_host(self, *, max_streams: int = 16) -> None: """The host half of ``_build`` (no device): usable alone by host tests and tools.""" from .tt.model import KNOBS self.config = BEVFusionConfig.from_weights(self.weights_path, self.variant) k = KNOBS.read() from .host import malloc as _malloc self.malloc = _malloc.apply(str(k.HOST_MALLOC)) # H7: process-wide, once # OPT_PLAN M0 (BEVFUSION_DEVICE_PLAN=1): the host builds the level sets only; the device builds the maps plan = "levels" if int(k.DEVICE_PLAN) else str(k.HOST_PLAN) self.pre = BEVFusionPreprocessor(self.config, max_streams=max_streams, plan=plan, voxelizer=str(k.HOST_VOXELIZER), densify=str(k.HOST_DENSIFY)) self.post = BEVFusionPostprocessor(self.config, nms=str(k.HOST_NMS)) concurrent = True if {k.HOST_PLAN, k.HOST_VOXELIZER, k.HOST_INDEX, k.HOST_NMS, k.HOST_DENSIFY} - {"numpy", 0, 1}: from .host import fast fast.THREADS = max(1, int(k.HOST_THREADS)) fast.warm_up() # compile the host twins at load, not on the first frame concurrent = fast.concurrent_ok() # OPT_PLAN T1 (BEVFUSION_PIPELINE=1): per-stage locks (see __call__). The pre-processing lock also guards the # staged-input writes in _forward when numba's threading layer cannot run two parallel twins at once self.pipeline = int(k.PIPELINE) self.pipelined = self.pipeline > 0 self.pipeline_decode = bool(int(k.PIPELINE_DECODE)) self._pre_lock = threading.Lock() self._stage_guard = contextlib.nullcontext() if concurrent else self._pre_lock # T1b (PIPELINE=2): the staged-input writes as their own stage (one writer at a time) into one of two # ping-pong staging sets; a set returns to the pool once its frame's device stage has finished self._stage_lock = threading.Lock() if concurrent else self._pre_lock self._slots: "queue.Queue[int]" = queue.Queue() for slot in range(2): self._slots.put(slot) def _warm_one(self, variant: Dict[str, Any]) -> None: """Warm up and capture the traces of one entry of ``default_warmup_variants()``. ``TraceRunner.capture`` runs every pending variant eagerly (kernel JIT, program cache) before any capture, then captures with program-cache misses forbidden; it is idempotent.""" self.runner.capture() # HOST_STAGING=1: the persistent host input buffers, at load (PIPELINE=2: both ping-pong sets) self.graph.prepare_staging(slots=2 if getattr(self, "pipeline", 0) == 2 else 1) def _prepare(self, points: Any = None, *, sweeps: Optional[Any] = None, stream: Optional[Any] = None, images: Any = None, calibration: Any = None, inputs: Any = None, **_: Any) -> Prepared: """Host pre-processing (``host/preprocess.py``): densification, hard voxelization, the VFE mean and the sparse rulebooks with the frame's capacity bucket. Client mistakes raise ``tio.InputError`` (HTTP 400).""" if points is None: raise tio.InputError("this model needs 'points'") if images or inputs is not None: raise tio.InputError("the lidar profile of BEVFusion takes a point cloud only (no 'images' / 'inputs')") prepared = self.pre(points, sweeps=sweeps, stream=stream) plan = prepared.plan if plan is not None and plan.overflow: # No clipping (VERIFY_PORT L2): a frame whose sparse level counts exceed the largest capacity bucket is # refused with an explicit error before the device runs. None of the 122 measured frames came close. caps = dict(zip(LEVEL_NAMES, CAPACITY_BUCKETS[plan.bucket])) over = {lv: f"{n:,} > {caps[lv]:,}" for lv, n in plan.counts().items() if n > caps[lv]} raise tio.InputError( f"frame exceeds the sparse encoder's largest capacity bucket '{plan.bucket}' " f"({prepared.num_voxels:,} voxels; active rows per level {over}); this port does not clip " "(README: Caveats)") return prepared def _forward(self, prepared: Prepared, inputs: Optional[Dict[str, Any]] = None) -> Any: """Upload the frame's voxel features and gather tables into the persistent device inputs of its capacity bucket, replay that trace, read the packed outputs back (one D2H) and unpack them into ``{bbox_pred, score, label_pred}``. A frame without voxels runs nothing. Runs under the model lock.""" from .tt.model import sparse_variant if prepared.skipped is not None: return None plan = prepared.plan if not getattr(self, "pipelined", False): inputs = self.graph.frame_inputs(prepared.vfe, plan, plan.bucket, staged=True) self.runner.run(sparse_variant(plan.bucket), inputs=inputs) # sparse segment -> the bev state raw = self.runner("dense") # dense segment, one packed readback self._device_pairs(plan) return self.graph.assemble(raw) import ttnn if inputs is None: with self._stage_guard: inputs = self.graph.frame_inputs(prepared.vfe, plan, plan.bucket, staged=True) self.runner.run(sparse_variant(plan.bucket), inputs=inputs) self.runner.run("dense") # T1: wait for the replays in synchronize_device, which releases the GIL (a blocking read does not), so the # next frame's pre-processing runs meanwhile; the read-back then finds the output ready ttnn.synchronize_device(self.device) raw = self.runner.read("dense") self._device_pairs(plan) return self.graph.assemble(raw) def _device_pairs(self, plan: Any) -> None: """M0: the frame's pair counts (``meta["sparse"]["pairs"]``) from the device rulebook (a small read-back after the replays).""" if getattr(self.graph, "device_plan", False): plan.pair_counts = self.graph.read_pairs(self.runner, plan.bucket) def __call__(self, points: Any = None, *, sweeps: Optional[Sequence] = None, images: Any = None, calibration: Optional[Any] = None, stream: Optional[Any] = None, inputs: Any = None, **params: Any): """One frame (``ModelBase.__call__``). With ``BEVFUSION_PIPELINE=1`` (OPT_PLAN T1) concurrent calls overlap: the host pre-processing runs under its own lock (stream densification state, the persistent host tables), the upload + replays + read-back under the model lock, the post-processing under none, so one frame's pre-processing runs while another frame is on the device. ``PIPELINE=2`` (T1b) adds a third stage between them: the staged-input writes into one of two ping-pong staging sets under their own lock. Same outputs as serialized calls; ``timing_ms`` adds ``wait`` (time spent queued for the locks and a staging set) and, with 2, ``stage``.""" if not getattr(self, "pipelined", False): return super().__call__(points, sweeps=sweeps, images=images, calibration=calibration, stream=stream, inputs=inputs, **params) self._check_open() extra = {k: params.pop(k) for k in list(params) if k in self.EXTRA_INPUTS} p = self.validate_params(params) if inputs is not None and self.INPUT_SCHEMA is not None: inputs = tio.load_named_arrays(inputs, self.INPUT_SCHEMA) t0 = time.perf_counter() decode_ms = 0.0 if getattr(self, "pipeline_decode", False) and points is not None: # T1c: the request decode is stateless; outside the stage lock concurrent calls decode in parallel points = tio.load_points(points) decode_ms = (time.perf_counter() - t0) * 1e3 t0 = time.perf_counter() with self._pre_lock: self._check_open() t1 = time.perf_counter() prepared = self._prepare(points=points, sweeps=sweeps, images=images, calibration=calibration, stream=stream, inputs=inputs, **extra) t2 = time.perf_counter() slot, staged, stage_ms = None, None, 0.0 try: if self.pipeline == 2 and prepared.skipped is None and prepared.plan is not None: slot = self._slots.get() with self._stage_lock: self._check_open() ts = time.perf_counter() plan = prepared.plan staged = self.graph.frame_inputs(prepared.vfe, plan, plan.bucket, staged=True, slot=slot) stage_ms = (time.perf_counter() - ts) * 1e3 with self._lock: self._check_open() # close() may have run while this call waited for the lock t3 = time.perf_counter() raw = self._forward(prepared, staged) t4 = time.perf_counter() finally: if slot is not None: self._slots.put(slot) out = self._postprocess(raw, prepared, p) t5 = time.perf_counter() out.timing_ms.update({"preprocess": (t2 - t1) * 1e3 + decode_ms, "device": (t4 - t3) * 1e3, "postprocess": (t5 - t4) * 1e3, "wait": ((t1 - t0) + (t3 - t2)) * 1e3 - stage_ms, "total": (t5 - t0) * 1e3 + decode_ms}) if self.pipeline == 2: out.timing_ms["stage"] = stage_ms return out predict = __call__ def _postprocess(self, raw: Any, prepared: Prepared, params: Dict[str, Any]) -> Output: """Host post-processing (``host/postprocess.py``): decode, distance-binned thresholds, stable sort, circle NMS, ROS conversion, IoU-BEV NMS, class remapper -> ``Detections3D``; ``raw`` = ``{bbox_pred, score, label_pred}`` or None for a skipped frame (no objects, as the node publishes nothing).""" meta = {"variant": self.variant, **prepared.stats()} if raw is None: return self.post.empty(model=self.MODEL_NAME, frame_id=prepared.frame_id, meta={**meta, "counts": {}}) return self.post(raw, params, model=self.MODEL_NAME, frame_id=prepared.frame_id, meta=meta) def _release(self) -> None: """Release the traces, persistent tensors and device weights; also called when ``from_pretrained`` fails half-way (tolerates a partial build).""" runner = getattr(self, "runner", None) if runner is not None: runner.release() self.runner = None graph = getattr(self, "graph", None) if graph is not None: graph.release() self.graph = None def extra_info(self) -> Dict[str, Any]: """Additions to ``model.info`` and ``/info``: the config, the host state and the runner.""" info: Dict[str, Any] = {} if getattr(self, "config", None) is not None: info["config"] = self.config.describe() if getattr(self, "pre", None) is not None: info["preprocess"] = self.pre.describe() runner = getattr(self, "runner", None) if runner is not None: info["trace"] = runner.describe() graph = getattr(self, "graph", None) if graph is not None: info["graph"] = graph.describe() return info