Function apply_on_physical_tensor

imperative/src/impl/ops/concatenate.cpp:107–143

}

SmallVector<TensorPtr> apply_on_physical_tensor(
        const OpDef& def, const SmallVector<TensorPtr>& inputs,
        SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) {
    auto&& op_def = def.cast_final_safe<Concat>();
    int axis = op_def.axis >= 0 ? op_def.axis : op_def.axis + inputs[0]->layout().ndim;

    CompNode& oup_cn = output_descs[0].comp_node;
    TensorLayout& oup_layout = output_descs[0].layout;
    if (validated) {
        if (op_def.comp_node.valid()) {
            mgb_assert(op_def.comp_node == oup_cn, "Concat compnode infer error");
        }
    } else {
        // prepare inputs and output layout
        oup_cn = inputs[0]->comp_node();
        SmallVector<const TensorLayout*> inputs_holder(inputs.size());
        for (size_t i = 0; i < inputs.size(); ++i) {
            inputs_holder[i] = &inputs[i]->layout();
        }
        oup_layout = concat_layout_deduce(inputs_holder, axis);
    }
    auto oup = Tensor::make(oup_layout, oup_cn);
    // because the dnn concat is very slow, we copy the slice code from
    // src/opr/impl/tensor_manip.cpp
    auto&& out = oup->dev_tensor();
    size_t end = 0;
    for (auto&& input : inputs) {
        auto&& in = input->dev_tensor();
        auto begin = end;
        end = begin + in.shape().shape[axis];
        if (!in.layout().is_empty()) {
            out.sub(Slice(begin, end).apply(out.layout(), axis))
                    .copy_from_fixlayout(in);
        }
    }
    return {oup};
}

OP_TRAIT_REG(Concat, Concat)
        .apply_on_var_node(apply_on_var_node)