| |
| """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 |
|
|
| def run(self, gold: Mapping[str, np.ndarray]) -> Dict[str, np.ndarray]: |
| imgs = np.asarray(gold["input.imgs"], np.uint8)[0] |
| 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} |
|
|
