Method update_args

src/jit/impl/executor_opr.cpp:204–279

}

void JITExecutor::update_args() {
    m_args.outputs.clear();
    for (auto out : output()) {
        m_args.outputs.push_back({out, out->layout(), -1});
    }
    m_args.inputs.resize(input().size());

    auto is_host_value_shape_input = [this](size_t idx) {
        return m_internal_graph->placeholders().at(idx)->is_host_value_shape_input();
    };

    for (size_t i = 0; i < input().size(); i++) {
        auto&& dst_data = m_args.inputs[i];
        dst_data.from = input(i);
        dst_data.idx = i;
        if (is_host_value_shape_input(i)) {
            auto&& mgr = owner_graph()->static_infer_manager();
            auto&& shpval_inp_val = &mgr.infer_value(input(i));
            cg::copy_tensor_value_to_shape(dst_data.layout, *shpval_inp_val);
            dst_data.layout.dtype = {};
            for (size_t i = 0; i < dst_data.layout.ndim; ++i) {
                dst_data.layout.stride[i] = 0;
            }
        } else {
            dst_data.layout = input(i)->layout();
        }
    }

    //! dimshuffle opr need to change the input.
    if (has_dimshuffle()) {
        do_dimshuffle();
    }

    if (m_compiler->property().contain_flag(CPFlag::NEED_INPUT_COLLAPSE)) {
        // collective collapse datum layout, try to reduce the output ndim
        opr::Elemwise::TensorLayoutPtrArray inp_layouts;
        inp_layouts.reserve(m_args.inputs.size());
        for (size_t i = 0; i < m_args.inputs.size(); i++) {
            if (!is_host_value_shape_input(i)) {
                inp_layouts.push_back(&m_args.inputs[i].layout);
            }
        }
        opr::Elemwise::broadcast_collective_collapse(
                inp_layouts, &m_args.outputs[0].layout);
    }

    // compute and update hash
    XXHash hstate;

    //  update layout info
    auto prop = m_compiler->property();
    if (prop.contain_flag(CPFlag::BIND_NDIM | CPFlag::BIND_SHAPE)) {
        mgb_assert(
                prop.contain_flag(CPFlag::BIND_NDIM),
                "BIND_NDIM must be set if bind_shape is set");
        std::vector<size_t> buf;
        buf.reserve(1024);
        buf.push_back(m_args.inputs.size());