MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / defrag

Method defrag

imperative/src/impl/blob_manager_impl.cpp:54–97  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

52}
53
54void BlobManagerImpl::defrag(const CompNode& cn) {
55 auto& blobs_set_ptr = ([&]() -> auto& {
56 MGB_LOCK_GUARD(m_mtx);
57 return m_comp2blobs_map[cn];
58 })();
59 MGB_LOCK_GUARD(blobs_set_ptr.mtx);
60 std::vector<BlobData> blob_data_arrary;
61 std::set<Blob::RawStorage> storage_set;
62
63 auto alignment = cn.get_mem_addr_alignment();
64 size_t tot_sz = 0;
65
66 // copy to HostTensorStorage, and release
67 for (auto i : blobs_set_ptr.blobs_set) {
68 // skip if blob do not have m_storage
69 if (!i->m_storage)
70 continue;
71
72 // skip if ues_count() > 1
73 if (i->m_storage.use_count() > 1)
74 continue;
75
76 // two blobs can't share same storage
77 mgb_assert(storage_set.insert(i->m_storage).second);
78
79 tot_sz += get_aligned_power2(i->m_size, alignment);
80 BlobData blob_data(i);
81 blob_data_arrary.push_back(blob_data);
82 i->m_storage.reset();
83 }
84 // clear all, make sure m_storage will be release
85 storage_set.clear();
86
87 // skip if no blob to defrag
88 if (!blob_data_arrary.size())
89 return;
90
91 // wait all other comp nodes to avoid moved var being read; note that
92 // ExecEnv has been paused, so no new task would not be dispatched
93 CompNode::sync_all();
94 CompNode::try_coalesce_all_free_memory();
95
96 // try free all
97 MGB_TRY { cn.free_device(cn.alloc_device(tot_sz)); }
98 MGB_CATCH(MemAllocError&, {})
99
100 // sort blobs by created time, may be helpful for reduce memory fragment

Callers 2

init_utilsFunction · 0.45

Calls 10

use_countMethod · 0.80
get_aligned_power2Function · 0.50
insertMethod · 0.45
push_backMethod · 0.45
resetMethod · 0.45
clearMethod · 0.45
sizeMethod · 0.45
free_deviceMethod · 0.45
alloc_deviceMethod · 0.45

Tested by

no test coverage detected