[JNI] rmm based pinned pool #15219
[JNI] rmm based pinned pool #15219
Conversation
Signed-off-by: Alessandro Bellina <abellina@nvidia.com>
abellina
added
feature request
|
/ok to test |
|
One thing that could improve the tests is to actually write something to the buffers that are returned. |
|
/ok to test |
@jbrennan333 I added a test in my latest commit. Thanks! |
| long available = PinnedMemoryPool.getTotalPoolSizeBytes(); | ||
| if (available < size) { |
There was a problem hiding this comment.
(nit) Is this size check really needed at all? Doesn't hurt anything though, so just a nit.
There was a problem hiding this comment.
yeah it's a little weird, but we can keep calling this function and it will shutdown/recreate it as needed, so I don't see a huge need to change it right now, plus it's pre-existing code that continues to work.
|
/ok to test |
|
/ok to test |
|
Unrelated python failure:
|
|
/merge |
Part of #14782.
This PR removes our old implementation of the java based pinned memory pool and replaces it with a jni layer on top of
rmm::pool_memory_resource<rmm::pinned_host_memory_resource>This PR does NOT set the default cuIO pinned host resource. That is happening after this PR goes in #15079. We'll need a follow on PR to change
PinnedMemoryPool.initializemethod to add an argument to set the cuIO pinned host resource.I have run with this and version of it that are shared with cuIO and I can't find regressions in NDS at SF3K.
Note that we don't align anymore on our side. RMM is doing the same alignment we were doing before, using
std::max_align_t.Note also that the rmm pool doesn't have a quick way to find out what the current size is. So we had some tests that were asserting for this, and I have removed the asserts. If we would like to get that back I am happy to work with RMM to figure out how to do that.