[DLPACK] Support from_dlpack with shared memory

cmake/external/dlpack.cmake

@@ -15,7 +15,7 @@
 include(ExternalProject)
 
 set(DLPACK_PREFIX_DIR ${THIRD_PARTY_PATH}/dlpack)
-set(DLPACK_TAG v0.4)
+set(DLPACK_TAG v0.8)
 set(DLPACK_INCLUDE_DIR ${THIRD_PARTY_PATH}/dlpack/src/extern_dlpack/include)
 set(SOURCE_DIR ${PADDLE_SOURCE_DIR}/third_party/dlpack)
 include_directories(${SOURCE_DIR}/include)

paddle/fluid/framework/dlpack_tensor.cc

@@ -33,6 +33,10 @@ static ::DLDataType GetDLDataTypeCode() {
   } else if (std::is_same<T, phi::dtype::float16>::value ||
              std::is_floating_point<T>::value) {
     dtype.code = kDLFloat;
+  } else if (std::is_same<T, bool>::value) {
+    // Since std::is_unsigned<bool>::value is True,
+    // it is necessary to evaluate bool before std::is_unsigned.
+    dtype.code = kDLBool;
   } else if (std::is_unsigned<T>::value) {
     dtype.code = kDLUInt;
   } else if (std::is_integral<T>::value) {
@@ -99,7 +103,7 @@ struct DLDeviceVisitor {
   inline ::DLDevice operator()(const phi::GPUPlace &place) const {
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
     ::DLDevice device;
-    device.device_type = kDLGPU;
+    device.device_type = kDLCUDA;
     device.device_id = place.device;  // NOLINT
     return device;
 #else
@@ -111,7 +115,7 @@ struct DLDeviceVisitor {
   inline ::DLDevice operator()(const phi::GPUPinnedPlace &place) const {
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
     ::DLDevice device;
-    device.device_type = kDLCPUPinned;
+    device.device_type = kDLCUDAHost;
     device.device_id = 0;
     return device;
 #else
@@ -125,52 +129,54 @@ struct DLDeviceVisitor {
 struct PaddleDLMTensor {
   phi::DenseTensor handle;
   DLManagedTensor tensor;
-  PaddleDLMTensor() : tensor() {}
 };
 
-void deleter(DLManagedTensor *arg) {
-  delete[] arg->dl_tensor.shape;
-  delete[] arg->dl_tensor.strides;
-  delete static_cast<PaddleDLMTensor *>(arg->manager_ctx);
+static void deleter(DLManagedTensor *self) {
+  if (self && self->manager_ctx) {
+    delete[] self->dl_tensor
+        .shape;  // delete shape allocated in toDLPack manually
+    delete[] self->dl_tensor
+        .strides;  // delete strides allocated in toDLPack manually
+    delete static_cast<PaddleDLMTensor *>(self->manager_ctx);
+  }
 }
 
 DLManagedTensor *toDLPack(const phi::DenseTensor &src) {
   PaddleDLMTensor *pdDLMTensor(new PaddleDLMTensor);
   pdDLMTensor->handle = const_cast<phi::DenseTensor &>(src);
   pdDLMTensor->tensor.manager_ctx = pdDLMTensor;
   pdDLMTensor->tensor.deleter = &deleter;
-  pdDLMTensor->tensor.dl_tensor.data = const_cast<void *>(src.data());
 
   // init ndim
-  using DimType = decltype(pdDLMTensor->tensor.dl_tensor.ndim);  // int
-  pdDLMTensor->tensor.dl_tensor.ndim = static_cast<DimType>(src.dims().size());
+  using DimType = decltype(pdDLMTensor->tensor.dl_tensor.ndim);  // int32_t
+  auto _shape = src.dims();
+  pdDLMTensor->tensor.dl_tensor.ndim = static_cast<DimType>(_shape.size());
   DimType ndim = pdDLMTensor->tensor.dl_tensor.ndim;
 
   // init shape
-  auto shape = new int64_t[ndim];
+  int64_t *shape = new int64_t[ndim];
   for (DimType i = 0; i < ndim; ++i) {
-    shape[i] = src.dims()[i];
+    shape[i] = _shape[i];
   }
   pdDLMTensor->tensor.dl_tensor.shape = shape;
 
-  // init stride
-  auto strides = new int64_t[ndim];
-  for (DimType i = 0; i < ndim; ++i) {
-    strides[i] = 1;
-  }
-  for (DimType i = ndim - 2; i >= 0; --i) {
-    strides[i] = shape[i + 1] * strides[i + 1];
+  // init strides
+  auto _strides = src.strides();
+  int64_t *strides = new int64_t[ndim];
+  for (int i = 0; i < src.dims().size(); i++) {
+    strides[i] = _strides[i];
+    if (shape[i] < 2) {
+      strides[i] = 1;
+    }
   }
   pdDLMTensor->tensor.dl_tensor.strides = strides;
 
-  // init device, DLDevice type with device_type and device_id
+  pdDLMTensor->tensor.dl_tensor.data = const_cast<void *>(src.data());
   auto place = src.place();
   pdDLMTensor->tensor.dl_tensor.device =
       phi::VisitPlace(place, internal::DLDeviceVisitor());
-
   pdDLMTensor->tensor.dl_tensor.dtype = internal::GetDLDataTypeFromTypeIndex(
       framework::TransToProtoVarType(src.dtype()));
-
   pdDLMTensor->tensor.dl_tensor.byte_offset = 0;
   return &(pdDLMTensor->tensor);
 }

paddle/fluid/framework/tensor_util.cc

@@ -15,6 +15,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/tensor_util.h"
 
 #include <algorithm>
+#include <functional>
 #include <limits>
 #include <memory>
 #include <string>
@@ -793,6 +794,158 @@ void* GetDstPtrByDLDataType(DLDataType type,
   }
 }
 
+// get Tensor data dtype from given DLDataType
+phi::DataType GetDstPtrByDLDataType(DLDataType type) {
+  // vector types not currently supported
+  PADDLE_ENFORCE_LE(
+      type.lanes,
+      1,
+      common::errors::Unimplemented("Vector type is not supported currently."));
+
+  switch (type.bits) {
+    case 8:
+      if (type.code == kDLBool) return phi::DataType::BOOL;
+      if (type.code == kDLInt) return phi::DataType::INT8;
+      if (type.code == kDLUInt) return phi::DataType::UINT8;
+      PADDLE_THROW(common::errors::Unimplemented(
+          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
+          type.code,
+          type.bits));
+    case 16:
+      if (type.code == kDLInt) return phi::DataType::INT16;
+      if (type.code == kDLFloat) return phi::DataType::FLOAT16;
+      if (type.code == kDLBfloat) return phi::DataType::BFLOAT16;
+      PADDLE_THROW(common::errors::Unimplemented(
+          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
+          type.code,
+          type.bits));
+    case 32:
+      if (type.code == kDLInt) return phi::DataType::INT32;
+      if (type.code == kDLFloat) return phi::DataType::FLOAT32;
+      PADDLE_THROW(common::errors::Unimplemented(
+          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
+          type.code,
+          type.bits));
+    case 64:
+      if (type.code == kDLInt) return phi::DataType::INT64;
+      if (type.code == kDLFloat) return phi::DataType::FLOAT64;
+      if (type.code == kDLComplex) return phi::DataType::COMPLEX64;
+      PADDLE_THROW(common::errors::Unimplemented(
+          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
+          type.code,
+          type.bits));
+    case 128:
+      if (type.code == kDLComplex) return phi::DataType::COMPLEX128;
+      PADDLE_THROW(common::errors::Unimplemented(
+          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
+          type.code,
+          type.bits));
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unsupported DLDataType.bits %d.", type.bits));
+  }
+}
+
+/*
+dlpack related code ref:
+https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/DLConvertor.cpp
+and paddle/phi/api/lib/tensor_utils.cc
+*/
+using Deleter = std::function<void(void*)>;
+
+std::unordered_map<void*, std::function<void(phi::Allocation*)>> ptr_to_deleter;
+std::mutex ptr_to_deleter_mutex;  // use mutex to keep thread safe
+
+void DeleterBridge(phi::Allocation* alloc) {
+  std::lock_guard<std::mutex> lock(ptr_to_deleter_mutex);
+  auto it = ptr_to_deleter.find(static_cast<void*>(alloc->ptr()));
+  if (it != ptr_to_deleter.end()) {
+    it->second(alloc);         // call the deleter
+    ptr_to_deleter.erase(it);  // remove the entry from the map safely
+  }
+}
+
+phi::DenseTensor from_blob(void* data,
+                           DLManagedTensor* src,
+                           const phi::DDim& shape,
+                           phi::DataType dtype,
+                           phi::DataLayout layout,
+                           const phi::Place& place,
+                           const Deleter& deleter) {
+  PADDLE_ENFORCE_NOT_NULL(
+      data, phi::errors::InvalidArgument("data can not be nullptr."));
+
+  auto meta = phi::DenseTensorMeta(dtype, shape, layout);
+  size_t size = SizeOf(dtype) * (meta.is_scalar ? 1 : product(meta.dims));
+  phi::Allocation::DeleterFnPtr f = nullptr;
+
+  if (deleter) {
+    auto g = [deleter, src](phi::Allocation* p) {
+      if (src->manager_ctx) {
+        deleter(src);
+      }
+    };
+
+    {
+      std::lock_guard<std::mutex> lock(ptr_to_deleter_mutex);
+      ptr_to_deleter[data] = g;
+    }
+
+    f = DeleterBridge;
+  }
+
+  auto alloc = std::make_shared<phi::Allocation>(data, size, f, place);
+  return phi::DenseTensor(alloc, meta);
+}
+
+phi::DenseTensor TensorFromDLPack(DLManagedTensor* src, Deleter deleter) {
+  std::vector<int64_t> vec;
+  std::copy(src->dl_tensor.shape,
+            src->dl_tensor.shape + src->dl_tensor.ndim,
+            std::back_inserter(vec));
+
+  phi::Place place;
+  if (src->dl_tensor.device.device_type == kDLCPU) {
+    place = phi::CPUPlace();
+  } else if (src->dl_tensor.device.device_type == kDLCUDA) {
+    place = phi::GPUPlace();
+  } else if (src->dl_tensor.device.device_type == kDLCUDAHost) {
+    place = phi::GPUPinnedPlace();
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented("Given Place is not supported"));
+  }
+
+  ::DLDataType type = src->dl_tensor.dtype;
+  auto dtype = GetDstPtrByDLDataType(type);
+  if (!src->dl_tensor.strides) {
+    return from_blob(src->dl_tensor.data,
+                     src,
+                     common::make_ddim(vec),
+                     dtype,
+                     phi::DataLayout::NCHW,
+                     place,
+                     std::move(deleter));
+  } else {
+    return from_blob(src->dl_tensor.data,
+                     src,
+                     common::make_ddim(vec),
+                     dtype,
+                     phi::DataLayout::NCHW,
+                     place,
+                     deleter);
+  }
+}
+
+phi::DenseTensor TensorFromDLPack(DLManagedTensor* src) {
+  auto deleter = [src](void* self [[maybe_unused]]) {
+    if (src->deleter) {
+      src->deleter(src);
+    }
+  };
+  return TensorFromDLPack(src, std::move(deleter));
+}
+
+// Keep the this overloaded version of the interface unchanged.
 void TensorFromDLPack(const ::DLTensor& dl_tensor, phi::DenseTensor* dst) {
   phi::CPUPlace dst_place = phi::CPUPlace();
   phi::CPUPlace src_place = phi::CPUPlace();
@@ -815,7 +968,7 @@ void TensorFromDLPack(const ::DLTensor& dl_tensor, phi::DenseTensor* dst) {
     memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
   }
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-  if (dl_tensor.device.device_type == kDLGPU) {
+  if (dl_tensor.device.device_type == kDLCUDA) {
     phi::GPUPlace dst_place = phi::GPUPlace(dl_tensor.device.device_id);
     phi::GPUPlace src_place = phi::GPUPlace(dl_tensor.device.device_id);
     dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
@@ -833,46 +986,6 @@ 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));
-
-  phi::DDim vddim = common::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 = common::product(vddim) * type.bits / 8;
-
-  if (src->dl_tensor.device.device_type == kDLCPU) {
-    phi::CPUPlace dst_place = phi::CPUPlace();
-    phi::CPUPlace src_place = phi::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) {
-    phi::GPUPlace dst_place = phi::GPUPlace(src->dl_tensor.device.device_id);
-    phi::GPUPlace src_place = phi::GPUPlace(src->dl_tensor.device.device_id);
-    void* dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
-    auto* ctx = phi::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(common::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.

paddle/fluid/framework/tensor_util.h

@@ -106,11 +106,16 @@ template <typename T>
 void TensorToVector(const phi::DenseTensor& src, std::vector<T>* dst);
 
 // convert dlpack's DLTensor to tensor
-
 TEST_API void TensorFromDLPack(const ::DLTensor& dl_tensor,
                                phi::DenseTensor* dst);
-void TensorFromDLPack(const DLManagedTensor* src, phi::DenseTensor* dst);
 
+TEST_API phi::DenseTensor TensorFromDLPack(DLManagedTensor* src);
+inline phi::DenseTensor TensorFromDLPack(const DLManagedTensor* src) {
+  return TensorFromDLPack(const_cast<DLManagedTensor*>(src));
+}
+
+phi::DenseTensor TensorFromDLPack(DLManagedTensor* src,
+                                  std::function<void(void*)> deleter);
 //
 // The implementation of template functions.
 //

paddle/fluid/pybind/pybind.cc

@@ -1264,29 +1264,23 @@ PYBIND11_MODULE(libpaddle, m) {
     phi::DeviceContextPool::Instance().Get(place)->Wait();
   });
 
-  m.def("from_dlpack", [](py::capsule *dltensor) {
-    DLManagedTensor *dmt = reinterpret_cast<DLManagedTensor *>(
-        PyCapsule_GetPointer(dltensor->ptr(), "dltensor"));
+  m.def("from_dlpack", [](py::object data) {
+    DLManagedTensor *dlMTensor = reinterpret_cast<DLManagedTensor *>(
+        PyCapsule_GetPointer(data.ptr(), "dltensor"));
 
     PADDLE_ENFORCE_NOT_NULL(
-        dmt,
-        common::errors::InvalidArgument(
+        dlMTensor,
+        phi::errors::InvalidArgument(
             "from_dlpack received an invalid capsule. "
-            "Note that a DLPack tensor can be consumed only once."));
+            "Note that DLTensor capsules can be consumed only once, "
+            "so you might have already constructed a tensor from it once."));
 
-    PyCapsule_SetName(dltensor->ptr(), "used_dltensor");
-    DLTensor dl = dmt->dl_tensor;
-    phi::DenseTensor tensor;
+    // NOTE: Might meet bugged numpy version, see:
+    // https://github.com/pytorch/pytorch/blob/main/torch/csrc/utils/tensor_new.cpp#L1636-L1638
+    auto ptensor = paddle::framework::TensorFromDLPack(dlMTensor);
 
-    if (dl.device.device_type == kDLCPU) {
-      paddle::framework::TensorFromDLPack(dmt, &tensor);
-    }
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-    if (dl.device.device_type == kDLGPU) {
-      paddle::framework::TensorFromDLPack(dmt, &tensor);
-    }
-#endif
-    return tensor;
+    PyCapsule_SetName(data.ptr(), "used_dltensor");
+    return ptensor;
   });
 
   m.def("_create_loaded_parameter",