Browse Source

fix(src/core): fix weight preprocess free combine memory UAF issue

GitOrigin-RevId: b779e05956
master
Megvii Engine Team 2 years ago
parent
commit
4c27522a2d
1 changed files with 7 additions and 0 deletions
  1. +7
    -0
      src/core/impl/graph/var_node_mem_mgr.cpp

+ 7
- 0
src/core/impl/graph/var_node_mem_mgr.cpp View File

@@ -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(


Loading…
Cancel
Save