| # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Device topology naming helpers shared by TTTv2 modules.""" | |
| from __future__ import annotations | |
| import ttnn | |
| def is_blackhole() -> bool: | |
| return "blackhole" in ttnn.get_arch_name() | |
| def get_device_name(mesh_device: ttnn.MeshDevice, num_devices: int | None = None) -> str: | |
| """Return the product/topology name for a TT mesh device. | |
| By default, the full mesh device count is used. CCL callers can pass a | |
| host-local device count when they need link-count tuning for the current | |
| process rather than for the full mesh. | |
| """ | |
| num_devices = mesh_device.get_num_devices() if num_devices is None else num_devices | |
| dram_grid_size = mesh_device.dram_grid_size() | |
| if ttnn.device.is_blackhole(mesh_device): | |
| device_names = { | |
| 1: "P100" if dram_grid_size and dram_grid_size.x == 7 else "P150", | |
| 2: "P300", | |
| 4: "P150x4", | |
| 8: "P150x8", | |
| 32: "BHGLX", | |
| } | |
| elif ttnn.device.is_wormhole_b0(mesh_device): | |
| device_names = { | |
| 1: "N150", | |
| 2: "N300", | |
| 4: "N150x4", | |
| 8: "T3K", | |
| 32: "TG", | |
| } | |
| else: | |
| raise ValueError(f"Unsupported architecture: {ttnn.get_arch_name()}") | |
| if num_devices in device_names: | |
| return device_names[num_devices] | |
| raise ValueError(f"Unsupported number of devices: {num_devices} for {ttnn.get_arch_name()}") | |