Method replace_dev_input_pass

lite/src/mge/network_impl.cpp:367–429

}

void NetworkImplDft::replace_dev_input_pass() {
    mgb::CompNode::Locator locator;
    m_load_config.comp_node_mapper(locator);
    //! CPU is not need use device input
    if (locator.type == mgb::CompNode::DeviceType::CPU) {
        return;
    }
    //! repalce the H2D with VolatileSharedDeviceTensor, and keep the dev tensor
    //! in m_network_io.input, user can directly change the dev tensor
    //! storage through m_network_io.input.lite_tensor->reset() befor forward
    using DeviceTensorMap =
            std::unordered_map<std::string, std::shared_ptr<mgb::DeviceTensorND>>;
    DeviceTensorMap name2dev_tensor;

    mgb::ThinHashMap<mgb::HostTensorND*, mgb::SymbolVar> host_val2var;

    //! construct host_val2var that maps from host tensor to corresponding var
    auto on_opr = [&](mgb::cg::OperatorNodeBase* opr) {
        if (opr->same_type<mgb::opr::Host2DeviceCopy>()) {
            mgb::HostTensorND* tensor =
                    opr->cast_final<mgb::opr::Host2DeviceCopy>().host_data().get();
            host_val2var[tensor] = opr->output(0);
        }
    };
    mgb::cg::DepOprIter dep_iter{on_opr};
    for (auto i : m_load_result.output_var_list) {
        dep_iter.add(i.node()->owner_opr());
    }

    mgb::ThinHashMap<mgb::SymbolVar, mgb::SymbolVar> inp_var_map, out_var_map;

    mgb::SmallVector<std::string> to_clear;
    for (auto&& config_in : m_network_io->inputs) {
        if (!config_in.is_host) {
            auto host_val = m_load_result.tensor_map[config_in.name];
            auto dev_val = TensorHelper::implement(config_in.lite_tensor)
                                   ->cast_final_safe<TensorImplDft>()
                                   .m_dev_tensor;
            auto dev_var = mgb::opr::VolatileSharedDeviceTensor::make(
                    *m_load_result.graph, dev_val, {config_in.name});
            inp_var_map[host_val2var.at(host_val.get())] = dev_var;
            name2dev_tensor[config_in.name] = dev_val;
        }
        //! reset lite_tensor in discrete mode
        if (config_in.name == m_user_config->discrete_input_name) {
            config_in.lite_tensor.reset();
        }
    }
    auto new_ovar = mgb::cg::replace_vars(m_load_result.output_var_list, inp_var_map);
    for (size_t i = 0; i < new_ovar.size(); ++i) {
        out_var_map[m_load_result.output_var_list[i]] = new_ovar[i];
    }
    for (auto&& i : m_load_result.output_var_map) {
        i.second = out_var_map.at(i.second);
    }
    for (auto&& i : m_load_result.output_var_map_id) {
        i.second = out_var_map.at(i.second);
    }