Merge pull request #128 from Project-HAMi/new_branch

archlitchi · web-flow · commit 3506e6222ca0 · 2025-09-15T17:49:50.000+08:00
Fix cuMemCreate not been properly counted
diff --git a/src/allocator/allocator.c b/src/allocator/allocator.c
@@ -20,19 +20,20 @@ allocated_list *device_allocasync;
 #define __CHUNK_SIZE__  CHUNK_SIZE
 
 extern size_t initial_offset;
-extern CUresult cuMemoryAllocate(CUdeviceptr* dptr, size_t bytesize, size_t* bytesallocated,void* data);
+extern CUresult
+    cuMemoryAllocate(CUdeviceptr* dptr, size_t bytesize, void* data);
 extern CUresult cuMemoryFree(CUdeviceptr dptr);
 
 pthread_once_t allocator_allocate_flag = PTHREAD_ONCE_INIT;
 pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
 
-size_t round_up(size_t size,size_t unit){
+size_t round_up(size_t size, size_t unit) {
     if (size & (unit-1))
         return ((size / unit) + 1 ) * unit;
     return size;
 }
 
-int oom_check(const int dev,size_t addon) {
+int oom_check(const int dev, size_t addon) {
     int count1=0;
     CUDA_OVERRIDE_CALL(cuda_library_entry,cuDeviceGetCount,&count1);
     CUdevice d;
@@ -59,7 +60,7 @@ int oom_check(const int dev,size_t addon) {
     return 0;
 }
 
-CUresult view_vgpu_allocator(){
+CUresult view_vgpu_allocator() {
     allocated_list_entry *al;
     size_t total;
     total=0;
@@ -74,7 +75,7 @@ CUresult view_vgpu_allocator(){
     return 0;
 }
 
-CUresult get_listsize(allocated_list *al,size_t *size){
+CUresult get_listsize(allocated_list *al, size_t *size) {
     if (al->length == 0){
         *size = 0;
         return CUDA_SUCCESS;
@@ -88,7 +89,7 @@ CUresult get_listsize(allocated_list *al,size_t *size){
     return CUDA_SUCCESS;
 }
 
-void allocator_init(){
+void allocator_init() {
     LOG_DEBUG("Allocator_init\n");
     
     device_overallocated = malloc(sizeof(allocated_list));
@@ -99,7 +100,7 @@ void allocator_init(){
     pthread_mutex_init(&mutex,NULL);
 }
 
-int add_chunk(CUdeviceptr *address,size_t size){
+int add_chunk(CUdeviceptr *address, size_t size) {
     size_t addr=0;
     size_t allocsize;
     CUresult res = CUDA_SUCCESS;
@@ -113,9 +114,8 @@ int add_chunk(CUdeviceptr *address,size_t size){
     if (size <= IPCSIZE)
         res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAlloc_v2,&e->entry->address,size);
     else{
-        //size = round_up(size,ALIGN);
         e->entry->length = size;
-        res = cuMemoryAllocate(&e->entry->address,size,&e->entry->length,e->entry->allocHandle);
+        res = cuMemoryAllocate(&e->entry->address, size, e->entry->allocHandle);
     }
     if (res!=CUDA_SUCCESS){
         LOG_ERROR("cuMemoryAllocate failed res=%d",res);
@@ -126,11 +126,11 @@ int add_chunk(CUdeviceptr *address,size_t size){
     *address = e->entry->address;
     allocsize = size;
     cuCtxGetDevice(&dev);
-    add_gpu_device_memory_usage(getpid(),dev,allocsize,2);
+    add_gpu_device_memory_usage(getpid(), dev, allocsize, 2);
     return 0;
 }
 
-int add_chunk_only(CUdeviceptr address,size_t size){
+int add_chunk_only(CUdeviceptr address, size_t size) {
     pthread_mutex_lock(&mutex);
     size_t addr=0;
     size_t allocsize;
@@ -147,7 +147,7 @@ int add_chunk_only(CUdeviceptr address,size_t size){
     e->entry->address=address;
     allocsize = size;
     cuCtxGetDevice(&dev);
-    add_gpu_device_memory_usage(getpid(),dev,allocsize,2);
+    add_gpu_device_memory_usage(getpid(), dev, allocsize, 2);
     pthread_mutex_unlock(&mutex);
     return 0;
 }
@@ -162,50 +162,70 @@ int check_memory_type(CUdeviceptr address) {
     return CU_MEMORYTYPE_HOST;
 }
 
-int remove_chunk(allocated_list *a_list, CUdeviceptr dptr){
+int remove_chunk(allocated_list *a_list, CUdeviceptr dptr) {
     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){
+        if (val->entry->address == dptr) {
             t_size=val->entry->length;
             cuMemoryFree(dptr);
             LIST_REMOVE(a_list,val);
-        
             CUdevice dev;
             cuCtxGetDevice(&dev);
-            rm_gpu_device_memory_usage(getpid(),dev,t_size,2);
+            rm_gpu_device_memory_usage(getpid(), dev, t_size, 2);
             return 0;
         }
     }
     return -1;
 }
 
-int allocate_raw(CUdeviceptr *dptr, size_t size){
+int remove_chunk_only(CUdeviceptr dptr) {
+    allocated_list *a_list = device_overallocated;
+    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) {
+            t_size = val->entry->length;
+            LIST_REMOVE(a_list, val);
+            CUdevice dev;
+            cuCtxGetDevice(&dev);
+            rm_gpu_device_memory_usage(getpid(), dev, t_size, 2);
+            return 0;
+        }
+    }
+    return -1;
+}
+
+int allocate_raw(CUdeviceptr *dptr, size_t size) {
     int tmp;
     pthread_mutex_lock(&mutex);
-    tmp = add_chunk(dptr,size);
+    tmp = add_chunk(dptr, size);
     pthread_mutex_unlock(&mutex);
     return tmp;
 }
 
-int free_raw(CUdeviceptr dptr){
+int free_raw(CUdeviceptr dptr) {
     pthread_mutex_lock(&mutex);
-    unsigned int tmp = remove_chunk(device_overallocated,dptr);
+    unsigned int tmp = remove_chunk(device_overallocated, dptr);
     pthread_mutex_unlock(&mutex);
     return tmp;
 }
 
-int remove_chunk_async(allocated_list *a_list, CUdeviceptr dptr, CUstream hStream){
+int remove_chunk_async(
+    allocated_list *a_list, CUdeviceptr dptr, CUstream hStream) {
     size_t t_size;
-    if (a_list->length==0) {
+    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){
+    for (val = a_list->head; val != NULL; val = val->next) {
+        if (val->entry->address == dptr) {
             t_size=val->entry->length;
             CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemFreeAsync,dptr,hStream);
             LIST_REMOVE(a_list,val);
@@ -219,14 +239,14 @@ int remove_chunk_async(allocated_list *a_list, CUdeviceptr dptr, CUstream hStrea
     return -1;
 }
 
-int free_raw_async(CUdeviceptr dptr, CUstream hStream){
+int free_raw_async(CUdeviceptr dptr, CUstream hStream) {
     pthread_mutex_lock(&mutex);
-    unsigned int tmp = remove_chunk_async(device_allocasync,dptr,hStream);
+    unsigned int tmp = remove_chunk_async(device_allocasync, dptr, hStream);
     pthread_mutex_unlock(&mutex);
     return tmp;
 }
 
-int add_chunk_async(CUdeviceptr *address,size_t size, CUstream hStream){
+int add_chunk_async(CUdeviceptr *address, size_t size, CUstream hStream) {
     size_t addr=0;
     size_t allocsize;
     CUresult res = CUDA_SUCCESS;
@@ -238,28 +258,28 @@ int add_chunk_async(CUdeviceptr *address,size_t size, CUstream hStream){
     allocated_list_entry *e;
     INIT_ALLOCATED_LIST_ENTRY(e,addr,size);
     res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAllocAsync,&e->entry->address,size,hStream);
-    if (res!=CUDA_SUCCESS){
+    if (res != CUDA_SUCCESS) {
         LOG_ERROR("cuMemoryAllocate failed res=%d",res);
         return res;
     }
     *address = e->entry->address;
     CUmemoryPool pool;
     res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuDeviceGetMemPool,&pool,dev);
-    if (res!=CUDA_SUCCESS){
+    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);
-    if (res!=CUDA_SUCCESS) {
+    if (res != CUDA_SUCCESS) {
         LOG_ERROR("cuMemPoolGetAttribute failed res=%d",res);
         return res;
     }
-    if (poollimit!=0) {
+    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);
+            add_gpu_device_memory_usage(getpid(), dev, allocsize, 2);
             device_allocasync->limit=device_allocasync->limit+allocsize;
             e->entry->length=allocsize;
         }else{
@@ -270,7 +290,7 @@ int add_chunk_async(CUdeviceptr *address,size_t size, CUstream hStream){
     return 0;
 }
 
-int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream){
+int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream) {
     int tmp;
     pthread_mutex_lock(&mutex);
     tmp = add_chunk_async(dptr,size,hStream);
diff --git a/src/allocator/allocator.h b/src/allocator/allocator.h
@@ -7,6 +7,9 @@
 #include <pthread.h>
 #include <errno.h>
 
+#define CUMALLOC 0
+#define CUCREATE 1
+
 struct allocated_device_memory_struct{
     CUdeviceptr address;
     size_t length;
@@ -155,6 +158,7 @@ int oom_check(const int dev,size_t addon);
 int allocate_raw(CUdeviceptr *dptr, size_t size);
 int free_raw(CUdeviceptr dptr);
 int add_chunk_only(CUdeviceptr address,size_t size);
+int remove_chunk_only(CUdeviceptr address);
 int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream);
 int free_raw_async(CUdeviceptr dptr, CUstream hStream);
 
diff --git a/src/cuda/hook.c b/src/cuda/hook.c
@@ -132,7 +132,9 @@ cuda_entry_t cuda_library_entry[] = {
     /* Virtual Memory Part */
     {.name = "cuMemAddressReserve"},
     {.name = "cuMemCreate"},
+    {.name = "cuMemRelease"},
     {.name = "cuMemMap"},
+    {.name = "cucuMemImportFromShareableHandle"},
     {.name = "cuMemAllocAsync"},
     {.name = "cuMemFreeAsync"},
     /* cuda11.7 new api memory part */
diff --git a/src/cuda/memory.c b/src/cuda/memory.c
@@ -126,10 +126,8 @@ CUresult cuArrayDestroy(CUarray arr) {
     return res;
 }
 
-CUresult cuMemoryAllocate(CUdeviceptr* dptr, size_t bytesize, size_t* bytesallocated,void* data){
+CUresult cuMemoryAllocate(CUdeviceptr* dptr, size_t bytesize, void* data) {
     CUresult res;
-    if (bytesallocated!=NULL)
-        *bytesallocated = bytesize;
     res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAlloc_v2,dptr,bytesize);
     return res;
 }
@@ -322,8 +320,8 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr){
     return CUDA_OVERRIDE_CALL(cuda_library_entry,cuIpcCloseMemHandle,dptr);
 }
 
-CUresult cuIpcGetMemHandle ( CUipcMemHandle* pHandle, CUdeviceptr dptr ){
-    LOG_DEBUG("cuIpcGetMemHandle dptr=%llx",dptr);
+CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) {
+    LOG_MSG("cuIpcGetMemHandle dptr=%llx", dptr);
     ENSURE_RUNNING();
     return CUDA_OVERRIDE_CALL(cuda_library_entry,cuIpcGetMemHandle,pHandle,dptr);
 }
@@ -499,21 +497,23 @@ CUresult cuMemGetInfo_v2(size_t* free, size_t* total) {
     size_t limit = get_current_device_memory_limit(dev);
     if (limit == 0) {
         CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemGetInfo_v2, free, total);
-        LOG_MSG("orig free=%ld total=%ld",*free,*total);
+        LOG_INFO("orig free=%ld total=%ld", *free, *total);
         *free = *total - usage;
-        LOG_MSG("after free=%ld total=%ld",*free,*total); 
+        LOG_INFO("after free=%ld total=%ld", *free, *total);
         return CUDA_SUCCESS;
     } else if (limit < usage) {
-        LOG_WARN("limit < usage; usage=%ld, limit=%ld",usage,limit);
+        LOG_WARN("limit < usage; usage=%ld, limit=%ld", usage, limit);
         return CUDA_ERROR_INVALID_VALUE;
     } else {
         CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemGetInfo_v2, free, total);
-        LOG_MSG("orig free=%ld total=%ld limit=%ld usage=%ld",*free,*total,limit,usage);
+        LOG_INFO("orig free=%ld total=%ld limit=%ld usage=%ld",
+            *free, *total, limit, usage);
         // Ensure total memory does not exceed the physical or imposed limit.
         size_t actual_limit = (limit > *total) ? *total : limit;
         *free = (actual_limit > usage) ? (actual_limit - usage) : 0;
         *total = actual_limit;
-        LOG_MSG("after free=%ld total=%ld limit=%ld usage=%ld",*free,*total,limit,usage);
+        LOG_INFO("after free=%ld total=%ld limit=%ld usage=%ld",
+            *free, *total, limit, usage);
         return CUDA_SUCCESS;
     }
 }
@@ -566,24 +566,53 @@ CUresult cuMemoryFree(CUdeviceptr dptr) {
     return res;
 }
 
-CUresult cuMemAddressReserve ( CUdeviceptr* ptr, size_t size, size_t alignment, CUdeviceptr addr, unsigned long long flags ) {
-    LOG_INFO("cuMemAddressReserve:%lx %lld",size,addr);
-    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAddressReserve,ptr,size,alignment,addr,flags);
+CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size,
+    size_t alignment, CUdeviceptr addr, unsigned long long flags ) {
+    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,
+        cuMemAddressReserve, ptr, size, alignment, addr, flags);
+    LOG_INFO("cuMemAddressReserve:%lx %llx", size, *ptr);
     return res;
 }
 
 CUresult cuMemCreate ( CUmemGenericAllocationHandle* handle, size_t size, const CUmemAllocationProp* prop, unsigned long long flags ) {
-    LOG_INFO("cuMemCreate:");
-    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemCreate,handle,size,prop,flags);
+    LOG_INFO("cuMemCreate:%lld:%d", size, prop->location.id);
+    ENSURE_RUNNING();
+    CUdevice dev;
+    CUDA_OVERRIDE_CALL(cuda_library_entry, cuCtxGetDevice, &dev);
+    if (oom_check(dev, size)) {
+        return CUDA_ERROR_OUT_OF_MEMORY;
+    }
+    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,
+        cuMemCreate, handle, size, prop, flags);
+    if (res == CUDA_SUCCESS) {
+        add_chunk_only(*handle, size);
+    }
+    return res;
+}
+
+CUresult cuMemRelease(CUmemGenericAllocationHandle handle) {
+    LOG_INFO("cuMemRelease:%llx", handle);
+    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemRelease, handle);
+    if (res == CUDA_SUCCESS) {
+        remove_chunk_only(handle);
+    }
     return res;
 }
 
 CUresult cuMemMap( CUdeviceptr ptr, size_t size, size_t offset, CUmemGenericAllocationHandle handle, unsigned long long flags ) {
-    LOG_INFO("cuMemMap");
+    LOG_INFO("cuMemMap:%lld(%llx,%llx)", size, ptr, offset);
     CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemMap,ptr,size,offset,handle,flags);
     return res;
 }
 
+CUresult cuMemImportFromShareableHandle(CUmemGenericAllocationHandle* handle,
+    void* osHandle, CUmemAllocationHandleType shHandleType) {
+    LOG_INFO("cuMemImportFromSharableHandle");
+    CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,
+        cuMemImportFromShareableHandle, handle, osHandle, shHandleType);
+    return res;
+}
+
 CUresult cuMemAllocAsync(CUdeviceptr *dptr, size_t bytesize, CUstream hStream) {
     LOG_DEBUG("cuMemAllocAsync:%ld",bytesize);
     return allocate_async_raw(dptr,bytesize,hStream);
diff --git a/src/include/libcuda_hook.h b/src/include/libcuda_hook.h
@@ -164,7 +164,9 @@ typedef enum {
     /* Virtual Memory Part */
     CUDA_OVERRIDE_ENUM(cuMemAddressReserve),
     CUDA_OVERRIDE_ENUM(cuMemCreate),
+    CUDA_OVERRIDE_ENUM(cuMemRelease),
     CUDA_OVERRIDE_ENUM(cuMemMap),
+    CUDA_OVERRIDE_ENUM(cuMemImportFromShareableHandle),
     CUDA_OVERRIDE_ENUM(cuMemAllocAsync),
     CUDA_OVERRIDE_ENUM(cuMemFreeAsync),
     /* cuda11.7 new api memory part */
diff --git a/src/libvgpu.c b/src/libvgpu.c
@@ -252,7 +252,9 @@ void* __dlsym_hook_section(void* handle, const char* symbol) {
     DLSYM_HOOK_FUNC(cuLinkDestroy);
     DLSYM_HOOK_FUNC(cuMemAddressReserve);
     DLSYM_HOOK_FUNC(cuMemCreate);
+    DLSYM_HOOK_FUNC(cuMemRelease);
     DLSYM_HOOK_FUNC(cuMemMap);
+    DLSYM_HOOK_FUNC(cuMemImportFromShareableHandle);
     DLSYM_HOOK_FUNC(cuMemAllocAsync);
     // cuda 11.7 new memory ops
     DLSYM_HOOK_FUNC(cuMemHostGetDevicePointer_v2);
