File size: 2,034 Bytes
ba7051a 13b4736 2a74422 13b4736 ba7051a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | # SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Minimal pytest fixtures for tt-MoGe.
Provides a single-device ``device`` fixture so the tt_moge tests run without the
tt-metal monorepo's own conftest.py. The device is selected with ``--device-id N``
on the pytest CLI or the ``MOGE_DEVICE`` environment variable (default 0). The
fixture opens with a trace region + 2 command queues so the perf tests exercise
the same deployment config as the benchmark. Set ``TT_VISIBLE_DEVICES`` to
constrain which PCIe chips ttnn enumerates at all (e.g. to pin one p150a).
"""
from __future__ import annotations
import gc
import os
import pytest
def pytest_addoption(parser):
parser.addoption("--device-id", action="store", default=None, help="ttnn device id")
@pytest.fixture(autouse=True)
def _gc_between_tests():
gc.collect()
@pytest.fixture(scope="session")
def device(request):
import ttnn
opt = request.config.getoption("--device-id")
device_id = int(opt) if opt is not None else int(os.environ.get("MOGE_DEVICE", "0"))
two_cq = bool(int(os.environ.get("MOGE_2CQ", "0")))
kw = dict(device_id=device_id, l1_small_size=32768, trace_region_size=1500000000,
num_command_queues=2 if two_cq else 1)
# Default = the p150 configuration of the published numbers: ETH dispatch, 1 CQ, 12x10 compute grid
# (tt_moge.device.open_device; falls back to worker dispatch with a warning without the ETH-dispatch patch).
# MOGE_DISPATCH=worker [MOGE_2CQ=1] = the Galaxy-only worker-dispatch mode; MOGE_GRID=12x10|11x10|full.
from tt_moge.device import open_device
dispatch = os.environ.get("MOGE_DISPATCH", "eth")
kw.pop("device_id")
dev, _info = open_device(device_id, grid=os.environ.get("MOGE_GRID", "12x10"), dispatch=dispatch, **kw)
dispatch = _info["dispatch"]
g = dev.compute_with_storage_grid_size()
print(f"\n[conftest] dispatch={dispatch or 'default'} grid={g.x}x{g.y}")
yield dev
ttnn.close_device(dev)
|