zeroshade commented on code in PR #42118:
URL: https://github.com/apache/arrow/pull/42118#discussion_r2047495330
##########
cpp/src/arrow/c/dlpack.cc:
##########
@@ -130,4 +131,71 @@ Result<DLDevice> ExportDevice(const
std::shared_ptr<Array>& arr) {
}
}
+struct TensorManagerCtx {
+ std::shared_ptr<Tensor> t;
+ std::vector<int64_t> strides;
+ std::vector<int64_t> shape;
+ DLManagedTensor tensor;
+};
+
+Result<DLManagedTensor*> ExportTensor(const std::shared_ptr<Tensor>& t) {
+ // Define the DLDataType struct
+ const DataType& type = *t->type();
+ ARROW_ASSIGN_OR_RAISE(DLDataType dlpack_type, GetDLDataType(type));
+
+ // Define DLDevice struct
+ ARROW_ASSIGN_OR_RAISE(DLDevice device, ExportDevice(t))
+
+ // Create TensorManagerCtx that will serve as the owner of the
DLManagedTensor
+ std::unique_ptr<TensorManagerCtx> ctx(new TensorManagerCtx);
+
+ // Define the data pointer to the DLTensor
+ // If tensor is of length 0, data pointer should be NULL
+ if (t->size() == 0) {
+ ctx->tensor.dl_tensor.data = NULL;
+ } else {
+ ctx->tensor.dl_tensor.data = t->raw_mutable_data();
+ }
+
+ ctx->tensor.dl_tensor.device = device;
+ ctx->tensor.dl_tensor.ndim = t->ndim();
+ ctx->tensor.dl_tensor.dtype = dlpack_type;
+ ctx->tensor.dl_tensor.byte_offset = 0;
+
+ std::vector<int64_t>& shape_arr = ctx->shape;
+ shape_arr.reserve(t->ndim());
+ for (auto i : t->shape()) {
+ shape_arr.emplace_back(i);
+ }
+ ctx->tensor.dl_tensor.shape = shape_arr.data();
+
+ std::vector<int64_t>& strides_arr = ctx->strides;
+ strides_arr.reserve(t->ndim());
+ auto byte_width = t->type()->byte_width();
+ for (auto i : t->strides()) {
+ strides_arr.emplace_back(i / byte_width);
+ }
+ ctx->tensor.dl_tensor.strides = strides_arr.data();
+
+ ctx->t = std::move(t);
+ ctx->tensor.manager_ctx = ctx.get();
+ ctx->tensor.deleter = [](struct DLManagedTensor* self) {
+ delete reinterpret_cast<TensorManagerCtx*>(self->manager_ctx);
+ };
+ return &ctx.release()->tensor;
+}
+
+Result<DLDevice> ExportDevice(const std::shared_ptr<Tensor>& t) {
+ // Define DLDevice struct
+ DLDevice device;
+ if (t->data()->device_type() == DeviceAllocationType::kCPU) {
+ device.device_id = 0;
+ device.device_type = DLDeviceType::kDLCPU;
+ return device;
+ } else {
+ return Status::NotImplemented(
+ "DLPack support is implemented only for buffers on CPU device.");
Review Comment:
any reason we aren't implementing the non-cpu cases here since we support it
in the generic gpu package etc.?
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
