Allow data casting along with cloning. (pytorch#15510)

Summary:

.

Differential Revision: D86070966

Author: shoumikhin

extension/tensor/tensor_ptr.cpp

@@ -164,7 +164,9 @@ TensorPtr make_tensor_ptr(
       [data = std::move(data)](void*) {});
 }
 
-TensorPtr clone_tensor_ptr(const executorch::aten::Tensor& tensor) {
+TensorPtr clone_tensor_ptr(
+    const executorch::aten::Tensor& tensor,
+    executorch::aten::ScalarType type) {
   std::vector<executorch::aten::SizesType> sizes(
       tensor.sizes().begin(), tensor.sizes().end());
   std::vector<executorch::aten::DimOrderType> dim_order{
@@ -178,23 +180,63 @@ TensorPtr clone_tensor_ptr(const executorch::aten::Tensor& tensor) {
 #ifndef USE_ATEN_LIB
   dynamism = tensor.shape_dynamism();
 #endif // USE_ATEN_LIB
-  return tensor.const_data_ptr()
-      ? make_tensor_ptr(
-            std::move(sizes),
-            std::vector<uint8_t>(
-                (uint8_t*)tensor.const_data_ptr(),
-                (uint8_t*)tensor.const_data_ptr() + tensor.nbytes()),
-            std::move(dim_order),
-            std::move(strides),
-            tensor.scalar_type(),
-            dynamism)
-      : make_tensor_ptr(
-            std::move(sizes),
-            nullptr,
-            std::move(dim_order),
-            std::move(strides),
-            tensor.scalar_type(),
-            dynamism);
+  const auto* tensor_data = tensor.const_data_ptr();
+  if (!tensor_data) {
+    return make_tensor_ptr(
+        std::move(sizes),
+        nullptr,
+        std::move(dim_order),
+        std::move(strides),
+        type,
+        dynamism);
+  }
+  const auto tensor_type = tensor.scalar_type();
+  if (tensor_type == type) {
+    return make_tensor_ptr(
+        std::move(sizes),
+        std::vector<uint8_t>(
+            (uint8_t*)tensor_data, (uint8_t*)tensor_data + tensor.nbytes()),
+        std::move(dim_order),
+        std::move(strides),
+        tensor_type,
+        dynamism);
+  }
+  ET_CHECK_MSG(
+      runtime::canCast(tensor_type, type),
+      "Cannot cast tensor type to desired type.");
+  const auto tensor_numel = static_cast<size_t>(tensor.numel());
+  std::vector<uint8_t> data(tensor_numel * aten::elementSize(type));
+
+  // Create a minimal context for error handling in ET_SWITCH
+  struct {
+    [[noreturn]] void fail(torch::executor::Error /* error */) {
+      ET_CHECK_MSG(false, "Unsupported dtype in clone_tensor_ptr");
+    }
+  } ctx;
+
+  ET_SWITCH_REALHBBF16_AND_UINT_TYPES(
+      tensor_type, ctx, "clone_tensor_ptr_from", CTYPE_FROM, [&] {
+        const CTYPE_FROM* tensor_data_ptr =
+            static_cast<const CTYPE_FROM*>(tensor_data);
+        ET_SWITCH_REALHBBF16_AND_UINT_TYPES(
+            type, ctx, "clone_tensor_ptr_to", CTYPE_TO, [&] {
+              CTYPE_TO* data_ptr = reinterpret_cast<CTYPE_TO*>(data.data());
+              std::transform(
+                  tensor_data_ptr,
+                  tensor_data_ptr + tensor_numel,
+                  data_ptr,
+                  [](const CTYPE_FROM& val) {
+                    return static_cast<CTYPE_TO>(val);
+                  });
+            });
+      });
+  return make_tensor_ptr(
+      std::move(sizes),
+      std::move(data),
+      std::move(dim_order),
+      std::move(strides),
+      type,
+      dynamism);
 }
 
 runtime::Error resize_tensor_ptr(

extension/tensor/tensor_ptr.h

@@ -114,7 +114,7 @@ inline TensorPtr make_tensor_ptr(
     ET_CHECK_MSG(
         runtime::canCast(deduced_type, type),
         "Cannot cast deduced type to specified type.");
-    std::vector<uint8_t> casted_data(data.size() * runtime::elementSize(type));
+    std::vector<uint8_t> casted_data(data.size() * aten::elementSize(type));
 
     // Create a minimal context for error handling in ET_SWITCH
     struct {
@@ -408,6 +408,21 @@ inline TensorPtr make_tensor_ptr(
       [tensor_ptr](void*) {});
 }
 
+/**
+ * Creates a TensorPtr that manages a new Tensor with the same properties
+ * as the given Tensor, but with a copy of the data owned by the returned
+ * TensorPtr, or nullptr if the original data is null.
+ *
+ * @param tensor The Tensor to clone.
+ * @param type The data type for the cloned tensor. The data will be cast
+ * from the source tensor's type.
+ * @return A new TensorPtr that manages a Tensor with the specified type
+ * and copied/cast data.
+ */
+TensorPtr clone_tensor_ptr(
+    const executorch::aten::Tensor& tensor,
+    executorch::aten::ScalarType type);
+
 /**
  * Creates a TensorPtr that manages a new Tensor with the same properties
  * as the given Tensor, but with a copy of the data owned by the returned
@@ -417,7 +432,25 @@ inline TensorPtr make_tensor_ptr(
  * @return A new TensorPtr that manages a Tensor with the same properties as the
  * original but with copied data.
  */
-TensorPtr clone_tensor_ptr(const executorch::aten::Tensor& tensor);
+inline TensorPtr clone_tensor_ptr(const executorch::aten::Tensor& tensor) {
+  return clone_tensor_ptr(tensor, tensor.scalar_type());
+}
+
+/**
+ * Creates a new TensorPtr by cloning the given TensorPtr, copying the
+ * underlying data.
+ *
+ * @param tensor The TensorPtr to clone.
+ * @param type The data type for the cloned tensor. The data will be cast
+ * from the source tensor's type.
+ * @return A new TensorPtr that manages a Tensor with the specified type
+ * and copied/cast data.
+ */
+inline TensorPtr clone_tensor_ptr(
+    const TensorPtr& tensor,
+    executorch::aten::ScalarType type) {
+  return clone_tensor_ptr(*tensor, type);
+}
 
 /**
  * Creates a new TensorPtr by cloning the given TensorPtr, copying the
@@ -428,7 +461,7 @@ TensorPtr clone_tensor_ptr(const executorch::aten::Tensor& tensor);
  * original but with copied data.
  */
 inline TensorPtr clone_tensor_ptr(const TensorPtr& tensor) {
-  return clone_tensor_ptr(*tensor);
+  return clone_tensor_ptr(*tensor, tensor->scalar_type());
 }
 
 /**

extension/tensor/test/tensor_ptr_test.cpp

@@ -571,6 +571,82 @@ TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorInt32) {
   EXPECT_EQ(cloned_tensor->scalar_type(), executorch::aten::ScalarType::Int);
 }
 
+TEST_F(TensorPtrTest, CloneTensorPtrCastInt32ToFloat) {
+  std::vector<int32_t> data = {1, 2, 3, 4};
+  auto tensor = make_tensor_ptr({2, 2}, std::move(data));
+  auto cloned_tensor =
+      clone_tensor_ptr(*tensor, executorch::aten::ScalarType::Float);
+
+  EXPECT_EQ(cloned_tensor->dim(), 2);
+  EXPECT_EQ(cloned_tensor->size(0), 2);
+  EXPECT_EQ(cloned_tensor->size(1), 2);
+  EXPECT_EQ(cloned_tensor->scalar_type(), executorch::aten::ScalarType::Float);
+  auto ptr = cloned_tensor->const_data_ptr<float>();
+  EXPECT_FLOAT_EQ(ptr[0], 1.0f);
+  EXPECT_FLOAT_EQ(ptr[1], 2.0f);
+  EXPECT_FLOAT_EQ(ptr[2], 3.0f);
+  EXPECT_FLOAT_EQ(ptr[3], 4.0f);
+}
+
+TEST_F(TensorPtrTest, CloneTensorPtrCastFloatToBFloat16) {
+  std::vector<float> data = {1.0f, 2.0f, 3.5f};
+  auto tensor = make_tensor_ptr({3}, std::move(data));
+  auto cloned_tensor =
+      clone_tensor_ptr(*tensor, executorch::aten::ScalarType::BFloat16);
+
+  EXPECT_EQ(cloned_tensor->dim(), 1);
+  EXPECT_EQ(cloned_tensor->size(0), 3);
+  EXPECT_EQ(
+      cloned_tensor->scalar_type(), executorch::aten::ScalarType::BFloat16);
+  auto ptr = cloned_tensor->const_data_ptr<executorch::aten::BFloat16>();
+  EXPECT_NEAR(static_cast<float>(ptr[0]), 1.0f, 0.01f);
+  EXPECT_NEAR(static_cast<float>(ptr[1]), 2.0f, 0.01f);
+  EXPECT_NEAR(static_cast<float>(ptr[2]), 3.5f, 0.01f);
+}
+
+TEST_F(TensorPtrTest, CloneTensorPtrCastKeepsMetadata) {
+  std::vector<uint8_t> data(
+      6 * executorch::aten::elementSize(executorch::aten::ScalarType::Float));
+  auto tensor = make_tensor_ptr({2, 3}, std::move(data));
+  auto cloned_tensor =
+      clone_tensor_ptr(*tensor, executorch::aten::ScalarType::Float);
+
+  EXPECT_EQ(cloned_tensor->dim(), 2);
+  EXPECT_EQ(cloned_tensor->size(0), 2);
+  EXPECT_EQ(cloned_tensor->size(1), 3);
+  EXPECT_EQ(cloned_tensor->strides()[0], 3);
+  EXPECT_EQ(cloned_tensor->strides()[1], 1);
+  EXPECT_EQ(cloned_tensor->scalar_type(), executorch::aten::ScalarType::Float);
+}
+
+TEST_F(TensorPtrTest, CloneTensorPtrCastNullData) {
+  auto tensor = make_tensor_ptr(
+      {2, 2},
+      nullptr,
+      {},
+      {},
+      executorch::aten::ScalarType::Float,
+      executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND);
+  auto cloned_tensor =
+      clone_tensor_ptr(*tensor, executorch::aten::ScalarType::Int);
+
+  EXPECT_EQ(cloned_tensor->dim(), 2);
+  EXPECT_EQ(cloned_tensor->size(0), 2);
+  EXPECT_EQ(cloned_tensor->size(1), 2);
+  EXPECT_EQ(cloned_tensor->const_data_ptr(), nullptr);
+  EXPECT_EQ(cloned_tensor->scalar_type(), executorch::aten::ScalarType::Int);
+}
+
+TEST_F(TensorPtrTest, CloneTensorPtrCastInvalidExpectDeath) {
+  std::vector<float> data = {1.0f, 2.0f};
+  auto tensor = make_tensor_ptr({2}, std::move(data));
+  ET_EXPECT_DEATH(
+      {
+        auto _ = clone_tensor_ptr(*tensor, executorch::aten::ScalarType::Int);
+      },
+      "");
+}
+
 TEST_F(TensorPtrTest, MakeTensorPtrFromTensorPtrInt32) {
   std::vector<int32_t> data = {1, 2, 3, 4};
   auto tensor = make_tensor_ptr({2, 2}, data);