Function test_param_pack_concat

src/opr/test/tensor_manip.cpp:2083–2137

namespace {
void test_param_pack_concat(const TensorShapeArray& shapes, DType type) {
    auto cn = CompNode::load("xpu0");
    auto graph = ComputingGraph::make();
    auto align = cn.get_mem_addr_alignment() / type.size();

    size_t size = 0;
    std::vector<size_t> begins;
    for (auto&& shape : shapes) {
        size = get_aligned_power2(size, align);
        begins.push_back(size);
        size += shape.total_nr_elems();
    }

    SmallVector<SymbolVar> srcs;
    for (size_t i = 0; i < shapes.size(); i++) {
        auto data = std::make_shared<HostTensorND>();
        data->comp_node(cn).dtype(dtype::Int32()).resize(shapes[i]);
        auto ptr = data->ptr<dt_int32>();
        for (size_t j = 0; j < shapes[i].total_nr_elems(); j++) {
            ptr[j] = j;
        }
        auto nd = opr::Host2DeviceCopy::make(*graph, data);
        srcs.push_back(nd);
    }

    auto host_offsets_gen = megdnn::ParamPackConcat::gen_offsets(
            shapes, cn.get_mem_addr_alignment(), 4);
    ASSERT_EQ(host_offsets_gen.back(), size);
    auto host_offsets = std::make_shared<HostTensorND>();
    host_offsets->comp_node(cn).dtype(dtype::Int32{}).resize({srcs.size() * 2});
    memcpy(host_offsets->raw_ptr(), host_offsets_gen.data(), srcs.size() * 8);
    auto offsets = opr::Host2DeviceCopy::make(*graph, host_offsets);

    auto z = opr::ParamPackConcat::make(srcs, offsets, host_offsets_gen);
    HostTensorND host_z;

    auto func = graph->compile({make_callback_copy(z, host_z)});
    func->execute();

    HostTensorND expected;
    expected.comp_node(cn).dtype(dtype::Int32()).resize({size});
    {
        auto ptr = expected.ptr<dt_int32>();

        memset(ptr, 0, sizeof(int32_t) * size);
        for (size_t i = 0; i < begins.size(); i++) {
            auto begin = begins[i];
            auto shape = shapes[i];
            for (size_t j = 0; j < shape.total_nr_elems(); j++) {
                ptr[begin + j] = j;
            }
        }
    }
    MGB_ASSERT_TENSOR_EQ(expected, host_z);
}

template <size_t nr_out>
void test_param_pack_split(const TensorShapeArray& shapes) {