[Paddle Tensor No.12][BUPT] 新增 Tensor.__dlpack_device__ (#69632)

ZHOU05030 · web-flow · commit eae4c8334de2 · 2024-11-27T14:58:59.000+08:00
* 新增__dlpack_device__

* 修改为DLDeviceType

* 修改cpu的device_id

* 添加单侧文件

* 添加Place判断逻辑

* 修改单侧文件

* 修改device_id
diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py
@@ -29,8 +29,10 @@
     _PADDLE_DTYPE_2_NUMPY_DTYPE,
     convert_uint16_to_float,
 )
+from paddle.base.libpaddle import Place
 from paddle.profiler.utils import in_profiler_mode
 from paddle.utils import deprecated
+from paddle.utils.dlpack import DLDeviceType
 
 from .. import core, framework, unique_name
 from ..framework import (
@@ -127,6 +129,7 @@ def _to_static_var(self, to_parameter=False, **kwargs):
             'strides',
             'offset',
             '__cuda_array_interface__',
+            '__dlpack_device__',
         ]
         param_keys = ['stop_gradient', 'trainable']
         if isinstance(self, EagerParamBase):
@@ -1257,6 +1260,39 @@ def coalesce(self: Tensor, name: str | None = None) -> Tensor:
         """
         return _C_ops.sparse_coalesce(self)
 
+    @framework.dygraph_only
+    def __dlpack_device__(self):
+        """
+        Extract the DLPack device type and device ID for the current tensor.
+
+        Returns:
+            tuple: A tuple containing the DLPack device type and device ID.
+                - device_type (DLDeviceType): The type of device (e.g., kDLCPU, kDLCUDA, etc.).
+                - device_id (int): The device ID.
+        """
+        place = self.place
+        if isinstance(place, Place):
+            if place.is_gpu_place():
+                return DLDeviceType.kDLCUDA, place.gpu_device_id()
+            elif place.is_cpu_place():
+                return DLDeviceType.kDLCPU, None
+            elif place.is_cuda_pinned_place():
+                return DLDeviceType.kDLCUDAHost, None
+            elif place.is_xpu_place():
+                return DLDeviceType.kDLOneAPI, place.xpu_device_id()
+            else:
+                raise RuntimeError(f"Unsupported Paddle device type {place}")
+        elif place.is_cpu_place():
+            return DLDeviceType.kDLCPU, None
+        elif place.is_cuda_pinned_place():
+            return DLDeviceType.kDLCUDAHost, None
+        elif place.is_gpu_place():
+            return DLDeviceType.kDLCUDA, place.get_device_id()
+        elif place.is_xpu_place():
+            return DLDeviceType.kDLOneAPI, place.get_device_id()
+        else:
+            raise ValueError(f"Unsupported tensor place: {place}")
+
     @property
     def __cuda_array_interface__(self):
         """Array view description for cuda tensors.
@@ -1374,6 +1410,7 @@ def __cuda_array_interface__(self):
         ("_use_gpudnn", _use_gpudnn),
         ("_md5sum", _md5sum),
         ("__cuda_array_interface__", __cuda_array_interface__),
+        ("__dlpack_device__", __dlpack_device__),
     ):
         setattr(core.eager.Tensor, method_name, method)
 
diff --git a/test/dygraph_to_static/test_tensor_attr_consistency.py b/test/dygraph_to_static/test_tensor_attr_consistency.py
@@ -78,6 +78,7 @@
         'value',
         'zero_',
         "__cuda_array_interface__",
+        "__dlpack_device__",
     ]
 )
 STATIC_ONLY_TENSOR_ATTRS_ALLOW_LIST = OrderedSet(
diff --git a/test/legacy_test/test_dlpack.py b/test/legacy_test/test_dlpack.py
@@ -326,6 +326,86 @@ def test_to_dlpack_from_zero_size(self):
                     np.testing.assert_array_equal(x.numpy(), y2.numpy())
 
 
+from paddle.utils.dlpack import DLDeviceType
+
+
+class TestDLPackDevice(unittest.TestCase):
+    def test_dlpack_device(self):
+        with dygraph_guard():
+
+            tensor_cpu = paddle.to_tensor([1, 2, 3], place=base.CPUPlace())
+            device_type, device_id = tensor_cpu.__dlpack_device__()
+            self.assertEqual(device_type, DLDeviceType.kDLCPU)
+            self.assertEqual(device_id, None)
+
+            if paddle.is_compiled_with_cuda():
+                tensor_cuda = paddle.to_tensor(
+                    [1, 2, 3], place=base.CUDAPlace(0)
+                )
+                device_type, device_id = tensor_cuda.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDA)
+                self.assertEqual(device_id, 0)
+
+            if paddle.is_compiled_with_cuda():
+                tensor_pinned = paddle.to_tensor(
+                    [1, 2, 3], place=base.CUDAPinnedPlace()
+                )
+                device_type, device_id = tensor_pinned.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDAHost)
+                self.assertEqual(device_id, None)
+
+            if paddle.is_compiled_with_xpu():
+                tensor_xpu = paddle.to_tensor([1, 2, 3], place=base.XPUPlace(0))
+                device_type, device_id = tensor_xpu.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLOneAPI)
+                self.assertEqual(device_id, 0)
+
+    def test_dlpack_device_zero_dim(self):
+        with dygraph_guard():
+
+            tensor = paddle.to_tensor(5.0, place=base.CPUPlace())
+            device_type, device_id = tensor.__dlpack_device__()
+            self.assertEqual(device_type, DLDeviceType.kDLCPU)
+            self.assertEqual(device_id, None)
+
+            if paddle.is_compiled_with_cuda():
+                tensor_cuda = paddle.to_tensor(5.0, place=base.CUDAPlace(0))
+                device_type, device_id = tensor_cuda.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDA)
+                self.assertEqual(device_id, 0)
+
+            if paddle.is_compiled_with_xpu():
+                tensor_xpu = paddle.to_tensor(5.0, place=base.XPUPlace(0))
+                device_type, device_id = tensor_xpu.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLOneAPI)
+                self.assertEqual(device_id, 0)
+
+    def test_dlpack_device_zero_size(self):
+        with dygraph_guard():
+            tensor = paddle.to_tensor(
+                paddle.zeros([0, 10]), place=base.CPUPlace()
+            )
+            device_type, device_id = tensor.__dlpack_device__()
+            self.assertEqual(device_type, DLDeviceType.kDLCPU)
+            self.assertEqual(device_id, None)
+
+            if paddle.is_compiled_with_cuda():
+                tensor_cuda = paddle.to_tensor(
+                    paddle.zeros([0, 10]), place=base.CUDAPlace(0)
+                )
+                device_type, device_id = tensor_cuda.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDA)
+                self.assertEqual(device_id, 0)
+
+            if paddle.is_compiled_with_xpu():
+                tensor_xpu = paddle.to_tensor(
+                    paddle.zeros([0, 10]), place=base.XPUPlace(0)
+                )
+                device_type, device_id = tensor_xpu.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLOneAPI)
+                self.assertEqual(device_id, 0)
+
+
 class TestRaiseError(unittest.TestCase):
     def test_to_dlpack_raise_type_error(self):
         self.assertRaises(TypeError, paddle.utils.dlpack.to_dlpack, np.zeros(5))
diff --git a/test/legacy_test/test_eager_tensor.py b/test/legacy_test/test_eager_tensor.py
@@ -23,6 +23,7 @@
 import paddle.nn.functional as F
 from paddle import base
 from paddle.base import core
+from paddle.utils.dlpack import DLDeviceType
 
 
 class TestEagerTensor(unittest.TestCase):
@@ -1292,6 +1293,88 @@ def test___cuda_array_interface__(self):
                     self.assertIn("version", interface)
                     self.assertEqual(interface["version"], 2)
 
+    def test_dlpack_device(self):
+        """test Tensor.__dlpack_device__"""
+        with dygraph_guard():
+            # test CPU
+            tensor_cpu = paddle.to_tensor([1, 2, 3], place=base.CPUPlace())
+            device_type, device_id = tensor_cpu.__dlpack_device__()
+            self.assertEqual(device_type, DLDeviceType.kDLCPU)
+            self.assertEqual(device_id, None)
+
+            # test CUDA
+            if paddle.is_compiled_with_cuda():
+                tensor_cuda = paddle.to_tensor(
+                    [1, 2, 3], place=base.CUDAPlace(0)
+                )
+                device_type, device_id = tensor_cuda.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDA)
+                self.assertEqual(device_id, 0)
+
+            # test CUDA Pinned
+            if paddle.is_compiled_with_cuda():
+                tensor_pinned = paddle.to_tensor(
+                    [1, 2, 3], place=base.CUDAPinnedPlace()
+                )
+                device_type, device_id = tensor_pinned.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDAHost)
+                self.assertEqual(device_id, None)
+
+            # test XPU
+            if paddle.is_compiled_with_xpu():
+                tensor_xpu = paddle.to_tensor([1, 2, 3], place=base.XPUPlace(0))
+                device_type, device_id = tensor_xpu.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLOneAPI)
+                self.assertEqual(device_id, 0)
+
+            # zero_dim
+            # test CPU
+            tensor = paddle.to_tensor(5.0, place=base.CPUPlace())
+            device_type, device_id = tensor.__dlpack_device__()
+            self.assertEqual(device_type, DLDeviceType.kDLCPU)
+            self.assertEqual(device_id, None)
+
+            # test CUDA
+            if paddle.is_compiled_with_cuda():
+                tensor_cuda = paddle.to_tensor(5.0, place=base.CUDAPlace(0))
+                device_type, device_id = tensor_cuda.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDA)
+                self.assertEqual(device_id, 0)
+
+            # test XPU
+            if paddle.is_compiled_with_xpu():
+                tensor_xpu = paddle.to_tensor(5.0, place=base.XPUPlace(0))
+                device_type, device_id = tensor_xpu.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLOneAPI)
+                self.assertEqual(device_id, 0)
+
+            # zero_size
+            # test CPU
+            tensor = paddle.to_tensor(
+                paddle.zeros([0, 10]), place=base.CPUPlace()
+            )
+            device_type, device_id = tensor.__dlpack_device__()
+            self.assertEqual(device_type, DLDeviceType.kDLCPU)
+            self.assertEqual(device_id, None)
+
+            # test CUDA
+            if paddle.is_compiled_with_cuda():
+                tensor_cuda = paddle.to_tensor(
+                    paddle.zeros([0, 10]), place=base.CUDAPlace(0)
+                )
+                device_type, device_id = tensor_cuda.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLCUDA)
+                self.assertEqual(device_id, 0)
+
+            # test XPU
+            if paddle.is_compiled_with_xpu():
+                tensor_xpu = paddle.to_tensor(
+                    paddle.zeros([0, 10]), place=base.XPUPlace(0)
+                )
+                device_type, device_id = tensor_xpu.__dlpack_device__()
+                self.assertEqual(device_type, DLDeviceType.kDLOneAPI)
+                self.assertEqual(device_id, 0)
+
     def test_tensor__format__(self):
         # test for floating point scalar
         for width in range(0, 5):

Original file line number	Diff line number	Diff line change
`@@ -78,6 +78,7 @@`
`78`	`78`	`'value',`
`79`	`79`	`'zero_',`
`80`	`80`	`"__cuda_array_interface__",`
	`81`	`+ "__dlpack_device__",`
`81`	`82`	`]`
`82`	`83`	`)`
`83`	`84`	`STATIC_ONLY_TENSOR_ATTRS_ALLOW_LIST = OrderedSet(`