# SPDX-License-Identifier: Apache-2.0 """Shared helpers for the ``mb2_*`` p150a micro-benchmarks (GR00T port, 01-architecture-synthesis §6 items 2/3/7/8/10/17). Every ``mb2_*.py`` script is a stand-alone argparse program that opens the device itself (always launched through ``bin/with-device.sh``). This module only holds the helpers they share so that all numbers come from one code path: * device open/close with program cache (``open_device``), compute-kernel configs (``compute_config``), * host<->device tensor helpers (``to_tt`` / ``to_torch``), * accuracy metrics (``compare``: PCC in float64, max|Δ|, relative L2) — the PCC definition matches ``golden-tap-map-and-thresholds.md`` §4 (Pearson correlation of the flattened tensors), * wall-clock timing around ``ttnn.synchronize_device`` (``time_calls``), * a ``Report`` container that records every sub-test (including failures with the exception text) and writes JSON. Nothing here imports ttnn at module import time; the ``ttnn`` module object is passed in by the caller. """ from __future__ import annotations import json import math import os import platform import statistics import subprocess import sys import time import traceback from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional import torch TILE = 32 RESULTS_DIR_DEFAULT = "/home/deepgadget/experiments/gr00t/docs/research/device-microbench-trace-semantics" TT_METAL_HOME_DEFAULT = "/home/deepgadget/experiments/gr00t/tt-metal" # -------------------------------------------------------------------------------------------------------------------- # environment # -------------------------------------------------------------------------------------------------------------------- def env_info(ttnn=None) -> Dict[str, Any]: """Record the software stack the numbers were produced with.""" info: Dict[str, Any] = { "python": sys.version.split()[0], "torch": torch.__version__, "platform": platform.platform(), "tt_metal_home": os.environ.get("TT_METAL_HOME", TT_METAL_HOME_DEFAULT), "time": time.strftime("%Y-%m-%d %H:%M:%S"), } try: info["tt_metal_commit"] = ( subprocess.check_output( ["git", "-C", info["tt_metal_home"], "rev-parse", "HEAD"], stderr=subprocess.DEVNULL, timeout=10 ) .decode() .strip() ) except Exception as e: # noqa: BLE001 - informational only info["tt_metal_commit"] = f"unknown ({e})" if ttnn is not None: info["ttnn_version"] = getattr(ttnn, "__version__", "unknown") return info def open_device(ttnn, l1_small_size: int = 32768, trace_region_size: int = 0, num_cqs: int = 1): """Open device 0 (``ttnn.open_device``, ttnn-nanobind/device.cpp:71-80) with the program cache enabled. ``trace_region_size=0`` selects dynamic (top-down DRAM) trace storage (tt_metal/distributed/mesh_trace.cpp:66-70); a positive size reserves a dedicated region and ``end_trace_capture`` raises the FATAL "Creating trace buffers of size {}B ... but only {}B is allocated for trace region" when it is too small (mesh_trace.cpp:77-82) — the ``mb2_trace_replay.py --probe-trace-size`` mode exploits that message. """ kwargs = dict(device_id=0, l1_small_size=l1_small_size, trace_region_size=trace_region_size) if num_cqs != 1: kwargs["num_command_queues"] = num_cqs device = ttnn.open_device(**kwargs) device.enable_program_cache() return device def close_device(ttnn, device) -> None: try: ttnn.synchronize_device(device) except Exception: # noqa: BLE001 pass ttnn.close_device(device) def compute_config( ttnn, fidelity: str = "HiFi2", fp32_acc: bool = False, approx: bool = False, packer_l1_acc: bool = False ): """``ttnn.WormholeComputeKernelConfig`` (alias for Blackhole, ttnn-nanobind/operations/core.cpp:58-84).""" fid = { "LoFi": ttnn.MathFidelity.LoFi, "HiFi2": ttnn.MathFidelity.HiFi2, "HiFi3": ttnn.MathFidelity.HiFi3, "HiFi4": ttnn.MathFidelity.HiFi4, }[fidelity] return ttnn.WormholeComputeKernelConfig( math_fidelity=fid, math_approx_mode=approx, fp32_dest_acc_en=fp32_acc, packer_l1_acc=packer_l1_acc, ) def tt_dtype(ttnn, name: str): return {"bf16": ttnn.bfloat16, "bfp8": ttnn.bfloat8_b, "bfp4": ttnn.bfloat4_b, "fp32": ttnn.float32}[name] # -------------------------------------------------------------------------------------------------------------------- # tensors # -------------------------------------------------------------------------------------------------------------------- def to_tt(ttnn, device, t: torch.Tensor, dtype=None, layout=None, mem=None, pad_value: float = 0.0): """Host torch -> device TILE tensor (logical shape kept, tile padding filled with ``pad_value``).""" dtype = ttnn.bfloat16 if dtype is None else dtype layout = ttnn.TILE_LAYOUT if layout is None else layout mem = ttnn.DRAM_MEMORY_CONFIG if mem is None else mem return ttnn.from_torch(t, dtype=dtype, layout=layout, device=device, memory_config=mem, pad_value=pad_value) def to_host_tt(ttnn, t: torch.Tensor, dtype=None, layout=None, pad_value: float = 0.0): """Host-side ttnn tensor (for ``ttnn.copy_host_to_device_tensor``).""" dtype = ttnn.bfloat16 if dtype is None else dtype layout = ttnn.TILE_LAYOUT if layout is None else layout return ttnn.from_torch(t, dtype=dtype, layout=layout, pad_value=pad_value) def to_torch(ttnn, t) -> torch.Tensor: return ttnn.to_torch(t).float() def bf16_round(t: torch.Tensor) -> torch.Tensor: return t.to(torch.bfloat16).float() # -------------------------------------------------------------------------------------------------------------------- # metrics # -------------------------------------------------------------------------------------------------------------------- def pcc(a: torch.Tensor, b: torch.Tensor) -> float: """Pearson correlation of the flattened tensors in float64 (golden-tap-map-and-thresholds §4 definition).""" a = a.detach().double().flatten() b = b.detach().double().flatten() if a.numel() != b.numel(): raise ValueError(f"pcc: numel mismatch {a.numel()} vs {b.numel()}") if torch.isnan(a).any() or torch.isnan(b).any(): return float("nan") a = a - a.mean() b = b - b.mean() den = a.norm() * b.norm() if den == 0: return 1.0 if torch.equal(a, b) else 0.0 return float((a @ b) / den) def compare(ref: torch.Tensor, out: torch.Tensor) -> Dict[str, float]: ref = ref.detach().float() out = out.detach().float() if ref.shape != out.shape: raise ValueError(f"compare: shape mismatch ref {tuple(ref.shape)} vs out {tuple(out.shape)}") d = (ref - out).double() rn = ref.double().norm() return { "pcc": pcc(ref, out), "max_abs": float(d.abs().max()) if d.numel() else 0.0, "mean_abs": float(d.abs().mean()) if d.numel() else 0.0, "rel_l2": float(d.norm() / rn) if rn > 0 else float("nan"), "ref_absmax": float(ref.abs().max()) if ref.numel() else 0.0, "out_absmax": float(out.abs().max()) if out.numel() else 0.0, "nan_count": int(torch.isnan(out).sum()), } def bitwise_equal_bf16(a: torch.Tensor, b: torch.Tensor) -> bool: """Exact equality of two bf16-valued float tensors (compared as int16 bit patterns).""" return bool(torch.equal(a.to(torch.bfloat16).view(torch.int16), b.to(torch.bfloat16).view(torch.int16))) # -------------------------------------------------------------------------------------------------------------------- # timing # -------------------------------------------------------------------------------------------------------------------- def time_calls(ttnn, device, fn: Callable[[], Any], reps: int, warmup: int = 2) -> Dict[str, Any]: """Wall time (µs) of ``fn()`` followed by ``ttnn.synchronize_device`` — ``reps`` samples, min/median/mean.""" for _ in range(warmup): fn() ttnn.synchronize_device(device) samples: List[float] = [] for _ in range(reps): t0 = time.perf_counter() fn() ttnn.synchronize_device(device) samples.append((time.perf_counter() - t0) * 1e6) return { "min_us": min(samples), "median_us": statistics.median(samples), "mean_us": statistics.fmean(samples), "max_us": max(samples), "reps": reps, "samples_us": [round(s, 1) for s in samples], } # -------------------------------------------------------------------------------------------------------------------- # reporting # -------------------------------------------------------------------------------------------------------------------- @dataclass class Report: name: str env: Dict[str, Any] = field(default_factory=dict) results: Dict[str, Any] = field(default_factory=dict) errors: Dict[str, str] = field(default_factory=dict) notes: List[str] = field(default_factory=list) def run(self, key: str, fn: Callable[[], Any]) -> Any: """Run one sub-test; store its dict result or the exception text (never abort the whole script).""" print(f"\n===== [{self.name}] {key} =====", flush=True) t0 = time.perf_counter() try: res = fn() if isinstance(res, dict): res = dict(res) res.setdefault("wall_s", round(time.perf_counter() - t0, 2)) self.results[key] = res print(json.dumps(res, indent=1, default=_json_default)[:6000], flush=True) return res except Exception as e: # noqa: BLE001 - we want the exact failure text in the report tb = traceback.format_exc() msg = f"{type(e).__name__}: {e}" self.errors[key] = msg self.results[key] = {"error": msg, "wall_s": round(time.perf_counter() - t0, 2)} print(f"!! FAILED {key}: {msg}\n{tb[-3000:]}", flush=True) return None def note(self, s: str) -> None: print(f"NOTE: {s}", flush=True) self.notes.append(s) def write(self, path: str) -> None: os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) with open(path, "w") as f: json.dump( { "name": self.name, "env": self.env, "results": self.results, "errors": self.errors, "notes": self.notes, }, f, indent=1, default=_json_default, ) print(f"\nwrote {path}", flush=True) def _json_default(o): if isinstance(o, (torch.Tensor,)): return o.tolist() if isinstance(o, float) and (math.isnan(o) or math.isinf(o)): return str(o) return str(o) def default_out(name: str) -> str: return os.path.join(RESULTS_DIR_DEFAULT, f"{name}.json") def gelu_tanh_ref(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.gelu(x.float(), approximate="tanh") def gelu_erf_ref(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.gelu(x.float(), approximate="none")