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

Method determine_swap_edge

src/core/impl/graph/swap/memory_swap.cpp:140–413  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

138MemorySwap::~MemorySwap() noexcept = default;
139
140void MemorySwap::determine_swap_edge(
141 PIPSet& heap, size_t loss_idx, const cg::OprNodeArray& opr_seq,
142 std::vector<std::vector<size_t>>& g, std::vector<std::vector<size_t>>& tg) {
143 auto&& infer_mgr = m_owner_graph->static_infer_manager();
144 static_cast<void>(infer_mgr);
145
146 size_t fin = opr_seq.size() + 10;
147 auto segT = new SegmentTree(fin);
148
149 for (auto x : m_var_map) {
150 auto v = x.second;
151 mgb_assert(tg[v->id()].size() == 1);
152 std::vector<size_t> a;
153 std::vector<int> m_consume_opr_set;
154 std::vector<PPI> segment_set;
155 auto sz = m_var_map[x.first]->dtype().size(
156 infer_mgr.infer_shape(m_var_map[x.first]).total_nr_elems());
157 ++m_segment_race_id;
158 size_t last = x.first;
159 std::vector<size_t> b;
160 for (auto y : g[x.first]) {
161 if (m_opr_seq_dist.find(y) == m_opr_seq_dist.end())
162 continue;
163 auto s = m_opr_map[y]->node_prop().dep_map().find(m_var_map[x.first]);
164 if (s == m_opr_map[y]->node_prop().dep_map().end())
165 continue;
166 if (s->second != cg::OperatorNodeBase::NodeProp::DepType::DEV_VALUE)
167 continue;
168 b.push_back(y);
169 }
170 sort(b.begin(), b.end(), [&](const int& lhs, const int& rhs) {
171 return m_opr_seq_dist[lhs] < m_opr_seq_dist[rhs];
172 });
173 for (auto y : b) {
174 if (m_opr_seq_dist.find(y) == m_opr_seq_dist.end())
175 continue;
176 bool flag = true;
177 a.push_back(m_opr_seq_dist[y]);
178 m_consume_opr_set.push_back(y);
179 if (m_opr_seq_dist[last] + 1 <= m_opr_seq_dist[y] - 1) {
180 if (flag) {
181 m_all_valid_segments.push_back(
182 PPI(PII(m_opr_seq_dist[last] + 1, m_opr_seq_dist[y] - 1),
183 m_segment_race_id));
184 }
185 int seg_id = flag ? (int)m_all_valid_segments.size() - 1 : -1;
186 segment_set.push_back(PPI(
187 PII(m_opr_seq_dist[last] + 1, m_opr_seq_dist[y] - 1), seg_id));
188 segT->insert(
189 1, 1, fin, m_opr_seq_dist[last] + 1, m_opr_seq_dist[y] - 1,
190 PII(sz, seg_id));
191 }
192 last = y;
193 }
194 if (a.empty())
195 continue;
196 sort(a.begin(), a.end());
197 auto s = new SegmentRace(

Callers

nothing calls this directly

Calls 15

sortFunction · 0.85
maxFunction · 0.85
query_maxMethod · 0.80
queryMethod · 0.80
sizeMethod · 0.45
idMethod · 0.45
dtypeMethod · 0.45
total_nr_elemsMethod · 0.45
infer_shapeMethod · 0.45
findMethod · 0.45
endMethod · 0.45
push_backMethod · 0.45

Tested by

no test coverage detected