Method infer_output_comp_node

src/opr/impl/custom_opnode.cpp:10–29

MGB_DYN_TYPE_OBJ_FINAL_IMPL(CustomOpNode);

void CustomOpNode::infer_output_comp_node(void) {
    SmallVector<CompNode> input_comp_nodes(input_num());
    for (size_t i = 0; i < input_num(); ++i) {
        input_comp_nodes[i] = input(i)->comp_node();
    }

    SmallVector<CompNode> output_comp_nodes =
            custom::to_builtin<CompNode, custom::Device>(m_op->infer_output_device(
                    custom::to_custom<CompNode, custom::Device>(input_comp_nodes),
                    m_param));

    for (size_t i = 0; i < output_num(); ++i) {
        mgb_assert(
                output_comp_nodes[i] == output_comp_nodes[0],
                "only single comp node operator is supported");
        output(i)->comp_node(output_comp_nodes[i]);
    }

    m_comp_node = output_comp_nodes[0];
}

void CustomOpNode::infer_output_dtype(void) {
    SmallVector<DType> input_dtypes(input_num());