Add paddle.utils.dlpack APIs (#35067)

* add dlpack api and fix a from_dlpack 

paddle/fluid/pybind/pybind.cc

@@ -524,6 +524,12 @@ PYBIND11_MODULE(core_noavx, m) {
   m.def("from_dlpack", [](py::capsule *dltensor) {
     DLManagedTensor *dmt = reinterpret_cast<DLManagedTensor *>(
         PyCapsule_GetPointer(dltensor->ptr(), "dltensor"));
+
+    PADDLE_ENFORCE_NOT_NULL(
+        dmt, platform::errors::InvalidArgument(
+                 "from_dlpack received an invalid capsule. "
+                 "Note that a DLPack tensor can be consumed only once."));
+
     PyCapsule_SetName(dltensor->ptr(), "used_dltensor");
     DLTensor dl = dmt->dl_tensor;
     framework::Tensor tensor;

python/paddle/tests/test_dlpack.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import unittest
+import numpy as np
+
+import paddle
+import paddle.fluid as fluid
+import paddle.fluid.core as core
+
+
+class TestDLPack(unittest.TestCase):
+    def test_dlpack_dygraph(self):
+        tensor = paddle.to_tensor(np.array([1, 2, 3, 4]).astype('int'))
+        dlpack = paddle.utils.dlpack.to_dlpack(tensor)
+        out_from_dlpack = paddle.utils.dlpack.from_dlpack(dlpack)
+        self.assertTrue(isinstance(out_from_dlpack, paddle.Tensor))
+        self.assertTrue(
+            np.array_equal(
+                np.array(out_from_dlpack), np.array([1, 2, 3, 4]).astype(
+                    'int')))
+
+    def test_dlpack_static(self):
+        paddle.enable_static()
+        tensor = fluid.create_lod_tensor(
+            np.array([[1], [2], [3], [4]]).astype('int'), [[1, 3]],
+            fluid.CPUPlace())
+        dlpack = paddle.utils.dlpack.to_dlpack(tensor)
+        out_from_dlpack = paddle.utils.dlpack.from_dlpack(dlpack)
+        self.assertTrue(isinstance(out_from_dlpack, fluid.core.Tensor))
+        self.assertTrue(
+            np.array_equal(
+                np.array(out_from_dlpack),
+                np.array([[1], [2], [3], [4]]).astype('int')))
+
+        # when build with cuda
+        if core.is_compiled_with_cuda():
+            gtensor = fluid.create_lod_tensor(
+                np.array([[1], [2], [3], [4]]).astype('int'), [[1, 3]],
+                fluid.CUDAPlace(0))
+            gdlpack = paddle.utils.dlpack.to_dlpack(gtensor)
+            gout_from_dlpack = paddle.utils.dlpack.from_dlpack(gdlpack)
+            self.assertTrue(isinstance(gout_from_dlpack, fluid.core.Tensor))
+            self.assertTrue(
+                np.array_equal(
+                    np.array(gout_from_dlpack),
+                    np.array([[1], [2], [3], [4]]).astype('int')))
+
+
+class TestRaiseError(unittest.TestCase):
+    def test_from_dlpack_raise_type_error(self):
+        self.assertRaises(TypeError, paddle.utils.dlpack.from_dlpack,
+                          np.zeros(5))
+
+    def test_to_dlpack_raise_type_error(self):
+        self.assertRaises(TypeError, paddle.utils.dlpack.to_dlpack, np.zeros(5))
+
+
+if __name__ == '__main__':
+    unittest.main()

python/paddle/utils/__init__.py

@@ -26,6 +26,7 @@ from ..fluid.framework import require_version  # noqa: F401
 from . import download  # noqa: F401
 from . import image_util  # noqa: F401
 from . import cpp_extension  # noqa: F401
+from . import dlpack
 
 __all__ = [  #noqa
     'deprecated', 'run_check', 'require_version', 'try_import'

python/paddle/utils/dlpack.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import paddle
+from ..fluid.core import LoDTensor
+from ..fluid.framework import in_dygraph_mode
+from ..fluid.data_feeder import check_type, check_dtype, convert_dtype
+
+
+def to_dlpack(x):
+    """
+    Encodes a tensor to DLPack.
+
+    Args:
+        x (Tensor): A tensor, and the data type is bool, float32, float64, int32, int64.
+
+    Returns:
+        dltensor, and the data type is PyCapsule.
+    
+    Examples:
+        .. code-block:: python
+
+            import paddle
+            # x is a tensor with shape [2, 4]
+            x = paddle.to_tensor([[0.2, 0.3, 0.5, 0.9],
+                                  [0.1, 0.2, 0.6, 0.7]])
+            dlpack = paddle.utils.dlpack.to_dlpack(x)
+            print(dlpack)
+            # <capsule object "dltensor" at 0x7f6103c681b0>
+    """
+
+    if in_dygraph_mode():
+        if not isinstance(x, paddle.Tensor):
+            raise TypeError(
+                "The type of 'x' in to_dlpack must be paddle.Tensor,"
+                " but received {}.".format(type(x)))
+
+        dtype = convert_dtype(x.dtype)
+
+        if dtype not in ['bool', 'int32', 'int64', 'float32', 'float64']:
+            raise TypeError(
+                "the dtype of 'x' in to_dlpack must be any of [bool, int32, int64, "
+                "float32, float64], but received {}.".format(dtype))
+
+        return x.value().get_tensor()._to_dlpack()
+
+    check_type(x, 'x', (LoDTensor), 'to_dlpack')
+    check_dtype(x._dtype(), 'x',
+                ['bool', 'int32', 'int64', 'float32', 'float64'], 'to_dlpack')
+
+    return x._to_dlpack()
+
+
+def from_dlpack(dlpack):
+    """Decodes a DLPack to a tensor.
+    
+    Args:
+        dlpack (PyCapsule): a PyCapsule object with the dltensor.
+
+    Returns:
+        out (Tensor): a tensor decoded from DLPack.
+
+    Examples:
+        .. code-block:: python
+
+            import paddle
+            # x is a tensor with shape [2, 4]
+            x = paddle.to_tensor([[0.2, 0.3, 0.5, 0.9],
+                                  [0.1, 0.2, 0.6, 0.7]])
+            dlpack = paddle.utils.dlpack.to_dlpack(x)
+            x = paddle.utils.dlpack.from_dlpack(dlpack)
+            print(x)
+            # Tensor(shape=[2, 4], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
+              [[0.20000000, 0.30000001, 0.50000000, 0.89999998],
+              [0.10000000, 0.20000000, 0.60000002, 0.69999999]]) 
+    """
+
+    t = type(dlpack)
+    dlpack_flag = (t.__module__ == 'builtins' and t.__name__ == 'PyCapsule')
+    if not dlpack_flag:
+        raise TypeError(
+            "The type of 'dlpack' in from_dlpack must be PyCapsule object,"
+            " but received {}.".format(type(dlpack)))
+
+    if in_dygraph_mode():
+        out = paddle.fluid.core.from_dlpack(dlpack)
+        out = paddle.to_tensor(out)
+        return out
+
+    out = paddle.fluid.core.from_dlpack(dlpack)
+    return out