Use cudaStream_t instead of cuda_stream_view in pylibcudf Cython#22368
Use cudaStream_t instead of cuda_stream_view in pylibcudf Cython#22368vyasr wants to merge 18 commits intorapidsai:mainfrom
Conversation
vyasr
added
improvement
github-actions
Bot
added
Python
vyasr
force-pushed
the
feat/cy_raw_cudastream
branch
from
14dc2de to
4de0682
Compare
github-actions
Bot
added
libcudf
vyasr
force-pushed
the
feat/cy_raw_cudastream
branch
from
4de0682 to
508f091
Compare
There was a problem hiding this comment.
Just commenting here as a signpost, this file has actual C++ in it. Looks fine to me.
| if isinstance(stream, Stream): | ||
| return <Stream>stream | ||
| return Stream(stream) # Handles __cuda_stream__ protocol |
There was a problem hiding this comment.
If there's something here that you think should be upstreamed (ha) to RMM, I'm open to it.
There was a problem hiding this comment.
Nah nothing worth upstreaming. We could simplify this code further by removing the second if and just delegating to the the Stream constructor, at which point we'd have the same semantics at the small extra cost of creating a new Stream object pointing to the same underlying cudaStream_t and owner. I don't think it really matters either way though.
… inline C++ wrapper
vyasr
force-pushed
the
feat/cy_raw_cudastream
branch
from
51bc52f to
db5faa4
Compare
…nd stream test marker - Fix self.stream -> self._stream in ChunkedParquetReader.read_chunk() - Mark test_join_streams_type_error with uses_custom_stream to skip in stream-testing mode (uses protocol object with default stream pointer)
wence-
left a comment
wence-
left a comment
There was a problem hiding this comment.
Question: I note all the Stream stream parameters in the cpdef functions change to object stream. Is that actually a type change? If yes, should we update the typestubs appropriately?
| #include <cudf/detail/utilities/stream_pool.hpp> | ||
| #include <cudf/utilities/span.hpp> | ||
| #include <rmm/cuda_stream_view.hpp> | ||
| #include <vector> | ||
|
|
||
| void join_streams_wrapper( | ||
| cudf::host_span<cudaStream_t const> streams, | ||
| cudaStream_t stream | ||
| ) { | ||
| std::vector<rmm::cuda_stream_view> stream_views(streams.begin(), streams.end()); | ||
| cudf::detail::join_streams(stream_views, stream); | ||
| } |
There was a problem hiding this comment.
| #include <cudf/detail/utilities/stream_pool.hpp> | |
| #include <cudf/utilities/span.hpp> | |
| #include <rmm/cuda_stream_view.hpp> | |
| #include <vector> | |
| void join_streams_wrapper( | |
| cudf::host_span<cudaStream_t const> streams, | |
| cudaStream_t stream | |
| ) { | |
| std::vector<rmm::cuda_stream_view> stream_views(streams.begin(), streams.end()); | |
| cudf::detail::join_streams(stream_views, stream); | |
| } | |
| #include <cudf/detail/utilities/stream_pool.hpp> | |
| #include <cudf/utilities/span.hpp> | |
| #include <rmm/cuda_stream_view.hpp> | |
| #include <vector> | |
| namespace { | |
| void join_streams_wrapper( | |
| cudf::host_span<cudaStream_t const> streams, | |
| cudaStream_t stream | |
| ) { | |
| std::vector<rmm::cuda_stream_view> stream_views(streams.begin(), streams.end()); | |
| cudf::detail::join_streams(stream_views, stream); | |
| } | |
| } |
I think,
Introduce HasCudaStream Protocol and CudaStreamLike type alias in utils.pyi to reflect that stream parameters now accept any object implementing the __cuda_stream__ protocol, not just rmm Stream objects. Update all .pyi files to use CudaStreamLike | None instead of Stream | None for stream parameters.
Good call. Yes it is a meaningful change. I've added a new protocol for the type. |
The pylibcudf join_streams stub now expects list[CudaStreamLike] due to the __cuda_stream__ protocol support. Update join_cuda_streams and get_joined_cuda_stream parameter types from Sequence[Stream] to Sequence[CudaStreamLike] to satisfy mypy's list invariance check.
4 participants
Description
Contributes to rapidsai/rmm#2359
Checklist