Set memory pool through RMM

[viclafargue]

This PR aims at updating the device_resources_snmg utility to make sure that the method in charge of setting up the RMM pools on the GPUs effectively work in multi-GPU context. This method used to set a pool workspace on the device resource to no effect. We instead now set the per device memory resource of each GPU to a pool through RMM.

Benchmark of search throughput for 100 rows queries on 2 GPUs CAGRA indexes built with NN Descent :

[achirkin]

Hi @viclafargue , could you please update the description of the PR to let the reviewers know what is the problem it addresses?
I suspect, this is about perf bottlenecks related to locking the cuda context in the multi-gpu setup. If so, please attach the benchmarks. I'm also concerned a little about the lifetime of the objects allocated with the per-device pools. Could there be a situation when an object outlives its memory pool and segfaults the program?

[viclafargue]

I'm also concerned a little about the lifetime of the objects allocated with the per-device pools. Could there be a situation when an object outlives its memory pool and segfaults the program?

Yes, that is a real issue. Once a user sets a memory pool (e.g. handle.set_memory_pool(80)). Anything allocated afterward should be released before the handle is released. That is why a management from the RAFT side would be preferable.

[jinsolp]

Thanks Victor!

[bdice]

Could there be a situation when an object outlives its memory pool and segfaults the program?

Yes, this is a tricky problem with RMM today. I am working to improve this for 26.02 or 26.04.

This class of problems will be eliminated by RMM migrating to the CCCL 3.2 memory resource design. There is a new form of memory resource ownership in any_resource, which is maybe-shared. Stateful resources will have shared (refcounted) ownership but stateless resources will not, to avoid overhead.

This issue comment (rapidsai/rmm#2011 (comment)) tracks adoption of the new memory resources, which is my primary active project right now. I am hopeful that 26.02 will entirely resolve this class of issues, though deprecation cycles may force it out to 26.04.

[viclafargue]

@achirkin @jinsolp Do you think we can merge the PR as-is for now and come back with a follow-up once the new RMM feature is available or would you like to wait for it?

[jinsolp]

I'm fine either way!