diff --git a/cpp/include/cugraph/mtmg/resource_manager.hpp b/cpp/include/cugraph/mtmg/resource_manager.hpp index 7e86f6e260e..e1e1d7ffc9d 100644 --- a/cpp/include/cugraph/mtmg/resource_manager.hpp +++ b/cpp/include/cugraph/mtmg/resource_manager.hpp @@ -121,6 +121,8 @@ class resource_manager_t { * @param instance_manager_id a ncclUniqueId that is shared by all processes participating * in this instance. All processes must use the same ID in this call, it is up * to the calling code to share this ID properly before the call. + * @param n_streams The number of streams to create in a stream pool for + * each GPU. Defaults to 16. * * @return unique pointer to instance manager */ diff --git a/cpp/tests/mtmg/threaded_test.cu b/cpp/tests/mtmg/threaded_test.cu index 3a677eccf43..bc4d8cfef6a 100644 --- a/cpp/tests/mtmg/threaded_test.cu +++ b/cpp/tests/mtmg/threaded_test.cu @@ -94,8 +94,9 @@ class Tests_Multithreaded size_t device_buffer_size{64 * 1024 * 1024}; size_t thread_buffer_size{4 * 1024 * 1024}; + const int num_threads_per_gpu{4}; int num_gpus = gpu_list.size(); - int num_threads = num_gpus * 4; + int num_threads = num_gpus * num_threads_per_gpu; cugraph::mtmg::resource_manager_t resource_manager; @@ -106,8 +107,10 @@ class Tests_Multithreaded ncclUniqueId instance_manager_id; ncclGetUniqueId(&instance_manager_id); + // Currently the only uses for multiple streams for each CPU threads + // associated with a particular GPU, which is a constant set above auto instance_manager = resource_manager.create_instance_manager( - resource_manager.registered_ranks(), instance_manager_id, 4); + resource_manager.registered_ranks(), instance_manager_id, num_threads_per_gpu); cugraph::mtmg::edgelist_t edgelist; cugraph::mtmg::graph_t graph;