Method flush_record_oprs

src/core/impl/graph/eager_eval.cpp:380–440

}

void EagerEvalManager::flush_record_oprs(const VarNodeArray& dest_vars) {
    if (!enabled()) {
        mgb_assert(m_record_oprs.empty());
        return;
    }
    m_record_mode = false;
    using NodeProp = OperatorNodeBase::NodeProp;
    ThinHashSet<OperatorNodeBase*> need_exec_oprs;
    ThinHashSet<OperatorNodeBase*> dest_oprs;
    std::function<void(OperatorNodeBase*)> visit = [&](OperatorNodeBase* opr) {
        if (!m_record_oprs.count(opr) || need_exec_oprs.count(opr))
            return;
        need_exec_oprs.insert(opr);
        for (auto inp : opr->input()) {
            visit(inp->owner_opr());
        }
    };
    for (auto var : dest_vars) {
        dest_oprs.insert(var->owner_opr());
        visit(var->owner_opr());
    }
    for (auto opr : need_exec_oprs) {
        auto&& node_prop = opr->node_prop();
        for (auto&& pair : node_prop.dep_map()) {
            if (NodeProp::is_device_value_dep(pair.second) &&
                need_exec_oprs.count(pair.first->owner_opr())) {
                if (!dest_oprs.count(pair.first->owner_opr()))
                    m_var2nr_readers[pair.first] += opr->output().size();
            }
        }
    }
    SmallVector<std::pair<OperatorNodeBase*, size_t>> stack;
    ThinHashSet<OperatorNodeBase*> instack;
    auto push_stack = [&](OperatorNodeBase* opr) {
        if (need_exec_oprs.erase(opr)) {
            stack.push_back({opr, 0});
            instack.insert(opr);
        } else {
            mgb_assert(instack.count(opr) || m_opr2version.count(opr));
        }
    };
    for (auto&& var : dest_vars) {
        push_stack(var->owner_opr());
    }
    while (!stack.empty()) {
        auto&& frame = stack.back();
        auto opr = frame.first;
        if (frame.second < opr->input().size()) {
            auto var = opr->input()[frame.second++];
            push_stack(var->owner_opr());
        } else {
            MGB_TRY { do_on_opr_insert(opr); }
            MGB_CATCH(MegBrainError & exc, {
                if (!exc.extra_info()) {
                    OperatorNodeExcExtraInfo::record(opr, exc);
                }
                throw;
            });