问从CUDA中可能重复的三重奏列表中组装COO矩阵

EN

#include <thrust/device_ptr.h>
#include <thrust/device_vector.h>
#include <thrust/distance.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/reduce.h>
#include <thrust/sort.h>
#include <thrust/tuple.h>

// Transforms triplets into COO sparse marix format in-place and
// returns the number of non-zero entries (nnz).
auto triplets2coo(int *d_rows,
                  int *d_cols,
                  float *d_vals,
                  int num_triplets) {
    // thrust::device_ptr is a lightweight wrapper that lets Thrust algorithms
    // know which memory the pointer points to. thrust::device_pointer_cast is
    // the preferred way of creating them. While an API directly taking
    // wrapped pointers as arguments would be preferable, I chose it this way 
    // to show how to wrap the pointers.
    auto const coord_iter = thrust::make_zip_iterator(
        thrust::make_tuple(
            thrust::device_pointer_cast(d_rows),
            thrust::device_pointer_cast(d_cols)
        ));
    auto const vals = thrust::device_pointer_cast(d_vals);

    // Sort in-place.
    thrust::sort_by_key(coord_iter, coord_iter + num_triplets,
                        vals);
    // Segmented reduction can only be done out-of-place,
    // so allocate more memory.
    thrust::device_vector<int> tmp_rows(num_triplets);
    thrust::device_vector<int> tmp_cols(num_triplets);
    thrust::device_vector<float> tmp_vals(num_triplets);
    // Caution: These vectors are unnecessarily initialized to zero!

    auto const tmp_coord_iter = thrust::make_zip_iterator(
        thrust::make_tuple(
            tmp_rows.begin(),
            tmp_cols.begin()
        ));

    auto const new_ends =
        thrust::reduce_by_key(coord_iter, coord_iter + num_triplets,
                              vals,
                              tmp_coord_iter,
                              tmp_vals.begin());
    auto const nnz = thrust::distance(tmp_vals.begin(), new_ends.second);

    thrust::copy_n(tmp_coord_iter, nnz, coord_iter);
    thrust::copy_n(tmp_vals.cbegin(), nnz, vals);

    return nnz;
}