Method update_graph

src/tensorrt/impl/opr_replace.cpp:1432–1600

}

void TensorRTReplacePass::Impl::update_graph() {
    using GpuAllocator = opr::TensorRTOpr::GpuAllocator;
    using TensorRTOpr = opr::TensorRTOpr;

    std::shared_ptr<GpuAllocator> gpu_allocator;
    std::shared_ptr<ExtraDep> extra_dep = std::make_shared<ExtraDep>();

    // construct trt network
    auto construct_network = [this, &gpu_allocator, &extra_dep](OperatorNodeBase* opr) {
        if (!has_fail_msg(opr).valid()) {
            auto cn = opr->output(0)->comp_node();
            auto trt_graph = m_tensorrt_graphs[m_graph_map[opr] - 1];
            auto b = trt_graph->builder;
            mgb_assert(b != nullptr);
            if (!gpu_allocator) {
                gpu_allocator = std::make_shared<GpuAllocator>(cn);
                b->setGpuAllocator(gpu_allocator.get());
            } else {
                auto cn0 = gpu_allocator->comp_node();
                mgb_assert(
                        cn0 == cn,
                        "multiple comp nodes for trt graph are not "
                        "supported: %s %s",
                        cn0.to_string().c_str(), cn.to_string().c_str());
            }

            if (!trt_graph->network) {
#if NV_TENSOR_RT_VERSION >= 6001
                nvinfer1::NetworkDefinitionCreationFlags flags;
                flags = 1 << static_cast<int>(nvinfer1::NetworkDefinitionCreationFlag::
                                                      kEXPLICIT_BATCH);
                trt_graph->network = b->createNetworkV2(flags);
#else
                trt_graph->network = b->createNetwork();
#endif
            }
            // make extra dep
            for (auto&& inp : trt_graph->inputs) {
                extra_dep->operator[](opr).push_back(inp);
            }

            auto iter = m_opr_trait.find(opr->dyn_typeinfo());
            if (iter != m_opr_trait.end()) {
                if (iter->second.add_to_nvinfer) {
                    iter->second.add_to_nvinfer(trt_graph->network, opr);
                }
            }
        }
    };
    m_opt_state.graph().iter(construct_network);

    // trt network markOutput
    for (auto trt_graph : m_tensorrt_graphs) {
        // record traverse order
        size_t idx = 0;
        auto&& varnode2itensor = trt_graph->varnode2itensor;
        for (auto output : trt_graph->outputs) {
            trt_graph->output2idx[output] = idx++;