Function param_merge

src/gopt/impl/inference.cpp:51–95

template <typename SharedDeviceTensor, typename MultipleDeviceTensorHolder>
void param_merge(OptState& opt_state) {
    auto rewriter = opt_state.graph().make_rewriter();
    ThinHashMap<OperatorNodeBase*, size_t> opr2idx;
    std::vector<OperatorNodeBase*> all_oprs;
    typename MultipleDeviceTensorHolder::ValueArray all_values;

    auto cb_find_opr = [&](cg::OperatorNodeBase* opr) {
        if (opr->same_type<SharedDeviceTensor>()) {
            auto p = &opr->cast_final<SharedDeviceTensor>();
            // ShredD may be manu
            opr2idx[p] = all_values.size();
            all_values.push_back(p->dev_data());
            all_oprs.push_back(p);
        }
    };
    opt_state.graph().iter(cb_find_opr);
    SymbolVarArray new_vars;
    auto cb_replace = [&](cg::OperatorNodeBase* opr) {
        auto iter = opr2idx.find(opr);
        if (iter == opr2idx.end()) {
            rewriter.auto_replace_outputs(opr);
        } else {
            if (new_vars.empty()) {
                // new oprs must be created in iter callback; so we populate
                // new_vars lazily
                new_vars = MultipleDeviceTensorHolder::make(
                        *opt_state.graph().comp_graph(), std::move(all_values),
                        {ssprintf("merged%zu", all_values.size())});
                for (size_t i = 0; i < new_vars.size(); ++i) {
                    auto src = all_oprs[i]->output(0);
                    if (src->has_name_set()) {
                        new_vars[i].rename(src->name());
                    }
                }
            }
            rewriter.replace_var(
                    opr->output(0), new_vars.at(iter->second).node(),
                    mgb_cstr_log("replace multi SharedDeviceTensor(Format) to "
                                 "MultipleDeviceTensorHolder(Format)"));
        }
    };
    opt_state.graph().iter(cb_replace);

    rewriter.apply_inplace();
}

}  // namespace