| 69 | } |
| 70 | |
| 71 | void VarDevMemDefragmenter::defrag_impl( |
| 72 | VarNode* req_var, const CompNodeInfo& cn_info, size_t extra_size) { |
| 73 | ThinHashMap<MemAllocPlan::Chunk*, ChunkInfo> chunkinfo; |
| 74 | VarNodeSet non_movable_vars; |
| 75 | if (!m_move_safe_oprs.count(req_var->owner_opr())) { |
| 76 | // input and output vars of current opr can not be moved |
| 77 | auto opr = req_var->owner_opr(); |
| 78 | for (auto i : opr->node_prop().dep_map()) { |
| 79 | if (OperatorNodeBase::NodeProp::is_device_value_dep(i.second)) { |
| 80 | non_movable_vars.insert(i.first); |
| 81 | } |
| 82 | } |
| 83 | for (auto i : opr->output()) { |
| 84 | non_movable_vars.insert(i); |
| 85 | } |
| 86 | } |
| 87 | for (auto i : cn_info.vars) { |
| 88 | if (i->dev_tensor_valid() && !non_movable_vars.count(i)) { |
| 89 | auto chk = &i->mem_plan().chunk(); |
| 90 | chunkinfo[chk].readers.push_back(i); |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | auto cn = req_var->comp_node(); |
| 95 | |
| 96 | // here we do not need to handle exceptions and restore vars, since |
| 97 | // allocation failure requires the whole graph to be re-executed and all |
| 98 | // vars would be re-allocated |
| 99 | |
| 100 | // release all memory |
| 101 | size_t tot_size = extra_size, nr_refcnt_mismatch = 0, nr_var = 0; |
| 102 | auto alignment = cn.get_mem_addr_alignment(); |
| 103 | for (decltype(chunkinfo.begin()) iter = chunkinfo.begin(), inext; |
| 104 | iter != chunkinfo.end(); iter = inext) { |
| 105 | inext = iter; |
| 106 | ++inext; |
| 107 | |
| 108 | auto refcnt = iter->first->m_refcnt.load(std::memory_order_relaxed); |
| 109 | if (refcnt == iter->second.readers.size()) { |
| 110 | tot_size += get_aligned_power2(iter->first->size(), alignment); |
| 111 | nr_var += iter->second.readers.size(); |
| 112 | auto owner_var = iter->first->owner_var; |
| 113 | auto&& tensor = owner_var->m_dev_tensor; |
| 114 | iter->second.value.comp_node(cn) |
| 115 | .ensure_size(iter->first->size()) |
| 116 | .copy_from(tensor.storage(), iter->first->size()); |
| 117 | |
| 118 | // release memory of all readers |
| 119 | for (auto var : iter->second.readers) { |
| 120 | const_cast<DeviceTensorND&>(var->dev_tensor()).storage({}); |
| 121 | } |
| 122 | // release memory of owner_var |
| 123 | auto&& mem_plan = owner_var->mem_plan(); |
| 124 | if (!mem_plan.valid()) { |
| 125 | // mem_plan of owner_var was invalid here if all reader oprs |
| 126 | // of owner_var have already been executed, but its tensor |
| 127 | // storage should not be released until the refcnt of chunk |
| 128 | // decreasing to zero (see release_chunk() for more details) |
nothing calls this directly
no test coverage detected