Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Add pinned memory resource to C API #311

Merged
merged 4 commits into from
Dec 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions cpp/include/cuvs/core/c_api.h
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,22 @@ cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
*/
cuvsError_t cuvsRMMMemoryResourceReset();

/**
* @brief Allocates pinned memory on the host using RMM
* @param[out] ptr Pointer to allocated host memory
* @param[in] bytes Size in bytes to allocate
* @return cuvsError_t
*/
cuvsError_t cuvsRMMHostAlloc(void** ptr, size_t bytes);

/**
* @brief Deallocates pinned memory on the host using RMM
* @param[in] ptr Pointer to allocated host memory to free
* @param[in] bytes Size in bytes to deallocate
* @return cuvsError_t
*/
cuvsError_t cuvsRMMHostFree(void* ptr, size_t bytes);

/** @} */

#ifdef __cplusplus
Expand Down
16 changes: 16 additions & 0 deletions cpp/src/core/c_api.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
#include <rmm/mr/device/owning_wrapper.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/mr/device/pool_memory_resource.hpp>
#include <rmm/mr/host/pinned_memory_resource.hpp>
#include <thread>

extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res)
Expand Down Expand Up @@ -130,6 +131,21 @@ extern "C" cuvsError_t cuvsRMMMemoryResourceReset()
});
}

thread_local std::unique_ptr<rmm::mr::pinned_memory_resource> pinned_mr;

extern "C" cuvsError_t cuvsRMMHostAlloc(void** ptr, size_t bytes)
{
return cuvs::core::translate_exceptions([=] {
if (pinned_mr == nullptr) { pinned_mr = std::make_unique<rmm::mr::pinned_memory_resource>(); }
*ptr = pinned_mr->allocate(bytes);
});
}

extern "C" cuvsError_t cuvsRMMHostFree(void* ptr, size_t bytes)
{
return cuvs::core::translate_exceptions([=] { pinned_mr->deallocate(ptr, bytes); });
}

thread_local std::string last_error_text = "";

extern "C" const char* cuvsGetLastErrorText()
Expand Down
9 changes: 9 additions & 0 deletions cpp/test/core/c_api.c
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,15 @@ int main()
error = cuvsRMMMemoryResourceReset();
if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }

// Alloc memory on host (pinned)
void* ptr3;
cuvsError_t alloc_error_pinned = cuvsRMMHostAlloc(&ptr3, 1024);
if (alloc_error_pinned == CUVS_ERROR) { exit(EXIT_FAILURE); }

// Free memory
cuvsError_t free_error_pinned = cuvsRMMHostFree(ptr3, 1024);
if (free_error_pinned == CUVS_ERROR) { exit(EXIT_FAILURE); }

// Destroy resources
error = cuvsResourcesDestroy(res);
if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
Expand Down
Loading