Skip to content

Commit ac2033c

Browse files
committed
try to free for not hook func
Signed-off-by: root <lihuang@alauda.io>
1 parent b94215b commit ac2033c

File tree

6 files changed

+23
-3
lines changed

6 files changed

+23
-3
lines changed

.gitignore

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
build
2-
cmake-build-debug
2+
cmake-build-debug
3+
.idea

src/allocator/allocator.c

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -163,21 +163,31 @@ int check_memory_type(CUdeviceptr address) {
163163
int remove_chunk(allocated_list *a_list, CUdeviceptr dptr){
164164
size_t t_size;
165165
if (a_list->length==0) {
166-
return -1;
166+
LOG_ERROR("remove_chunk a_list length is 0 res=%d",a_list->length);
167+
if (dptr != NULL){
168+
//maybe some graph func, we can not hook the memory alloc
169+
//try to free raw
170+
return cuMemoryFree(dptr);
171+
}
167172
}
168173
allocated_list_entry *val;
169174
for (val=a_list->head;val!=NULL;val=val->next){
175+
LOG_ERROR("remove_chunk remove list entry ");
170176
if (val->entry->address==dptr){
171177
t_size=val->entry->length;
172178
cuMemoryFree(dptr);
173179
LIST_REMOVE(a_list,val);
174-
175180
CUdevice dev;
176181
cuCtxGetDevice(&dev);
177182
rm_gpu_device_memory_usage(getpid(),dev,t_size,2);
178183
return 0;
179184
}
180185
}
186+
if (dptr != NULL){
187+
//maybe some graph func, we can not hook the memory alloc
188+
//try to free raw
189+
return cuMemoryFree(dptr);
190+
}
181191
return -1;
182192
}
183193

src/cuda/graph.c

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@ CUresult cuGraphAddMemcpyNode(CUgraphNode *phGraphNode, CUgraph hGraph, const CU
2525
return CUDA_OVERRIDE_CALL(cuda_library_entry,cuGraphAddMemcpyNode,phGraphNode,hGraph,dependencies,numDependencies,copyParams,ctx);
2626
}
2727

28+
CUresult cuGraphAddMemAllocNode(CUgraphNode *phGraphNode, CUgraph hGraph, const CUgraphNode *dependencies,
29+
size_t numDependencies, CUDA_MEM_ALLOC_NODE_PARAMS *nodeParams) {
30+
LOG_INFO("call in cuGraphAddMemAllocNode");
31+
return CUDA_OVERRIDE_CALL(cuda_library_entry,cuGraphAddMemAllocNode,phGraphNode,hGraph,dependencies,numDependencies,nodeParams);
32+
}
33+
2834
CUresult cuGraphMemcpyNodeGetParams(CUgraphNode hNode, CUDA_MEMCPY3D *nodeParams) {
2935
LOG_DEBUG("cuGraphMemcpyNodeGetParams");
3036
return CUDA_OVERRIDE_CALL(cuda_library_entry,cuGraphMemcpyNodeGetParams,hNode,nodeParams);

src/cuda/hook.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,7 @@ cuda_entry_t cuda_library_entry[] = {
176176
{.name = "cuGraphMemcpyNodeGetParams"},
177177
{.name = "cuGraphMemcpyNodeSetParams"},
178178
{.name = "cuGraphAddMemsetNode"},
179+
{.name = "cuGraphAddMemAllocNode"},
179180
{.name = "cuGraphMemsetNodeGetParams"},
180181
{.name = "cuGraphMemsetNodeSetParams"},
181182
{.name = "cuGraphAddHostNode"},

src/include/libcuda_hook.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ typedef enum {
208208
CUDA_OVERRIDE_ENUM(cuGraphMemcpyNodeGetParams),
209209
CUDA_OVERRIDE_ENUM(cuGraphMemcpyNodeSetParams),
210210
CUDA_OVERRIDE_ENUM(cuGraphAddMemsetNode),
211+
CUDA_OVERRIDE_ENUM(cuGraphAddMemAllocNode),
211212
CUDA_OVERRIDE_ENUM(cuGraphMemsetNodeGetParams),
212213
CUDA_OVERRIDE_ENUM(cuGraphMemsetNodeSetParams),
213214
CUDA_OVERRIDE_ENUM(cuGraphAddHostNode),

src/libvgpu.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,7 @@ void* __dlsym_hook_section(void* handle, const char* symbol) {
300300
DLSYM_HOOK_FUNC(cuGraphMemcpyNodeGetParams);
301301
DLSYM_HOOK_FUNC(cuGraphMemcpyNodeSetParams);
302302
DLSYM_HOOK_FUNC(cuGraphAddMemsetNode);
303+
DLSYM_HOOK_FUNC(cuGraphAddMemAllocNode);
303304
DLSYM_HOOK_FUNC(cuGraphMemsetNodeGetParams);
304305
DLSYM_HOOK_FUNC(cuGraphMemsetNodeSetParams);
305306
DLSYM_HOOK_FUNC(cuGraphAddHostNode);

0 commit comments

Comments
 (0)