From 8dc050d8a6a586581e7fd484ec39977a3ce51ee6 Mon Sep 17 00:00:00 2001 From: Siming Dai <908660116@qq.com> Date: Thu, 26 Aug 2021 23:08:26 +0800 Subject: [PATCH] Add paddle.utils.dlpack APIs (#35067) * add dlpack api and fix a from_dlpack --- paddle/fluid/pybind/pybind.cc | 6 ++ python/paddle/tests/test_dlpack.py | 71 ++++++++++++++++++++ python/paddle/utils/__init__.py | 1 + python/paddle/utils/dlpack.py | 102 +++++++++++++++++++++++++++++ 4 files changed, 180 insertions(+) create mode 100644 python/paddle/tests/test_dlpack.py create mode 100644 python/paddle/utils/dlpack.py diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index a53bd1f535..677da35b41 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/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( 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; diff --git a/python/paddle/tests/test_dlpack.py b/python/paddle/tests/test_dlpack.py new file mode 100644 index 0000000000..2880901d1a --- /dev/null +++ b/python/paddle/tests/test_dlpack.py @@ -0,0 +1,71 @@ +# 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() diff --git a/python/paddle/utils/__init__.py b/python/paddle/utils/__init__.py index 2c7bca7169..6994fd139b 100644 --- a/python/paddle/utils/__init__.py +++ b/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' diff --git a/python/paddle/utils/dlpack.py b/python/paddle/utils/dlpack.py new file mode 100644 index 0000000000..d3a7e7918f --- /dev/null +++ b/python/paddle/utils/dlpack.py @@ -0,0 +1,102 @@ +# 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) + # + """ + + 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 -- GitLab