File size: 1,421 Bytes
d18366a | 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 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import os
import pytest
import ttnn
@pytest.fixture
def device_params(request, galaxy_type):
# Get param dict passed in from test parametrize (or default to empty dict).
# Any TRACE_MODEL_KEY_PARAM is left in place; the mesh_device fixture resolves it
# to trace_region_size using the logical submesh SKU.
params = getattr(request, "param", {}).copy()
mesh_device = {"N150": (1, 1), "N300": (1, 2), "N150x4": (1, 4), "T3K": (1, 8), "TG": (8, 4), "P150x8": (1, 8)}.get(
os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
)
is_single_device = (mesh_device == (1, 1)) if isinstance(mesh_device, tuple) else (mesh_device == 1)
if "fabric_config" in params:
if is_single_device:
params["fabric_config"] = None
elif params["fabric_config"] == True:
cluster_type = ttnn.cluster.get_cluster_type()
if cluster_type == ttnn.cluster.ClusterType.BLACKHOLE_GALAXY:
# The 8x4 decode path uses Ring collectives along both mesh axes.
params["fabric_config"] = ttnn.FabricConfig.FABRIC_2D_TORUS_XY
else:
params["fabric_config"] = (
ttnn.FabricConfig.FABRIC_1D_RING if galaxy_type == "6U" else ttnn.FabricConfig.FABRIC_1D
)
return params
|