Skip to content

Commit

Permalink
#0: disable and enable tests based on to_layout hang + new support fo…
Browse files Browse the repository at this point in the history 
…r N-d
  • Loading branch information
@sjameelTT
sjameelTT committed Dec 18, 2024
1 parent ed94ee6 commit dca2fe3
Show file tree
Hide file tree
Showing 5 changed files with 58 additions and 32 deletions.
40 changes: 22 additions & 18 deletions tests/tt_eager/python_api_testing/unit_testing/misc/test_transpose.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -672,7 +672,7 @@ def test_transpose_2D(dtype, shape, layout, device):
tt_input = ttnn.from_torch(torch_input, dtype=ttnn.DataType.BFLOAT16, layout=layout, device=device)
tt_output = ttnn.transpose(tt_input, 0, 1)
tt_output = ttnn.to_torch(tt_output)
assert_with_pcc(torch_output, tt_output, 0.9999)
assert_with_pcc(torch_output, tt_output, 0.99)


@pytest.mark.parametrize(
Expand Down Expand Up @@ -745,17 +745,20 @@ def test_transpose_4d_wh_tile(shape, device):
assert_with_pcc(torch_output, tt_output, 0.9999)


@pytest.mark.skip("Skipping due to hang on to_layout to tile where input shape has 1 in it")
@pytest.mark.parametrize(
"config",
[
[[1, 1370, 1, 3, 1280], [0, -2], ttnn.TILE_LAYOUT], # untilize doesn't work with 4D
[[1, 50, 1, 3, 768], [0, -2], ttnn.TILE_LAYOUT], # untilize doesn't work with 4D
[[21843, 768], [0, 1], ttnn.ROW_MAJOR_LAYOUT], # circular buffer overflow
[[1, 1370, 1, 3, 1280], [0, -2], ttnn.TILE_LAYOUT], # hang
[[1, 50, 1, 3, 768], [0, -2], ttnn.TILE_LAYOUT], # hang
[[1, 50, 1, 3, 1024], [0, -2], ttnn.TILE_LAYOUT], # hang
[[1, 197, 1, 3, 768], [0, -2], ttnn.TILE_LAYOUT], # hang
[[1, 197, 1, 3, 1024], [0, -2], ttnn.TILE_LAYOUT], # hang
[[2, 7, 2, 7, 384], [-4, -3], ttnn.TILE_LAYOUT], # hang
],
)
@pytest.mark.parametrize("memory_config", [ttnn.L1_MEMORY_CONFIG, ttnn.DRAM_MEMORY_CONFIG])
@pytest.mark.parametrize("memory_config", [ttnn.DRAM_MEMORY_CONFIG])
def test_transpose_failures(config, memory_config, device):
pytest.skip("Failing pytorch 2.0 trace sweeps")
torch.manual_seed(2005)
torch_input = torch.randn(config[0], dtype=torch.bfloat16)
torch_output = torch_input.transpose(config[1][0], config[1][1])
Expand Down Expand Up @@ -787,23 +790,24 @@ def test_transpose_failures(config, memory_config, device):
[1, 2],
ttnn.ROW_MAJOR_LAYOUT,
],
[[1, 9, 8, 18], [1, 2], ttnn.ROW_MAJOR_LAYOUT], # unaligned RM that fallsback to tiled
[[1, 9, 8, 14], [1, 2], ttnn.ROW_MAJOR_LAYOUT], # unaligned RM that fallsback to tiled
[[1, 9, 8, 2], [1, 2], ttnn.ROW_MAJOR_LAYOUT], # unaligned RM that fallsback to tiled
[[1, 2, 8, 2], [1, 2], ttnn.ROW_MAJOR_LAYOUT], # unaligned RM that fallsback to tiled
[[64, 4, 49, 32], [-2, -1], ttnn.ROW_MAJOR_LAYOUT], # Page size must be divisible by sizeof(uint32_t)
[[12, 3], [0, 1], ttnn.ROW_MAJOR_LAYOUT], # need tensor for this one
[[1, 9, 8, 18], [1, 2], ttnn.ROW_MAJOR_LAYOUT],
[[1, 9, 8, 14], [1, 2], ttnn.ROW_MAJOR_LAYOUT],
[[1, 9, 8, 2], [1, 2], ttnn.ROW_MAJOR_LAYOUT],
[[1, 2, 8, 2], [1, 2], ttnn.ROW_MAJOR_LAYOUT],
[[64, 4, 49, 32], [-2, -1], ttnn.ROW_MAJOR_LAYOUT],
[[12, 3], [0, 1], ttnn.ROW_MAJOR_LAYOUT],
[
[1, 8, 4096, 40],
[1, 2],
ttnn.ROW_MAJOR_LAYOUT,
], # RM that fallsback to tiled only when reading from DRAM (32B alignment requirement on DRAM, 16B on L1)
[[1, 9, 8, 40], [1, 2], ttnn.ROW_MAJOR_LAYOUT], # RM that fallsback to tiled only when reading from DRAM
[[1, 8, 8, 8], [1, 2], ttnn.ROW_MAJOR_LAYOUT], # RM that fallsback to tiled only when reading from DRAM
],
[[1, 9, 8, 40], [1, 2], ttnn.ROW_MAJOR_LAYOUT],
[[1, 8, 8, 8], [1, 2], ttnn.ROW_MAJOR_LAYOUT],
[[21843, 768], [0, 1], ttnn.ROW_MAJOR_LAYOUT],
],
)
@pytest.mark.parametrize("memory_config", [ttnn.L1_MEMORY_CONFIG, ttnn.DRAM_MEMORY_CONFIG])
def test_transpose_unaligned(config, memory_config, device):
@pytest.mark.parametrize("memory_config", [ttnn.DRAM_MEMORY_CONFIG])
def test_transpose_former_failures(config, memory_config, device):
torch.manual_seed(2005)
# this will convert to tiled for now
torch_input = torch.randn(config[0], dtype=torch.bfloat16)
Expand All @@ -813,7 +817,7 @@ def test_transpose_unaligned(config, memory_config, device):
)
tt_output = ttnn.transpose(tt_input, config[1][0], config[1][1])
tt_output = ttnn.to_torch(tt_output)
assert_with_pcc(torch_output, tt_output, 0.9999)
assert_with_pcc(torch_output, tt_output, 0.99)


@pytest.mark.parametrize(
Expand Down
25 changes: 25 additions & 0 deletions tests/ttnn/unit_tests/test_to_layout.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,3 +164,28 @@ def test_to_layout_6D(shape, input_layout, output_layout, device):
output_tensor = ttnn.to_layout(input_tensor, output_layout)
output_tensor = ttnn.to_torch(output_tensor)
assert_with_pcc(input_a, output_tensor)


@pytest.mark.skip("Skipping due to hang on to_layout to tile where input shape has 1 in it")
@pytest.mark.parametrize(
"config",
[
[[3, 1370, 1, 1, 1280], ttnn.ROW_MAJOR_LAYOUT], # hang
[[3, 50, 1, 1, 768], ttnn.ROW_MAJOR_LAYOUT], # hang
[[3, 50, 1, 1, 1024], ttnn.ROW_MAJOR_LAYOUT], # hang
[[3, 197, 1, 1, 768], ttnn.ROW_MAJOR_LAYOUT], # hang
[[3, 197, 1, 1, 1024], ttnn.ROW_MAJOR_LAYOUT], # hang
],
)
@pytest.mark.parametrize("memory_config", [ttnn.DRAM_MEMORY_CONFIG])
def test_to_layout_hangs(config, memory_config, device):
torch.manual_seed(2005)
torch_input = torch.randn(config[0], dtype=torch.bfloat16)

tt_input = ttnn.from_torch(
torch_input, dtype=ttnn.DataType.BFLOAT16, layout=config[1], device=device, memory_config=memory_config
)
tt_output = ttnn.to_layout(tt_input, ttnn.TILE_LAYOUT)
tt_output = ttnn.to_torch(tt_output)

assert_with_pcc(torch_input, tt_output, 0.9999)
13 changes: 6 additions & 7 deletions ttnn/cpp/ttnn/operations/data_movement/permute/permute.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,7 @@ ttnn::Tensor permute_impl(
TT_FATAL(
!(pad_value.has_value() && pad_value.value() != 0.0f),
"Non-zero padding is not supported for permute on tensors with rank > 4.");
SmallVector<uint32_t> permute_dims(dims.begin(), dims.end());
input = ttnn::prim::permute(input, permute_dims, output_mem_config, std::nullopt);
input = ttnn::prim::permute(input, dims, output_mem_config, std::nullopt);
return ttnn::to_layout(input, a.get_layout(), std::nullopt, std::nullopt, (Device*)nullptr);
}

Expand Down Expand Up @@ -150,7 +149,7 @@ ttnn::Tensor permute_launch(
return output_tensors.at(0);
}

bool is_permute_nop(const ttnn::Tensor& a, tt::stl::Span<const uint32_t> dims) {
bool is_permute_nop(const ttnn::Tensor& a, const ttnn::SmallVector<uint32_t>& dims) {
if (a.get_shape().rank() <= 1) {
return true;
}
Expand All @@ -165,7 +164,7 @@ bool is_permute_nop(const ttnn::Tensor& a, tt::stl::Span<const uint32_t> dims) {
ttnn::Tensor ExecutePermute::invoke(
uint8_t queue_id,
const ttnn::Tensor& input_tensor,
tt::stl::Span<const int64_t> dims,
const ttnn::SmallVector<int64_t>& dims,
const std::optional<MemoryConfig>& memory_config,
const std::optional<float>& pad_value) {
const auto input_rank = input_tensor.get_logical_shape().rank();
Expand All @@ -182,7 +181,7 @@ ttnn::Tensor ExecutePermute::invoke(
return ttnn::to_memory_config(input_tensor, memory_config.value_or(input_tensor.memory_config()));
}

auto adjust_order = [](tt::stl::Span<const uint32_t> dims) {
auto adjust_order = [](const ttnn::SmallVector<uint32_t>& dims) {
ttnn::SmallVector<uint32_t> new_order;
TT_FATAL(dims.size() <= 4, "Minimum rank of tensor required is 4");
int additional_ranks = 4 - dims.size();
Expand Down Expand Up @@ -211,14 +210,14 @@ ttnn::Tensor ExecutePermute::invoke(

ttnn::Tensor ExecutePermute::invoke(
const ttnn::Tensor& input_tensor,
tt::stl::Span<const int64_t> dims,
const ttnn::SmallVector<int64_t>& dims,
const std::optional<MemoryConfig>& memory_config,
const std::optional<float>& pad_value) {
return invoke(DefaultQueueId, input_tensor, dims, memory_config, pad_value);
}

ttnn::Tensor ExecutePermute::invoke(
const ttnn::Tensor& input_tensor, tt::stl::Span<const int64_t> dims, const std::optional<float>& pad_value) {
const ttnn::Tensor& input_tensor, const ttnn::SmallVector<int64_t>& dims, const std::optional<float>& pad_value) {
return invoke(input_tensor, dims, std::nullopt, pad_value);
}

Expand Down
6 changes: 3 additions & 3 deletions ttnn/cpp/ttnn/operations/data_movement/permute/permute.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,19 @@ struct ExecutePermute {
static ttnn::Tensor invoke(
uint8_t queue_id,
const ttnn::Tensor& input_tensor,
tt::stl::Span<const int64_t> dims,
const SmallVector<int64_t>& dims,
const std::optional<MemoryConfig>& memory_config,
const std::optional<float>& pad_value = 0.0f);

static ttnn::Tensor invoke(
const ttnn::Tensor& input_tensor,
tt::stl::Span<const int64_t> dims,
const SmallVector<int64_t>& dims,
const std::optional<MemoryConfig>& memory_config,
const std::optional<float>& pad_value = 0.0f);

static ttnn::Tensor invoke(
const ttnn::Tensor& input_tensor,
tt::stl::Span<const int64_t> dims,
const SmallVector<int64_t>& dims,
const std::optional<float>& pad_value = 0.0f);
};

Expand Down
6 changes: 2 additions & 4 deletions ttnn/cpp/ttnn/operations/data_movement/transpose/transpose.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,7 @@ inline Tensor transpose_(
if (a.device()->arch() == tt::ARCH::GRAYSKULL) {
tiled_only = a.shape()[-2] > 256; // hangs right now past this dimension, #13660 will turn it from a
// hang into a PCC issue for GS and improve perf for WH
} else if (!a.is_sharded() && a.layout() == Layout::ROW_MAJOR) { // rm is L1 intensive, if it overflows we
// can do tiled which allocates much
// smaller CBs
} else if (!a.is_sharded() && a.layout() == Layout::ROW_MAJOR) {
return ttnn::prim::permute(
a, ttnn::SmallVector<uint32_t>({0, 1, 3, 2}), output_mem_config, std::nullopt);
}
Expand Down Expand Up @@ -91,7 +89,7 @@ ttnn::Tensor transpose_nd(
const uint32_t dim2,
const std::optional<MemoryConfig>& memory_config_arg,
const std::optional<float>& pad_value) {
std::vector<int64_t> permutation;
ttnn::SmallVector<int64_t> permutation;
permutation.reserve(input_tensor.get_shape().rank());
for (uint32_t i = 0; i < input_tensor.get_shape().rank(); ++i) {
permutation.push_back(i);
Expand Down

0 comments on commit dca2fe3

Please sign in to comment.