Method auto_replace_outputs

src/gopt/impl/framework.cpp:31–90

/* ================ SubGraph ================ */

OperatorNodeBase* SubGraph::Rewriter::auto_replace_outputs(OperatorNodeBase* opr) {
    auto&& new_inp = m_opr_new_inp_cache;
    new_inp.clear();
    new_inp.reserve(opr->input().size());
    bool has_replaced_inp = false;

    for (auto i : opr->input()) {
        auto new_var = get_var(i);
        if (new_var != i) {
            has_replaced_inp = true;
            new_inp.push_back(new_var);
        } else {
            new_inp.push_back(i);
        }
    }

    if (has_replaced_inp) {
        auto new_opr = serialization::copy_opr_shallow(*opr, new_inp, opr->config());
        auto &&out0 = opr->output(), &&out1 = new_opr->output();
        size_t i = 0;
        auto err_msg = [opr, new_opr] {
            return ssprintf(
                    "bad opr copy: src=%s{%s} dst=%s{%s}", opr->cname(),
                    opr->dyn_typeinfo()->name, new_opr->cname(),
                    new_opr->dyn_typeinfo()->name);
        };
        MGB_MARK_USED_VAR(err_msg);
        // opr output size mismatch may be caused by:
        //     0) inplace arith optimization (e.g. PowC need an extra workspace)
        //     1) other post-insert optimization (e.g. const folding)
        // we can't handle only usable_output here, since some output var with
        // volatile flag could be the graph's endpoint (e.g. RemoteSend)
        for (; i < std::min(out0.size(), out1.size()); ++i) {
            bool v0 = out0[i]->contain_flag(VarNode::Flag::VOLATILE_CONTENT),
                 v1 = out1[i]->contain_flag(VarNode::Flag::VOLATILE_CONTENT);
            mgb_assert(v0 == v1, "%s", err_msg().c_str());
            //! rename new var
            out1[i]->name(out0[i]->cname());
            auto&& ins = m_varmap.insert({out0[i], {true, nullptr}});
            mgb_assert(
                    ins.second || ins.first->second.first,
                    "opr output already replaced");
            // handle repeated call on the same opr
            ins.first->second.second = out1[i];
            on_var_replaced(out0[i], out1[i], nullptr);
        }
        for (; i < out0.size(); ++i) {
            mgb_assert(
                    out0[i]->contain_flag(VarNode::Flag::VOLATILE_CONTENT), "%s",
                    err_msg().c_str());
        }
        for (; i < out1.size(); ++i) {
            mgb_assert(
                    out1[i]->contain_flag(VarNode::Flag::VOLATILE_CONTENT), "%s",
                    err_msg().c_str());
        }
        return new_opr;
    }