dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix: thread_local reentrant guard — prevent cudaMalloc infinite recursion

Browse Source 
CUB CachingDeviceAllocator::DeviceAllocate calls cudaMalloc internally
on cache miss. Without a guard, our intercepted cudaMalloc recurses
into DeviceAllocate → cudaMalloc → DeviceAllocate → segfault.

thread_local g_in_allocator flag detects reentrant calls and forwards
them directly to the real cudaMalloc/cudaFree via dlsym(RTLD_NEXT).
This commit is contained in:
Claude
2026-08-13 10:37:32 +00:00
parent 3ce5bff10f
commit 45161610f0
1 changed files with 16 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
qwen3_6_scripts/cccl_preload/cccl_allocator_preload.cu
View File

@@ -123,20 +123,32 @@ static void cccl_preload_init() {
/* ======================================================================== /* ========================================================================
* cudaMalloc / cudaFree intercepts * cudaMalloc / cudaFree intercepts
*
* CUB's DeviceAllocate internally calls cudaMalloc on cache miss.
* We must detect this reentrant call and forward to the real function,
* otherwise we get infinite recursion → segfault.
* ======================================================================== */ * ======================================================================== */
static thread_local bool g_in_allocator = false;
extern "C" cudaError_t cudaMalloc(void** devPtr, size_t size) extern "C" cudaError_t cudaMalloc(void** devPtr, size_t size)
{ {
if (!g_preload_active) { if (!g_preload_active || g_in_allocator) {
return get_real_malloc()(devPtr, size); return get_real_malloc()(devPtr, size);
} }
return get_allocator().DeviceAllocate(devPtr, size); g_in_allocator = true;
cudaError_t err = get_allocator().DeviceAllocate(devPtr, size);
g_in_allocator = false;
return err;
} }
extern "C" cudaError_t cudaFree(void* devPtr) extern "C" cudaError_t cudaFree(void* devPtr)
{ {
if (!g_preload_active || devPtr == nullptr) { if (!g_preload_active || devPtr == nullptr || g_in_allocator) {
return get_real_free()(devPtr); return get_real_free()(devPtr);
} }
return get_allocator().DeviceFree(devPtr); g_in_allocator = true;
cudaError_t err = get_allocator().DeviceFree(devPtr);
g_in_allocator = false;
return err;
} }