Copilot commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3966858460
##########
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
+ return std::reduce(shape.begin(), shape.end(), int64_t{1},
std::multiplies{});
+ }
+
+ /// Number of bytes needed to store this tensor data.
+ Result<int64_t> ComputeNumBytes() const {
+ ARROW_ASSIGN_OR_RAISE(const auto nelements, ComputeNumElements());
+ int64_t nbytes = 0;
+ if (ARROW_PREDICT_FALSE(internal::MultiplyWithOverflow(
+ nelements, static_cast<int64_t>(byte_width()), &nbytes))) {
+ return Status::Invalid("Overflow computing DLPack tensor size in
bytes.");
+ }
+ return nbytes;
+ }
+
+ 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) {
+ ARROW_ASSIGN_OR_RAISE(const int64_t nbytes, dl.ComputeNumBytes());
+
+ // 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;
Review Comment:
`ImportBuffer()` does pointer arithmetic on `dl.tensor().data` when `nbytes
> 0`, but doesn’t validate that the DLPack producer actually provided a
non-null data pointer. If a producer returns `data == nullptr` for a non-empty
tensor, `dl.data_as<uint8_t>() + byte_offset` is undefined behavior. Add an
explicit check and fail fast with a clear error.
--
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]
