Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for streams in CuPy allocator #654

Merged
merged 19 commits into from
Mar 4, 2021
Merged

Add support for streams in CuPy allocator #654

merged 19 commits into from
Mar 4, 2021

Conversation

pentschev
Copy link
Member

@pentschev pentschev commented Dec 11, 2020 •
edited by kkraus14
Loading

This PR adds support for CuPy streams in rmm_cupy_allocator. It works by getting CuPy's current stream and passing that to the DeviceBuffer constructor.

There's also a fix for the casting of CuPy/Numba streams to cudaStream_t, it needs to be cast to uintptr_t first, without that the resulting pointer would be wrong and result in a segfault.

Depends on #661

@pentschev pentschev requested a review from a team as a code owner December 11, 2020 00:05
kkraus14
kkraus14 approved these changes Dec 11, 2020
View reviewed changes
Copy link
Contributor

@kkraus14 kkraus14 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor fix, otherwise LGTM

python/rmm/tests/test_rmm.py Outdated Show resolved Hide resolved
@kkraus14 kkraus14 added breaking Breaking change improvement Improvement / enhancement to an existing function Python Related to RMM Python API labels Dec 11, 2020
@kkraus14 kkraus14 added non-breaking Non-breaking change and removed breaking Breaking change labels Dec 11, 2020
shwina
shwina reviewed Dec 11, 2020
View reviewed changes
python/rmm/_cuda/stream.pyx Outdated Show resolved Hide resolved
kkraus14
kkraus14 previously requested changes Dec 11, 2020
View reviewed changes
Copy link
Contributor

@kkraus14 kkraus14 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requesting changes to block merging until the stream reference is added to maintain lifetime

@pentschev
Copy link
Member Author

While looking for stream lifetime management in Numba I stumbled upon https://github.com/numba/numba/blob/0bac18af44d08e913cd512babb9f9b7f6386d30a/numba/cuda/cudadrv/driver.py#L1834-L1836 which allowed me to simplify _init_from_numba_stream considerably in 079d24d.

@jakirkham
Copy link
Member

cc @gmarkall

@rapidsai rapidsai deleted a comment from harrism Dec 14, 2020
pentschev and others added 2 commits December 16, 2020 21:59
@pentschev @harrism
Improve docs for DeviceBuffer's stream kwarg
fd7e004 
Co-authored-by: Mark Harris <mharris@nvidia.com>
@pentschev
Fix flake8 errors
f83e9cf
@pentschev
Copy link
Member Author

Sorry everyone for the delay here, I was moving over the past several days and had very limited access to PC. I now updated the PR with the suggested changes, please take some time to review when possible.

@kkraus14
Copy link
Contributor

kkraus14 commented Jan 7, 2021

Generally things look good, but would suggest we merge #661 first which handles the Stream lifetime management, and then merge this in on top of it.

@kkraus14 kkraus14 added the 5 - Merge After Dependencies Depends on another PR: do not merge out of order label Jan 7, 2021
@harrism
Copy link
Member

harrism commented Feb 3, 2021

#661 hasn't progressed. So I guess we need to move this to 0.19.

@harrism harrism changed the base branch from branch-0.18 to branch-0.19 February 3, 2021 02:57
@pentschev
Copy link
Member Author

Agreed, thanks @harrism for keeping track.

@kkraus14 kkraus14 removed the 5 - Merge After Dependencies Depends on another PR: do not merge out of order label Mar 1, 2021
kkraus14
kkraus14 reviewed Mar 1, 2021
View reviewed changes
python/rmm/_lib/device_buffer.pxd Outdated Show resolved Hide resolved
python/rmm/_lib/device_buffer.pyx Outdated Show resolved Hide resolved
jakirkham
jakirkham reviewed Mar 1, 2021
View reviewed changes
Copy link
Member

@jakirkham jakirkham left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small suggestion below to simplify some of the casting logic here

python/rmm/_cuda/stream.pyx Show resolved Hide resolved
python/rmm/_cuda/stream.pyx Show resolved Hide resolved
@kkraus14 kkraus14 dismissed their stale review March 2, 2021 01:34

Took over latest changes

@kkraus14
Copy link
Contributor

kkraus14 commented Mar 2, 2021

rerun tests

@kkraus14
Copy link
Contributor

kkraus14 commented Mar 2, 2021

This should be ready for review

@kkraus14 kkraus14 added 3 - Ready for review Ready for review by team 4 - Needs Reviewer Waiting for reviewer to review or respond labels Mar 2, 2021
@kkraus14
Copy link
Contributor

kkraus14 commented Mar 4, 2021

@gpucibot merge

@rapids-bot rapids-bot bot merged commit 0d6d2ad into rapidsai:branch-0.19 Mar 4, 2021
@jakirkham
Copy link
Member

Thanks Peter for the PR and everyone for the reviews! 😄

@pentschev
Copy link
Member Author

Thanks everyone for reviews!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@leofang leofang leofang left review comments

@kkraus14 kkraus14 kkraus14 left review comments

@shwina shwina shwina left review comments

@jakirkham jakirkham jakirkham approved these changes

@harrism harrism harrism approved these changes

Assignees
No one assigned
Labels
3 - Ready for review Ready for review by team 4 - Needs Reviewer Waiting for reviewer to review or respond improvement Improvement / enhancement to an existing function non-breaking Non-breaking change Python Related to RMM Python API
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@pentschev @jakirkham @kkraus14 @harrism @shwina @leofang