pitrou commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3979362920
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +262,265 @@ Result<DLDevice> ExportDevice(const
std::shared_ptr<Tensor>& t) {
return ExportDeviceImpl(t);
}
+/***************
+ * Consumers *
+ ***************/
+
+namespace {
+
+class CppDLTensor {
+ public:
+ using value_type = DLManagedTensorVersioned;
+ using pointer_type = value_type*;
+
+ static Result<CppDLTensor> TakeOwnership(pointer_type ptr) {
+ if (ARROW_PREDICT_FALSE(ptr == nullptr)) {
+ return Status::Invalid("Received null pointer.");
+ }
+ // Create the wrapper before checking the version as the spec mandates
that the
+ // deleter MUST be called on version major mismatch.
+ auto out = CppDLTensor(ptr);
+ if (ARROW_PREDICT_FALSE(out.ptr_->version.major != kVersion.major)) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", kVersion.major);
+ }
+ if (ARROW_PREDICT_FALSE(out.tensor().ndim < 0)) {
+ return Status::Invalid("Invalid DLPack tensor: ndim must be >= 0");
+ }
+ if (ARROW_PREDICT_FALSE(out.tensor().ndim != 0 && out.tensor().shape ==
nullptr)) {
+ return Status::Invalid(
+ "Invalid DLPack tensor: shape must be non-null when ndim != 0");
+ }
+ // Null strides are handled as row major
+ return out;
+ }
+
+ const DLTensor& tensor() const { return ptr_->dl_tensor; }
+
+ int64_t ndim() const {
+ DCHECK_GE(tensor().ndim, 0);
+ return tensor().ndim;
+ }
+
+ template <typename T>
+ T* data_as() {
+ return static_cast<T*>(tensor().data);
+ }
+
+ std::span<const int64_t> shape() const {
+ return {tensor().shape, static_cast<std::size_t>(ndim())};
+ }
+
+ /// Strides or empty span for old DLPack row-major convention.
+ std::span<const int64_t> strides() const {
+ if (auto strides = tensor().strides; strides != nullptr) {
+ return {strides, static_cast<std::size_t>(ndim())};
+ }
+ return {};
+ }
+
+ bool flag_is_set(uint8_t bits) const { return (ptr_->flags & bits) == bits; }
+
+ bool is_readonly() const { return
flag_is_set(DLPACK_FLAG_BITMASK_READ_ONLY); }
+
+ int32_t byte_width() const { return tensor().dtype.bits / 8; }
+
+ /// Number of element in this tensor's buffer.
+ ///
+ /// Possibly more elements than represented in the tensor for non-contiguous
tensors.
+ ///
+ /// A zero dimensional tensor is a scalar, it holds a single element.
+ Result<int64_t> ComputeNumElements() const {
+ const auto strides = this->strides();
+ const auto shape = this->shape();
+ if (strides.size() > 0) {
+ // DLPack strides are in number of elements, so is the size we compute
from them.
+ return internal::ComputeTensorSize(shape, strides, 1);
+ }
+ // DLPack <1.3 my set strides == nullptr for row major
Review Comment:
```suggestion
// DLPack <1.3 may set strides == nullptr for row major
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
