From 4c27522a2d27c2b599e8590757aba9c198b73f3b Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Mon, 5 Sep 2022 12:51:08 +0800 Subject: [PATCH] fix(src/core): fix weight preprocess free combine memory UAF issue GitOrigin-RevId: b779e059560fb91ffd9d2d770bda097353b3ba5b --- src/core/impl/graph/var_node_mem_mgr.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/core/impl/graph/var_node_mem_mgr.cpp b/src/core/impl/graph/var_node_mem_mgr.cpp index caff321c..0ed7ddc8 100644 --- a/src/core/impl/graph/var_node_mem_mgr.cpp +++ b/src/core/impl/graph/var_node_mem_mgr.cpp @@ -546,6 +546,13 @@ bool VarNodeMemManager::free_combine_memory_no_need_var() { size_t size = var->layout().span().dist_byte(); storage.ensure_size(size); storage.copy_from(var->m_dev_tensor.storage(), size); + auto ref_cn = + var->m_dev_tensor.storage().raw_storage().use_count(); + //! var->m_dev_tensor.storage() will free now, + //! insure do not trigger UAF issue for copy_from. + if (ref_cn <= 2) { + comp_node.sync(); + } var->m_dev_tensor.reset(storage, var->layout()); opr_base->mutable_values()[index]->reset(