From e4003776b10a7604217209ebdca52449fc927261 Mon Sep 17 00:00:00 2001
From: Brian Liu <109366241+TT-BrianLiu@users.noreply.github.com>
Date: Fri, 4 Oct 2024 16:14:08 -0700
Subject: [PATCH] #13127: Add simple tensor creation gtest (#13500)

---
 tests/ttnn/unit_tests/gtests/CMakeLists.txt   |  5 ++
 .../gtests/tensor/test_create_tensor.cpp      | 70 +++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 tests/ttnn/unit_tests/gtests/tensor/test_create_tensor.cpp

diff --git a/tests/ttnn/unit_tests/gtests/CMakeLists.txt b/tests/ttnn/unit_tests/gtests/CMakeLists.txt
index 25a5bb3fde4..3937aa52f13 100644
--- a/tests/ttnn/unit_tests/gtests/CMakeLists.txt
+++ b/tests/ttnn/unit_tests/gtests/CMakeLists.txt
@@ -12,9 +12,13 @@ set(TTNN_UNIT_TESTS_SRC
 set(TTNN_CCL_UNIT_TESTS_SRC
     ${CMAKE_CURRENT_SOURCE_DIR}/ccl/test_erisc_data_mover_with_workers.cpp
 )
+set(TTNN_TENSOR_UNIT_TESTS_SRC
+    ${CMAKE_CURRENT_SOURCE_DIR}/tensor/test_create_tensor.cpp
+)
 
 add_executable(unit_tests_ttnn ${TTNN_UNIT_TESTS_SRC})
 add_executable(unit_tests_ttnn_ccl ${TTNN_CCL_UNIT_TESTS_SRC})
+add_executable(unit_tests_ttnn_tensor ${TTNN_TENSOR_UNIT_TESTS_SRC})
 add_executable(test_multi_device ${CMAKE_CURRENT_SOURCE_DIR}/test_multi_device.cpp)
 add_executable(galaxy_unit_tests_ttnn ${CMAKE_CURRENT_SOURCE_DIR}/test_ccl_on_galaxy.cpp)
 
@@ -34,5 +38,6 @@ endfunction()
 # Set up properties for both targets
 setup_ttnn_test_target(unit_tests_ttnn)
 setup_ttnn_test_target(unit_tests_ttnn_ccl)
+setup_ttnn_test_target(unit_tests_ttnn_tensor)
 setup_ttnn_test_target(test_multi_device)
 setup_ttnn_test_target(galaxy_unit_tests_ttnn)
diff --git a/tests/ttnn/unit_tests/gtests/tensor/test_create_tensor.cpp b/tests/ttnn/unit_tests/gtests/tensor/test_create_tensor.cpp
new file mode 100644
index 00000000000..a0545180de0
--- /dev/null
+++ b/tests/ttnn/unit_tests/gtests/tensor/test_create_tensor.cpp
@@ -0,0 +1,70 @@
+// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "gtest/gtest.h"
+
+#include "tt_metal/common/bfloat16.hpp"
+#include "ttnn/device.hpp"
+#include "ttnn/operations/core/core.hpp"
+#include "ttnn/async_runtime.hpp"
+#include "ttnn/operations/numpy/functions.hpp"
+#include "tt_metal/common/logger.hpp"
+
+#include "ttnn_test_fixtures.hpp"
+
+void run_create_tensor_test(tt::tt_metal::Device* device, ttnn::Shape input_shape) {
+    MemoryConfig mem_cfg = MemoryConfig{
+        .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED,
+        .buffer_type = BufferType::DRAM,
+        .shard_spec = std::nullopt};
+
+    const uint32_t io_cq = 0;
+    constexpr DataType dtype = DataType::BFLOAT16;
+    constexpr uint32_t datum_size_bytes = 2;
+
+    auto input_buf_size_datums = input_shape.volume();
+
+    auto host_data = std::shared_ptr<uint16_t[]>(new uint16_t[input_buf_size_datums]);
+    auto readback_data = std::shared_ptr<uint16_t[]>(new uint16_t[input_buf_size_datums]);
+
+    for (int i = 0; i < input_buf_size_datums; i++) {
+        host_data[i] = 1;
+    }
+
+    auto input_buffer = ttnn::allocate_buffer_on_device(input_buf_size_datums * datum_size_bytes, device, input_shape, dtype, Layout::TILE, mem_cfg);
+
+    auto input_storage = tt::tt_metal::DeviceStorage{input_buffer};
+
+    Tensor input_tensor = Tensor(input_storage, input_shape, dtype, Layout::TILE);
+    tt::log_debug("input_data: \n {}", input_tensor.write_to_string());
+
+    ttnn::write_buffer(io_cq, input_tensor, {host_data});
+
+    ttnn::read_buffer(io_cq, input_tensor, {readback_data});
+
+    for (int i = 0; i < input_buf_size_datums; i++) {
+        EXPECT_EQ(host_data[i], readback_data[i]);
+    }
+
+    input_tensor.deallocate();
+}
+
+struct CreateTensorParams {
+    ttnn::Shape shape;
+};
+
+class CreateTensorTest : public ttnn::TTNNFixtureWithDevice, public ::testing::WithParamInterface<CreateTensorParams> {};
+
+TEST_P(CreateTensorTest, Tile) {
+    CreateTensorParams params = GetParam();
+    run_create_tensor_test(device_, params.shape);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CreateTensorTestWithShape,
+    CreateTensorTest,
+    ::testing::Values(
+        CreateTensorParams{.shape=ttnn::Shape({1, 1, 32, 32})}
+    )
+);