diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 2e215ba414..398b983663 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -347,19 +347,24 @@ extern "C" cuvsError_t cuvsRMMMemoryResourceReset() }); } -thread_local std::unique_ptr pinned_mr; +namespace { +// One process-wide resource: pinned_host_memory_resource is stateless, so a buffer allocated on +// one thread can be freed on another (a thread_local resource left it null on other threads). +rmm::mr::pinned_host_memory_resource& pinned_mr() +{ + static rmm::mr::pinned_host_memory_resource mr; + return mr; +} +} // namespace extern "C" cuvsError_t cuvsRMMHostAlloc(void** ptr, size_t bytes) { - return cuvs::core::translate_exceptions([=] { - if (pinned_mr == nullptr) { pinned_mr = std::make_unique(); } - *ptr = pinned_mr->allocate_sync(bytes); - }); + return cuvs::core::translate_exceptions([=] { *ptr = pinned_mr().allocate_sync(bytes); }); } extern "C" cuvsError_t cuvsRMMHostFree(void* ptr, size_t bytes) { - return cuvs::core::translate_exceptions([=] { pinned_mr->deallocate_sync(ptr, bytes); }); + return cuvs::core::translate_exceptions([=] { pinned_mr().deallocate_sync(ptr, bytes); }); } thread_local std::string last_error_text = ""; diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index cfceb58f33..b53b2d75b4 100644 --- a/c/tests/CMakeLists.txt +++ b/c/tests/CMakeLists.txt @@ -109,6 +109,7 @@ ConfigureTest(NAME cuvs_c_headers PATH core/headers.c) ConfigureTest(NAME cuvs_c_test PATH core/c_api.c) ConfigureTest(NAME cuvs_c_resource_pool_test PATH core/resource_pool.cpp) ConfigureTest(NAME cuvs_c_resource_stream_pool_test PATH core/resource_stream_pool.cpp) +ConfigureTest(NAME cuvs_c_host_alloc_test PATH core/host_alloc.cpp) target_link_libraries(cuvs_c_test PRIVATE CUDA::cudart_static) target_link_libraries(cuvs_c_resource_pool_test PRIVATE CUDA::cudart_static) target_link_libraries(cuvs_c_resource_stream_pool_test PRIVATE CUDA::cudart_static) diff --git a/c/tests/core/host_alloc.cpp b/c/tests/core/host_alloc.cpp new file mode 100644 index 0000000000..507e11fd34 --- /dev/null +++ b/c/tests/core/host_alloc.cpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include + +// Pinned host buffers may be freed on a different thread than the one that allocated them, e.g. +// when a binding moves the owning object across threads. +TEST(RMMHostAlloc, FreeOnAnotherThread) +{ + constexpr std::size_t kBytes = 1024; + void* ptr = nullptr; + cuvsError_t alloc_status = CUVS_ERROR; + std::thread([&] { alloc_status = cuvsRMMHostAlloc(&ptr, kBytes); }).join(); + ASSERT_EQ(alloc_status, CUVS_SUCCESS); + ASSERT_NE(ptr, nullptr); + + cuvsError_t free_status = CUVS_ERROR; + std::thread([&] { free_status = cuvsRMMHostFree(ptr, kBytes); }).join(); + EXPECT_EQ(free_status, CUVS_SUCCESS); +} + +TEST(RMMHostAlloc, ConcurrentAllocAndFree) +{ + constexpr int kThreads = 8; + constexpr int kIters = 50; + constexpr std::size_t kBytes = 4096; + std::vector ptrs(kThreads * kIters, nullptr); + std::atomic failures{0}; + + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + for (int i = 0; i < kIters; ++i) { + if (cuvsRMMHostAlloc(&ptrs[t * kIters + i], kBytes) != CUVS_SUCCESS) { failures++; } + } + }); + } + for (auto& th : threads) { + th.join(); + } + threads.clear(); + ASSERT_EQ(failures.load(), 0); + + // Free every buffer from a different thread than the one that allocated it. + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + int const owner = (t + 1) % kThreads; + for (int i = 0; i < kIters; ++i) { + if (cuvsRMMHostFree(ptrs[owner * kIters + i], kBytes) != CUVS_SUCCESS) { failures++; } + } + }); + } + for (auto& th : threads) { + th.join(); + } + EXPECT_EQ(failures.load(), 0); +}