[DLPACK] Fix size of allocation in from_dlpack for strided Tensor (#68559)

HydrogenSulfate · web-flow · commit c16ee765fe1e · 2024-10-09T15:19:27.000+08:00
corresponding unittest
diff --git a/paddle/fluid/framework/tensor_util.cc b/paddle/fluid/framework/tensor_util.cc
@@ -872,13 +872,9 @@ phi::DenseTensor from_blob(void* data,
                            phi::DataType dtype,
                            const phi::Place& place,
                            const Deleter& deleter) {
-  PADDLE_ENFORCE_NOT_NULL(
-      data, common::errors::InvalidArgument("data can not be nullptr."));
-
   auto meta = phi::DenseTensorMeta(dtype, shape, strides);
-  size_t size = SizeOf(dtype) * (meta.is_scalar ? 1 : product(meta.dims));
-  phi::Allocation::DeleterFnPtr f = nullptr;
 
+  phi::Allocation::DeleterFnPtr f = nullptr;
   if (deleter) {
     auto g = [deleter, src](phi::Allocation* p) {
       if (src->manager_ctx) {
@@ -894,7 +890,18 @@ phi::DenseTensor from_blob(void* data,
     f = DeleterBridge;
   }
 
-  auto alloc = std::make_shared<phi::Allocation>(data, size, f, place);
+  // Calculate the number of elements of underlying storage
+  size_t size = 1;
+  for (auto i = 0; i < shape.size(); ++i) {
+    if (shape[i] == 0) {
+      size = 0;
+      break;
+    }
+    size += strides[i] * (shape[i] - 1);
+  }
+
+  auto alloc =
+      std::make_shared<phi::Allocation>(data, size * SizeOf(dtype), f, place);
   return phi::DenseTensor(alloc, meta);
 }
 
diff --git a/test/legacy_test/test_dlpack.py b/test/legacy_test/test_dlpack.py
@@ -31,6 +31,7 @@ def test_dlpack_dygraph(self):
             self.assertTrue(
                 isinstance(out_from_dlpack, paddle.base.core.eager.Tensor)
             )
+            self.assertEqual(str(tensor.place), str(out_from_dlpack.place))
             np.testing.assert_array_equal(
                 out_from_dlpack.numpy(), np.array([1, 2, 3, 4]).astype("int")
             )
@@ -41,6 +42,7 @@ def test_dlpack_tensor_larger_than_2dim(self):
             t = paddle.to_tensor(numpy_data)
             dlpack = paddle.utils.dlpack.to_dlpack(t)
             out = paddle.utils.dlpack.from_dlpack(dlpack)
+            self.assertEqual(str(t.place), str(out.place))
             np.testing.assert_allclose(numpy_data, out.numpy(), rtol=1e-05)
 
     def test_dlpack_static(self):
@@ -100,7 +102,7 @@ def test_dlpack_dtype_and_place_consistency(self):
                     o = paddle.utils.dlpack.from_dlpack(dlpack)
                     self.assertEqual(x.dtype, o.dtype)
                     np.testing.assert_allclose(x.numpy(), o.numpy(), rtol=1e-05)
-                    self.assertEqual(type(x.place), type(o.place))
+                    self.assertEqual(str(x.place), str(o.place))
 
             complex_dtypes = ["complex64", "complex128"]
             for place in places:
@@ -114,7 +116,7 @@ def test_dlpack_dtype_and_place_consistency(self):
                     o = paddle.utils.dlpack.from_dlpack(dlpack)
                     self.assertEqual(x.dtype, o.dtype)
                     np.testing.assert_allclose(x.numpy(), o.numpy(), rtol=1e-05)
-                    self.assertEqual(type(x.place), type(o.place))
+                    self.assertEqual(str(x.place), str(o.place))
 
     def test_dlpack_deletion(self):
         # See Paddle issue 47171
@@ -129,6 +131,7 @@ def test_dlpack_deletion(self):
                     )
                     dlpack = paddle.utils.dlpack.to_dlpack(a)
                     b = paddle.utils.dlpack.from_dlpack(dlpack)
+                    self.assertEqual(str(a.place), str(b.place))
 
     def test_to_dlpack_for_loop(self):
         # See Paddle issue 50120
@@ -154,6 +157,7 @@ def test_to_dlpack_modification(self):
                     y = paddle.utils.dlpack.from_dlpack(dlpack)
                     y[1:2, 2:5] = 2.0
                     np.testing.assert_allclose(x.numpy(), y.numpy())
+                    self.assertEqual(str(x.place), str(y.place))
 
     def test_to_dlpack_data_ptr_consistency(self):
         # See Paddle issue 50120
@@ -168,6 +172,7 @@ def test_to_dlpack_data_ptr_consistency(self):
                     y = paddle.utils.dlpack.from_dlpack(dlpack)
 
                     self.assertEqual(x.data_ptr(), y.data_ptr())
+                    self.assertEqual(str(x.place), str(y.place))
 
     def test_to_dlpack_strides_consistency(self):
         with dygraph_guard():
@@ -182,6 +187,8 @@ def test_to_dlpack_strides_consistency(self):
                     y = paddle.utils.dlpack.from_dlpack(dlpack)
 
                     self.assertEqual(x_strided.strides, y.strides)
+                    self.assertEqual(str(x_strided.place), str(y.place))
+                    np.testing.assert_equal(x_strided.numpy(), y.numpy())
 
     def test_to_dlpack_from_ext_tensor(self):
         with dygraph_guard():
@@ -192,6 +199,38 @@ def test_to_dlpack_from_ext_tensor(self):
                 self.assertEqual(x.__array_interface__['data'][0], y.data_ptr())
                 np.testing.assert_allclose(x, y.numpy())
 
+    def test_to_dlpack_from_zero_dim(self):
+        with dygraph_guard():
+            places = [base.CPUPlace()]
+            if paddle.is_compiled_with_cuda():
+                places.append(base.CUDAPlace(0))
+            for place in places:
+                for _ in range(4):
+                    x = paddle.to_tensor(1.0, place=place)
+                    dlpack = paddle.utils.dlpack.to_dlpack(x)
+                    y = paddle.utils.dlpack.from_dlpack(dlpack)
+                    self.assertEqual(x.data_ptr(), y.data_ptr())
+                    self.assertEqual(str(x.place), str(y.place))
+                    self.assertEqual(y.shape, [])
+                    self.assertEqual(y.numel().item(), 1)
+                    np.testing.assert_array_equal(x.numpy(), y.numpy())
+
+    def test_to_dlpack_from_zero_size(self):
+        with dygraph_guard():
+            places = [base.CPUPlace()]
+            if paddle.is_compiled_with_cuda():
+                places.append(base.CUDAPlace(0))
+            for place in places:
+                for _ in range(4):
+                    x = paddle.zeros([0, 10]).to(device=place)
+                    dlpack = paddle.utils.dlpack.to_dlpack(x)
+                    y = paddle.utils.dlpack.from_dlpack(dlpack)
+                    self.assertEqual(x.data_ptr(), y.data_ptr())
+                    self.assertEqual(str(x.place), str(y.place))
+                    self.assertEqual(y.shape, [0, 10])
+                    self.assertEqual(y.numel().item(), 0)
+                    np.testing.assert_array_equal(x.numpy(), y.numpy())
+
 
 class TestRaiseError(unittest.TestCase):
     def test_to_dlpack_raise_type_error(self):
