Hello RMM team,

I'm currently using the RMM device memory pool rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> as described in the README. My program operates in with multistream and enabled PTDS.
When profiling, I've noticed that the mutex operations in do_allocate and do_deallocate become a major bottleneck, for example when using 32 streams.

Bottleneck Example

Let me illustrate this with a simple profiling result. I'm using an RMM device buffer rmm::device_buffer as temp buffer for calling cub::reduce_by_key. The execution time of the CUB kernel (~40μs) is only a tiny fraction of the total time. However, both do_allocate and do_deallocate each take over 1ms, due to concurrent threads holding the mutex during allocation or deallocation. This issue prevents the CPU from efficiently launching more kernels, and as a result, the GPU isn't fully utilized.

The first pthread_mutex_lock comes from do_allocate:

The last pthread_mutex_lock comes from do_deallocate:

In RMM: stream_ordered_memory_resource

I've examined the code in RMM and believe that the critical sections in do_allocate and do_deallocate are quite large. Given that my program uses PTDS, I think some parts of these functions could be moved outside the critical section. However, implementing such a fix isn't straightforward, and I haven't found an obvious solution yet :(

My Questions

My main question is how to work around the bottlenecks caused by do_allocate and do_deallocate. Should I consider switching to another memory resource inside RMM? Or is this overhead an inevitable cost of using a pooled memory system?
Ultimately, I aim to fully optimize the memory allocation process from the CPU mutex side without being limited by these bottlenecks.