| 328 | |
| 329 | template <typename TensorType> |
| 330 | void FillTensorFromCudaArray(const py::object &object, |
| 331 | TensorType *batch, |
| 332 | int device_id, |
| 333 | const std::optional<std::string> &layout) { |
| 334 | auto cu_a_interface_val = getattr(object, "__cuda_array_interface__", py::none()); |
| 335 | if (cu_a_interface_val.is_none()) { |
| 336 | DALI_FAIL("Provided object doesn't support cuda array interface protocol.") |
| 337 | } |
| 338 | py::dict cu_a_interface = py::cast<py::dict>(cu_a_interface_val); |
| 339 | |
| 340 | DALI_ENFORCE(cu_a_interface.contains("typestr") && |
| 341 | // see detail::PyUnicode_Check_Permissive implementation |
| 342 | (PyUnicode_Check(cu_a_interface["typestr"].ptr()) || |
| 343 | PYBIND11_BYTES_CHECK(cu_a_interface["typestr"].ptr())) && |
| 344 | cu_a_interface.contains("shape") && |
| 345 | PyTuple_Check(cu_a_interface["shape"].ptr()) && |
| 346 | cu_a_interface.contains("data") && |
| 347 | PyTuple_Check(cu_a_interface["data"].ptr()) && |
| 348 | cu_a_interface["data"].cast<py::tuple>().size() >= 2 && |
| 349 | cu_a_interface.contains("version"), |
| 350 | "Provided object doesn't have required cuda array interface " |
| 351 | "protocol fields of necessary type."); |
| 352 | DALI_ENFORCE(!cu_a_interface.contains("mask") || cu_a_interface["mask"].is_none(), |
| 353 | "Masked tensors are not supported"); |
| 354 | |
| 355 | // Create the Tensor and wrap the data |
| 356 | TensorShape<> shape = shape_from_py(cu_a_interface["shape"].cast<py::tuple>()); |
| 357 | |
| 358 | std::string typestr = cu_a_interface["typestr"].cast<py::str>(); |
| 359 | const TypeInfo &type = TypeFromFormatStr(typestr); |
| 360 | size_t bytes = volume(shape) * type.size(); |
| 361 | |
| 362 | if (cu_a_interface.contains("strides") && !cu_a_interface["strides"].is_none()) { |
| 363 | TensorShape<> strides = shape_from_py(cu_a_interface["strides"].cast<py::tuple>()); |
| 364 | CheckContiguousTensor(strides, shape, type.size()); |
| 365 | } |
| 366 | |
| 367 | const auto &typed_shape = ConvertShape(shape, batch); |
| 368 | auto *ptr = PyLong_AsVoidPtr(cu_a_interface["data"].cast<py::tuple>()[0].ptr()); |
| 369 | |
| 370 | // it is for __cuda_array_interface__ so device_id < 0 is not a valid value |
| 371 | if (device_id < 0) { |
| 372 | CUDA_CALL(cudaGetDevice(&device_id)); |
| 373 | } |
| 374 | |
| 375 | batch->Reset(); |
| 376 | |
| 377 | if (cu_a_interface.contains("stream")) { |
| 378 | const auto &stream_obj = cu_a_interface["stream"]; |
| 379 | if (!stream_obj.is_none()) { |
| 380 | auto stream_long_value = cu_a_interface["stream"].cast<int64_t>(); |
| 381 | auto stream_value = PyLong_AsVoidPtr(cu_a_interface["stream"].ptr()); |
| 382 | DALI_ENFORCE(stream_value != 0, make_string("Provided stream is not a valid CUDA stream ", |
| 383 | "based on CUDA Array Interface v3. `0` value is ambiguous and disallowed")); |
| 384 | if (stream_long_value == 1) stream_value = 0; |
| 385 | if (stream_long_value == 2) stream_value = CU_STREAM_PER_THREAD; |
| 386 | auto order = AccessOrder(cudaStream_t(stream_value)); |
| 387 | batch->set_order(order); |
no test coverage detected