support CUDA async memory resource in JNI

[rongou]

CUDA 11.2 introduced stream ordered memory allocator that can potentially resolve memory fragmentation issues. See https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-1/

[jlowe]

It would be nice to have at least a smoke test of the new allocator type in RmmTest that sets up the allocator, allocates and frees memory to exercise it. Bonus points if it also sets up the allocator with a small limit and verifies it gets an OOM if it tries to allocate just beyond that size.

[rongou]

Added smoke test, which will be skipped if cuda < 11.2.

[abellina]

I just filed another pretty much exact PR to this one: #9208. The main difference is that I am wrapping the async allocator with limiting_resource_adaptor.

I did not know @rongou's PR was up, for some reason. I closed mine in favor of @rongou's PR, since folks already spent time reviewing his.

[jlowe]

Note that we now support Java in the CI, so Java PRs should not skip ci.

[jlowe]

rerun tests

[rongou]

rerun tests

[rongou]

rerun tests

[jlowe]

rerun tests

[jrhemstad]

fyi, you could also experiment with using cuda_async_resource as the upstream for arena. That might give you best of both worlds of what you're looking for.

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (branch-21.10@4349232). Click here to learn what that means.
The diff coverage is n/a.

❗ Current head 6e8cf91 differs from pull request most recent head 5d91944. Consider uploading reports for the commit 5d91944 to get more accurate results

@@               Coverage Diff               @@
##             branch-21.10    #9201   +/-   ##
===============================================
  Coverage                ?   10.82%           
===============================================
  Files                   ?      115           
  Lines                   ?    19166           
  Branches                ?        0           
===============================================
  Hits                    ?     2074           
  Misses                  ?    17092           
  Partials                ?        0           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 4349232...5d91944. Read the comment docs.

[rongou]

@gpucibot merge

[rongou]

fyi, you could also experiment with using cuda_async_resource as the upstream for arena. That might give you best of both worlds of what you're looking for.

@jrhemstad Yeah that's something we can try if it turns out small allocations are too expensive with async.

[abellina]

@jrhemstad filed this: rapidsai/rmm#868, we need to fix this before we start using the async allocator. He thought it was a quick fix, and that it could be included in 21.10. FYI @sameerz

[jlowe]

you could also experiment with using cuda_async_resource as the upstream for arena

It seems that circumvents the fragmentation-solving feature we want from the async allocator. If arena only allocates large chunks from the async allocator, won't we still have fragmentation within the arena blocks that the async allocator cannot solve since the async allocator will be unaware of the sub-utilization of the allocations it sees?

[rongou]

The per-thread arenas are just caches for small allocations. If cuda async proves to be slow for small allocations, we can use the arena allocator to speed up these, as in a typical job there are tons of small allocations. The number of free blocks are now capped in each per-thread arena, so in theory it shouldn't cause too much additional fragmentation. If/when we decide to try this, we can probably further tweak the algorithm.