[FEA] Max workspace size

[tfeher]

Is your feature request related to a problem? Please describe.

PR #1137 added a workspace resource, which can be used for temporary memory allocations. I would like an interface to

• set / get max workspace size,
• query the currently allocated workspace.

There are several algorithms, where the workspace size that we actually use depends on an internal batch_size parameter:

• ivf-pq:

  raft/cpp/include/raft/neighbors/detail/ivf_pq_search.cuh

  Lines 1608 to 1611 in ae15be7

  	// Maximum number of query vectors to search at the same time.
  	const auto max_queries = std::min<uint32_t>(std::max<uint32_t>(n_queries, 1), 4096);
  	auto max_batch_size = get_max_batch_size(k, n_probes, max_queries, max_samples);

• radix-topk:

  raft/cpp/include/raft/matrix/detail/select_radix.cuh

  Lines 598 to 603 in ae15be7

  	rmm::device_uvector<Counter<T, IdxT>> counters(max_chunk_size, stream, mr);
  	rmm::device_uvector<IdxT> histograms(max_chunk_size * num_buckets, stream, mr);
  	rmm::device_uvector<T> buf1(max_chunk_size * len, stream, mr_buf);
  	rmm::device_uvector<IdxT> idx_buf1(max_chunk_size * len, stream, mr_buf);
  	rmm::device_uvector<T> buf2(max_chunk_size * len, stream, mr_buf);
  	rmm::device_uvector<IdxT> idx_buf2(max_chunk_size * len, stream, mr_buf);

• kmeans balanced:

  raft/cpp/include/raft/cluster/detail/kmeans_balanced.cuh

  Lines 385 to 388 in ae15be7

  	auto [max_minibatch_size, _mem_per_row] =
  	calc_minibatch_size<MathT>(n_clusters, n_rows, dim, params.metric, std::is_same_v<T, MathT>);
  	rmm::device_uvector<MathT> cur_dataset(
  	std::is_same_v<T, MathT> ? 0 : max_minibatch_size * dim, stream, mr);

• cuML svm predict https://github.com/rapidsai/cuml/blob/branch-23.04/cpp/src/svm/svc_impl.cuh#L113-L123

In some cases above, we just hardcode a reasonable parameter. But the performance of an algorithm usually depends on the batch size, so we would like to allow the user to specify the workspace size, and the algorithms can adapt their batch size accordingly.

Describe the solution you'd like

• Add an extra max_size attribute to the workspace resource,
• track currently allocated size
• print warning if we try to allocate larger than max workspace size.

Describe alternatives you've considered

• Simply use a pool allocator, with initial_size = max_size = workspace_size.
• This allows to specify the max workspace size during initialization.
• We can query the pool size.
• We don't have access of the currently allocated workspace.

The last point is problematic if we have nested algorithms that need workspace. Concrete example is IVF-PQ search can use radix-top-k.

Additional context

• While the difference between max_workspace_size - currently_allocated_size gives the theoretical free space, in practice it is not guaranteed that we can allocate that much memory when a pooling allocator is used (due to fragmentation of the pool).
• The workspace is always freed when an algorithm returns, therefore fragmentation might not be an issue.
• Even if max_workspace_size - currrently_allocated_size is not accurate, it is an estimate that is useful for setting batch size.

[tfeher]

cc @achirkin

[achirkin]

A reasonable alternative: https://github.com/rapidsai/rmm/blob/branch-23.04/include/rmm/mr/device/limiting_resource_adaptor.hpp
One can use the pool allocator as an upstream for this one for better performance.
Though there's some inconvenience to this: to check the available memory amount inside raft's code, we'd have to try dynamic cast the allocator to limiting_resource_adaptor before querying get_allocation_limit , but that would only work if we know exactly what's the upstream resource (template parameter).
On the other hand, we can just force get_workspace_resource to be of type limiting_resource_adaptor<device_memory_resource>, so that the upstream can be anything, and set the default limit to be equal to the current amount of free or total memory (or just int max).