Function test_param_pack_split

src/opr/test/tensor_manip.cpp:2140–2188

template <size_t nr_out>
void test_param_pack_split(const TensorShapeArray& shapes) {
    auto cn = CompNode::load("xpu0");
    auto align = std::max<size_t>(cn.get_mem_addr_alignment() / 4, 1);
    size_t concat_size = 0;
    mgb_assert(shapes.size() == nr_out);
    for (auto&& i : shapes) {
        concat_size = get_aligned_power2(concat_size, align) + i.total_nr_elems();
    }

    using Checker = AutoOprChecker<1, nr_out>;

    auto make_graph = [&](const typename Checker::SymInpArray& inputs) ->
            typename Checker::SymOutArray {
                auto offsets_val = megdnn::ParamPackConcat::gen_offsets(
                        shapes, cn.get_mem_addr_alignment(), 4);
                HostTensorND offsets;
                std::copy_n(
                        offsets_val.data(), offsets_val.size(),
                        offsets.dtype(dtype::Int32{})
                                .comp_node(cn)
                                .resize({offsets_val.size()})
                                .ptr<dt_int32>());
                auto out = opr::ParamPackSplit::make(inputs[0], offsets_val, shapes);
                mgb_assert(out.size() == nr_out);
                typename Checker::SymOutArray ret;
                for (size_t i = 0; i < nr_out; ++i) {
                    ret[i] = out[i];
                }
                return ret;
            };

    auto fwd = [&](typename Checker::NumOutArray& dest,
                   typename Checker::NumInpArray inp) {
        size_t offset = 0;
        auto ptr = inp[0]->template ptr<float>();
        for (size_t i = 0; i < nr_out; ++i) {
            dest[i].resize(shapes[i]);
            offset = get_aligned_power2(offset, align);
            auto nr_elem = shapes[i].total_nr_elems();
            memcpy(dest[i].template ptr<float>(), ptr + offset, nr_elem * 4);
            offset += nr_elem;
        }
    };

    Checker{make_graph, fwd}
            .run({TensorShape{concat_size}})
            .run({TensorShape{concat_size}})
            .run({TensorShape{concat_size}});
}

}  // anonymous namespace