Method apply_action

src/core/impl/graph/seq_sublinear_memory.cpp:975–1027

}

void SeqModifierForSublinearMemory::apply_action(
        SeqModifyAction& action, const OprNodeArray& oprseq) {
    auto cur_priority = std::numeric_limits<
            decltype(OperatorNodeBase::NodeProp::Attribute::priority)>::min();

    ThinHashSet<OperatorNodeBase*> modified_opr;

    // each operator should be set no more than once
    auto set_priority = [&](OperatorNodeBase* opr) {
        mgb_assert(modified_opr.insert(opr).second);
        mem_opt().set_priority(opr, cur_priority++);
    };

    auto on_opr_visited = [&](OperatorNodeBase* opr) {
        if (replace_vars(opr->input())) {
            auto&& repl_info = m_opr2replace_info[opr];
            mgb_assert(
                    !repl_info.recomp, "input of operator %s{%s} already replaced",
                    opr->cname(), opr->dyn_typeinfo()->name);
            opr = copy_opr_from_new_inputs(opr, true);
            repl_info.recomp = opr;
        }
        set_priority(opr);
    };

    // use a DepOprIter rather than directly iterate on oprseq because shape-dep
    // oprs would be omitted in the opr_seq generated by topo sorter; but they
    // should be replaced too
    DepOprIter dep_iter{on_opr_visited};

    // setup m_var_map and priority
    for (auto opr : oprseq) {
        auto iter = action.find(opr);

        if (iter != action.end()) {
            // insert duplicated oprs
            for (auto i : iter->second) {
                replace_vars(i->input());
                auto&& repl_info = m_opr2replace_info[i];
                mgb_assert(
                        !repl_info.dup, "operator %s{%s} already duplicated",
                        i->cname(), i->dyn_typeinfo()->name);
                auto opr_new = copy_opr_from_new_inputs(i, false);
                repl_info.dup = opr_new;
                set_priority(opr_new);
            }
            action.erase(iter);
        }

        dep_iter.add(opr);
    }
    mgb_assert(action.empty());
}

void SeqModifierForSublinearMemory::modify_endpoint_vars(VarNodeArray& endpoints) {
    auto comp_seq = MemoryOptimizerHelper::CompSeq(owner_graph(), endpoints);