tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
Apache License 2.0
438 stars 63 forks source link

Transpose fails for unaligned shapes #13749

Open yieldthought opened 1 week ago

yieldthought commented 1 week ago

Describe the bug WH transpose fails if W is an unaligned value such as 5.

To Reproduce Can be trivially reproduced by adding [[1, 1024, 5, 1280]], # Non page-aligned to shape_wh in tests/tt_eager/python_api_testing/sweep_tests/pytests/tt_dnn/test_transpose.py or just running this code:

import pytest
import torch
from functools import partial

from tests.tt_eager.python_api_testing.sweep_tests import comparison_funcs, generation_funcs
from tests.tt_eager.python_api_testing.sweep_tests.run_pytorch_ci_tests import run_single_pytorch_test
from models.utility_functions import skip_for_grayskull, skip_for_blackhole

shape_wh = [
    [[1, 1, 32, 32]],  # Single core
    [[3, 1, 320, 384]],  # Multi core
    [[1, 1024, 5, 1280]], # Non page-aligned
]

@skip_for_blackhole("Mismatching on BH, see #12349")
@pytest.mark.parametrize("input_shapes", shape_wh)
def test_run_transpose_wh_test(input_shapes, device, function_level_defaults):
    datagen_func = [
        generation_funcs.gen_func_with_cast(partial(generation_funcs.gen_rand, low=-100, high=100), torch.bfloat16)
    ]
    test_args = generation_funcs.gen_default_dtype_layout_device(input_shapes)[0]
    test_args.update(
        {
            "dim0": -2,
            "dim1": -1,
        }
    )
    comparison_func = partial(comparison_funcs.comp_equal)
    run_single_pytorch_test("transpose", input_shapes, datagen_func, comparison_func, device, test_args)

Expected behavior Matches torch behaviour.

Screenshots RuntimeError: TT_FATAL @ ../tt_metal/impl/buffers/buffer.cpp:41: page_size % sizeof(uint32_t) == 0

Please complete the following environment information: IRD-provided N150

Additional context Found whilst bringing up Llama 3.2

sjameelTT commented 1 week ago

Row Major doesn't support inner dimension 5 at the moment though we're working on doing a big refactor to support that.