| 278 | |
| 279 | template<typename SrcBackend, template<typename> class SourceDataType> |
| 280 | void FillTensorFromDlPack( |
| 281 | py::capsule capsule, |
| 282 | SourceDataType<SrcBackend> *batch, |
| 283 | const std::optional<std::string> &layout) { |
| 284 | auto dlm_tensor_ptr = DLMTensorPtrFromCapsule(capsule); |
| 285 | const auto &dl_tensor = dlm_tensor_ptr->dl_tensor; |
| 286 | DALI_ENFORCE((std::is_same<SrcBackend, GPUBackend>::value && |
| 287 | dl_tensor.device.device_type == kDLCUDA) || |
| 288 | (std::is_same<SrcBackend, CPUBackend>::value && |
| 289 | (dl_tensor.device.device_type == kDLCPU || |
| 290 | dl_tensor.device.device_type == kDLCUDAHost)), |
| 291 | "DLPack device type doesn't match Tensor type"); |
| 292 | |
| 293 | const TypeInfo &dali_type = TypeTable::GetTypeInfo(ToDALIType(dl_tensor.dtype)); |
| 294 | TensorShape<> shape; |
| 295 | shape.resize(dl_tensor.ndim); |
| 296 | for (ssize_t i = 0; i < dl_tensor.ndim; ++i) { |
| 297 | shape[i] = dl_tensor.shape[i]; |
| 298 | } |
| 299 | |
| 300 | if (dl_tensor.strides) |
| 301 | CheckContiguousTensor(dl_tensor.strides, dl_tensor.ndim, dl_tensor.shape, dl_tensor.ndim, 1); |
| 302 | |
| 303 | size_t bytes = volume(shape) * dali_type.size(); |
| 304 | |
| 305 | const auto &typed_shape = ConvertShape(shape, batch); |
| 306 | bool is_pinned = dl_tensor.device.device_type == kDLCUDAHost; |
| 307 | int device_id = CPU_ONLY_DEVICE_ID; |
| 308 | // according to the docs kDLCUDAHost = kDLCPU | kDLCUDA so test it as a the first option |
| 309 | if (dl_tensor.device.device_type == kDLCUDAHost) { |
| 310 | device_id = CPU_ONLY_DEVICE_ID; |
| 311 | } else if (dl_tensor.device.device_type == kDLCPU) { |
| 312 | device_id = CPU_ONLY_DEVICE_ID; |
| 313 | } else if (dl_tensor.device.device_type == kDLCUDA) { |
| 314 | device_id = dl_tensor.device.device_id; |
| 315 | } else { |
| 316 | DALI_FAIL(make_string("Not supported DLPack device type: ", dl_tensor.device.device_type, ".")); |
| 317 | } |
| 318 | |
| 319 | // empty lambda that just captures dlm_tensor_ptr unique ptr that would be destructed when |
| 320 | // shared ptr is destroyed |
| 321 | batch->ShareData(shared_ptr<void>(dl_tensor.data, |
| 322 | [dlm_tensor_ptr = std::move(dlm_tensor_ptr)](void*) {}), |
| 323 | bytes, is_pinned, typed_shape, dali_type.id(), device_id); |
| 324 | |
| 325 | |
| 326 | SetLayout(*batch, layout); |
| 327 | } |
| 328 | |
| 329 | template <typename TensorType> |
| 330 | void FillTensorFromCudaArray(const py::object &object, |
no test coverage detected