未验证 提交 1f3ff41c 编写于 作者: S Siming Dai 提交者: GitHub

Fix dlpack deletion (#47310)

* fix dlpack deletion

* add unittest

* fix unittest
上级 d8314ff5
......@@ -1108,6 +1108,48 @@ void TensorFromDLPack(const ::DLTensor& dl_tensor, phi::DenseTensor* dst) {
#endif
}
void TensorFromDLPack(const DLManagedTensor* src, phi::DenseTensor* dst) {
std::vector<int64_t> vec;
std::copy(src->dl_tensor.shape,
src->dl_tensor.shape + src->dl_tensor.ndim,
std::back_inserter(vec));
framework::DDim vddim = phi::make_ddim(vec);
dst->Resize(vddim);
::DLDataType type = src->dl_tensor.dtype;
auto src_ptr = static_cast<const void*>(src->dl_tensor.data);
auto size = phi::product(vddim) * type.bits / 8;
if (src->dl_tensor.device.device_type == kDLCPU) {
platform::CPUPlace dst_place = platform::CPUPlace();
platform::CPUPlace src_place = platform::CPUPlace();
void* dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
}
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
if (src->dl_tensor.device.device_type == kDLGPU) {
platform::CUDAPlace dst_place =
platform::CUDAPlace(src->dl_tensor.device.device_id);
platform::CUDAPlace src_place =
platform::CUDAPlace(src->dl_tensor.device.device_id);
void* dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
auto* ctx = platform::DeviceContextPool::Instance().GetByPlace(dst_place);
// Fix copy by share allocation.
memory::Copy(dst_place,
dst_ptr,
src_place,
src_ptr,
size,
reinterpret_cast<const phi::GPUContext&>(*ctx).stream());
}
#endif
src->deleter(const_cast<DLManagedTensor*>(src));
#ifdef PADDLE_WITH_XPU
PADDLE_THROW(platform::errors::Unimplemented("XPUPlace is not supported"));
#endif
}
template <typename T>
std::string format_tensor(const phi::DenseTensor& tensor) {
// TODO(zhiqiu): use the print option to format tensor.
......
......@@ -112,7 +112,9 @@ template <typename T>
void TesnorToVector(const phi::DenseTensor& src, std::vector<T>* dst);
// convert dlpack's DLTensor to tensor
void TensorFromDLPack(const ::DLTensor& dl_tensor, phi::DenseTensor* dst);
void TensorFromDLPack(const DLManagedTensor* src, phi::DenseTensor* dst);
//
// The implementation of template functions.
......
......@@ -703,11 +703,11 @@ PYBIND11_MODULE(libpaddle, m) {
phi::DenseTensor tensor;
if (dl.device.device_type == kDLCPU) {
paddle::framework::TensorFromDLPack(dl, &tensor);
paddle::framework::TensorFromDLPack(dmt, &tensor);
}
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
if (dl.device.device_type == kDLGPU) {
paddle::framework::TensorFromDLPack(dl, &tensor);
paddle::framework::TensorFromDLPack(dmt, &tensor);
}
#endif
return tensor;
......
......@@ -123,6 +123,14 @@ class TestDLPack(unittest.TestCase):
self.func_test_dlpack_dtype_conversion()
self.func_test_dlpack_dtype_conversion()
def test_dlpack_deletion(self):
# See Paddle issue 47171
if paddle.is_compiled_with_cuda():
for i in range(80):
a = paddle.rand(shape=[1024 * 128, 1024], dtype="float32")
dlpack = paddle.utils.dlpack.to_dlpack(a)
b = paddle.utils.dlpack.from_dlpack(dlpack)
class TestRaiseError(unittest.TestCase):
def func_test_from_dlpack_raise_type_error(self):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册