Enable clear cache upon GC.

glianyi · Oct 25, 2022 · 095263b · 095263b
1 parent f5ef19c
commit 095263b
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/lib/nnc/ccv_nnc_dynamic_graph_alloc.c b/lib/nnc/ccv_nnc_dynamic_graph_alloc.c
@@ -4,10 +4,16 @@
 #include "ccv_nnc_easy.h"
 #include "ccv_internal.h"
 #include "_ccv_nnc_dynamic_graph.h"
+#ifdef HAVE_MPS
+#include "mps/ccv_nnc_mps.h"
+#endif
 
 void ccv_nnc_dynamic_graph_gc(ccv_nnc_dynamic_graph_t* const graph)
 {
 	ccv_nnc_xpu_gc(-1, &graph->xpu_alloc);
+#ifdef HAVE_MPS
+	ccv_nnc_mps_clear_graph_executable_cache();
+#endif
 }
 
 ccv_nnc_compilation_artifact_t* ccv_nnc_compilation_artifact_new(ccv_nnc_graph_t* const graph, ccv_nnc_tensor_arena_t* const tensor_arena, ccv_nnc_graph_exec_arena_t* const exec_arena)

diff --git a/lib/nnc/mps/ccv_nnc_mps.h b/lib/nnc/mps/ccv_nnc_mps.h
@@ -29,6 +29,7 @@ void ccv_nnc_stream_compat_wait_signal(const ccv_nnc_stream_context_t* const str
 void ccv_nnc_deinit_stream_signal(ccv_nnc_stream_signal_t* const signal);
 CCV_WARN_UNUSED(int) ccv_nnc_gpu_device_count(void);
 void ccv_nnc_mps_unbounded_command_buffers(int state);
+void ccv_nnc_mps_clear_graph_executable_cache(void);
 
 #ifdef __OBJC__
 

diff --git a/lib/nnc/mps/ccv_nnc_mps.m b/lib/nnc/mps/ccv_nnc_mps.m
@@ -99,8 +99,6 @@ void mpunregmp(const int slot)
 	pthread_mutex_unlock(&g_mp_mutex);
 }
 
-static void mps_graph_cache_pressure(void);
-
 static void mptrigmp(void)
 {
 	ccv_nnc_synchronize_stream_context(0);
@@ -113,7 +111,7 @@ static void mptrigmp(void)
 			mp->func(0, mp->ctx);
 	}
 	pthread_mutex_unlock(&g_mp_mutex);
-	mps_graph_cache_pressure();
+	ccv_nnc_mps_clear_graph_executable_cache();
 }
 
 void* mpmalloc(int device, size_t size)
@@ -307,7 +305,7 @@ static inline void ccv_nnc_mps_graph_key_free(ccv_nnc_mps_graph_key_t key)
 		ccfree(key.inputs);
 }
 
-static void mps_graph_cache_pressure(void)
+void ccv_nnc_mps_clear_graph_executable_cache(void)
 {
 	if (!g_graph_executable_cache)
 		return;
@@ -319,6 +317,7 @@ static void mps_graph_cache_pressure(void)
 		ccv_nnc_mps_graph_key_free(kh_key(g_graph_executable_cache, k));
 		if (kh_val(g_graph_executable_cache, k).indices)
 			ccfree(kh_val(g_graph_executable_cache, k).indices);
+		[kh_val(g_graph_executable_cache, k).exec release];
 		kh_del(graph_executable_cache, g_graph_executable_cache, k);
 	}
 }