code/tt_diffusion_planner/tests/test_grid_fit_host.py

# SPDX-License-Identifier: Apache-2.0
"""The swept 2-D multicast configs fit the 12x10 ETH grid and fall back to the auto config on the 11x10 WORKER grid
(republish 2026-10-11: ``dispatch="worker"`` failed at capture with "Num output blocks along x (12) must be smaller
than or equal to the number of columns in compute grid (11)")."""
from types import SimpleNamespace

from tt_diffusion_planner.tt.layers import DEC_MM_CONFIGS, DEC_ROWS, KCAT_CONFIGS, fit_pcm, mcast_fits

ETH, WORKER = SimpleNamespace(x=12, y=10), SimpleNamespace(x=11, y=10)


def test_every_swept_config_fits_the_eth_grid():
    for rows in (DEC_ROWS, 32, 64, 96, 128, 192):
        for (k, n, _ps), (tr, pcm, pcn, _kb, _sw) in DEC_MM_CONFIGS.items():
            assert mcast_fits(tr, fit_pcm(tr, pcm, rows, ETH), pcn, rows, n, ETH), (rows, k, n)
    for (rows, _ktp, n), (tr, pcm, pcn, _kb, _sw) in KCAT_CONFIGS.items():
        assert mcast_fits(tr, pcm, pcn, rows, n, ETH), (rows, n)


def test_the_qkv_config_does_not_fit_the_worker_grid():
    tr, pcm, pcn, _kb, _sw = KCAT_CONFIGS[(352, 25, 768)]
    assert not mcast_fits(tr, pcm, pcn, 352, 768, WORKER)        # 24 N tiles / 2 = 12 blocks on x > 11
    assert mcast_fits(False, 2, 1, 352, 256, WORKER)