rmm
rmm copied to clipboard
Accept stream argument in DeviceMemoryResource allocate/deallocate
Description
In #980, the DeviceMemoryResource class in Python gained allocation and deallocation routines. This was to facility writing Python allocate/deallocate callbacks for the CallbackMemoryResource.
These routines should, to match the C++ API, accept a stream parameter such that one can use them for stream-ordered allocation. Although we recommend that users allocate on the Python side using the DeviceBuffer interface, exposing these routines implicitly makes them public.
To fix this, add an optional stream argument defaulting to the default stream.
- Closes #1493
Checklist
- [x] I am familiar with the Contributing Guidelines.
- [x] New or existing tests cover these changes.
- [x] The documentation is up to date with these changes.
This is a breaking change since any user of the CallbackMemoryResource
must adapt their callback signatures. That said, I can't find any...
/merge