Fix deleter initialization in from_blob api (#68859)

huangjiyi · web-flow · commit ce51e82e84fc · 2024-10-22T19:31:33.000+08:00
diff --git a/paddle/phi/api/include/tensor_utils.h b/paddle/phi/api/include/tensor_utils.h
@@ -35,6 +35,7 @@ using Deleter = std::function<void(void*)>;
  *
  * @param data The pointer to the memory buffer.
  * @param shape The dims of the tensor.
+ * @param strides The strides of the tensor.
  * @param dtype The data type of the tensor, should correspond to data type of
  *              `data`. See PD_FOR_EACH_DATA_TYPE in `phi/common/data_type.h`
  * @param layout The data layout of the tensor.
@@ -54,6 +55,13 @@ PADDLE_API Tensor from_blob(void* data,
                             phi::DataLayout layout = phi::DataLayout::NCHW,
                             const phi::Place& place = phi::Place(),
                             const Deleter& deleter = nullptr);
+PADDLE_API Tensor from_blob(void* data,
+                            const phi::IntArray& shape,
+                            const phi::IntArray& strides,
+                            phi::DataType dtype,
+                            phi::DataLayout layout = phi::DataLayout::NCHW,
+                            const phi::Place& place = phi::Place(),
+                            const Deleter& deleter = nullptr);
 
 #ifdef PADDLE_WITH_DISTRIBUTE
 /**
diff --git a/paddle/phi/api/lib/tensor_utils.cc b/paddle/phi/api/lib/tensor_utils.cc
@@ -58,22 +58,44 @@ phi::Place GetPlaceFromPtr(void* data) {
   return phi::CPUPlace();
 }
 
+struct DeleterManeger {
+  static DeleterManeger* Instance() {
+    static DeleterManeger instance;
+    return &instance;
+  }
+  DeleterManeger() = default;
+
+  void DeletePtr(void* ptr) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto it = ptr2deleter_.find(ptr);
+    if (it != ptr2deleter_.end()) {
+      it->second(ptr);
+      ptr2deleter_.erase(it);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "The deleter of the pointer is not found."));
+    }
+  }
+
+  void RegisterPtr(void* ptr, const Deleter& deleter) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    ptr2deleter_[ptr] = deleter;
+  }
+
+ private:
+  std::unordered_map<void*, Deleter> ptr2deleter_;
+  std::mutex mutex_;
+};
+
 using AllocationDeleter = void (*)(phi::Allocation*);
 
-PADDLE_API Tensor from_blob(void* data,
-                            const phi::IntArray& shape,
-                            phi::DataType dtype,
-                            phi::DataLayout layout,
-                            const phi::Place& place,
-                            const Deleter& deleter) {
+Tensor FromBlobImpl(void* data,
+                    const phi::DenseTensorMeta& meta,
+                    const phi::Place& place,
+                    const Deleter& deleter) {
   PADDLE_ENFORCE_NOT_NULL(
       data, common::errors::InvalidArgument("data can not be nullptr."));
 
-  PADDLE_ENFORCE_EQ(shape.FromTensor(),
-                    false,
-                    common::errors::InvalidArgument(
-                        "shape cannot be constructed from a Tensor."));
-
   phi::Place data_place;
   if (place.GetType() == phi::AllocationType::UNDEFINED ||
       place.GetType() == phi::AllocationType::CPU ||
@@ -92,23 +114,54 @@ PADDLE_API Tensor from_blob(void* data,
     data_place = place;
   }
 
-  auto meta =
-      phi::DenseTensorMeta(dtype, common::make_ddim(shape.GetData()), layout);
-
-  size_t size = SizeOf(dtype) * (meta.is_scalar ? 1 : product(meta.dims));
+  // Calculate the number of elements of underlying storage
+  size_t size = 1;
+  for (auto i = 0; i < meta.dims.size(); ++i) {
+    if (meta.dims[i] == 0) {
+      size = 0;
+      break;
+    }
+    size += meta.strides[i] * (meta.dims[i] - 1);
+  }
 
   AllocationDeleter alloc_deleter = nullptr;
   if (deleter) {
-    static thread_local Deleter g_deleter = deleter;
-    alloc_deleter = [](phi::Allocation* p) { g_deleter(p->ptr()); };
+    DeleterManeger::Instance()->RegisterPtr(data, deleter);
+    alloc_deleter = [](phi::Allocation* p) {
+      DeleterManeger::Instance()->DeletePtr(p->ptr());
+    };
   }
 
-  auto alloc =
-      std::make_shared<phi::Allocation>(data, size, alloc_deleter, data_place);
+  auto alloc = std::make_shared<phi::Allocation>(
+      data, size * SizeOf(meta.dtype), alloc_deleter, data_place);
 
   return Tensor(std::make_shared<phi::DenseTensor>(alloc, meta));
 }
 
+PADDLE_API Tensor from_blob(void* data,
+                            const phi::IntArray& shape,
+                            phi::DataType dtype,
+                            phi::DataLayout layout,
+                            const phi::Place& place,
+                            const Deleter& deleter) {
+  auto meta =
+      phi::DenseTensorMeta(dtype, common::make_ddim(shape.GetData()), layout);
+  return FromBlobImpl(data, meta, place, deleter);
+}
+
+PADDLE_API Tensor from_blob(void* data,
+                            const phi::IntArray& shape,
+                            const phi::IntArray& strides,
+                            phi::DataType dtype,
+                            phi::DataLayout layout,
+                            const phi::Place& place,
+                            const Deleter& deleter) {
+  auto meta = phi::DenseTensorMeta(dtype,
+                                   common::make_ddim(shape.GetData()),
+                                   common::make_ddim(strides.GetData()));
+  return FromBlobImpl(data, meta, place, deleter);
+}
+
 #ifdef PADDLE_WITH_DISTRIBUTE
 PD_REGISTER_API(reshard)
 
diff --git a/test/cpp/phi/api/test_from_blob.cc b/test/cpp/phi/api/test_from_blob.cc
@@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#include <glog/logging.h>
 #include <gtest/gtest.h>
 
 #include "paddle/phi/api/include/api.h"
@@ -172,31 +173,52 @@ TEST(from_blob, GPU) {
 #endif
 
 TEST(from_blob, Option) {
-  // 1. create data
-  auto data = new int64_t[8];
-  for (int64_t i = 0; i < 8; i++) {
-    data[i] = i;
-  }
-
-  // 2. test Deleter and Layout
-  int isdelete = 0;
-  auto deleter = [&isdelete](void* data) {
+  int delete_count = 0, f_delete_count = 0;
+  auto deleter = [&delete_count](void* data) {
     delete[] static_cast<int64_t*>(data);
-    isdelete++;
+    delete_count++;
+  };
+  auto f_deleter = [&f_delete_count](void* ptr) {
+    delete[] static_cast<float*>(ptr);
+    f_delete_count++;
   };
   {
+    auto data = new int64_t[8];
+    for (int64_t i = 0; i < 8; i++) {
+      data[i] = i;
+    }
     auto test_tensor = from_blob(data,
-                                 {1, 2, 2, 1},
+                                 {1, 2, 2, 2},
                                  DataType::INT64,
                                  phi::DataLayout::NHWC,
                                  phi::CPUPlace(),
                                  deleter);
-
-    // check tensor attributes
-    ASSERT_EQ(test_tensor.layout(), phi::DataLayout::NHWC);  // check layout
-
-    // check deleter
-    ASSERT_EQ(isdelete, 0);
+    ASSERT_EQ(test_tensor.layout(), phi::DataLayout::NHWC);
+    ASSERT_EQ(delete_count, 0);
+
+    auto f_data = new float[8];
+    for (int i = 0; i < 8; i++) {
+      f_data[i] = static_cast<float>(i);
+    }
+    auto test_tensor_f = from_blob(f_data,
+                                   {1, 2, 2, 2},
+                                   DataType::FLOAT32,
+                                   common::DataLayout::NHWC,
+                                   phi::CPUPlace(),
+                                   f_deleter);
+    ASSERT_EQ(test_tensor_f.layout(), phi::DataLayout::NHWC);
+    ASSERT_EQ(f_delete_count, 0);
   }
-  ASSERT_EQ(isdelete, 1);
+  ASSERT_EQ(delete_count, 1);
+  ASSERT_EQ(f_delete_count, 1);
+}
+
+TEST(from_blob, Strides) {
+  int64_t data[8] = {0, 1, 2, 3, 4, 5, 6, 7};
+  auto test_tensor =
+      from_blob(data, {1, 2, 2, 1}, {0, 4, 2, 0}, DataType::INT64);
+  ASSERT_EQ(test_tensor.shape()[1], 2);
+  ASSERT_EQ(test_tensor.shape()[2], 2);
+  ASSERT_EQ(test_tensor.strides()[1], 4);
+  ASSERT_EQ(test_tensor.strides()[2], 2);
 }
