Add async view memory resource bindings to Python. #1864
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
github-actions
bot
added
Python
|
/ok to test |
|
Currently tests are failing locally and I am struggling to identify why. It seems like something is breaking when passing the pool from Python to C++? |
|
/ok to test |
bdice
added
feature request
|
This PR is out of date with the commits I've pushed. I will close and reopen to retrigger it. |
| @@ -72,7 +72,8 @@ class cuda_async_view_memory_resource final : public device_memory_resource { | |||
| */ | |||
| [[nodiscard]] cudaMemPool_t pool_handle() const noexcept { return cuda_pool_handle_; } | |||
|
|
|||
| cuda_async_view_memory_resource() = default; | |||
| cuda_async_view_memory_resource() = default; | |||
| ~cuda_async_view_memory_resource() = default; | |||
There was a problem hiding this comment.
Fixed a missing rule-of-6 destructor declaration here, which raised a warning in my IDE.
|
|
||
| Parameters | ||
| ---------- | ||
| valid_pool_handle : cudaMemPool_t |
There was a problem hiding this comment.
I don't like the name valid_pool_handle but I copied it from C++. Should we consider renaming it?
There was a problem hiding this comment.
Sure, though please mention somewhere in this docstring that the user is responsible for keeping the mempool alive
There was a problem hiding this comment.
Done in b7e899d. The new sentence says,
The memory pool passed in must not be destroyed during the lifetime of this memory resource.
| raw_pool_handle = int(valid_pool_handle) | ||
| c_valid_pool_handle = <cyruntime.cudaMemPool_t><uintptr_t>raw_pool_handle | ||
| else: | ||
| raw_pool_handle = int(runtime.cudaMemPool_t(valid_pool_handle)) |
There was a problem hiding this comment.
Is there any typing we can apply to the constructor argument that would prevent someone passing something bad in here (or does this constructor do that for hs?)
There was a problem hiding this comment.
I narrowed the constructor a bit in b7e899d. I am now forcing users to pass a cuda.bindings.runtime.cudaMemPool_t or cuda.bindings.driver.CUmemoryPool instance. The constructor no longer accepts raw ints as pointers to the corresponding types. I think that is much safer.
| props.allocType = runtime.cudaMemAllocationType.cudaMemAllocationTypePinned | ||
| props.location.id = rmm._cuda.gpu.getDevice() | ||
| props.location.type = runtime.cudaMemLocationType.cudaMemLocationTypeDevice | ||
| err, pool = runtime.cudaMemPoolCreate(props) |
There was a problem hiding this comment.
This test should also destroy the pool
There was a problem hiding this comment.
Good idea! Tested in b7e899d. It also makes sure that the memory resource raises a MemoryError when asked to allocate after the pool is destroyed.
| def test_cuda_async_view_memory_resource_default_pool(dtype, nelem, alloc): | ||
| # Get the default memory pool handle | ||
| current_device = rmm._cuda.gpu.getDevice() | ||
| err, pool = runtime.cudaDeviceGetDefaultMemPool(current_device) |
There was a problem hiding this comment.
Does this pool need to be released? I don't know how the default mempool behaves
There was a problem hiding this comment.
I don't think so. It cannot be destroyed, per the docs. Attempting to destroy it returns an error code. 😄
Note:
A device's default memory pool cannot be destroyed.
|
/merge |
Description
Closes #1611.
Checklist