code/tt_diffusion_planner/tests/test_tt_graph_host.py

# SPDX-License-Identifier: Apache-2.0
"""The whole ttnn graph of ``tt.model.TtDiffusionPlanner`` on the FAKE ttnn (numpy numerics: bf16 tensors hold
bf16-rounded values, fp32 tensors float32), against the fp32 CPU reference's goldens: catches wiring mistakes
(shapes, transposes, token order, masks, solver indices, packing) without a device. The device numbers are gated
in ``test_pcc_device.py`` / ``test_e2e_device.py``.

The fake op set is ``common/tests/host/fake_ttnn.py`` + ``fake_ttnn_cnn.py`` + ``fake_ttnn_attention.py`` (loaded
by path like ``fake_ttnn_plugin``) plus LayerNorm / mean / subtract / tanh-GELU defined here; skipped when the
workspace's common tree is absent (an installed package).

    TT_VISIBLE_DEVICES=none python -m pytest -q -p fake_ttnn_plugin code/tt_diffusion_planner/tests/test_tt_graph_host.py
"""
from __future__ import annotations

import importlib.util
import math
import sys
from pathlib import Path

import numpy as np
import pytest

from tt_diffusion_planner.host import pipeline as hp
from tt_diffusion_planner.reference import config as C
from tt_diffusion_planner.ttaw.metrics import pcc

PKG = Path(__file__).resolve().parents[1]
FAKES = PKG.parents[3] / "common" / "tests" / "host"
GOLDENS = PKG.parents[3] / "research" / "diffusion-planner" / "goldens"
SCENE = "kashiwanoha_dense"


def _load(name: str):
    path = FAKES / f"{name}.py"
    if not path.is_file():
        pytest.skip(f"{path} not found (workspace fake ttnn op sets)")
    if name in sys.modules:
        return sys.modules[name]
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


def _install_planner_ops(fake):
    """LayerNorm, mean, subtract, rsqrt, scalar add and a linear with "gelu_tanh" on top of the CNN / attention
    fakes."""
    names = ("linear", "layer_norm", "mean", "subtract", "rsqrt", "add", "gelu", "GeluVariant")
    saved = {n: getattr(fake, n) for n in names if hasattr(fake, n)}
    base_linear, base_add = fake.linear, fake.add
    variants = type("GeluVariant", (), {"Accurate": "accurate", "FastLut": "fast_lut", "Tanh": "tanh"})

    def gelu_tanh(x):
        return 0.5 * x * (1.0 + np.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x ** 3)))

    def linear(a, b, *, bias=None, activation=None, dtype=None, compute_kernel_config=None, **kw):
        if activation != "gelu_tanh":
            return base_linear(a, b, bias=bias, activation=activation, dtype=dtype,
                               compute_kernel_config=compute_kernel_config, **kw)
        inputs = [a, b] + ([bias] if bias is not None else [])

        def fn(x, y, *bb):
            z = np.matmul(np.asarray(x, np.float64), np.asarray(y, np.float64))
            return gelu_tanh(z + (bb[0].reshape(-1) if bb else 0.0))

        return fake._op("linear_gelu_tanh", inputs, fn, dtype or a.dtype, fake.TILE_LAYOUT)

    def layer_norm(x, *, epsilon=1e-12, weight=None, bias=None, compute_kernel_config=None, **kw):
        inputs = [x] + [t for t in (weight, bias) if t is not None]

        def fn(a, *gb):
            a = np.asarray(a, np.float64)
            y = (a - a.mean(-1, keepdims=True)) / np.sqrt(a.var(-1, keepdims=True) + epsilon)
            if weight is not None:
                y = y * np.asarray(gb[0], np.float64).reshape(-1)
            if bias is not None:
                y = y + np.asarray(gb[-1], np.float64).reshape(-1)
            return y

        return fake._op("layer_norm", inputs, fn, x.dtype, x.layout, (float(epsilon),))

    def mean(x, dim, keepdim=False, **kw):
        return fake._op("mean", [x], lambda a: np.asarray(a, np.float64).mean(axis=dim, keepdims=keepdim), x.dtype,
                        x.layout, (dim, keepdim))

    def subtract(a, b, **kw):
        return fake._op("subtract", [a, b], lambda x, y: np.asarray(x, np.float64) - y, a.dtype, a.layout)

    def rsqrt(x, *, fast_and_approximate_mode=True, **kw):
        return fake._op("rsqrt", [x], lambda a: 1.0 / np.sqrt(np.asarray(a, np.float64)), x.dtype, x.layout)

    def add(a, b, **kw):
        if isinstance(b, (int, float)):
            value = float(b)
            return fake._op("add_scalar", [a], lambda x: np.asarray(x, np.float64) + value, a.dtype, a.layout,
                            (value,))
        return base_add(a, b, **kw)

    def gelu(x, *, variant=None, fast_and_approximate_mode=False, **kw):
        if variant == variants.Tanh:
            return fake._op("gelu_tanh", [x], lambda a: gelu_tanh(np.asarray(a, np.float64)), x.dtype, x.layout)
        erf = np.vectorize(math.erf)
        return fake._op("gelu", [x], lambda a: 0.5 * a * (1.0 + erf(np.asarray(a, np.float64) / math.sqrt(2.0))),
                        x.dtype, x.layout)

    for name, fn in (("linear", linear), ("layer_norm", layer_norm), ("mean", mean), ("subtract", subtract),
                     ("rsqrt", rsqrt), ("add", add), ("gelu", gelu), ("GeluVariant", variants)):
        setattr(fake, name, fn)

    def restore():
        for name in names:
            if name in saved:
                setattr(fake, name, saved[name])
            elif hasattr(fake, name):
                delattr(fake, name)
    return restore


@pytest.fixture(scope="module")
def fake_planner():
    import ttnn  # the fake (fake_ttnn_plugin)

    if not hasattr(ttnn, "_op"):
        pytest.skip("needs the fake ttnn (-p fake_ttnn_plugin)")
    from tt_diffusion_planner.reference.weights import find_weights_dir, load_weights

    wd = find_weights_dir()
    if wd is None:
        pytest.skip("weights not found")
    ttnn.reset()
    undo = [_load("fake_ttnn_cnn").install(ttnn), _load("fake_ttnn_attention").install(ttnn)]
    undo.append(_install_planner_ops(ttnn))
    from tt_diffusion_planner.ttaw.device import close_device, open_device
    from tt_diffusion_planner.tt.model import TtDiffusionPlanner

    weights = load_weights(wd)
    dev = open_device(0, dispatch="eth", num_command_queues=1)
    tt = TtDiffusionPlanner(dev, weights, debug=True)
    yield tt, weights
    tt.release()
    close_device(dev)
    for fn in reversed(undo):
        fn()


@pytest.fixture(scope="module")
def scene(fake_planner):
    tt, weights = fake_planner
    path = GOLDENS / f"{SCENE}.npz"
    if not path.is_file():
        pytest.skip(f"research goldens {path} not found")
    with np.load(path, allow_pickle=False) as z:
        gold = {k: z[k] for k in z.files if k != "__meta__"}
    raw = {k: gold[f"in.{k}"] for k in C.INPUT_NAMES}
    return hp.prepare(raw, weights.normalization.observation), gold


def test_encoder_taps_match_the_reference(fake_planner, scene):
    tt, _ = fake_planner
    prep, gold = scene
    taps = tt.encoder_taps(prep, eager=True)
    for name, _ in C.TOKEN_LAYOUT:
        rows = np.flatnonzero(gold[f"host.valid.{name}"])
        if rows.size:
            assert pcc(taps[f"enc.{name}"][rows], gold[f"enc.{name}"][rows]) > 0.999, name
    for c in ("ego", "neighbor", "lane", "route", "line_string"):
        rows = gold[f"enc.{c}.pre.rows"]
        assert pcc(taps[f"enc.{c}.pre"][rows], gold[f"enc.{c}.pre"]) > 0.9999, c
        assert pcc(taps[f"enc.{c}.mixer"][rows], gold[f"enc.{c}.mixer"]) > 0.999, c
    tok = np.flatnonzero(gold["host.token_valid"])
    assert pcc(taps["enc.encoding"][tok], gold["enc.encoding"][tok]) > 0.999


def test_decode_once_matches_the_reference(fake_planner, scene):
    tt, _ = fake_planner
    prep, gold = scene
    rows = gold["dec.rows"]
    for k in (0, 10):
        got = tt.decode_once(prep, gold["dec.x_in"][k], float(gold["dec.t"][k]), encoding=gold["enc.encoding"],
                             eager=True)
        assert not got[:, 0].any()                                       # the masked t = 0 columns
        assert pcc(got[rows][:, 1:], gold["dec.out"][k][:, 1:]) > 0.999, k


def test_plan_and_trace(fake_planner, scene):
    """The plan eagerly vs the reference's final x0 / logits, then captured: replay == eager."""
    tt, weights = fake_planner
    prep, gold = scene
    eager = tt.forward(prep, eager=True)
    rows = gold["dec.rows"]
    assert pcc(eager["final_x0"][rows], gold["final_x0"][rows]) > 0.999
    np.testing.assert_array_equal(eager["final_x0"][:, 0], prep.decoder.current_states)   # prefix constraint
    assert np.abs(eager["logit"] - gold["turn.logit"]).max() < 0.05
    assert len(eager["denoising_steps"]) == C.DPM_SOLVER_STEPS + 1
    np.testing.assert_array_equal(eager["denoising_steps"][-1][0], eager["final_x0"][0])
    tt.capture()
    replay = tt.forward(prep)
    for k in ("final_x0", "logit"):
        np.testing.assert_array_equal(replay[k], eager[k])
    out = hp.make_output(replay["final_x0"], replay["logit"], prep, weights.normalization,
                         {k: v[3] for k, v in hp.RUNTIME_PARAMS.items()})
    assert int(out.turn_indicator["command"]) == int(gold["out.turn_command"])