Method concat_and_prepare

src/core/impl/graph/cg_impl_partial.cpp:584–694

}

const OprNodeArray* ComputingGraphImpl::MultiPartCompiler::concat_and_prepare() {
    // no callback in out_spec_concat, so CallbackCaller would not be inserted
    OutputSpec out_spec_concat;
    std::vector<bool> out_spec_concat_from_original;

    // init out_spec_concat
    {
        SymbolVarArray part_vars;
        ExtraDependencyMerger dep_merger;
        for (size_t part = 0; part < m_out_specs.size(); ++part) {
            part_vars.clear();
            for (auto&& i : m_out_specs[part]) {
                part_vars.push_back(i.first);
            }
            auto&& dest_vars = dep_merger.add(part_vars);
            for (size_t i = 0; i < dest_vars.size(); ++i) {
                out_spec_concat.push_back({dest_vars[i].node(), {}});
                out_spec_concat_from_original.push_back(i < part_vars.size());
            }
            dest_vars.clear();
        }
    }

    auto remap_priority = [&](const VarNodeArray& dest_vars,
                              const TopoSorter::PriorityItem* items, size_t nr_item) {
        const size_t nr_part = m_out_specs.size();
        mgb_assert(
                nr_item <= static_cast<size_t>(std::numeric_limits<int>::max()),
                "too many oprs");

        ThinHashMap<const OperatorNodeBase*, size_t> endpoint2part;

        // remap optimized vars to specs and init endpoint2part
        mgb_assert(dest_vars.size() == out_spec_concat.size());
        size_t dest_var_idx = 0;
        auto skip_non_orig_vars = [&](size_t part) {
            while (!out_spec_concat_from_original[dest_var_idx]) {
                ++dest_var_idx;
                if (dest_var_idx == dest_vars.size()) {
                    mgb_assert(part == m_out_specs.size() - 1);
                    break;
                } else {
                    // add extra vars to out spec so we do not need to handle
                    // extra_vardeps in graph copy
                    m_out_specs[part].push_back({dest_vars[dest_var_idx - 1], {}});
                }
            }
        };
        for (size_t part = 0; part < nr_part; ++part) {
            int begin = dest_var_idx;
            for (auto&& i : m_out_specs[part]) {
                mgb_assert(out_spec_concat_from_original[dest_var_idx]);
                i.first = dest_vars[dest_var_idx++];
            }

            if (dest_var_idx < dest_vars.size()) {
                skip_non_orig_vars(part);
            }