tenstorrent / tt-metal

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

Transpose op doesn't work when dimensions are not divisible by tile #13005

Open dgolubovicTT opened 1 week ago

dgolubovicTT commented 1 week ago

Describe the bug

for tensor shape (1,32,12,100) transpose doesn't work for -2, -3 dimensions. Namely, it throws an error that shape of the tensor must be divisible by tile dim. It seems that padding is not done in this case.

As I looked into implementation of transpose I have noticed the difference: In ttnn/cpp/ttnn/operations/data_movement/transpose/transpose.cpp when padding is HC (-2,-3) padding is only done on C.

To Reproduce

import pytest
import torch
import numpy as np

import ttnn

from loguru import logger
from models.utility_functions import is_grayskull
from tests.tt_eager.python_api_testing.sweep_tests.comparison_funcs import comp_pcc, comp_equal
from models.utility_functions import skip_for_grayskull, skip_for_blackhole
from tests.ttnn.utils_for_testing import assert_with_pcc

def test_transpose_hc(dtype, device):
    if is_grayskull() and dtype == ttnn.float32:
        pytest.skip("Skipping float32 tests on Grayskull")

    logger.info("transpose on C H dim")

    N = 1
    H = 32
    C = 12
    W = 100
    input_shape = (N, C, H, W)
    transpose(input_shape, device, dim0=1, dim1=-2, input_dtype=dtype)

tt-metal commit: 96fd0df449931e9c4958e761f40e0b8550c3e2c1

Expected behaviour Transpose op should work regardless of the divisibility by tile dimensions.

To give a little context, we need this transpose case for Llama bringup on forge. This transpose case appears in self attention graph, and we need to support it.

dgolubovicTT commented 1 week ago

fyi @tt-mpantic

nvukobratTT commented 1 week ago

@dgolubovicTT for quicker and easier repro, can you provide the TTNN test code as part of the issue description (example).

Also, TTNN folks have their templates for opening up bug issues. In order to make their life a bit easier, let's try to use those template fields for information we currently posses (not all are mandatory).