kkraus14 commented on code in PR #34972:
URL: https://github.com/apache/arrow/pull/34972#discussion_r1163346998
##########
cpp/src/arrow/c/abi.h:
##########
@@ -106,6 +212,98 @@ struct ArrowArrayStream {
#endif // ARROW_C_STREAM_INTERFACE
+#ifndef ARROW_C_DEVICE_STREAM_INTERFACE
+#define ARROW_C_DEVICE_STREAM_INTERFACE
+
+/// \brief Equivalent to ArrowArrayStream, but for ArrowDeviceArrays.
+///
+/// This stream is intended to provide a stream of data on a single
+/// device, if a producer wants data to be produced on multiple devices
+/// then multiple streams should be provided. One per device.
+struct ArrowDeviceArrayStream {
+ /// \brief The device that this stream produces data on.
+ ///
+ /// All ArrowDeviceArrays that are produced by this
+ /// stream should have the same device_type as set
+ /// here. The device_type needs to be provided here
+ /// so that consumers can provide the correct type
+ /// of queue_ptr when calling get_next.
+ ArrowDeviceType device_type;
+
+ /// \brief Callback to get the stream schema
+ /// (will be the same for all arrays in the stream).
+ ///
+ /// If successful, the ArrowSchema must be released independantly from the
stream.
+ /// The schema should be accessible via CPU memory.
+ ///
+ /// \param[in] self The ArrowDeviceArrayStream object itself
+ /// \param[out] out C struct to export the schema to
+ /// \return 0 if successful, an `errno`-compatible error code otherwise.
+ int (*get_schema)(struct ArrowDeviceArrayStream* self, struct ArrowSchema*
out);
+
+ /// \brief Callback to get the device id for the next array.
+ ///
+ /// This is necessary so that the proper/correct stream pointer can be
provided
+ /// to get_next.
+ ///
+ /// The next call to `get_next` should provide an ArrowDeviceArray whose
+ /// device_id matches what is provided here, and whose device_type is the
+ /// same as the device_type member of this stream.
Review Comment:
There's a lot of discussion in that issue related to internal stream
handling under contexts and schedulers and similar terms. It all boils down to
the same discussion of many producers being unable to release or share
ownership of their streams.
I think this comment does a good job of summarizing the options that were
considered: https://github.com/dmlc/dlpack/issues/57#issuecomment-753220425
And then this comment summarizes discussion of those options:
https://github.com/dmlc/dlpack/issues/57#issuecomment-774482973
The lifetime management of streams as defined by the Numba documentation for
`__cuda_array_interface__`
(https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html#streams)
requires that keeping the object (typically array) that produces the
`__cuda_array_interface__` alive also keeps the stream alive. In most cases
libraries don't associate a stream with the object since it's valid to use
multiple streams with a single object.
Here's the current state of things across a handful of projects:
- Numba
- Handles stream lifetime properly for `__cuda_array_interface__` since
they store a stream object as part of their array:
https://github.com/numba/numba/blob/008077553b558bd183668ecd581d4d0bc54bd32c/numba/cuda/cudadrv/devicearray.py#L119-L139
- Doesn't support `__dlpack__`
- CuPy
- Doesn't implement stream lifetime properly for
`__cuda_array_interface__` as far as I can tell, where they just get a stream
ptr as an integer and if it's not the default stream someone could change the
current stream ptr and end up having it be destroyed out from underneath it:
https://github.com/cupy/cupy/blob/c92d5bc16293300297b843b4ebb364125697c131/cupy/_core/core.pyx#L258-L262
(cc @leofang)
- Handles `__dlpack__` properly:
https://github.com/cupy/cupy/blob/c92d5bc16293300297b843b4ebb364125697c131/cupy/_core/core.pyx#L285-L327
- PyTorch
- Doesn't support passing the stream in `__cuda_array_interface__`
currently:
https://github.com/pytorch/pytorch/blob/def50d253401540cfdc6c0fffa444d0ee643cc11/torch/_tensor.py#L1005-L1066
- Handles `__dlpack__` properly:
https://github.com/pytorch/pytorch/blob/def50d253401540cfdc6c0fffa444d0ee643cc11/torch/_tensor.py#L1345
- Tensorflow and JAX
- Doesn't support passing the stream in `__cuda_array_interface__`
currently:
https://github.com/tensorflow/tensorflow/blob/ea6a0f282d2b7ce20891dfc24ec8fe107eeaf22d/tensorflow/compiler/xla/python/py_buffer.cc#L254-L300
- Doesn't support `__dlpack__`, but has explicit `to` and `from` functions
for using dlpack pycapsules but does not handle streams:
https://github.com/tensorflow/tensorflow/blob/6a050b6c15ed2a545693bc171f5e95dacbe05839/tensorflow/compiler/xla/python/dlpack.cc#L285-L363
- cuDF
- Doesn't support passing the stream in `__cuda_array_interface__`
currently:
https://github.com/rapidsai/cudf/blob/50718e673ff53b18706cf66c6e02cda8e30681fe/python/cudf/cudf/core/column/numerical.py#L169-L191
- Doesn't support `__dlpack__` but has explicit `to` and `from` functions
for using dlpack pycapsules but does not handle streams:
https://github.com/rapidsai/cudf/blob/50718e673ff53b18706cf66c6e02cda8e30681fe/cpp/src/interop/dlpack.cpp#L218-L294
--
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]
