Change create_tt_tensor_from_py_data to use from_vector

tests/ttnn/unit_tests/test_convert_python_tensor.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import pathlib
+import pytest
+
+import torch
+
+import ttnn
+
+
+@pytest.mark.parametrize("size", [64])
+@pytest.mark.parametrize("mode", [ttnn.graph.RunMode.NO_DISPATCH, ttnn.graph.RunMode.NORMAL])
+@pytest.mark.parametrize("dtype", [torch.int32, torch.float, torch.bfloat16])
+def test_convert_python_tensor(device, size, mode, dtype):
+    torch.manual_seed(0)
+
+    ttnn.graph.begin_graph_capture(mode)
+    torch_input_tensor = torch.rand((size,), (dtype))
+    input_tensor = ttnn.from_torch(torch_input_tensor, layout=ttnn.TILE_LAYOUT, device=device)
+    output_tensor = ttnn.to_torch(input_tensor, torch_rank=1)
+    captured_graph = ttnn.graph.end_graph_capture()
+    calltrace = ttnn.graph.extract_calltrace(captured_graph)
+
+    assert output_tensor == input_tensor
+    assert "tt::tt_metal::detail::convert_python_tensor_to_tt_tensor" in calltrace
+    assert captured_graph[0]["node_type"] == "capture_start"
+    assert captured_graph[1]["node_type"] == "function_start"
+    assert captured_graph[1]["params"]["name"] == "tt::tt_metal::detail::convert_python_tensor_to_tt_tensor"
+    assert captured_graph[-2]["node_type"] == "buffer_deallocate"
+    assert captured_graph[-1]["node_type"] == "capture_end"
+
+
+@pytest.mark.parametrize("size", [64])
+@pytest.mark.parametrize("mode", [ttnn.graph.RunMode.NO_DISPATCH, ttnn.graph.RunMode.NORMAL])
+@pytest.mark.parametrize("dtype", [ttnn.bfloat4_b, ttnn.bfloat8_b])
+def test_convert_python_tensor_bfp_b(device, size, mode, dtype):
+    torch.manual_seed(0)
+
+    ttnn.graph.begin_graph_capture(mode)
+    torch_input_tensor = torch.rand((size,), torch.float)
+    input_tensor = ttnn.from_torch(torch_input_tensor, layout=ttnn.TILE_LAYOUT, device=device, dtype=(dtype))
+    output_tensor = ttnn.to_torch(input_tensor, torch_rank=1)
+    captured_graph = ttnn.graph.end_graph_capture()
+    calltrace = ttnn.graph.extract_calltrace(captured_graph)
+
+    assert output_tensor == input_tensor
+    assert "tt::tt_metal::detail::convert_python_tensor_to_tt_tensor" in calltrace
+    assert captured_graph[0]["node_type"] == "capture_start"
+    assert captured_graph[1]["node_type"] == "function_start"
+    assert captured_graph[1]["params"]["name"] == "tt::tt_metal::detail::convert_python_tensor_to_tt_tensor"
+    assert captured_graph[-2]["node_type"] == "buffer_deallocate"
+    assert captured_graph[-1]["node_type"] == "capture_end"

ttnn/cpp/pybind11/pytensor.cpp

@@ -85,6 +85,45 @@ Tensor create_owned_tensor(T* data_ptr, const ttnn::TensorSpec& tensor_spec) {
     return Tensor(std::move(storage), tensor_spec);
 }
 
+template <typename T>
+Tensor create_typed_tt_tensor_from_py_data(
+    std::size_t py_data_ptr,
+    const TensorSpec& tensor_spec,
+    IDevice* device,
+    const std::function<void()>& on_creation_callback,
+    const std::function<void()>& on_destruction_callback,
+    const bool force_disable_borrow) {
+    const bool requires_padding = tensor_spec.logical_2d_shape() != tensor_spec.physical_shape();
+    const bool requires_tilization = layout != Layout::ROW_MAJOR;
+    const bool enable_borrow = !requires_padding and !requires_tilization and !force_disable_borrow;
+
+    TT_FATAL(
+        !tensor_spec.memory_config().is_sharded() or tensor_spec.memory_config().shard_spec.has_value(),
+        "Sharded tensors must have a shard spec when converting to tt tensors!");
+
+    // Use template type for generic function - TODO find better way, maybe decltype or variants w/ array or map?
+    auto data_ptr = reinterpret_cast<T*>(py_data_ptr);
+
+    std::size_t num_elements = tensor_spec.logical_shape().volume();
+
+    // never enable_borrow for bfloat8 and bfloat4 since they're tt specific types
+    if (enable_borrow and
+        !(tensor_spec.data_type == DataType::BFLOAT8_B || tensor_spec.data_type == DataType::BFLOAT4_B)) {
+        auto storage = BorrowedStorage(
+            borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
+        return Tensor(std::move(storage), tensor_spec);
+    } else {
+        std::size_t num_elements = tensor_spec.logical_shape().volume();
+        auto logical_data = std::vector<T>(data_ptr, data_ptr + num_elements);
+
+        // Abstract away handling by calling from_vector which calls from_span which handles bfloats
+        return Tensor::from_vector(
+            std::move(logical_data),
+            tensor_spec,
+            device == nullptr ? std::nullopt : std::optional<ttnn::AnyDevice>(device));
+    }
+}
+
 Tensor create_tt_tensor_from_py_data(
     std::size_t py_data_ptr,
     const TensorSpec& tensor_spec,
@@ -94,97 +133,43 @@ Tensor create_tt_tensor_from_py_data(
     const std::function<void()>& on_destruction_callback) {
     auto layout = tensor_spec.layout();
 
-    const bool requires_padding = tensor_spec.logical_2d_shape() != tensor_spec.physical_shape();
-    const bool requires_tilization = layout != Layout::ROW_MAJOR;
-    const bool enable_borrow = !requires_padding and !requires_tilization and !force_disable_borrow;
-
     auto data_type = tensor_spec.data_type();
     std::size_t num_elements = tensor_spec.logical_shape().volume();
     switch (data_type) {
         case DataType::UINT8: {
-            auto data_ptr = reinterpret_cast<uint8_t*>(py_data_ptr);
-            if (enable_borrow) {
-                auto storage = BorrowedStorage(
-                    borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
-                return Tensor(std::move(storage), tensor_spec);
-            } else {
-                return create_owned_tensor(data_ptr, tensor_spec);
-            }
+            return create_typed_tt_tensor_from_py_data<uint8_t>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
         case DataType::UINT16: {
-            auto data_ptr = reinterpret_cast<uint16_t*>(py_data_ptr);
-            if (enable_borrow) {
-                auto storage = BorrowedStorage(
-                    borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
-                return Tensor(std::move(storage), tensor_spec);
-            } else {
-                return create_owned_tensor(data_ptr, tensor_spec);
-            }
+            return create_typed_tt_tensor_from_py_data<uint16_t>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
         case DataType::INT32: {
-            auto data_ptr = reinterpret_cast<int32_t*>(py_data_ptr);
-            if (enable_borrow) {
-                auto storage = BorrowedStorage(
-                    borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
-                return Tensor(std::move(storage), tensor_spec);
-            } else {
-                return create_owned_tensor(data_ptr, tensor_spec);
-            }
+            return create_typed_tt_tensor_from_py_data<int32_t>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
         case DataType::UINT32: {
-            auto data_ptr = reinterpret_cast<uint32_t*>(py_data_ptr);
-            if (enable_borrow) {
-                auto storage = BorrowedStorage(
-                    borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
-                return Tensor(std::move(storage), tensor_spec);
-            } else {
-                return create_owned_tensor(data_ptr, tensor_spec);
-            }
+            return create_typed_tt_tensor_from_py_data<uint32_t>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
         case DataType::FLOAT32: {
-            auto data_ptr = reinterpret_cast<float*>(py_data_ptr);
-            if (enable_borrow) {
-                auto storage = BorrowedStorage(
-                    borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
-                return Tensor(std::move(storage), tensor_spec);
-            } else {
-                return create_owned_tensor(data_ptr, tensor_spec);
-            }
+            return create_typed_tt_tensor_from_py_data<float>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
         // TODO: This is not supported for numpy
         case DataType::BFLOAT16: {
-            auto data_ptr = reinterpret_cast<::bfloat16*>(py_data_ptr);
-            if (enable_borrow) {
-                auto storage = BorrowedStorage(
-                    borrowed_buffer::Buffer(data_ptr, num_elements), on_creation_callback, on_destruction_callback);
-                return Tensor(std::move(storage), tensor_spec);
-            } else {
-                return create_owned_tensor(data_ptr, tensor_spec);
-            }
+            return create_typed_tt_tensor_from_py_data<bfloat16>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
         case DataType::BFLOAT8_B:
         case DataType::BFLOAT4_B: {
-            auto data_ptr = reinterpret_cast<float*>(py_data_ptr);
-            auto float_tensor_spec = TensorSpec(
-                tensor_spec.logical_shape(),
-                TensorLayout(DataType::FLOAT32, tensor_spec.page_config(), tensor_spec.memory_config()));
-            auto float_tensor = create_owned_tensor(data_ptr, float_tensor_spec);
-
-            auto tile = tensor_spec.tensor_layout().get_page_config().get_tile();
-            auto output_float_data = owned_buffer::get_as<float>(float_tensor).get();
-            auto output_packed_data = data_type == DataType::BFLOAT8_B
-                                          ? pack_fp32_vec_as_bfp8_tiles(
-                                                output_float_data, /*row_major_input=*/false, /*is_exp_a=*/false, tile)
-                                          : pack_fp32_vec_as_bfp4_tiles(
-                                                output_float_data, /*row_major_input=*/false, /*is_exp_a=*/false, tile);
-            auto output_buffer = owned_buffer::create<uint32_t>(std::move(output_packed_data));
-            return Tensor(std::move(OwnedStorage{std::move(output_buffer)}), tensor_spec);
-        }
-        default: {
-            TT_THROW("Unsupported DataType: {}", data_type);
-            break;
+            return create_typed_tt_tensor_from_py_data<bfloat16>(
+                py_data_ptr, tensor_spec, device, on_creation_callback, on_destruction_callback, force_disable_borrow);
         }
     }
+
+    // remove default case in switch
+    TT_THROW("Unsupported DataType: {}", data_type);
 }
 
 Tensor convert_python_tensor_to_tt_tensor(

ttnn/cpp/ttnn/tensor/tensor.cpp

@@ -40,7 +40,9 @@ Tensor create_owned_tensor_from_row_major_data(
         spec.logical_shape(),
         TensorLayout(spec.data_type(), PageConfig(Layout::ROW_MAJOR, spec.tile()), MemoryConfig{}));
 
-    Tensor output(OwnedStorage{owned_buffer::create(std::move(data))}, result_cpu_spec);
+    auto physical_data = tensor_impl::encode_tensor_data(std::move(logical_data), tensor_spec);
+
+    Tensor output(OwnedStorage{owned_buffer::create(std::move(physical_data))}, result_cpu_spec);
 
     if (spec.layout() == Layout::TILE) {
         // TODO: whenever possible, perform tiliziation on device.