File size: 2,721 Bytes
b025706 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
import pytest
import ttnn
from models.common.modules.tt_ccl import default_topology
from models.common.tests.conftest import (
_allowed_req_shapes_for_system,
_default_fabric_config,
_is_physical_p150x4_cluster,
_pick_parent_shape_for_submesh,
)
@pytest.mark.parametrize(
("logical_shape", "expected"),
[
((1, 1), None),
((1, 2), ttnn.FabricConfig.FABRIC_1D),
((1, 8), ttnn.FabricConfig.FABRIC_1D_RING),
((1, 16), ttnn.FabricConfig.FABRIC_1D_RING),
((2, 4), ttnn.FabricConfig.FABRIC_1D),
((4, 8), ttnn.FabricConfig.FABRIC_1D),
((8, 4), ttnn.FabricConfig.FABRIC_1D),
],
)
def test_default_fabric_config(logical_shape, expected):
assert _default_fabric_config(logical_shape) == expected
@pytest.mark.parametrize(
("system_shape", "expected"),
[
((1, 1), {(1, 1)}),
((1, 2), {(1, 2), (1, 1)}),
((2, 1), {(1, 2), (2, 1), (1, 1)}),
((1, 4), {(1, 4), (1, 2), (1, 1)}),
((4, 1), {(1, 4), (4, 1), (1, 2), (1, 1)}),
((2, 2), {(2, 2), (1, 4), (1, 2), (1, 1)}),
((2, 4), {(2, 4), (1, 8), (1, 4), (1, 2), (1, 1)}),
((8, 4), {(8, 4), (4, 8), (1, 8), (1, 4), (1, 2), (1, 1)}),
],
)
def test_allowed_req_shapes_for_system(system_shape, expected):
assert _allowed_req_shapes_for_system(system_shape) == expected
def test_quietbox_square_system_uses_linear_full_device_view():
system_shape = (2, 2)
allowed = _allowed_req_shapes_for_system(system_shape, blackhole_selected=True)
assert (1, 4) in allowed
assert (2, 2) not in allowed
assert _pick_parent_shape_for_submesh(system_shape, (1, 4)) == (1, 4)
@pytest.mark.parametrize("cluster_type", [ttnn.cluster.ClusterType.P150_X4, ttnn.cluster.ClusterType.P300_X2])
def test_physical_four_die_bh_systems_admit_p150x4(cluster_type):
assert _is_physical_p150x4_cluster(cluster_type)
@pytest.mark.parametrize(
("cluster_type", "num_devices", "expected"),
[
(ttnn.cluster.ClusterType.T3K, 8, ttnn.Topology.Ring),
(ttnn.cluster.ClusterType.P150_X4, 4, ttnn.Topology.Ring),
(ttnn.cluster.ClusterType.GALAXY, 8, ttnn.Topology.Linear),
(ttnn.cluster.ClusterType.GALAXY, 4, ttnn.Topology.Linear),
(ttnn.cluster.ClusterType.GALAXY, 1, None),
],
)
def test_default_ccl_topology(monkeypatch, cluster_type, num_devices, expected):
class MeshDevice:
def get_num_devices(self):
return num_devices
monkeypatch.setattr(ttnn.cluster, "get_cluster_type", lambda: cluster_type)
assert default_topology(MeshDevice()) == expected
|