【SCU】【Paddle Tensor No.3】新增 Tensor.__dlpack__

paddle/fluid/pybind/tensor.cc

@@ -446,6 +446,28 @@ void BindTensor(pybind11::module &m) {  // NOLINT
                 });
             return capsule;
           })
+      .def(
+          "__dlpack__",
+          [](phi::DenseTensor &self) {
+              DLManagedTensor *dlMTensor = framework::toDLPack(self);
+              return pybind11::capsule(
+                  dlMTensor,
+                  "dltensor",
+                  [](PyObject *capsule) {
+                      DLManagedTensor *managedTensor =
+                          reinterpret_cast<DLManagedTensor *>(
+                              PyCapsule_GetPointer(capsule, "dltensor"));
+                      if (managedTensor && managedTensor->deleter) {
+                          managedTensor->deleter(managedTensor);
+                      }
+                  });
+          },
+          R"DOC(
+          Encode the tensor to a DLPack capsule.
+
+          Returns:
+              PyCapsule: The DLPack representation of the tensor.
+          )DOC")
       .def("_set_float_element", TensorSetElement<float>)
       .def("_get_float_element", TensorGetElement<float>)
       .def("_set_double_element", TensorSetElement<double>)

python/paddle/__init__.py

@@ -591,6 +591,7 @@
 from .utils.dlpack import (
     from_dlpack,
     to_dlpack,
+    __dlpack__,
 )
 
 # CINN has to set a flag to include a lib
@@ -1212,6 +1213,7 @@
     'positive',
     'from_dlpack',
     'to_dlpack',
+    '__dlpack__',
     'inf',
     'newaxis',
     'nan',

python/paddle/tensor/__init__.py

@@ -466,6 +466,7 @@
     var,
 )
 from .to_string import set_printoptions  # noqa: F401
+from paddle.utils.dlpack import __dlpack__
 
 # this list used in math_op_patch.py for _binary_creator_
 tensor_method_func = [

python/paddle/tensor/tensor.prototype.pyi

@@ -301,6 +301,7 @@ class AbstractTensor:
     def _grad_ivar(self) -> Tensor | None: ...
 
     # annotation: ${tensor_alias}
+    def __dlpack__(self) -> Any: ...
 
 class Tensor(AbstractTensor, AbstractEagerParamBase):
     # annotation: ${tensor_docstring}

python/paddle/utils/dlpack.py

@@ -33,6 +33,7 @@
 __all__ = [
     'to_dlpack',
     'from_dlpack',
+    '__dlpack__'
 ]
 
 _T_contra = TypeVar("_T_contra", contravariant=True)
@@ -221,3 +222,42 @@ def from_dlpack(
         out: Tensor = paddle.Tensor(out, place=out._place())
 
     return out
+
+
+def __dlpack__(tensor: paddle.Tensor, stream: int | None = None):
+    """
+    Export a Paddle Tensor to a DLPack capsule for data interchange.
+
+    Args:
+        tensor (paddle.Tensor): The input tensor.
+        stream (int | None): An optional stream for synchronization. 
+
+    Returns:
+        CapsuleType: A PyCapsule object representing the tensor in DLPack format.
+    """
+    if not isinstance(tensor, paddle.Tensor):
+        raise TypeError(f"Expected a paddle.Tensor, but got {type(tensor)}.")
+
+    if tensor.stop_gradient is False:
+        raise RuntimeError(
+            "Tensors that require gradients cannot be exported. "
+            "Call `detach()` before exporting."
+        )
+    if tensor.layout != paddle.strided:
+        raise RuntimeError(
+            "Only strided tensors can be exported to DLPack."
+        )
+
+    if stream is not None:
+        if not isinstance(stream, int):
+            raise TypeError("The `stream` argument must be an int or None.")
+        if tensor.place.is_gpu_place():
+            current_stream = paddle.device.cuda.current_stream()
+            if stream != 1 and stream != current_stream.cuda_stream:
+                event = paddle.device.cuda.Event()
+                event.record(current_stream)
+                sync_stream = paddle.device.cuda.ExternalStream(stream)
+                sync_stream.wait_event(event)
+
+    return paddle.utils.dlpack.to_dlpack(tensor)
+

test/legacy_test/test_dlpack.py

 def __cuda_array_interface__(self): 
 def __cuda_array_interface__(self): 
@@ -325,6 +325,62 @@ def test_to_dlpack_from_zero_size(self):
                     np.testing.assert_array_equal(x.numpy(), y1.numpy())
                     np.testing.assert_array_equal(x.numpy(), y2.numpy())
 
+    def test__dlpack_basic(self):
+        with dygraph_guard():
+            tensor = paddle.to_tensor(np.array([1, 2, 3, 4]).astype("int"))
+            dlpack_from_method = tensor.__dlpack__()
+            out_from_dlpack = paddle.utils.dlpack.from_dlpack(dlpack_from_method)
+
+            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")
+            )
+
+    def test__dlpack_consistency_with_to_dlpack(self):
+        with dygraph_guard():
+            tensor = paddle.to_tensor(np.random.rand(4, 5).astype("float32"))
+            dlpack_via_method = tensor.__dlpack__()
+            dlpack_via_func = paddle.to_dlpack(tensor)
+
+            self.assertEqual(dlpack_via_method.__dlpack_device__, dlpack_via_func.__dlpack_device__)
+            self.assertEqual(dlpack_via_method.__dlpack_dtype__, dlpack_via_func.__dlpack_dtype__)
+
+            out_from_method = paddle.from_dlpack(dlpack_via_method)
+            out_from_func = paddle.from_dlpack(dlpack_via_func)
+
+            np.testing.assert_allclose(out_from_method.numpy(), out_from_func.numpy(), rtol=1e-05)
+
+    def test__dlpack_on_special_cases(self):
+        with dygraph_guard():
+            zero_dim_tensor = paddle.to_tensor(1.0)
+            dlpack = zero_dim_tensor.__dlpack__()
+            out_from_dlpack = paddle.from_dlpack(dlpack)
+            self.assertEqual(out_from_dlpack.shape, [])
+            np.testing.assert_array_equal(zero_dim_tensor.numpy(), out_from_dlpack.numpy())
+
+            empty_tensor = paddle.zeros([0, 10])
+            dlpack = empty_tensor.__dlpack__()
+            out_from_dlpack = paddle.from_dlpack(dlpack)
+            self.assertEqual(out_from_dlpack.shape, [0, 10])
+            np.testing.assert_array_equal(empty_tensor.numpy(), out_from_dlpack.numpy())
+
+    def test__dlpack_on_different_devices_and_dtypes(self):
+        with dygraph_guard():
+            dtypes = ["float32", "int32", "bool"]
+            places = [base.CPUPlace()]
+            if paddle.is_compiled_with_cuda():
+                places.append(base.CUDAPlace(0))
+
+            for dtype in dtypes:
+                for place in places:
+                    tensor = paddle.ones([3, 4], dtype=dtype)
+                    dlpack = tensor.__dlpack__()
+                    out = paddle.from_dlpack(dlpack)
+                    self.assertEqual(str(tensor.place), str(out.place))
+                    self.assertEqual(tensor.dtype, out.dtype)
+                    np.testing.assert_array_equal(tensor.numpy(), out.numpy())
+
 
 class TestRaiseError(unittest.TestCase):
     def test_to_dlpack_raise_type_error(self):