Make TensorImplPtr custom deleter copyable.

Differential Revision: D62338334

Pull Request resolved: #5161

Author: shoumikhin

extension/tensor/tensor_impl_ptr.cpp

@@ -30,20 +30,23 @@ static void noop_deleter(void*) {}
  * TensorImpl is destroyed.
  */
 struct TensorImplPtrDeleter final {
-  std::unique_ptr<void, std::function<void(void*)>> data;
-  std::vector<exec_aten::SizesType> sizes;
-  std::vector<exec_aten::DimOrderType> dim_order;
-  std::vector<exec_aten::StridesType> strides;
+  // A custom deleter of the std::shared_ptr is required to be copyable until
+  // C++20, so any data it holds must be copyable too. Hence, we use shared_ptr
+  // to hold the data and metadata to avoid unnecessary copies.
+  std::shared_ptr<void> data;
+  std::shared_ptr<std::vector<exec_aten::SizesType>> sizes;
+  std::shared_ptr<std::vector<exec_aten::DimOrderType>> dim_order;
+  std::shared_ptr<std::vector<exec_aten::StridesType>> strides;
 
   void operator()(exec_aten::TensorImpl* pointer) {
     // Release all resources immediately since the data held by the
-    // TensorImplDeleter is tied to the managed object, not the smart pointer
+    // TensorImplPtrDeleter is tied to the managed object, not the smart pointer
     // itself. We need to free this memory when the object is destroyed, not
     // when the smart pointer (and deleter) are eventually destroyed or reset.
     data.reset();
-    sizes = {};
-    dim_order = {};
-    strides = {};
+    sizes.reset();
+    dim_order.reset();
+    strides.reset();
     delete pointer;
   }
 };
@@ -90,11 +93,13 @@ TensorImplPtr make_tensor_impl_ptr(
   return TensorImplPtr(
       tensor_impl.release(),
       TensorImplPtrDeleter{
-          std::unique_ptr<void, std::function<void(void*)>>(
+          std::shared_ptr<void>(
               data, deleter ? std::move(deleter) : noop_deleter),
-          std::move(sizes),
-          std::move(dim_order),
-          std::move(strides)});
+          std::make_shared<std::vector<exec_aten::SizesType>>(std::move(sizes)),
+          std::make_shared<std::vector<exec_aten::DimOrderType>>(
+              std::move(dim_order)),
+          std::make_shared<std::vector<exec_aten::StridesType>>(
+              std::move(strides))});
 #else
   auto options = c10::TensorOptions()
                      .dtype(c10::scalarTypeToTypeMeta(type))

extension/tensor/tensor_impl_ptr.h

@@ -94,15 +94,18 @@ TensorImplPtr make_tensor_impl_ptr(
     std::vector<exec_aten::StridesType> strides = {},
     exec_aten::TensorShapeDynamism dynamism =
         exec_aten::TensorShapeDynamism::STATIC) {
-  const auto data_ptr = data.data();
+  auto raw_data_ptr = data.data();
+  auto data_ptr = std::make_shared<
+      std::vector<typename runtime::ScalarTypeToCppType<T>::type>>(
+      std::move(data));
   return make_tensor_impl_ptr(
       T,
       std::move(sizes),
-      data_ptr,
+      raw_data_ptr,
       std::move(dim_order),
       std::move(strides),
       dynamism,
-      [data = std::move(data)](void*) {});
+      [data_ptr = std::move(data_ptr)](void*) {});
 }
 
 } // namespace extension

extension/tensor/test/tensor_impl_ptr_test.cpp

@@ -181,3 +181,46 @@ TEST_F(TensorImplPtrTest, TensorImplOwningEmptyData) {
   EXPECT_EQ(tensor_impl->strides()[1], 1);
   EXPECT_EQ(tensor_impl->data(), nullptr);
 }
+
+TEST_F(TensorImplPtrTest, SharedDataManagement) {
+  auto data = std::make_shared<std::vector<float>>(100, 1.0f);
+  auto tensor_impl1 = make_tensor_impl_ptr(
+      exec_aten::ScalarType::Float, {10, 10}, data->data());
+  auto tensor_impl2 = tensor_impl1;
+
+  EXPECT_EQ(tensor_impl1.get(), tensor_impl2.get());
+  EXPECT_EQ(tensor_impl1.use_count(), 2);
+  EXPECT_EQ(((float*)tensor_impl1->data())[0], 1.0f);
+
+  ((float*)tensor_impl1->mutable_data())[0] = 2.0f;
+  EXPECT_EQ(((float*)tensor_impl2->data())[0], 2.0f);
+
+  tensor_impl1.reset();
+  EXPECT_NE(tensor_impl2.get(), nullptr);
+  EXPECT_EQ(tensor_impl2.use_count(), 1);
+
+  EXPECT_EQ(((float*)tensor_impl2->data())[0], 2.0f);
+}
+
+TEST_F(TensorImplPtrTest, CustomDeleterWithSharedData) {
+  auto data = std::make_shared<std::vector<float>>(100, 1.0f);
+  bool deleter_called = false;
+  {
+    auto tensor_impl = make_tensor_impl_ptr(
+        exec_aten::ScalarType::Float,
+        {10, 10},
+        data->data(),
+        {},
+        {},
+        exec_aten::TensorShapeDynamism::STATIC,
+        [data, &deleter_called](void*) mutable {
+          deleter_called = true;
+          data.reset();
+        });
+
+    EXPECT_EQ(data.use_count(), 2);
+    EXPECT_FALSE(deleter_called);
+  }
+  EXPECT_TRUE(deleter_called);
+  EXPECT_EQ(data.use_count(), 1);
+}

extension/tensor/test/tensor_ptr_test.cpp

@@ -176,3 +176,36 @@ TEST_F(TensorPtrTest, TensorOwningEmptyData) {
   EXPECT_EQ(tensor->strides()[1], 1);
   EXPECT_EQ(tensor->data_ptr<float>(), nullptr);
 }
+
+TEST_F(TensorPtrTest, TensorSharingImplModifiesSharedDataVector) {
+  std::vector<float> data = {1, 2, 3, 4, 5, 6};
+
+  auto tensor1 = make_tensor_ptr({2, 3}, std::move(data));
+  auto tensor2 = make_tensor_ptr(tensor1);
+
+  tensor1->mutable_data_ptr<float>()[0] = 10;
+  EXPECT_EQ(tensor2->const_data_ptr<float>()[0], 10);
+
+  tensor2->mutable_data_ptr<float>()[5] = 20;
+  EXPECT_EQ(tensor1->const_data_ptr<float>()[5], 20);
+}
+
+TEST_F(TensorPtrTest, TensorSharingImplResizingAffectsBothVector) {
+  std::vector<float> data = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
+
+  auto tensor1 = make_tensor_ptr(
+      {3, 4},
+      std::move(data),
+      {},
+      {},
+      exec_aten::TensorShapeDynamism::DYNAMIC_UNBOUND);
+  auto tensor2 = make_tensor_ptr(tensor1);
+
+  EXPECT_EQ(resize_tensor_ptr(tensor1, {2, 6}), Error::Ok);
+  EXPECT_EQ(tensor2->size(0), 2);
+  EXPECT_EQ(tensor2->size(1), 6);
+
+  EXPECT_EQ(resize_tensor_ptr(tensor2, {4, 3}), Error::Ok);
+  EXPECT_EQ(tensor1->size(0), 4);
+  EXPECT_EQ(tensor1->size(1), 3);
+}