Removing more pools

Author: cjnolet

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

@@ -847,12 +847,6 @@ void radix_topk(const T* in,
   size_t req_buf = max_chunk_size * buf_len * 2 * (sizeof(T) + sizeof(IdxT));
   size_t mem_req = req_aux + req_buf + 256 * 6;  // might need extra memory for alignment
 
-  auto pool_guard = raft::get_pool_memory_resource(mr, mem_req);
-  if (pool_guard) {
-    RAFT_LOG_DEBUG("radix::select_k: using pool memory resource with initial size %zu bytes",
-                   mem_req);
-  }
-
   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 * buf_len, stream, mr);
@@ -1120,10 +1114,6 @@ void radix_topk_one_block(const T* in,
   const size_t max_chunk_size =
     calc_chunk_size<T, IdxT, BlockSize>(batch_size, len, sm_cnt, kernel, true);
 
-  auto pool_guard =
-    raft::get_pool_memory_resource(mr, max_chunk_size * buf_len * 2 * (sizeof(T) + sizeof(IdxT)));
-  if (pool_guard) { RAFT_LOG_DEBUG("radix::select_k: using pool memory resource"); }
-
   rmm::device_uvector<char> bufs(
     max_chunk_size * buf_len * 2 * (sizeof(T) + sizeof(IdxT)), stream, mr);
 

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

@@ -988,10 +988,6 @@ void select_k_(int num_of_block,
                rmm::cuda_stream_view stream,
                rmm::mr::device_memory_resource* mr = nullptr)
 {
-  auto pool_guard = raft::get_pool_memory_resource(
-    mr, num_of_block * k * batch_size * 2 * std::max(sizeof(T), sizeof(IdxT)));
-  if (pool_guard) { RAFT_LOG_DEBUG("warpsort::select_k: using pool memory resource"); }
-
   rmm::device_uvector<T> tmp_val(num_of_block * k * batch_size, stream, mr);
   rmm::device_uvector<IdxT> tmp_idx(num_of_block * k * batch_size, stream, mr);
 

cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh

@@ -124,10 +124,6 @@ void build_knn_graph(raft::resources const& res,
   bool first                    = true;
   const auto start_clock        = std::chrono::system_clock::now();
 
-  rmm::mr::device_memory_resource* device_memory = nullptr;
-  auto pool_guard = raft::get_pool_memory_resource(device_memory, 1024 * 1024);
-  if (pool_guard) { RAFT_LOG_DEBUG("ivf_pq using pool memory resource"); }
-
   raft::spatial::knn::detail::utils::batch_load_iterator<DataT> vec_batches(
     dataset.data_handle(),
     dataset.extent(0),

cpp/include/raft/neighbors/detail/ivf_flat_search-inl.cuh

@@ -233,12 +233,6 @@ inline void search(raft::resources const& handle,
                        raft::div_rounding_up_safe<uint64_t>(
                          kExpectedWsSize, 16ull * uint64_t{n_probes} * k + 4ull * index.dim()));
 
-  auto pool_guard = raft::get_pool_memory_resource(mr, max_queries * n_probes * k * 16);
-  if (pool_guard) {
-    RAFT_LOG_DEBUG("ivf_flat::search: using pool memory resource with initial size %zu bytes",
-                   n_queries * n_probes * k * 16ull);
-  }
-
   for (uint32_t offset_q = 0; offset_q < n_queries; offset_q += max_queries) {
     uint32_t queries_batch = min(max_queries, n_queries - offset_q);