From ac2033cb2751e8d82e969f8659047147e31c635e Mon Sep 17 00:00:00 2001
From: root <lihuang@alauda.io>
Date: Sun, 27 Apr 2025 16:15:00 +0800
Subject: [PATCH] try to free for not hook func

Signed-off-by: root <lihuang@alauda.io>
---
 .gitignore                 |  3 ++-
 src/allocator/allocator.c  | 14 ++++++++++++--
 src/cuda/graph.c           |  6 ++++++
 src/cuda/hook.c            |  1 +
 src/include/libcuda_hook.h |  1 +
 src/libvgpu.c              |  1 +
 6 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5ef4f87b..6d6ded84 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 build
-cmake-build-debug
\ No newline at end of file
+cmake-build-debug
+.idea
diff --git a/src/allocator/allocator.c b/src/allocator/allocator.c
index 99215b12..3c979a08 100755
--- a/src/allocator/allocator.c
+++ b/src/allocator/allocator.c
@@ -163,21 +163,31 @@ int check_memory_type(CUdeviceptr address) {
 int remove_chunk(allocated_list *a_list, CUdeviceptr dptr){
     size_t t_size;
     if (a_list->length==0) {
-        return -1;
+        LOG_ERROR("remove_chunk  a_list length is 0 res=%d",a_list->length);
+        if (dptr != NULL){
+          //maybe some graph func, we can not hook the memory alloc
+          //try to free raw
+          return cuMemoryFree(dptr);
+        }
     }
     allocated_list_entry *val;
     for (val=a_list->head;val!=NULL;val=val->next){
+        LOG_ERROR("remove_chunk  remove list entry ");
         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);
             return 0;
         }
     }
+    if (dptr != NULL){
+      //maybe some graph func, we can not hook the memory alloc
+      //try to free raw
+      return cuMemoryFree(dptr);
+    }
     return -1;
 }
 
diff --git a/src/cuda/graph.c b/src/cuda/graph.c
index 348993c1..a1878d56 100644
--- a/src/cuda/graph.c
+++ b/src/cuda/graph.c
@@ -25,6 +25,12 @@ CUresult cuGraphAddMemcpyNode(CUgraphNode *phGraphNode, CUgraph hGraph, const CU
 	return CUDA_OVERRIDE_CALL(cuda_library_entry,cuGraphAddMemcpyNode,phGraphNode,hGraph,dependencies,numDependencies,copyParams,ctx);
 }
 
+CUresult cuGraphAddMemAllocNode(CUgraphNode *phGraphNode, CUgraph hGraph, const CUgraphNode *dependencies,
+                                size_t numDependencies, CUDA_MEM_ALLOC_NODE_PARAMS *nodeParams) {
+        LOG_INFO("call in cuGraphAddMemAllocNode");
+        return CUDA_OVERRIDE_CALL(cuda_library_entry,cuGraphAddMemAllocNode,phGraphNode,hGraph,dependencies,numDependencies,nodeParams);
+}
+
 CUresult cuGraphMemcpyNodeGetParams(CUgraphNode hNode, CUDA_MEMCPY3D *nodeParams) {
 	LOG_DEBUG("cuGraphMemcpyNodeGetParams");
 	return CUDA_OVERRIDE_CALL(cuda_library_entry,cuGraphMemcpyNodeGetParams,hNode,nodeParams);
diff --git a/src/cuda/hook.c b/src/cuda/hook.c
index bccc12f1..c1970b6a 100644
--- a/src/cuda/hook.c
+++ b/src/cuda/hook.c
@@ -176,6 +176,7 @@ cuda_entry_t cuda_library_entry[] = {
     {.name = "cuGraphMemcpyNodeGetParams"},
     {.name = "cuGraphMemcpyNodeSetParams"},
     {.name = "cuGraphAddMemsetNode"},
+    {.name = "cuGraphAddMemAllocNode"},
     {.name = "cuGraphMemsetNodeGetParams"},
     {.name = "cuGraphMemsetNodeSetParams"},
     {.name = "cuGraphAddHostNode"},
diff --git a/src/include/libcuda_hook.h b/src/include/libcuda_hook.h
index f7bc43e3..98151ac0 100644
--- a/src/include/libcuda_hook.h
+++ b/src/include/libcuda_hook.h
@@ -208,6 +208,7 @@ typedef enum {
     CUDA_OVERRIDE_ENUM(cuGraphMemcpyNodeGetParams),
     CUDA_OVERRIDE_ENUM(cuGraphMemcpyNodeSetParams),
     CUDA_OVERRIDE_ENUM(cuGraphAddMemsetNode),
+    CUDA_OVERRIDE_ENUM(cuGraphAddMemAllocNode),
     CUDA_OVERRIDE_ENUM(cuGraphMemsetNodeGetParams),
     CUDA_OVERRIDE_ENUM(cuGraphMemsetNodeSetParams),
     CUDA_OVERRIDE_ENUM(cuGraphAddHostNode),
diff --git a/src/libvgpu.c b/src/libvgpu.c
index f645eafa..154e35eb 100644
--- a/src/libvgpu.c
+++ b/src/libvgpu.c
@@ -300,6 +300,7 @@ void* __dlsym_hook_section(void* handle, const char* symbol) {
     DLSYM_HOOK_FUNC(cuGraphMemcpyNodeGetParams);
     DLSYM_HOOK_FUNC(cuGraphMemcpyNodeSetParams);
     DLSYM_HOOK_FUNC(cuGraphAddMemsetNode);
+    DLSYM_HOOK_FUNC(cuGraphAddMemAllocNode);
     DLSYM_HOOK_FUNC(cuGraphMemsetNodeGetParams);
     DLSYM_HOOK_FUNC(cuGraphMemsetNodeSetParams);
     DLSYM_HOOK_FUNC(cuGraphAddHostNode);