Method update_graph

src/jit/impl/fusion_pass.cpp:166–213

}

void JITFusionPass::Impl::update_graph() {
    auto process = [this](OperatorNodeBase* opr) {
        if (!Compiler::is_supported_device(opr->output(0)->comp_node().device_type()))
            return;

        auto fuse_varnode = [this](VarNode* var) {
            auto ig_gen_iter = m_var2igraph_gen.find(var);
            if (ig_gen_iter == m_var2igraph_gen.end()) {
                return;
            }
            auto ig_gen = ig_gen_iter->second;
            if (m_endpoint_set.count(var) != 0 && ig_gen->opr_set().size() >= 2) {
                auto igraph = ig_gen->generate();
                auto&& inputs = ig_gen->orig_inps();
                if (m_after_grad || nr_non_const_vars(inputs) == 1) {
                    // in the forward pass, only fuse oprs with one non-const
                    // inp
                    VarNodeArray rewritten_inputs;
                    for (auto&& input : inputs) {
                        auto new_input = m_rewriter.get_var(input);
                        rewritten_inputs.push_back(new_input);
                    }
                    auto fusion_op = JITExecutor::make(igraph, rewritten_inputs);
                    m_rewriter.replace_var(
                            var, fusion_op.node(),
                            mgb_ssprintf_log(
                                    "fuse endpoint: %s", var->owner_opr()->cname())
                                    .c_str());
                }
            }
        };

        for (auto i : opr->input()) {
            if (!m_rewriter.has_manual_replace(i)) {
                // if input i is a endpoint, and number of oprs in this subgraph
                // is greater than 2
                m_opt_state.call_with_opr(i->owner_opr(), [&] { fuse_varnode(i); });
            }
        }
        m_rewriter.auto_replace_outputs(opr);
        if (m_opt_state.graph().endpoint_contain(opr->output(0))) {
            // process final endpoint
            fuse_varnode(opr->output(0));
        }
    };
    m_opt_state.graph().iter(process);
    m_rewriter.apply_inplace();
}

bool JITFusionPass::Impl::test_all_readers_in_the_graph(
        VarNode* var, InternalGraphGenerator* ig_gen) {