Proper support of padding / unpadding for block floats

tenstorrent · Dec 19, 2024 · d21d2c4 · d21d2c4
1 parent 108932d
commit d21d2c4
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 32 deletions.
diff --git 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<T> arange(int64_t start, int64_t end, int64_t step) {
     return result;
 }
 
+// Creates a vector with alternating ones and zeros.
+std::vector<float> ones_and_zeros(int total_size) {
+    std::vector<float> 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 <typename T>
 class VectorConversionTest : public ::testing::Test {};
 
@@ -130,35 +140,35 @@ TEST_P(BlockFloatVectorConversionTest, InvalidLayout) {
 
 TEST_P(BlockFloatVectorConversionTest, Roundtrip) {
     ttnn::SimpleShape shape{128, 128};
-    std::vector<float> 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<float> input = ones_and_zeros(shape.volume());
+
     auto output = Tensor::from_vector(input, get_tensor_spec(shape, GetParam(), Layout::TILE)).to_vector<float>();
     EXPECT_THAT(output, Pointwise(Eq(), input));
 }
 
-TEST_P(BlockFloatVectorConversionTest, AddPadding) {
+TEST_P(BlockFloatVectorConversionTest, RoundtripWithPadding) {
     ttnn::SimpleShape shape{14, 47};
-    std::vector<float> input(shape.volume(), 1.0f);
+    std::vector<float> 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<float>(), Pointwise(Eq(), input));
 }
 
-TEST_P(BlockFloatVectorConversionTest, AddPaddingWithCustomTile) {
+TEST_P(BlockFloatVectorConversionTest, RoundtripWithPaddingAndCustomTile) {
     ttnn::SimpleShape shape{14, 47};
-    std::vector<float> input(shape.volume(), 1.0f);
+    std::vector<float> 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<float>(), 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<float>();
-    EXPECT_THAT(output, Pointwise(Eq(), input));
+    EXPECT_THAT(output.to_vector<float>(), Pointwise(Eq(), input));
 }
 
 }  // namespace

diff --git a/ttnn/cpp/ttnn/tensor/tensor.cpp b/ttnn/cpp/ttnn/tensor/tensor.cpp
@@ -637,18 +637,18 @@ Tensor Tensor::from_span<float>(
                 "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<float>(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<float>(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<float> padded_row_major_data = owned_buffer::get_as<float>(result_cpu_tensor).get();
+            std::vector<float> padded_row_major_data = owned_buffer::get_as<float>(fp32_tensor).get();
             std::vector<uint32_t> packed_block_floats =
                 spec.data_type() == DataType::BFLOAT8_B
                     ? pack_fp32_vec_as_bfp8_tiles(
@@ -698,11 +698,18 @@ std::vector<float> Tensor::to_vector<float>() 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<float, float>(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<float>(std::get<OwnedStorage>(fp32_cpu_tensor.storage()).buffer).get();
         }
         default: {
             TT_THROW("Cannot convert tensor to vector for data type: {}", cpu_tensor.get_dtype());