| 272 | } |
| 273 | |
| 274 | std::optional<py::dict> ExternalBuffer::cudaArrayInterface() const |
| 275 | { |
| 276 | if (!m_cacheCudaArrayInterface) |
| 277 | { |
| 278 | if (!IsCudaAccessible(m_dlTensor->device.device_type)) |
| 279 | { |
| 280 | return std::nullopt; |
| 281 | } |
| 282 | |
| 283 | nvcv::DataType dataType = ToNVCVDataType(m_dlTensor->dtype); |
| 284 | |
| 285 | NVCV_ASSERT(m_dlTensor->dtype.bits % 8 == 0); |
| 286 | NVCV_ASSERT(dataType.strideBytes() * 8 == m_dlTensor->dtype.bits * m_dlTensor->dtype.lanes); |
| 287 | int elemStrideBytes = dataType.strideBytes(); |
| 288 | |
| 289 | py::object strides; |
| 290 | |
| 291 | if (m_dlTensor->strides == nullptr) |
| 292 | { |
| 293 | strides = py::none(); |
| 294 | } |
| 295 | else |
| 296 | { |
| 297 | py::tuple vStrides(m_dlTensor->ndim); |
| 298 | for (size_t i = 0; i < vStrides.size(); ++i) |
| 299 | { |
| 300 | vStrides[i] = m_dlTensor->strides[i] * elemStrideBytes; |
| 301 | } |
| 302 | strides = vStrides; |
| 303 | } |
| 304 | |
| 305 | std::string format = ToFormatString(m_dlTensor->dtype); |
| 306 | |
| 307 | // clang-format off |
| 308 | m_cacheCudaArrayInterface = py::dict |
| 309 | { |
| 310 | "shape"_a = this->shape(), |
| 311 | "strides"_a = strides, |
| 312 | "typestr"_a = format, |
| 313 | "data"_a = py::make_tuple(reinterpret_cast<long>(m_dlTensor->data), false /* read/write */), |
| 314 | "version"_a = 2 |
| 315 | }; |
| 316 | } |
| 317 | |
| 318 | return *m_cacheCudaArrayInterface; |
| 319 | } |
| 320 | |
| 321 | py::capsule ExternalBuffer::dlpack(py::object stream) const |
| 322 | { |
no test coverage detected