diff --git a/tests/ttnn/unit_tests/gtests/tensor/test_vector_conversion.cpp b/tests/ttnn/unit_tests/gtests/tensor/test_vector_conversion.cpp index ae2ec8bcb331..944aa6f0b623 100644 --- a/tests/ttnn/unit_tests/gtests/tensor/test_vector_conversion.cpp +++ b/tests/ttnn/unit_tests/gtests/tensor/test_vector_conversion.cpp @@ -56,6 +56,16 @@ std::vector arange(int64_t start, int64_t end, int64_t step) { return result; } +// Creates a vector with alternating ones and zeros. +std::vector ones_and_zeros(int total_size) { + std::vector result; + result.reserve(total_size); + for (int i = 0; i < total_size; ++i) { + result.push_back(i % 2 == 0 ? 0.0f : 1.0f); + } + return result; +} + template class VectorConversionTest : public ::testing::Test {}; @@ -130,35 +140,35 @@ TEST_P(BlockFloatVectorConversionTest, InvalidLayout) { TEST_P(BlockFloatVectorConversionTest, Roundtrip) { ttnn::SimpleShape shape{128, 128}; - std::vector input; - input.reserve(shape.volume()); - for (int i = 0; i < shape.volume(); ++i) { - input.push_back(i % 2 == 0 ? 0.0f : 1.0f); - } + std::vector input = ones_and_zeros(shape.volume()); + auto output = Tensor::from_vector(input, get_tensor_spec(shape, GetParam(), Layout::TILE)).to_vector(); EXPECT_THAT(output, Pointwise(Eq(), input)); } -TEST_P(BlockFloatVectorConversionTest, AddPadding) { +TEST_P(BlockFloatVectorConversionTest, RoundtripWithPadding) { ttnn::SimpleShape shape{14, 47}; - std::vector input(shape.volume(), 1.0f); + std::vector input = ones_and_zeros(shape.volume()); auto output = Tensor::from_vector(input, get_tensor_spec(shape, GetParam(), Layout::TILE)); EXPECT_THAT(output.get_logical_shape(), ShapeIs(14, 47)); EXPECT_THAT(output.get_padded_shape(), ShapeIs(32, 64)); + + EXPECT_THAT(output.to_vector(), Pointwise(Eq(), input)); } -TEST_P(BlockFloatVectorConversionTest, AddPaddingWithCustomTile) { +TEST_P(BlockFloatVectorConversionTest, RoundtripWithPaddingAndCustomTile) { ttnn::SimpleShape shape{14, 47}; - std::vector input(shape.volume(), 1.0f); + std::vector input = ones_and_zeros(shape.volume()); - TensorSpec spec( - shape, TensorLayout(DataType::FLOAT32, PageConfig(Layout::ROW_MAJOR, Tile({16, 16})), MemoryConfig{})); + TensorSpec spec(shape, TensorLayout(GetParam(), PageConfig(Layout::TILE, Tile({16, 16})), MemoryConfig{})); auto output = Tensor::from_vector(input, spec); EXPECT_THAT(output.get_logical_shape(), ShapeIs(14, 47)); - EXPECT_THAT(output.get_padded_shape(), ShapeIs(14, 47)); + EXPECT_THAT(output.get_padded_shape(), ShapeIs(16, 48)); + + EXPECT_THAT(output.to_vector(), Pointwise(Eq(), input)); } INSTANTIATE_TEST_SUITE_P( @@ -175,13 +185,12 @@ TEST_F(DeviceVectorConversionTest, RoundtripWithMemoryConfig) { TensorSpec spec( shape, TensorLayout(DataType::FLOAT32, Layout::ROW_MAJOR, MemoryConfig{.buffer_type = BufferType::L1})); - auto output_tensor = Tensor::from_vector(input, spec, device_); + auto output = Tensor::from_vector(input, spec, device_); - EXPECT_TRUE(is_device_tensor(output_tensor)); - EXPECT_TRUE(output_tensor.memory_config().is_l1()); + EXPECT_TRUE(is_device_tensor(output)); + EXPECT_TRUE(output.memory_config().is_l1()); - auto output = output_tensor.to_vector(); - EXPECT_THAT(output, Pointwise(Eq(), input)); + EXPECT_THAT(output.to_vector(), Pointwise(Eq(), input)); } } // namespace diff --git a/ttnn/cpp/ttnn/tensor/tensor.cpp b/ttnn/cpp/ttnn/tensor/tensor.cpp index c90d55d38a87..1ff9d70ed407 100644 --- a/ttnn/cpp/ttnn/tensor/tensor.cpp +++ b/ttnn/cpp/ttnn/tensor/tensor.cpp @@ -637,18 +637,18 @@ Tensor Tensor::from_span( "Tile layout is required for BFLOAT8_B and BFLOAT4_B"); const auto& tile = spec.tensor_layout().get_page_config().get_tile(); - // Create an intermediate tensor on CPU to perform padding. - TensorSpec result_cpu_spec( - spec.logical_shape(), - TensorLayout(DataType::FLOAT32, PageConfig(Layout::ROW_MAJOR, tile), MemoryConfig{})); - Tensor result_cpu_tensor = create_owned_tensor_from_row_major_data( - std::vector(buffer.begin(), buffer.end()), result_cpu_spec); - if (result_cpu_spec.logical_shape() != result_cpu_spec.padded_shape()) { - result_cpu_tensor = tensor_ops::tensor_pad( - result_cpu_tensor, result_cpu_spec.padded_shape(), ttnn::SimpleShape{0, 0, 0, 0}, 0); + // Create an intermediate tensor on CPU to perform padding and data packing. + Tensor fp32_tensor = create_owned_tensor_from_row_major_data( + std::vector(buffer.begin(), buffer.end()), + TensorSpec( + spec.logical_shape(), + TensorLayout(DataType::FLOAT32, PageConfig(Layout::ROW_MAJOR, tile), MemoryConfig{}))); + if (spec.logical_shape() != spec.padded_shape()) { + fp32_tensor = + tensor_ops::tensor_pad(fp32_tensor, spec.padded_shape(), ttnn::SimpleShape{0, 0, 0, 0}, 0); } - std::vector padded_row_major_data = owned_buffer::get_as(result_cpu_tensor).get(); + std::vector padded_row_major_data = owned_buffer::get_as(fp32_tensor).get(); std::vector packed_block_floats = spec.data_type() == DataType::BFLOAT8_B ? pack_fp32_vec_as_bfp8_tiles( @@ -698,11 +698,18 @@ std::vector Tensor::to_vector() const { : unpack_bfp4_tiles_into_float_vec( packed_data, /*row_major_output=*/true, /*is_exp_a=*/false, tile); - TensorSpec ft32_cpu_spec( - cpu_tensor.get_shape().logical_shape(), - TensorLayout(DataType::FLOAT32, PageConfig(Layout::ROW_MAJOR), MemoryConfig{})); - Tensor ft32_cpu_tensor = create_owned_tensor_from_row_major_data(std::move(unpacked_data), ft32_cpu_spec); - return unpad_tensor_to_vec(ft32_cpu_tensor); + Tensor fp32_cpu_tensor = create_owned_tensor_from_row_major_data( + std::move(unpacked_data), + TensorSpec( + cpu_tensor.get_shape().padded_shape(), + TensorLayout(DataType::FLOAT32, PageConfig(Layout::ROW_MAJOR, tile), MemoryConfig{}))); + + if (cpu_tensor.get_shape().logical_shape() != cpu_tensor.get_shape().padded_shape()) { + fp32_cpu_tensor = tensor_ops::tensor_unpad( + fp32_cpu_tensor, ttnn::SimpleShape{0, 0, 0, 0}, cpu_tensor.get_shape().logical_shape()); + } + + return owned_buffer::get_as(std::get(fp32_cpu_tensor.storage()).buffer).get(); } default: { TT_THROW("Cannot convert tensor to vector for data type: {}", cpu_tensor.get_dtype());