提交 4c27522a 编写于 作者: M Megvii Engine Team

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

GitOrigin-RevId: b779e059560fb91ffd9d2d770bda097353b3ba5b
上级 0a52b258
......@@ -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(
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册