python/dlpack.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file dlpack.cpp | ||
| 3 | /// @brief DLPack capsules, versioned and not, per the DLPack Python specification. | ||
| 4 | #include "ernie/python/dlpack.h" | ||
| 5 | |||
| 6 | #include <nanobind/stl/tuple.h> | ||
| 7 | |||
| 8 | #include <Python.h> | ||
| 9 | |||
| 10 | #include <cstdint> | ||
| 11 | #include <tuple> | ||
| 12 | |||
| 13 | namespace nb = nanobind; | ||
| 14 | |||
| 15 | namespace ernie::python { | ||
| 16 | namespace { | ||
| 17 | |||
| 18 | /// Names a capsule takes, per the DLPack Python specification. | ||
| 19 | constexpr const char* kVersioned = "dltensor_versioned"; | ||
| 20 | constexpr const char* kUnversioned = "dltensor"; | ||
| 21 | |||
| 22 | /// An unversioned DLPack tensor over a versioned one, for consumers that ask for no version. | ||
| 23 | ✗ | DLManagedTensor* Unversioned(DLManagedTensorVersioned* versioned) { | |
| 24 | ✗ | auto* managed = new DLManagedTensor{.dl_tensor = versioned->dl_tensor, | |
| 25 | .manager_ctx = versioned, | ||
| 26 | ✗ | .deleter = [](DLManagedTensor* self) { | |
| 27 | ✗ | auto* inner = static_cast<DLManagedTensorVersioned*>( | |
| 28 | self->manager_ctx); | ||
| 29 | ✗ | inner->deleter(inner); | |
| 30 | ✗ | delete self; | |
| 31 | ✗ | }}; | |
| 32 | ✗ | return managed; | |
| 33 | } | ||
| 34 | |||
| 35 | /// Takes the tensor from a DLPack capsule, marking the capsule used. | ||
| 36 | 4 | Tensor FromCapsule(const nb::object& capsule) { | |
| 37 | 4 | PyObject* object = capsule.ptr(); | |
| 38 | 4 | if (PyCapsule_IsValid(object, kVersioned)) { | |
| 39 | 4 | auto* managed = | |
| 40 | 4 | static_cast<DLManagedTensorVersioned*>(PyCapsule_GetPointer(object, kVersioned)); | |
| 41 | 4 | PyCapsule_SetName(object, "used_dltensor_versioned"); | |
| 42 | 4 | return Tensor::FromDlpack(managed); | |
| 43 | } | ||
| 44 | ✗ | if (PyCapsule_IsValid(object, kUnversioned)) { | |
| 45 | ✗ | auto* managed = static_cast<DLManagedTensor*>(PyCapsule_GetPointer(object, kUnversioned)); | |
| 46 | ✗ | PyCapsule_SetName(object, "used_dltensor"); | |
| 47 | ✗ | return Tensor::FromDlpack(managed); | |
| 48 | } | ||
| 49 | ✗ | throw nb::type_error("__dlpack__ returned no DLPack capsule"); | |
| 50 | } | ||
| 51 | |||
| 52 | } // namespace | ||
| 53 | |||
| 54 | 2 | nb::object Capsule(DLManagedTensorVersioned* managed, bool versioned) { | |
| 55 | 2 | PyObject* capsule = nullptr; | |
| 56 | 2 | if (versioned) { | |
| 57 | 2 | capsule = PyCapsule_New(managed, kVersioned, [](PyObject* self) { | |
| 58 | // A consumer renames the capsule when it takes the tensor. | ||
| 59 | 2 | if (PyCapsule_IsValid(self, kVersioned)) { | |
| 60 | ✗ | auto* tensor = | |
| 61 | ✗ | static_cast<DLManagedTensorVersioned*>(PyCapsule_GetPointer(self, kVersioned)); | |
| 62 | ✗ | tensor->deleter(tensor); | |
| 63 | } | ||
| 64 | 2 | }); | |
| 65 | } else { | ||
| 66 | ✗ | capsule = PyCapsule_New(Unversioned(managed), kUnversioned, [](PyObject* self) { | |
| 67 | ✗ | if (PyCapsule_IsValid(self, kUnversioned)) { | |
| 68 | ✗ | auto* tensor = | |
| 69 | ✗ | static_cast<DLManagedTensor*>(PyCapsule_GetPointer(self, kUnversioned)); | |
| 70 | ✗ | tensor->deleter(tensor); | |
| 71 | } | ||
| 72 | ✗ | }); | |
| 73 | } | ||
| 74 | 2 | if (capsule == nullptr) { | |
| 75 | ✗ | throw nb::python_error(); | |
| 76 | } | ||
| 77 | 2 | return nb::steal(capsule); | |
| 78 | } | ||
| 79 | |||
| 80 | 4 | Tensor TakeTensor(const nb::object& source, const nb::object& given) { | |
| 81 | // A stream handle of 0, as PyTorch gives its default stream, is spelled 1 in DLPack. | ||
| 82 | 4 | const nb::object stream = | |
| 83 | 4 | !given.is_none() && nb::cast<std::intptr_t>(given) == 0 ? nb::int_(1) : given; | |
| 84 | 4 | nb::object capsule; | |
| 85 | 4 | try { | |
| 86 | 20 | capsule = source.attr("__dlpack__")(nb::arg("stream") = stream, | |
| 87 | 8 | nb::arg("max_version") = | |
| 88 | 12 | std::make_tuple(DLPACK_MAJOR_VERSION, 0)); | |
| 89 | ✗ | } catch (const nb::python_error& error) { | |
| 90 | ✗ | if (!error.matches(PyExc_TypeError)) { | |
| 91 | ✗ | throw; | |
| 92 | } | ||
| 93 | ✗ | capsule = source.attr("__dlpack__")(nb::arg("stream") = stream); | |
| 94 | ✗ | } | |
| 95 | 4 | return FromCapsule(capsule); | |
| 96 | 4 | } | |
| 97 | |||
| 98 | } // namespace ernie::python | ||
| 99 |