[ENH] Consider switching to cuda::managed_memory_pool / cuda::pinned_memory_pool as the default managed/pinned resources

[achirkin]

Raft supports pinned_mdarray and managed_mdarray along with associated container policies since 24.02
These are hard-coded to use cudaMallocHost cudaMallocManaged under the hood and have host-synchronous semantics.

There are a few places in raft and cuvs where we explicitly use these mdarrays (raft::make_pinned_xxx / raft::make_managed_xxx) for host-device communication.

PR #2968 introduces first-class raft resource objects allowing the user to explicitly set/replace the memory resources associated with the raft::resources handle. These are used under the hood in the mdarray constructor helpers. Hence they give the user explicit control over the pinned/managed memory allocation. By default, we use cuda::mr::legacy_pinned_memory_resource and cuda::mr::legacy_managed_memory_resource, which preserve the original sync behavior.

Now the question: shall we switch the default behavior to use cuda::pinned_memory_pool and cuda::managed_memory_pool when available? They are directly pluggable in place of the cuda::mr::legacy_xxx_memory_resource, but should yield better performance and stream-ordered semantics.

[achirkin]

For context, with PR #2968, the user can write something like this:

  raft::resources res;
  raft::resource::set_pinned_memory_resource(res, cuda::pinned_memory_pool{});
  raft::resource::set_managed_memory_resource(res, cuda::managed_memory_pool{});
  
  // This buffer uses `cuda::pinned_memory_pool` under the hood
  auto buf = raft::make_pinned_vector<float, uint32_t>(res, 1024);

So, at the very least, we should try this in cuvs benchmarks when setting up the algorithms.

One problem I see is that the new cuda pools are stream-ordered while our interface is synchronous (and the cuvs internal code using these mdarrays assumes the synchronous behavior). Does that mean, by calling allocate_sync/deallocate_sync, the new pools would do some default stream/device sync under the hood? Do we get any perf improvement by using these pools via the sync api then?

[bdice]

One problem I see is that the new cuda pools are stream-ordered while our interface is synchronous (and the cuvs internal code using these mdarrays assumes the synchronous behavior). Does that mean, by calling allocate_sync/deallocate_sync, the new pools would do some default stream/device sync under the hood?

Yes, the [de]allocate_sync functions synchronize with a stream that is internal to CCCL:
https://github.com/NVIDIA/cccl/blob/d4e7fd71c404f0d4d31a319ed51943794f252c7f/libcudacxx/include/cuda/__memory_pool/memory_pool_base.h#L424-L425

Do we get any perf improvement by using these pools via the sync api then?

Yes, you should still use these. The benefits aren't strictly tied to stream-ordering, they're tied to lower-level parts of CUDA memory management. However, it would be a really good idea to make the interfaces in RAFT/cuVS stream-ordered wherever possible.

[achirkin]

Thanks for the insights!
I'm trying to think through how (and whether) do we change raft::make_managed_mdarray helper keep its behavior with the new pool resource. The raft::managed_container_policy passes the resource to the container under the hood, but can me adjusted to also pass a stream.

Yes, the [de]allocate_sync functions synchronize with a stream that is internal to CCCL

So as I read this, using allocate_sync inside raft::host_container is strictly better than passing any of raft::resources internal or user streams to allocate and then synchronizing on it - to avoid waiting on any other submitted work.

However, what if we could omit the stream sync and and just call allocate on raft::resources's active stream (same behavior as raft::make_device_mdarray)? The docs clearly state we cannot access the memory before waiting on the stream. Do I understand it right, that we cannot write to the allocated managed/pinned memory from the host too? A possible pattern: allocate and initialize a pinned_mdarray on host, then use it on both host and device. Such code will break, right?