From 28ee7d2184312b1313ab894d72eabaf2dbdf58fb Mon Sep 17 00:00:00 2001 From: Ajit Mistry <55892788+ajit283@users.noreply.github.com> Date: Fri, 27 Sep 2024 20:52:27 +0200 Subject: [PATCH] Add `managed` option for RMM Pool memory resource to C API (#305) Authors: - Ajit Mistry (https://github.com/ajit283) - Ray Douglass (https://github.com/raydouglass) - Corey J. Nolet (https://github.com/cjnolet) - Ben Frederickson (https://github.com/benfred) Approvers: - Ben Frederickson (https://github.com/benfred) - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cuvs/pull/305 --- cpp/include/cuvs/core/c_api.h | 5 +++- cpp/src/core/c_api.cpp | 27 +++++++++++++++++--- cpp/test/core/c_api.c | 47 +++++++++++++++++++++++------------ 3 files changed, 58 insertions(+), 21 deletions(-) diff --git a/cpp/include/cuvs/core/c_api.h b/cpp/include/cuvs/core/c_api.h index 4db7fd12c..c8c8d3934 100644 --- a/cpp/include/cuvs/core/c_api.h +++ b/cpp/include/cuvs/core/c_api.h @@ -17,6 +17,7 @@ #pragma once #include +#include #include #ifdef __cplusplus @@ -138,10 +139,12 @@ cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes); * available memory * @param[in] max_pool_size_percent The maximum pool size as a percentage of the total * available memory + * @param[in] managed Whether to use a managed memory resource as upstream resource or not * @return cuvsError_t */ cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent, - int max_pool_size_percent); + int max_pool_size_percent, + bool managed); /** * @brief Resets the memory resource to use the default memory resource (cuda_memory_resource) * @return cuvsError_t diff --git a/cpp/src/core/c_api.cpp b/cpp/src/core/c_api.cpp index a75e5a1dd..cfbeed2d5 100644 --- a/cpp/src/core/c_api.cpp +++ b/cpp/src/core/c_api.cpp @@ -21,6 +21,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -83,10 +86,14 @@ extern "C" cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes) }); } -thread_local std::unique_ptr> pool_mr; +thread_local std::shared_ptr< + rmm::mr::owning_wrapper, + rmm::mr::device_memory_resource>> + pool_mr; extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent, - int max_pool_size_percent) + int max_pool_size_percent, + bool managed) { return cuvs::core::translate_exceptions([=] { // Upstream memory resource needs to be a cuda_memory_resource @@ -95,10 +102,22 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per if (cuda_mr_casted == nullptr) { throw std::runtime_error("Current memory resource is not a cuda_memory_resource"); } + auto initial_size = rmm::percent_of_free_device_memory(initial_pool_size_percent); auto max_size = rmm::percent_of_free_device_memory(max_pool_size_percent); - pool_mr = std::make_unique>( - cuda_mr_casted, initial_size, max_size); + + auto mr = std::shared_ptr(); + if (managed) { + mr = std::static_pointer_cast( + std::make_shared()); + } else { + mr = std::static_pointer_cast( + std::make_shared()); + } + + pool_mr = + rmm::mr::make_owning_wrapper(mr, initial_size, max_size); + rmm::mr::set_current_device_resource(pool_mr.get()); }); } diff --git a/cpp/test/core/c_api.c b/cpp/test/core/c_api.c index 27973c2dd..a3dae6004 100644 --- a/cpp/test/core/c_api.c +++ b/cpp/test/core/c_api.c @@ -33,34 +33,49 @@ int main() // Allocate memory void* ptr; - size_t bytes = 1024; - cuvsError_t alloc_error = cuvsRMMAlloc(res, &ptr, bytes); - if (alloc_error == CUVS_ERROR) { exit(EXIT_FAILURE); } + size_t bytes = 1024; + cuvsError_t error = cuvsRMMAlloc(res, &ptr, bytes); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } // Free memory - cuvsError_t free_error = cuvsRMMFree(res, ptr, bytes); - if (free_error == CUVS_ERROR) { exit(EXIT_FAILURE); } + error = cuvsRMMFree(res, ptr, bytes); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } // Enable pool memory resource - cuvsError_t pool_error = cuvsRMMPoolMemoryResourceEnable(10, 100); - if (pool_error == CUVS_ERROR) { exit(EXIT_FAILURE); } + error = cuvsRMMPoolMemoryResourceEnable(10, 100, false); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } // Allocate memory again - void* ptr2; - cuvsError_t alloc_error_pool = cuvsRMMAlloc(res, &ptr2, 1024); - if (alloc_error_pool == CUVS_ERROR) { exit(EXIT_FAILURE); } + error = cuvsRMMAlloc(res, &ptr, 1024); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } // Free memory - cuvsError_t free_error_pool = cuvsRMMFree(res, ptr2, 1024); - if (free_error_pool == CUVS_ERROR) { exit(EXIT_FAILURE); } + error = cuvsRMMFree(res, ptr, 1024); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } // Reset pool memory resource - cuvsError_t reset_error = cuvsRMMMemoryResourceReset(); - if (reset_error == CUVS_ERROR) { exit(EXIT_FAILURE); } + error = cuvsRMMMemoryResourceReset(); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } + + // Enable pool memory resource (managed) + error = cuvsRMMPoolMemoryResourceEnable(10, 100, true); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } + + // Allocate memory again + error = cuvsRMMAlloc(res, &ptr, 1024); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } + + // Free memory + error = cuvsRMMFree(res, ptr, 1024); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } + + // Reset pool memory resource + error = cuvsRMMMemoryResourceReset(); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } // Destroy resources - cuvsError_t destroy_error = cuvsResourcesDestroy(res); - if (destroy_error == CUVS_ERROR) { exit(EXIT_FAILURE); } + error = cuvsResourcesDestroy(res); + if (error == CUVS_ERROR) { exit(EXIT_FAILURE); } return 0; }