code/models/common/tests/test_distribute_as.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Tests for distributing a torch.Tensor using the topology of an existing TTNN tensor.

Follows the same style as test_auto_compose.py, but in the opposite
direction: given a reference distributed tensor, verify that
`from_torch_dist_as(torch_tensor, ref_tensor)` produces a TTNN tensor
with the same topology and data that composes back to the original
torch input.
"""

import pytest
import torch

import ttnn
from models.common.auto_compose import extract_tensor_topology_info, to_torch_auto_compose
from models.common.distribute_as import from_torch_dist_as

# ======================================================================================
# Test Parameters
# ======================================================================================


pytestmark = [
    pytest.mark.parametrize(
        "ttnn_mesh_device",
        [
            (1, 1),  # single device
            (1, 2),  # 1D mesh, 2 devices
            (1, 8),  # 1D mesh, 8 devices
            (2, 4),  # 2D mesh, 8 devices
        ],
        ids=[
            "1x1",
            "1x2",
            "1x8",
            "2x4",
        ],
        indirect=True,
    ),
    pytest.mark.parametrize(
        "layout,dtype",
        [
            (ttnn.ROW_MAJOR_LAYOUT, ttnn.bfloat16),
            (ttnn.TILE_LAYOUT, ttnn.bfloat16),
            (ttnn.TILE_LAYOUT, ttnn.bfloat8_b),  # bfloat8_b only works with TILE_LAYOUT
            (ttnn.TILE_LAYOUT, ttnn.bfloat4_b),  # bfloat4_b only works with TILE_LAYOUT
        ],
        ids=["row_major_bf16", "tile_bf16", "tile_bf8b", "tile_bf4b"],
    ),
    pytest.mark.parametrize(
        "memory_config",
        [ttnn.DRAM_MEMORY_CONFIG, ttnn.L1_MEMORY_CONFIG],
        ids=["dram", "l1"],
    ),
]


# ======================================================================================
# Helpers
# ======================================================================================


def _make_arange_dtype(
    shape: tuple[int, ...], dtype: torch.dtype = torch.bfloat16, min_value: float = 0, max_value: float = 100
) -> torch.Tensor:
    """Create a deterministic tensor with arange data and specified dtype."""
    numel = 1
    for s in shape:
        numel *= s
    # Generate values from min_value to max_value with step of 1
    values = torch.arange(min_value, max_value + 1, dtype=dtype)
    # Randomly sample indices (with replacement) to fill the tensor
    indices = torch.randint(0, len(values), size=(numel,))
    data = values[indices].reshape(shape)
    return data


def _pos_dim(dim: int, rank: int) -> int:
    return dim % rank


def _topology_signature(tt: ttnn.Tensor) -> tuple[tuple[str, tuple[int, ...]], tuple[int, ...]]:
    placements, dist_shape = extract_tensor_topology_info(tt)
    sig_p = []
    for p in placements:
        if isinstance(p, ttnn.PlacementShard):
            sig_p.append(("shard", (p.dim,)))
        else:
            assert isinstance(p, ttnn.PlacementReplicate)
            sig_p.append(("replicate", ()))
    return tuple(sig_p), tuple(dist_shape)


# ======================================================================================
# Tests
# ======================================================================================


@pytest.mark.parametrize("storage", ["host", "device"])
def test_sharded_1d(
    ttnn_mesh_device: ttnn.MeshDevice, layout, dtype, memory_config: ttnn.MemoryConfig, storage: str
) -> None:
    is_host_ref = storage == "host"
    num_devices = ttnn_mesh_device.get_num_devices()

    # Reference topology: host-sharded along dim 0
    shape = (num_devices, 1, 3, 1)
    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
    max_value = 6  # allows for range of value to grow to 7 for new_input
    ref_input = _make_arange_dtype(shape, dtype=torch.float32, min_value=-7, max_value=max_value)
    ref_tt = ttnn.from_torch(
        ref_input,
        device=None if is_host_ref else ttnn_mesh_device,
        dtype=dtype,
        layout=layout,
        memory_config=memory_config,
        mesh_mapper=ttnn.ShardTensorToMesh(ttnn_mesh_device, dim=0),
    )

    # Distribute a different tensor using the same topology
    new_input = ref_input + 1  # change values to avoid accidental equality
    if is_host_ref:
        tt_as = from_torch_dist_as(new_input, ref_tt, device=ttnn_mesh_device)
    else:
        tt_as = from_torch_dist_as(new_input, ref_tt)

    # Verify topology matches reference
    expected_topology = _topology_signature(ref_tt)
    assert _topology_signature(tt_as) == expected_topology
    assert tt_as.dtype == dtype
    assert tt_as.memory_config() == memory_config

    # Verify round-trip via auto-composition matches the source
    if is_host_ref:
        torch_auto = to_torch_auto_compose(tt_as, device=ttnn_mesh_device)
    else:
        torch_auto = to_torch_auto_compose(tt_as)
    assert torch.equal(torch_auto, new_input)


@pytest.mark.parametrize("storage", ["host", "device"])
@pytest.mark.parametrize("dim", [0, 1, 2, -1])
def test_sharded_various_dims(
    ttnn_mesh_device: ttnn.MeshDevice,
    layout,
    dtype,
    memory_config: ttnn.MemoryConfig,
    dim: int,
    storage: str,
) -> None:
    is_host_ref = storage == "host"
    num_devices = ttnn_mesh_device.get_num_devices()

    rank = 4
    axis = _pos_dim(dim, rank)
    shape = [2, 3, 4, 1]
    shape[axis] = num_devices
    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
    max_value = 6  # allows for range of value to grow to 7 for new_input
    ref_input = _make_arange_dtype(shape, dtype=torch.float32, min_value=-7, max_value=max_value)

    ref_tt = ttnn.from_torch(
        ref_input,
        device=None if is_host_ref else ttnn_mesh_device,
        dtype=dtype,
        layout=layout,
        memory_config=memory_config,
        mesh_mapper=ttnn.ShardTensorToMesh(ttnn_mesh_device, dim=dim),
    )

    new_input = ref_input + 1  # change values to avoid accidental equality
    if is_host_ref:
        tt_as = from_torch_dist_as(new_input, ref_tt, device=ttnn_mesh_device)
    else:
        tt_as = from_torch_dist_as(new_input, ref_tt)

    assert _topology_signature(tt_as) == _topology_signature(ref_tt)
    assert tt_as.dtype == dtype
    assert tt_as.memory_config() == memory_config
    if is_host_ref:
        torch_auto = to_torch_auto_compose(tt_as, device=ttnn_mesh_device)
    else:
        torch_auto = to_torch_auto_compose(tt_as)
    assert torch.equal(torch_auto, new_input)


@pytest.mark.parametrize("storage", ["host", "device"])
@pytest.mark.parametrize("dims_pair", [(0, 1), (0, -1), (1, -1)])
def test_sharded_2d(
    ttnn_mesh_device: ttnn.MeshDevice,
    layout,
    dtype,
    memory_config: ttnn.MemoryConfig,
    dims_pair: tuple[int, int],
    storage: str,
) -> None:
    is_host_ref = storage == "host"
    mesh_shape = tuple(ttnn_mesh_device.shape)
    if len(mesh_shape) != 2 and torch.prod(torch.tensor(mesh_shape)).item() <= 1:
        pytest.skip("Requires a 2D mesh with both dims > 1")

    rank = 4
    d0 = _pos_dim(dims_pair[0], rank)
    d1 = _pos_dim(dims_pair[1], rank)
    assert d0 != d1

    shape = [2, 3, 4, 5]
    shape[d0] = mesh_shape[0]
    shape[d1] = mesh_shape[1]
    # get range of values for bfloat4_b quantization which has 4 bits for the mantissa and shared 8-bit exponent
    max_value = 6  # allows for range of value to grow to 7 for new_input
    ref_input = _make_arange_dtype(shape, dtype=torch.float32, min_value=-7, max_value=max_value)

    mapper = ttnn.ShardTensor2dMesh(ttnn_mesh_device, mesh_shape=mesh_shape, dims=(dims_pair[0], dims_pair[1]))
    ref_tt = ttnn.from_torch(
        ref_input,
        device=None if is_host_ref else ttnn_mesh_device,
        dtype=dtype,
        layout=layout,
        memory_config=memory_config,
        mesh_mapper=mapper,
    )

    new_input = ref_input + 1  # change values to avoid accidental equality
    if is_host_ref:
        tt_as = from_torch_dist_as(new_input, ref_tt, device=ttnn_mesh_device)
    else:
        tt_as = from_torch_dist_as(new_input, ref_tt)

    assert _topology_signature(tt_as) == _topology_signature(ref_tt)
    assert tt_as.dtype == dtype
    assert tt_as.memory_config() == memory_config
    if is_host_ref:
        torch_auto = to_torch_auto_compose(tt_as, device=ttnn_mesh_device)
    else:
        torch_auto = to_torch_auto_compose(tt_as)
    assert torch.equal(torch_auto, new_input)