diff --git a/src/allocator/allocator.c b/src/allocator/allocator.c index 22e0f587..37aa4f67 100755 --- a/src/allocator/allocator.c +++ b/src/allocator/allocator.c @@ -237,9 +237,6 @@ int free_raw(CUdeviceptr dptr) { int remove_chunk_async( allocated_list *a_list, CUdeviceptr dptr, CUstream hStream) { size_t t_size; - if (a_list->length == 0) { - return -1; - } allocated_list_entry *val; for (val = a_list->head; val != NULL; val = val->next) { if (val->entry->address == dptr) { @@ -253,7 +250,9 @@ int remove_chunk_async( return 0; } } - return -1; + /* Not tracked by libvgpu: free it for real instead of returning -1, which + * leaked and surfaced as an "unrecognized error code". */ + return CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFreeAsync, dptr, hStream); } int free_raw_async(CUdeviceptr dptr, CUstream hStream) { @@ -263,14 +262,40 @@ int free_raw_async(CUdeviceptr dptr, CUstream hStream) { return tmp; } +/* Track an allocated async chunk and account it against its pool's high-water + * limit. Call with mutex held. */ +static int account_async_chunk(allocated_list_entry *e, size_t size, CUmemoryPool pool) { + size_t poollimit; + CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemPoolGetAttribute, pool, + CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH, &poollimit); + if (res != CUDA_SUCCESS) { + LOG_ERROR("cuMemPoolGetAttribute failed res=%d", res); + return res; + } + if (poollimit != 0) { + if (poollimit > device_allocasync->limit) { + size_t allocsize = (poollimit - device_allocasync->limit < size) + ? poollimit - device_allocasync->limit : size; + add_gpu_device_memory_usage(getpid(), e->entry->dev, allocsize, 2); + device_allocasync->limit = device_allocasync->limit + allocsize; + e->entry->length = allocsize; + } else { + e->entry->length = 0; + } + } else { + e->entry->length = 0; + } + LIST_ADD(device_allocasync, e); + return 0; +} + int add_chunk_async(CUdeviceptr *address, size_t size, CUstream hStream) { - size_t addr=0; - size_t allocsize; + size_t addr = 0; CUresult res = CUDA_SUCCESS; CUdevice dev; cuCtxGetDevice(&dev); - if (oom_check(dev,size)) - return -1; + if (oom_check(dev, size)) + return CUDA_ERROR_OUT_OF_MEMORY; allocated_list_entry *e; INIT_ALLOCATED_LIST_ENTRY(e, addr, size, dev); @@ -280,31 +305,34 @@ int add_chunk_async(CUdeviceptr *address, size_t size, CUstream hStream) { return res; } *address = e->entry->address; + /* cuMemAllocAsync draws from the device's default pool. */ CUmemoryPool pool; res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuDeviceGetMemPool,&pool,dev); if (res != CUDA_SUCCESS) { LOG_ERROR("cuDeviceGetMemPool failed res=%d",res); return res; } - size_t poollimit; - res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemPoolGetAttribute,pool,CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,&poollimit); + return account_async_chunk(e, size, pool); +} + +int add_chunk_from_pool_async(CUdeviceptr *address, size_t size, CUmemoryPool pool, CUstream hStream) { + size_t addr = 0; + CUresult res = CUDA_SUCCESS; + CUdevice dev; + cuCtxGetDevice(&dev); + if (oom_check(dev, size)) + return CUDA_ERROR_OUT_OF_MEMORY; + + allocated_list_entry *e; + INIT_ALLOCATED_LIST_ENTRY(e, addr, size, dev); + res = CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemAllocFromPoolAsync, &e->entry->address, size, pool, hStream); if (res != CUDA_SUCCESS) { - LOG_ERROR("cuMemPoolGetAttribute failed res=%d",res); + LOG_ERROR("cuMemAllocFromPoolAsync failed res=%d", res); return res; } - if (poollimit != 0) { - if (poollimit> device_allocasync->limit) { - allocsize = (poollimit-device_allocasync->limit < size)? poollimit-device_allocasync->limit : size; - cuCtxGetDevice(&dev); - add_gpu_device_memory_usage(getpid(), dev, allocsize, 2); - device_allocasync->limit=device_allocasync->limit+allocsize; - e->entry->length=allocsize; - }else{ - e->entry->length=0; - } - } - LIST_ADD(device_allocasync,e); - return 0; + *address = e->entry->address; + /* Account against the caller's pool, not the default. */ + return account_async_chunk(e, size, pool); } int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream) { @@ -314,3 +342,11 @@ int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream) { pthread_mutex_unlock(&mutex); return tmp; } + +int allocate_from_pool_async_raw(CUdeviceptr *dptr, size_t size, CUmemoryPool pool, CUstream hStream) { + int tmp; + pthread_mutex_lock(&mutex); + tmp = add_chunk_from_pool_async(dptr, size, pool, hStream); + pthread_mutex_unlock(&mutex); + return tmp; +} diff --git a/src/allocator/allocator.h b/src/allocator/allocator.h index 38933287..2e0f6c60 100755 --- a/src/allocator/allocator.h +++ b/src/allocator/allocator.h @@ -162,6 +162,7 @@ int free_raw(CUdeviceptr dptr); int add_chunk_only(CUdeviceptr address, size_t size, CUdevice dev); int remove_chunk_only(CUdeviceptr address); int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream); +int allocate_from_pool_async_raw(CUdeviceptr *dptr, size_t size, CUmemoryPool pool, CUstream hStream); int free_raw_async(CUdeviceptr dptr, CUstream hStream); // Checks memory type diff --git a/src/cuda/memory.c b/src/cuda/memory.c index 364046a3..93c40363 100755 --- a/src/cuda/memory.c +++ b/src/cuda/memory.c @@ -687,8 +687,8 @@ CUresult cuMemPoolDestroy(CUmemoryPool pool) { } CUresult cuMemAllocFromPoolAsync(CUdeviceptr *dptr, size_t bytesize, CUmemoryPool pool, CUstream hStream) { - LOG_DEBUG("cuMemAllocFromPoolAsync"); - return CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAllocFromPoolAsync,dptr,bytesize,pool,hStream); + LOG_DEBUG("cuMemAllocFromPoolAsync:%ld", bytesize); + return allocate_from_pool_async_raw(dptr, bytesize, pool, hStream); } CUresult cuMemPoolExportToShareableHandle(void *handle_out, CUmemoryPool pool, CUmemAllocationHandleType handleType, unsigned long long flags) { diff --git a/test/test_alloc_pool_async.c b/test/test_alloc_pool_async.c new file mode 100644 index 00000000..d6cba804 --- /dev/null +++ b/test/test_alloc_pool_async.c @@ -0,0 +1,124 @@ +// Reproducer for Project-HAMi/HAMi-core issue #93 (the same untracked-free leak +// class reported in #67): cuMemAllocFromPoolAsync allocations are invisible to +// libvgpu, so the matching cuMemFreeAsync is rejected. +// +// Mechanism (src/cuda/memory.c + src/allocator/allocator.c on main): +// cuMemFreeAsync() -> free_raw_async() -> remove_chunk_async() walks the +// device_allocasync list for the pointer. cuMemAllocAsync() registers there +// (via add_chunk_async); cuMemAllocFromPoolAsync() is a bare pass-through and +// never does. For a pool-explicit allocation the pointer is not found, so +// remove_chunk_async() returns -1 WITHOUT calling the real cuMemFreeAsync: +// (a) -1 surfaces as a CUresult -> "unrecognized error code", and +// (b) the real free is skipped -> the device allocation leaks. +// +// Path [A] (default pool) is the control: it IS tracked, so it still frees +// cleanly under libvgpu. Path [B] (explicit pool) is the bug. On stock CUDA +// (no LD_PRELOAD) both paths return CUDA_SUCCESS. +// +// Build (standalone, no HAMi headers needed): +// gcc test_alloc_pool_async.c -o test_alloc_pool_async \ +// -I/usr/local/cuda/include -L/usr/local/cuda/lib64/stubs -lcuda +// Or drop into HAMi-core/test/ and `make build-in-docker` (auto-discovered). +// +// Run baseline (no libvgpu, expect PASS): +// ./test_alloc_pool_async +// Run under libvgpu (expect FAIL until fixed): +// LD_PRELOAD=/path/to/libvgpu.so CUDA_DEVICE_MEMORY_LIMIT=2g \ +// LIBCUDA_LOG_LEVEL=1 ./test_alloc_pool_async + +#include +#include +#include + +#ifndef TEST_DEVICE_ID +#define TEST_DEVICE_ID 0 +#endif + +#define ALLOC_BYTES (32ull * 1024 * 1024) /* 32 MiB */ + +/* Abort-on-error for setup calls that must succeed for the test to be valid. */ +#define CHECK_DRV(call) \ + do { \ + CUresult _e = (call); \ + if (_e != CUDA_SUCCESS) { \ + const char *_n = NULL; \ + cuGetErrorName(_e, &_n); \ + fprintf(stderr, "FATAL %s:%d: %s -> %d (%s)\n", __FILE__, \ + __LINE__, #call, (int)_e, _n ? _n : "?"); \ + return 2; \ + } \ + } while (0) + +/* Print a CUresult, flagging codes the driver itself does not recognise + (exactly what a caller sees as "unrecognized error code"). */ +static void describe(const char *label, CUresult res) { + const char *name = NULL; + CUresult q = cuGetErrorName(res, &name); + if (q != CUDA_SUCCESS || name == NULL) + printf(" %-34s -> %d \n", label, (int)res); + else + printf(" %-34s -> %d (%s)\n", label, (int)res, name); +} + +int main(void) { + CHECK_DRV(cuInit(0)); + + CUdevice dev; + CHECK_DRV(cuDeviceGet(&dev, TEST_DEVICE_ID)); + + CUcontext ctx; +#if CUDA_VERSION >= 13000 + CHECK_DRV(cuCtxCreate(&ctx, NULL, 0, dev)); +#else + CHECK_DRV(cuCtxCreate(&ctx, 0, dev)); +#endif + + CUstream stream; + CHECK_DRV(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING)); + + int failures = 0; + + /* [A] Control: default-pool async path. libvgpu tracks this one. */ + printf("[A] cuMemAllocAsync + cuMemFreeAsync (default pool)\n"); + { + CUdeviceptr d = 0; + CHECK_DRV(cuMemAllocAsync(&d, ALLOC_BYTES, stream)); + CUresult f = cuMemFreeAsync(d, stream); + describe("cuMemFreeAsync", f); + CHECK_DRV(cuStreamSynchronize(stream)); + if (f != CUDA_SUCCESS) + failures++; + } + + /* [B] Bug: explicit-pool async path. libvgpu never tracks this one. */ + printf("[B] cuMemAllocFromPoolAsync + cuMemFreeAsync (explicit pool)\n"); + { + CUmemPoolProps props; + memset(&props, 0, sizeof(props)); + props.allocType = CU_MEM_ALLOCATION_TYPE_PINNED; + props.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + props.location.id = dev; + + CUmemoryPool pool; + CHECK_DRV(cuMemPoolCreate(&pool, &props)); + + CUdeviceptr d = 0; + CHECK_DRV(cuMemAllocFromPoolAsync(&d, ALLOC_BYTES, pool, stream)); + CUresult f = cuMemFreeAsync(d, stream); + describe("cuMemFreeAsync", f); + if (f != CUDA_SUCCESS) + failures++; + + /* Best-effort teardown; a rejected free leaves the allocation live. */ + cuStreamSynchronize(stream); + cuMemPoolDestroy(pool); + } + + printf("\nResult: %s\n", + failures ? "FAIL - pool-async free rejected (issue #93 reproduced)" + : "PASS - both async free paths accepted"); + + cuStreamDestroy(stream); + cuCtxDestroy(ctx); + return failures ? 1 : 0; +}