Fix tensor cloning when data is null. (#5535)

Summary:
Pull Request resolved: pytorch/executorch#5535

.

Reviewed By: dltn

Differential Revision: D63201286

fbshipit-source-id: 1767a1c0cf876f7a3b6b4534a83c912c3de0eabf

Author: shoumikhin

extension/tensor/tensor_ptr.cpp

@@ -13,6 +13,39 @@
 namespace executorch {
 namespace extension {
 
+TensorPtr clone_tensor_ptr(const exec_aten::Tensor& tensor) {
+  std::vector<exec_aten::SizesType> sizes(
+      tensor.sizes().begin(), tensor.sizes().end());
+  std::vector<exec_aten::DimOrderType> dim_order{
+#ifndef USE_ATEN_LIB
+      tensor.dim_order().begin(), tensor.dim_order().end()
+#endif // USE_ATEN_LIB
+  };
+  std::vector<exec_aten::StridesType> strides(
+      tensor.strides().begin(), tensor.strides().end());
+  auto dynamism = exec_aten::TensorShapeDynamism::DYNAMIC_BOUND;
+#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);
+}
+
 runtime::Error resize_tensor_ptr(
     TensorPtr& tensor,
     const std::vector<exec_aten::SizesType>& sizes) {

extension/tensor/tensor_ptr.h

@@ -397,34 +397,13 @@ inline TensorPtr make_tensor_ptr(
 /**
  * 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.
+ * TensorPtr, or nullptr if the original data is null.
  *
  * @param tensor The Tensor to clone.
  * @return A new TensorPtr that manages a Tensor with the same properties as the
  * original but with copied data.
  */
-inline TensorPtr clone_tensor_ptr(const exec_aten::Tensor& tensor) {
-  return make_tensor_ptr(make_tensor_impl_ptr(
-      std::vector<exec_aten::SizesType>(
-          tensor.sizes().begin(), tensor.sizes().end()),
-      std::vector<uint8_t>(
-          (uint8_t*)tensor.const_data_ptr(),
-          (uint8_t*)tensor.const_data_ptr() + tensor.nbytes()),
-#ifndef USE_ATEN_LIB
-      std::vector<exec_aten::DimOrderType>(
-          tensor.dim_order().begin(), tensor.dim_order().end()),
-      std::vector<exec_aten::StridesType>(
-          tensor.strides().begin(), tensor.strides().end()),
-      tensor.scalar_type(),
-      tensor.shape_dynamism()
-#else // USE_ATEN_LIB
-      {},
-      std::vector<exec_aten::StridesType>(
-          tensor.strides().begin(), tensor.strides().end()),
-      tensor.scalar_type()
-#endif // USE_ATEN_LIB
-          ));
-}
+TensorPtr clone_tensor_ptr(const exec_aten::Tensor& tensor);
 
 /**
  * Creates a new TensorPtr by cloning the given TensorPtr, copying the

extension/tensor/test/tensor_ptr_test.cpp

@@ -478,6 +478,17 @@ TEST_F(TensorPtrTest, CloneTensorPtrFromTensorPtrInt64) {
   EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Long);
 }
 
+TEST_F(TensorPtrTest, CloneTensorPtrFromTensorPtrNull) {
+  auto tensor = make_tensor_ptr({2, 2}, nullptr);
+  auto cloned_tensor = clone_tensor_ptr(tensor);
+
+  EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
+  EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
+  EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
+  EXPECT_EQ(cloned_tensor->const_data_ptr(), tensor->const_data_ptr());
+  EXPECT_EQ(cloned_tensor->const_data_ptr(), nullptr);
+}
+
 TEST_F(TensorPtrTest, TensorDataCastingFromIntToFloat) {
   std::vector<int32_t> int_data = {1, 2, 3, 4, 5, 6};
   auto tensor = make_tensor_ptr(