| |
| |
|
|
| 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 |
|
|
