AntoinePrv commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3950793347
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -248,4 +261,238 @@ 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 != VERSION.major)) {
+ return Status::Invalid("Unsupported DLPack major version ",
out.ptr_->version.major,
+ ", expected ", VERSION.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");
+ }
+ if (ARROW_PREDICT_FALSE(out.tensor().ndim != 0 && out.tensor().strides ==
nullptr)) {
+ return Status::Invalid(
+ "Invalid DLPack tensor: strides must be non-null when ndim != 0");
+ }
+ 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())};
+ }
+
+ std::span<const int64_t> strides() const {
+ return {tensor().strides, static_cast<std::size_t>(ndim())};
+ }
+
+ 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; }
+
+ private:
+ struct Deleter {
+ void operator()(pointer_type ptr) {
+ // Null is valid in DLPack spec
+ if (auto del = ptr->deleter) {
+ del(ptr);
+ }
+ }
+ };
+
+ /// Make a safe wrapper that will delete the resource in case of exception.
+ std::unique_ptr<value_type, Deleter> ptr_;
+
+ explicit CppDLTensor(pointer_type ptr) : ptr_(ptr) {}
+};
+
+Result<std::shared_ptr<FixedWidthType>> DataTypeFromDLPack(DLDataType dtype) {
+ if (dtype.lanes != 1) {
+ return Status::TypeError("Only type with one lane are supported.");
+ }
+
+ auto constexpr as_fw = [](auto dt) {
+ return std::static_pointer_cast<FixedWidthType>(std::move(dt));
+ };
+
+ switch (dtype.code) {
+ case kDLInt: {
+ switch (dtype.bits) {
+ case 8:
+ return as_fw(int8());
+ case 16:
+ return as_fw(int16());
+ case 32:
+ return as_fw(int32());
+ case 64:
+ return as_fw(int64());
+ default:
+ return Status::Invalid("unsupported integer bit width ",
+ static_cast<int>(dtype.bits));
+ }
+ }
+ case kDLUInt: {
+ switch (dtype.bits) {
+ case 8:
+ return as_fw(uint8());
+ case 16:
+ return as_fw(uint16());
+ case 32:
+ return as_fw(uint32());
+ case 64:
+ return as_fw(uint64());
+ default:
+ return Status::Invalid("unsupported unsigned integer bit width ",
+ static_cast<int>(dtype.bits));
+ }
+ }
+ case kDLFloat: {
+ switch (dtype.bits) {
+ case 16:
+ return as_fw(float16());
+ case 32:
+ return as_fw(float32());
+ case 64:
+ return as_fw(float64());
+ default:
+ return Status::Invalid("unsupported float bit width ",
+ static_cast<int>(dtype.bits));
+ }
+ }
+ default: {
+ return Status::Invalid("unsupported DLPack type ",
static_cast<int>(dtype.code));
+ }
+ }
+}
+
+Result<std::vector<int64_t>> StridesInBytes(std::span<const int64_t> strides,
+ int64_t byte_width) {
+ std::vector<int64_t> out{};
+ out.reserve(strides.size());
+ for (const auto& s : strides) {
+ int64_t stride_bytes = 0;
+ if (ARROW_PREDICT_FALSE(
+ internal::MultiplyWithOverflow(s, byte_width, &stride_bytes))) {
+ return Status::Invalid("Overflow computing DLPack tensor stride in
bytes.");
+ }
+ out.push_back(stride_bytes);
+ }
+ return out;
+}
+
+Result<std::shared_ptr<Buffer>> ImportBuffer(CppDLTensor&& dl, bool copy) {
+ // DLPack strides are in number of elements, so is the size we compute from
them.
+ ARROW_ASSIGN_OR_RAISE(const auto nelements,
+ internal::ComputeTensorSize(dl.shape(), dl.strides(),
1));
+ int64_t nbytes = 0;
+ if (ARROW_PREDICT_FALSE(internal::MultiplyWithOverflow(
+ nelements, static_cast<int64_t>(dl.byte_width()), &nbytes))) {
+ return Status::Invalid("Overflow computing DLPack tensor size in bytes.");
+ }
+ // DLPack mandates a null data pointer when the tensor holds no element, so
there is
+ // neither anything to share nor to copy.
+ uint8_t* data =
+ (nbytes == 0) ? nullptr : dl.data_as<uint8_t>() +
dl.tensor().byte_offset;
+
+ std::shared_ptr<Buffer> buffer = nullptr;
+ if (nbytes == 0) {
+ // DLPack data pointer may be null on empty tensors
+ buffer = std::make_shared<Buffer>(data, nbytes);
+ } else if (copy) {
+ ARROW_ASSIGN_OR_RAISE(buffer, MutableBuffer::CopyNonOwned(
+ {data, nbytes},
default_cpu_memory_manager()));
+ } else {
Review Comment:
@pitrou I've been thinking with this and:
- First we do not need the copy parameter in `ImportArray` because we
already can pass `copy` to the producer (which is what we do in Python).
- Then without `copy` I believe we have no use for
`DLPACK_FLAG_BITMASK_IS_COPIED`.
What do you think, should we keep our `copy` parameter (currently unused) or
always rely on the producer to copy?
--
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]
