Avoid using cudaMemGetInfo and adjust the default workspace size

rapidsai · Feb 23, 2024 · 952c6b9 · 952c6b9
1 parent 4e5d842
commit 952c6b9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/cpp/include/raft/core/resource/device_memory_resource.hpp b/cpp/include/raft/core/resource/device_memory_resource.hpp
@@ -177,7 +177,7 @@ class workspace_resource_factory : public resource_factory {
     // Note, the workspace does not claim all this memory from the start, so it's still usable by
     // the main resource as well.
     // This limit is merely an order for algorithm internals to plan the batching accordingly.
-    return total_size / 2;
+    return total_size / 4;
   }
 };
 

diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
@@ -1586,8 +1586,7 @@ void extend(raft::resources const& handle,
     n_rows + (kIndexGroupSize - 1) * std::min<IdxT>(n_clusters, n_rows));
 
   // Available device memory
-  size_t free_mem, total_mem;
-  RAFT_CUDA_TRY(cudaMemGetInfo(&free_mem, &total_mem));
+  size_t free_mem = raft::resource::get_workspace_free_bytes(handle);
 
   // We try to use the workspace memory by default here.
   // If the workspace limit is too small, we change the resource for batch data to the