try to free for not hook func

Author: huanghuangzym

src/allocator/allocator.c

@@ -163,21 +163,33 @@ 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;
+    int has_found = 0;
     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);
-        
+            has_found = 1;
             CUdevice dev;
             cuCtxGetDevice(&dev);
             rm_gpu_device_memory_usage(getpid(),dev,t_size,2);
             return 0;
         }
     }
+    if (has_found == 0  && dptr != NULL){
+      //maybe some graph func, we can not hook the memory alloc
+      //try to free raw
+      return cuMemoryFree(dptr);
+    }
     return -1;
 }
 

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

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"},

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

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