# SPDX-License-Identifier: Apache-2.0 """Eager bring-up runners of the device stages, teacher-forced from the CPU-reference goldens (``code/scripts/device_stage_check.py``; PLAN 4.4: every new stage eager x2 before any trace). Production and the gated tests run the same stage modules through ``tt.model.TtMETEOR``'s traces.""" from __future__ import annotations from typing import Any, Dict, Mapping import numpy as np from ..reference import config as C from ..ttaw.metrics import pcc from ..ttaw.ops.conv import FeatureMap from .layers import Consts from .unpack import split_det2d, to_reference __all__ = ["STAGES", "compare_taps_to_golden", "read", "upload_rows", "golden_rows"] def read(t: Any) -> np.ndarray: import ttnn a = ttnn.to_torch(t.tensor if isinstance(t, FeatureMap) else t) return a.float().numpy() if a.is_floating_point() else a.numpy() def golden_rows(nchw: np.ndarray) -> np.ndarray: """NCHW -> NHWC rows [1, 1, N*H*W, C] float32.""" a = np.asarray(nchw, np.float32) n, c, h, w = a.shape return np.ascontiguousarray(a.transpose(0, 2, 3, 1).reshape(1, 1, n * h * w, c)) def upload_rows(dev: Any, rows: np.ndarray, *, dtype: str = "bfloat16", layout: str = "tile", shape=None) -> Any: import torch import ttnn from ..ttaw.tensors import ttnn_dtype a = np.ascontiguousarray(rows if shape is None else rows.reshape(shape)) return ttnn.from_torch(torch.from_numpy(a), dtype=ttnn_dtype(dtype), layout=ttnn.TILE_LAYOUT if layout == "tile" else ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) def compare_taps_to_golden(taps: Mapping[str, np.ndarray], gold: Mapping[str, np.ndarray]) -> Dict[str, Any]: """PCC (float) / agreement (class maps) of every tap that has a golden; det2d merged maps are split.""" rows: Dict[str, Any] = {} for k, v in taps.items(): if k.startswith("det2d_s"): s = int(k[-1]) hm, reg = split_det2d(v, s) for name, a in ((f"hm2d_s{s}", hm), (f"reg2d_s{s}", reg)): rows[name] = {"pcc": round(float(pcc(a, gold[name])), 7), "max_abs": float(np.abs(a - gold[name]).max())} continue if k not in gold: continue g = np.asarray(gold[k]) a = np.asarray(v) if a.shape != g.shape and a.size == g.size: a = a.reshape(g.shape) if a.shape != g.shape: rows[k] = {"error": f"shape {a.shape} vs golden {g.shape}"} continue if g.dtype == np.uint8 or g.dtype == np.bool_: rows[k] = {"agreement": round(float((a == g).mean()), 7)} else: a64, g64 = a.astype(np.float64), g.astype(np.float64) rows[k] = {"pcc": round(float(pcc(a64, g64)), 7), "max_abs": float(np.abs(a64 - g64).max()), "ref_max": float(np.abs(g64).max())} return rows class _Stage: def __init__(self, dev: Any, params: Any): self.dev = dev self.params = params self.consts = Consts() def release(self) -> None: self.consts.release() class ImageStage(_Stage): """uint8 cameras -> encoder + image heads (every tap).""" def __init__(self, dev: Any, params: Any): super().__init__(dev, params) from .image import TtImageBranch self.branch = TtImageBranch(params, self.consts) self.shape = self.branch.img_shape # IMG_SHAPE, or whole rows with IMG_IN=rows def run(self, gold: Mapping[str, np.ndarray]) -> Dict[str, np.ndarray]: imgs = np.asarray(gold["input.imgs"], np.uint8)[0] # [8, 3, 432, 768] rows = np.ascontiguousarray(imgs.transpose(0, 2, 3, 1).reshape(self.shape)) x = upload_rows(self.dev, rows, dtype="uint8", layout="row_major") taps: Dict[str, Any] = {} out = self.branch(x, taps=taps) res = {k: to_reference(k, read(v)) for k, v in taps.items()} for k in ("seg2d", "depth", "depth_mean", "det2d_s0", "det2d_s1", "det2d_s2", "tl"): res[k] = to_reference(k, read(out[k])) return res def release(self) -> None: self.branch.release() super().release() def nhwc_rows(nchw: np.ndarray, pad_to: int = 0) -> np.ndarray: rows = golden_rows(nchw) if pad_to and rows.shape[-1] < pad_to: rows = np.concatenate([rows, np.zeros(rows.shape[:-1] + (pad_to - rows.shape[-1],), np.float32)], axis=-1) return rows class LiftStage(_Stage): """golden ctx.painted / depth.prob + the host lift tables -> lift.weight, lift.bev, bev.raw.""" def __init__(self, dev: Any, params: Any): super().__init__(dev, params) from .lift import TtLift self.lift = TtLift(self.consts) def run(self, gold: Mapping[str, np.ndarray]) -> Dict[str, np.ndarray]: from ..host.calib import lift_geometry from .lift import TABLE_NAMES, lift_tables geom = lift_geometry(gold["input.K"], gold["input.T_cam_ego"], points=self.params.ground_points()) tb = lift_tables(geom) ctx = np.asarray(gold["ctx.painted"], np.float32).transpose(0, 2, 3, 1) prob = np.asarray(gold["depth.prob"], np.float32).transpose(0, 2, 3, 1) d = self.dev args = [upload_rows(d, ctx, layout="row_major"), upload_rows(d, prob, layout="row_major"), upload_rows(d, tb["lift_grid"], dtype="float32", layout="row_major"), [upload_rows(d, tb[name], dtype="float32") for name in TABLE_NAMES]] taps: Dict[str, Any] = {} raw = self.lift(*args, taps=taps) taps["bev.raw"] = raw return {k: to_reference(k, read(v)) for k, v in taps.items()} def release(self) -> None: self.lift.release() super().release() class BevStage(_Stage): """golden bev.raw -> BEV trunk and heads (every tap).""" def __init__(self, dev: Any, params: Any): super().__init__(dev, params) from .bev import TtBev self.bev = TtBev(params, self.consts) def run(self, gold: Mapping[str, np.ndarray]) -> Dict[str, np.ndarray]: raw = FeatureMap(upload_rows(self.dev, nhwc_rows(gold["bev.raw"])), 1, C.BEV_H, C.BEV_W, C.BEV_C) taps: Dict[str, Any] = {} out = self.bev(raw, taps=taps) res = {k: to_reference(k, read(v)) for k, v in taps.items()} res["bev.fused"] = to_reference("bev.fused", read(out["fused"])) res["lane.pre"] = to_reference("lane.pre", read(out["lane_pre"])) hr = to_reference("hm", read(out["hm_reg"]).reshape(-1, 8)[:, :2]) res["det.hm_pre"] = hr res["det.reg_pre"] = to_reference("reg", read(out["hm_reg"]).reshape(-1, 8)[:, 2:8]) res["occ"] = to_reference("occ", read(out["occ"])) res["traj"] = to_reference("traj", read(out["traj"])) res["stationary"] = to_reference("stationary", read(out["stationary"])) res["risk"] = to_reference("risk", read(out["risk32"]).reshape(-1, 32)[:, :1]) return res def release(self) -> None: self.bev.release() super().release() class HeadStage(_Stage): """golden bev.fused, lane.pre, det.hm_pre / reg_pre, risk, v0 -> planner, refiners, tail (every tap).""" def __init__(self, dev: Any, params: Any): super().__init__(dev, params) from .head import TtHead self.head = TtHead(params, self.consts) def run(self, gold: Mapping[str, np.ndarray]) -> Dict[str, np.ndarray]: d = self.dev fused = FeatureMap(upload_rows(d, nhwc_rows(gold["bev.fused"])), 1, C.BEV_H, C.BEV_W, C.BEV_C) lane_pre = FeatureMap(upload_rows(d, nhwc_rows(gold["lane.pre"]), dtype="float32"), 1, C.BEV_H, C.BEV_W, 9) hr = np.concatenate([gold["det.hm_pre"], gold["det.reg_pre"]], axis=1) hm_reg = FeatureMap(upload_rows(d, nhwc_rows(hr), dtype="float32"), 1, C.DET_H, C.DET_W, 8) risk = FeatureMap(upload_rows(d, nhwc_rows(gold["risk"], 32), dtype="float32"), 1, C.DET_H, C.DET_W, 32) v0 = upload_rows(d, np.asarray(gold["input.v0"], np.float32).reshape(1, 1, 1, 1), dtype="float32") taps: Dict[str, Any] = {} out = self.head(fused, lane_pre, hm_reg, risk, v0, taps=taps) res = {k: (to_reference(k, read(v)) if not k.startswith("ego.") or k == "ego.sem_in" else read(v)) for k, v in taps.items()} res["ego"] = to_reference("ego", read(out["ego"])) res["lane_logit"] = to_reference("lane_logit", read(out["lane_logit"])) res["lane"] = to_reference("lane", read(out["lane"])) hrr = read(out["hm_reg"]).reshape(-1, 8) res["hm"] = to_reference("hm", hrr[:, :2]) res["reg"] = to_reference("reg", hrr[:, 2:8]) return res def release(self) -> None: self.head.release() super().release() STAGES = {"image": ImageStage, "lift": LiftStage, "bev": BevStage, "head": HeadStage}