Cuda Thrust获得分段中的前一个元素

Cuda Thrust get previous element in a segment

本文关键字：元素 一个 Thrust 分段 Cuda      更新时间：2023-10-16

我有一个值向量和键向量(表示段)

v = [1, 1, 1, 2, 3, 5, 6]
k = [0, 0, 1, 1, 0, 2, 2]

对于每个元素，我想知道它的前一个元素(在同一段)。它可以是值，也可以是原向量的下标。

所以结果应该是(对于value)

r = [nan, 1, nan, 1, 1, nan, 5]

你可以使用任何元素代替nan，对于算法的剩余部分没有关系。

也许我可以存档它与独家分段扫描和max操作，而不是sum。两个问题:

    我的进近正确吗?
  1. 有更优雅或更有效的解决方案吗?

所需的功能可以通过以下步骤实现:

  1. 按k排序v以获得相邻相等的键值;这必须通过stable_sort_by_key来完成，因为您想要检索"前一个"元素，因此必须保留具有相等键的元素之间的顺序。

  2. 对排序后的数据应用以下转换:

if (previous element has the same key) then return value of previous element else return -1

以下代码实现了这些步骤: 
#include <cstdint>
#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>  
#include <thrust/sort.h>
#include <thrust/transform.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#define PRINTER(name) print(#name, (name))
template <template <typename...> class V, typename T, typename ...Args>
void print(const char* name, const V<T,Args...> & v)
{
    std::cout << name << ":t";
    thrust::copy(v.begin(), v.end(), std::ostream_iterator<T>(std::cout, "t"));
    std::cout << std::endl;
}
template<typename... Iterators>
__host__ __device__
thrust::zip_iterator<thrust::tuple<Iterators...>> zip(Iterators... its)
{
    return thrust::make_zip_iterator(thrust::make_tuple(its...));
}
template <typename IteratorType, typename Integer>
struct prev_value
{
    prev_value(IteratorType first) : first(first){}
   template <typename Tuple>
   __host__ __device__
   Integer operator()(const Tuple& t)
   {
      const auto& index = thrust::get<0>(t);
      const auto& previousValue = thrust::get<1>(t);
      Integer result = -1;
      const auto& currentKey = *(first+index);
      const auto& previousKey = *(first+index-1);
      if(currentKey == previousKey)
      {
          result = previousValue;
      }
      return result;
   }
   IteratorType first;
};
template <typename Integer, typename IteratorType>
prev_value<IteratorType, Integer> make_prev_value(IteratorType first)
{
  return prev_value<IteratorType, Integer>(first);
}

int main(int argc, char** argv)
{
    using Integer = std::int32_t;
    using HostVec = thrust::host_vector<Integer>;
    using DeviceVec = thrust::device_vector<Integer>;
    Integer v[] = {1, 1, 1, 2, 3, 5, 6};
    Integer k[] = {0, 0, 1, 1, 0, 2, 2};
    Integer size = sizeof(k)/sizeof(k[0]);
    HostVec h_k(k, k+size);
    HostVec h_v(v, v+size);
    // copy data to device
    DeviceVec d_k = h_k;
    DeviceVec d_v = h_v;
    std::cout << "---- input data ----" << std::endl;
    PRINTER(d_k);    
    PRINTER(d_v);
    thrust::stable_sort_by_key(d_k.begin(), d_k.end(), d_v.begin());
    std::cout << "---- after sorting ----" << std::endl;
    PRINTER(d_k);    
    PRINTER(d_v);
    DeviceVec d_r(size, -1);
    auto op = make_prev_value<Integer>(d_k.begin());
    thrust::transform(zip(thrust::make_counting_iterator(Integer(1)), d_v.begin()),
                      zip(thrust::make_counting_iterator(size), d_v.end()),
                      d_r.begin()+1,
                      op);
    std::cout << "---- result ----" << std::endl;
    PRINTER(d_r);
    return 0;
}
输出:

---- input data ----
d_k:    0   0   1   1   0   2   2   
d_v:    1   1   1   2   3   5   6   
---- after sorting ----
d_k:    0   0   0   1   1   2   2   
d_v:    1   1   3   1   2   5   6   
---- result ----
d_r:    -1  1   1   -1  1   -1  5
相关文章：