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

Method postprocess_trace_result

imperative/src/impl/transformations/trace.cpp:307–377  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

305}
306
307void TracingTransformation::postprocess_trace_result() {
308 std::unordered_map<size_t, size_t> identity_oi_map, identity_io_map;
309 for (auto&& op : m_seq) {
310 if (op.op == nullptr && op.inputs.size() == 1 && op.outputs.size() == 1) {
311 identity_oi_map[op.outputs[0]] = op.inputs[0];
312 identity_io_map[op.inputs[0]] = op.outputs[0];
313 }
314 }
315
316 for (auto&& op : m_seq) {
317 if (op.kind == OpKind::IOMarkVar) {
318 auto&& inpvar = m_vars[op.inputs[0]];
319 auto&& outvar = m_vars[op.outputs[0]];
320 if (inpvar.inp_marker.size() > 0) {
321 auto id = inpvar.id;
322 if (inpvar.kind != VarKind::External) {
323 while (identity_oi_map.find(id) != identity_oi_map.end()) {
324 id = identity_oi_map[id];
325 }
326 if (m_vars[id].kind == VarKind::External) {
327 for (auto mark : inpvar.inp_marker) {
328 mgb_assert(
329 inpmark_to_id.find(mark) == inpmark_to_id.end() ||
330 inpmark_to_id[mark] == id,
331 "two nodes have same mark");
332 inpmark_to_id[mark] = id;
333 m_vars[id].inp_marker.insert(mark);
334 }
335 inpvar.inp_marker.clear();
336 }
337 } else {
338 for (auto mark : inpvar.inp_marker) {
339 mgb_assert(
340 inpmark_to_id.find(mark) == inpmark_to_id.end() ||
341 inpmark_to_id[mark] == id,
342 "two nodes have same mark");
343 inpmark_to_id[mark] = id;
344 }
345 }
346 } else {
347 mgb_assert(outvar.out_marker.size() > 0);
348 auto id = outvar.id;
349 if (!outvar.data_required) {
350 while (identity_io_map.find(id) != identity_io_map.end()) {
351 id = identity_io_map[id];
352 }
353
354 if (m_vars[id].data_required) {
355 for (auto mark : outvar.out_marker) {
356 mgb_assert(
357 outmark_to_id.find(mark) == outmark_to_id.end() ||
358 outmark_to_id[mark] == id,
359 "two nodes have same mark");
360 outmark_to_id[mark] = id;
361 m_vars[id].out_marker.insert(mark);
362 }
363 outvar.out_marker.clear();
364 }

Callers 1

exitMethod · 0.80

Calls 5

sizeMethod · 0.45
findMethod · 0.45
endMethod · 0.45
insertMethod · 0.45
clearMethod · 0.45

Tested by

no test coverage detected