From 88be96376347b67750adb2919ba65b4445e093d8 Mon Sep 17 00:00:00 2001 From: Akhmed Rakhmati Date: Thu, 7 Dec 2023 23:00:24 +0000 Subject: [PATCH] #3003: updated ttnn tests --- .../ttnn/unit_tests/experimental/test_exp.py | 2 +- .../experimental/test_layer_norm.py | 8 +++---- tests/ttnn/unit_tests/test_add.py | 10 ++++---- .../{test_free.py => test_deallocate.py} | 6 ++--- tests/ttnn/unit_tests/test_dump_and_load.py | 4 ++-- tests/ttnn/unit_tests/test_slicing.py | 24 ------------------- tests/ttnn/unit_tests/test_softmax.py | 9 +++---- .../ttnn/unit_tests/test_to_and_from_torch.py | 2 +- ttnn/core.py | 13 ++++++---- ttnn/experimental.py | 8 +++++-- 10 files changed, 36 insertions(+), 50 deletions(-) rename tests/ttnn/unit_tests/{test_free.py => test_deallocate.py} (84%) delete mode 100644 tests/ttnn/unit_tests/test_slicing.py diff --git a/tests/ttnn/unit_tests/experimental/test_exp.py b/tests/ttnn/unit_tests/experimental/test_exp.py index a211b07c58d..e8c7f77b86e 100644 --- a/tests/ttnn/unit_tests/experimental/test_exp.py +++ b/tests/ttnn/unit_tests/experimental/test_exp.py @@ -16,7 +16,7 @@ def test_exp(device, h, w): torch.manual_seed(0) - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) torch_output_tensor = torch.exp(torch_input_tensor) input_tensor = ttnn.from_torch(torch_input_tensor) diff --git a/tests/ttnn/unit_tests/experimental/test_layer_norm.py b/tests/ttnn/unit_tests/experimental/test_layer_norm.py index 8be17eda0b9..fba5a6eb479 100644 --- a/tests/ttnn/unit_tests/experimental/test_layer_norm.py +++ b/tests/ttnn/unit_tests/experimental/test_layer_norm.py @@ -18,7 +18,7 @@ def test_layer_norm(device, h, w): torch.manual_seed(0) - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) torch_output_tensor = torch.nn.functional.layer_norm(torch_input_tensor, normalized_shape=[w]) input_tensor = ttnn.from_torch(torch_input_tensor) @@ -37,7 +37,7 @@ def test_layer_norm(device, h, w): def test_layer_norm_with_weight_and_bias(device, h, w): torch.manual_seed(0) - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) torch_weight = torch.rand((w,), dtype=torch.bfloat16) torch_bias = torch.rand((w,), dtype=torch.bfloat16) torch_output_tensor = torch.nn.functional.layer_norm( @@ -66,8 +66,8 @@ def test_layer_norm_with_weight_and_bias(device, h, w): def test_layer_norm_with_weight_bias_and_residual_input(device, h, w): torch.manual_seed(0) - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) - torch_residual_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) + torch_residual_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) torch_weight = torch.rand((w,), dtype=torch.bfloat16) torch_bias = torch.rand((w,), dtype=torch.bfloat16) torch_output_tensor = torch.nn.functional.layer_norm( diff --git a/tests/ttnn/unit_tests/test_add.py b/tests/ttnn/unit_tests/test_add.py index 7cbed67eb59..0adef993d7b 100644 --- a/tests/ttnn/unit_tests/test_add.py +++ b/tests/ttnn/unit_tests/test_add.py @@ -31,7 +31,7 @@ def test_add_1D_tensor_and_scalar(device, scalar, size): @pytest.mark.parametrize("h", [2 * 32]) @pytest.mark.parametrize("w", [4 * 32]) def test_add_scalar(device, s, h, w): - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) torch_output_tensor = torch_input_tensor + s input_tensor = ttnn.from_torch(torch_input_tensor) @@ -49,7 +49,7 @@ def test_add_scalar(device, s, h, w): @pytest.mark.parametrize("h", [1]) @pytest.mark.parametrize("w", [4]) def test_add_scalar_and_alpha(device, alpha, scalar_input_tensor_b, h, w): - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) torch_output_tensor = torch.add(torch_input_tensor, scalar_input_tensor_b, alpha=alpha) input_tensor = ttnn.from_torch(torch_input_tensor) @@ -65,8 +65,8 @@ def test_add_scalar_and_alpha(device, alpha, scalar_input_tensor_b, h, w): @pytest.mark.parametrize("h", [32]) @pytest.mark.parametrize("w", [2 * 32]) def test_add(device, h, w): - torch_a = torch.rand((1, 1, h, w), dtype=torch.bfloat16) - torch_b = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_a = torch.rand((h, w), dtype=torch.bfloat16) + torch_b = torch.rand((h, w), dtype=torch.bfloat16) torch_output = torch.add(torch_a, torch_b) a = ttnn.from_torch(torch_a) @@ -106,7 +106,7 @@ def test_add_4D(device, n, c, h, w): @pytest.mark.parametrize("w", [2 * 32]) @pytest.mark.parametrize("scalar", [0.42]) def test_add_scalar(device, h, w, scalar): - torch_a = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_a = torch.rand((h, w), dtype=torch.bfloat16) torch_output = scalar + torch_a a = ttnn.from_torch(torch_a) diff --git a/tests/ttnn/unit_tests/test_free.py b/tests/ttnn/unit_tests/test_deallocate.py similarity index 84% rename from tests/ttnn/unit_tests/test_free.py rename to tests/ttnn/unit_tests/test_deallocate.py index 209076725d0..a2818393c8d 100644 --- a/tests/ttnn/unit_tests/test_free.py +++ b/tests/ttnn/unit_tests/test_deallocate.py @@ -11,8 +11,8 @@ @pytest.mark.parametrize("h", [32]) @pytest.mark.parametrize("w", [2 * 32]) -def test_free(device, h, w): - torch_input_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) +def test_deallocate(device, h, w): + torch_input_tensor = torch.rand((h, w), dtype=torch.bfloat16) input_tensor = ttnn.from_torch(torch_input_tensor) @@ -25,7 +25,7 @@ def test_free(device, h, w): # Create a reference to the same storage by using reshape which will create a new flyweight # (If reshape operation changes, then this test might need to be updated) - output_tensor_reference = ttnn.reshape(output_tensor, (1, 1, h, w)) + output_tensor_reference = ttnn.reshape(output_tensor, (h, w)) ttnn.deallocate(output_tensor) with pytest.raises(RuntimeError) as exception: diff --git a/tests/ttnn/unit_tests/test_dump_and_load.py b/tests/ttnn/unit_tests/test_dump_and_load.py index f038ddf0f5f..33cb8496318 100644 --- a/tests/ttnn/unit_tests/test_dump_and_load.py +++ b/tests/ttnn/unit_tests/test_dump_and_load.py @@ -16,7 +16,7 @@ def test_dump_and_load(tmp_path, h, w): file_name = tmp_path / pathlib.Path("tensor.bin") - torch_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_tensor = torch.rand((h, w), dtype=torch.bfloat16) tt_tensor = ttnn.from_torch(torch_tensor) ttnn.dump_tensor(file_name, tt_tensor) @@ -30,7 +30,7 @@ def test_dump_and_load(tmp_path, h, w): def test_dump_and_load_tilized(tmp_path, h, w): file_name = tmp_path / pathlib.Path("tensor.bin") - torch_tensor = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_tensor = torch.rand((h, w), dtype=torch.bfloat16) tt_tensor = ttnn.from_torch(torch_tensor) tt_tensor = ttnn.to_layout(tt_tensor, ttnn.TILE_LAYOUT) ttnn.dump_tensor(file_name, tt_tensor) diff --git a/tests/ttnn/unit_tests/test_slicing.py b/tests/ttnn/unit_tests/test_slicing.py deleted file mode 100644 index 792a4817314..00000000000 --- a/tests/ttnn/unit_tests/test_slicing.py +++ /dev/null @@ -1,24 +0,0 @@ -# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. - -# SPDX-License-Identifier: Apache-2.0 - -import pytest - -import torch - -import ttnn - -from tests.ttnn.utils_for_testing import assert_with_pcc - - -@pytest.mark.skip(reason="Unable to slice large tensors") -def test_large_slicing(device): - torch_a = torch.rand((1, 1, 42, 250880), dtype=torch.bfloat16) - torch_output = torch_a[:, :, -1, :] - a = ttnn.from_torch(torch_a) - shape = a.shape - a = ttnn.to_device(a, device) - tt_output = a[:, :, -1, :] - tt_output = ttnn.from_device(tt_output) - tt_output = ttnn.to_torch(tt_output) - assert_with_pcc(torch_output, tt_output, 0.9999) diff --git a/tests/ttnn/unit_tests/test_softmax.py b/tests/ttnn/unit_tests/test_softmax.py index 4f66cb8841a..65290da8b61 100644 --- a/tests/ttnn/unit_tests/test_softmax.py +++ b/tests/ttnn/unit_tests/test_softmax.py @@ -14,12 +14,13 @@ @skip_for_wormhole_b0() -@pytest.mark.parametrize("h", [32]) -@pytest.mark.parametrize("w", [2 * 32]) -def test_softmax(device, h, w): +@pytest.mark.parametrize("batch_size", [1, 16]) +@pytest.mark.parametrize("h", [32, 64]) +@pytest.mark.parametrize("w", [32, 64]) +def test_softmax(device, batch_size, h, w): torch.manual_seed(0) - torch_input_tensor = torch_random((1, 16, 4, 4), -10, 10, dtype=torch.bfloat16) + torch_input_tensor = torch_random((batch_size, h, w), -10, 10, dtype=torch.bfloat16) torch_output_tensor = F.softmax(torch_input_tensor, dim=-1, dtype=torch.bfloat16) input_tensor = ttnn.from_torch(torch_input_tensor) input_tensor = ttnn.to_device(input_tensor, device) diff --git a/tests/ttnn/unit_tests/test_to_and_from_torch.py b/tests/ttnn/unit_tests/test_to_and_from_torch.py index 90df5beb2fd..066cd3d7940 100644 --- a/tests/ttnn/unit_tests/test_to_and_from_torch.py +++ b/tests/ttnn/unit_tests/test_to_and_from_torch.py @@ -12,7 +12,7 @@ @pytest.mark.parametrize("h", [7]) @pytest.mark.parametrize("w", [3]) def test_to_and_from_4D(h, w): - torch_input = torch.rand((1, 1, h, w), dtype=torch.bfloat16) + torch_input = torch.rand((h, w), dtype=torch.bfloat16) tt_output = ttnn.from_torch(torch_input) torch_output = ttnn.to_torch(tt_output) assert torch.allclose(torch_output, torch_input) diff --git a/ttnn/core.py b/ttnn/core.py index 2305e28a480..be11c4360d4 100644 --- a/ttnn/core.py +++ b/ttnn/core.py @@ -1009,7 +1009,7 @@ def ttnn_reshape(ttl_input_tensor, shape): ttl_input_tensor, shape ) - if len(input_tensor.shape) == 4 and len(shape) == 4: + if input_tensor.is_on_device and len(input_tensor.shape) == 4 and len(shape) == 4: w, z, y, x = shape return Tensor(ttl.tensor.reshape(ttl_input_tensor, w, z, y, x)) else: @@ -1063,7 +1063,7 @@ def permute(input_tensor: Tensor, order: Tuple[int, ...]) -> Tensor: ttl_input_tensor = input_tensor._tensor - if len(input_tensor.shape) == 4: + if input_tensor.is_on_device and len(input_tensor.shape) == 4: return Tensor(ttl.tensor.permute(ttl_input_tensor, order)) else: @@ -1099,15 +1099,20 @@ def softmax(input_tensor: Tensor, dim: int, memory_config: MemoryConfig = DRAM_M """ - rank = len(input_tensor.shape) + input_shape = tuple(input_tensor.shape) + rank = len(input_shape) if dim < 0: dim = rank + dim if dim != rank - 1: raise RuntimeError("Softmax can only operate on the last dimension.") + input_tensor = _reshape_to_4D(input_tensor) + ttl_input_tensor = input_tensor._tensor ttl_output_tensor = ttl.tensor.softmax(ttl_input_tensor, output_mem_config=memory_config) - return Tensor(ttl_output_tensor) + output_tensor = Tensor(ttl_output_tensor) + output_tensor = reshape(output_tensor, input_shape) + return output_tensor def embedding( diff --git a/ttnn/experimental.py b/ttnn/experimental.py index 0fb48451457..cb3c38cda9a 100644 --- a/ttnn/experimental.py +++ b/ttnn/experimental.py @@ -15,9 +15,13 @@ def exp(input_tensor: Tensor) -> Tensor: + original_shape = tuple(input_tensor.shape) + input_tensor = _reshape_to_4D(input_tensor) ttl_input_tensor = input_tensor._tensor - output_tensor = ttl.tensor.exp(ttl_input_tensor) - return Tensor(output_tensor) + ttl_output_tensor = ttl.tensor.exp(ttl_input_tensor) + output_tensor = Tensor(ttl_output_tensor) + output_tensor = reshape(output_tensor, original_shape) + return output_tensor def gelu(input_tensor: Tensor, fast_and_approx=True) -> Tensor: