Skip to content

Commit

Permalink
#13127: Add simple tensor creation gtest (#13500)
Browse files Browse the repository at this point in the history
  • Loading branch information
TT-BrianLiu authored Oct 4, 2024
1 parent fb60e69 commit e400377
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tests/ttnn/unit_tests/gtests/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)
70 changes: 70 additions & 0 deletions tests/ttnn/unit_tests/gtests/tensor/test_create_tensor.cpp
Original file line number Diff line number Diff line change
@@ -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})}
)
);

0 comments on commit e400377

Please sign in to comment.