Track cuMemAllocFromPoolAsync allocations and free untracked async pointers

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;
+}

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

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) {

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 <cuda.h>
+#include <stdio.h>
+#include <string.h>
+
+#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  <unrecognized error code>\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;
+}