From 7f564aa27186eb6bf52158ffc3b361104c1982ba Mon Sep 17 00:00:00 2001
From: Oleg Milyutin <omilyutin@tenstorrent.com>
Date: Tue, 19 Nov 2024 17:41:08 +0000
Subject: [PATCH] Add C++ multi-device tensor creation test, fix Python test

---
 tests/ttnn/unit_tests/gtests/CMakeLists.txt   |  1 +
 .../test_create_tensor_multi_device.cpp       | 44 +++++++++++++++++++
 .../unit_tests/operations/test_creation.py    |  2 +
 3 files changed, 47 insertions(+)
 create mode 100644 tests/ttnn/unit_tests/gtests/tensor/test_create_tensor_multi_device.cpp

diff --git a/tests/ttnn/unit_tests/gtests/CMakeLists.txt b/tests/ttnn/unit_tests/gtests/CMakeLists.txt
index 865b29daefc5..9861a843dc38 100644
--- a/tests/ttnn/unit_tests/gtests/CMakeLists.txt
+++ b/tests/ttnn/unit_tests/gtests/CMakeLists.txt
@@ -17,6 +17,7 @@ set(TTNN_TENSOR_UNIT_TESTS_SRC
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor/common_tensor_test_utils.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_create_tensor.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_tensor_layout.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_create_tensor_multi_device.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_create_tensor_with_layout.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_shape_base.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_sharding_with_alignment.cpp
diff --git a/tests/ttnn/unit_tests/gtests/tensor/test_create_tensor_multi_device.cpp b/tests/ttnn/unit_tests/gtests/tensor/test_create_tensor_multi_device.cpp
new file mode 100644
index 000000000000..86652d29b416
--- /dev/null
+++ b/tests/ttnn/unit_tests/gtests/tensor/test_create_tensor_multi_device.cpp
@@ -0,0 +1,44 @@
+
+// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <variant>
+
+#include "buffers/buffer_constants.hpp"
+#include "gtest/gtest.h"
+#include "ttnn/cpp/ttnn/operations/creation.hpp"
+#include "ttnn/cpp/ttnn/tensor/types.hpp"
+#include "ttnn/distributed/api.hpp"
+#include "ttnn/tensor/enum_types.hpp"
+#include "ttnn_test_fixtures.hpp"
+
+namespace ttnn::distributed::test {
+namespace {
+
+using ::tt::tt_metal::BufferType;
+using ::tt::tt_metal::Layout;
+using ::tt::tt_metal::MemoryConfig;
+using ::tt::tt_metal::StorageType;
+using ::tt::tt_metal::TensorMemoryLayout;
+
+TEST_F(T3kMultiDeviceFixture, CreateEmpty) {
+    MeshDevice* mesh_device = this->mesh_device_.get();
+
+    const auto mesh_replicated_tensor = ttnn::empty(
+        ttnn::Shape(std::array<uint32_t, 2>{32, 32}),
+        DataType::BFLOAT16,
+        Layout::ROW_MAJOR,
+        mesh_device,
+        MemoryConfig{TensorMemoryLayout::INTERLEAVED, BufferType::DRAM, std::nullopt});
+
+    EXPECT_EQ(mesh_replicated_tensor.storage_type(), StorageType::MULTI_DEVICE);
+    EXPECT_EQ(mesh_replicated_tensor.get_workers().size(), mesh_device->num_devices());
+
+    const auto distributed_tensor_config = get_distributed_tensor_config_from_tensor(mesh_replicated_tensor);
+
+    EXPECT_TRUE(std::holds_alternative<ReplicateTensor>(distributed_tensor_config));
+}
+
+}  // namespace
+}  // namespace ttnn::distributed::test
diff --git a/tests/ttnn/unit_tests/operations/test_creation.py b/tests/ttnn/unit_tests/operations/test_creation.py
index b2159aff4af8..5798298dce9e 100644
--- a/tests/ttnn/unit_tests/operations/test_creation.py
+++ b/tests/ttnn/unit_tests/operations/test_creation.py
@@ -347,6 +347,7 @@ def test_arange(device, start, end, step):
         [1, 1, 32, 32],
         [1, 1, 320, 384],
         [1, 3, 320, 384],
+        [1, 3, 180, 64],
         [2, 640, 64, 64],
         [2, 1280, 64, 64],
     ],
@@ -368,6 +369,7 @@ def test_empty(device, input_shapes):
         [1, 1, 32, 32],
         [1, 1, 320, 384],
         [1, 3, 320, 384],
+        [1, 3, 180, 64],
         [2, 640, 64, 64],
         [2, 1280, 64, 64],
     ],