Function test_param_pack_concat

dnn/test/rocm/param_pack.cpp:72–117

template <class T>
void test_param_pack_concat(
        Handle* handle, const TensorShapeArray& shapes, DType type) {
    auto concat = handle->create_operator<ParamPackConcat>();
    size_t nr_params = shapes.size();

    std::vector<T*> param_ptrs;
    std::vector<std::vector<T>> params = create_params<T>(nr_params, shapes);
    for (size_t i = 0; i < nr_params; ++i) {
        param_ptrs.push_back(create_device_data<T>(
                handle, params[i].data(), shapes[i].total_nr_elems()));
    }
    std::vector<int32_t> offsets =
            create_offsets<T>(shapes, handle->alignment_requirement());
    size_t pack_size = offsets.back();
    int32_t* offsets_gpu =
            create_device_data<int32_t>(handle, offsets.data(), offsets.size());

    std::vector<T> expected_pack = create_pack<T>(pack_size, offsets, params);
    T* pack_gpu = create_device_data<T>(handle, nullptr, expected_pack.size());

    TensorLayout dst_layout({pack_size}, type);
    TensorND dst_tensor(pack_gpu, dst_layout);

    TensorLayout offsets_layout({offsets.size()}, dtype::Int32());
    TensorND offsets_tensor(offsets_gpu, offsets_layout);

    test::WorkspaceWrapper workspace(
            handle,
            concat->get_workspace_in_bytes({nr_params}, offsets_layout, {pack_size}));
    TensorND src_tensor(param_ptrs.data(), TensorLayout({nr_params}, dtype::Int32()));

    concat->exec(src_tensor, offsets_tensor, dst_tensor, workspace.workspace());

    // check
    T* actual_pack = static_cast<T*>(malloc(pack_size * sizeof(T)));
    test::megdnn_memcpy_D2H(handle, actual_pack, pack_gpu, sizeof(T) * pack_size);
    for (size_t i = 0; i < pack_size; ++i) {
        ASSERT_EQ(actual_pack[i], expected_pack[i]);
    }
    free(actual_pack);
    test::megdnn_free(handle, pack_gpu);
    test::megdnn_free(handle, offsets_gpu);
    for (auto ptr : param_ptrs) {
        test::megdnn_free(handle, ptr);
    }
}

}  // namespace