这是一些示例代码,使用我在您的问题下方的评论中概述的方法来完成您所追求的事情。事实上,我们想使用 4 元组来获取您的键值。在此处复制经过适当修改的评论:
您可以制作一个 zip 迭代器,将您的 3 行压缩在一起加上键“行”并将 4 元组传递给一个特殊的函子。然后,您的特殊仿函数将对 3 元组数组进行归约(也使用键)并返回一个 4 元组的结果。推力点积示例可能会给您一些想法。
这是一种可能的方法:
#include <thrust/host_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/sequence.h>
#include <thrust/fill.h>
#include <thrust/tuple.h>
#define N 30 // make this evenly divisible by 3 for this example
typedef thrust::tuple<int, int, int, int> tpl4int;
typedef thrust::host_vector<int>::iterator intiter;
typedef thrust::tuple<intiter, intiter, intiter, intiter> tpl4intiter;
typedef thrust::zip_iterator<tpl4intiter> int4zip;
struct r3key_unary_op : public thrust::unary_function<tpl4int, tpl4int>
{
__host__ __device__
tpl4int operator()(const tpl4int& x) const
{
tpl4int result;
thrust::get<0>(result) = x.get<0>()*x.get<3>();
thrust::get<1>(result) = x.get<1>()*x.get<3>();
thrust::get<2>(result) = x.get<2>()*x.get<3>();
thrust::get<3>(result) = 1;
return result;
}
};
struct r3key_binary_op : public thrust::binary_function<tpl4int, tpl4int, tpl4int>
{
__host__ __device__
tpl4int operator()(const tpl4int& x, const tpl4int& y) const
{
tpl4int result;
thrust::get<0>(result) = x.get<0>()*x.get<3>() + y.get<0>()*y.get<3>();
thrust::get<1>(result) = x.get<1>()*x.get<3>() + y.get<1>()*y.get<3>();
thrust::get<2>(result) = x.get<2>()*x.get<3>() + y.get<2>()*y.get<3>();
thrust::get<3>(result) = 1;
return result;
}
};
int main() {
thrust::host_vector<int> A(N); // values, in 3 "rows" flattened
thrust::sequence(A.begin(), A.end());
thrust::host_vector<int> K(N/3); // keys in one row
thrust::fill(K.begin(), K.end(), 1); // set some keys to 1
K[9] = 0; // set some keys to zero
int4zip first = thrust::make_zip_iterator(thrust::make_tuple(A.begin(), A.begin() + N/3, A.begin() + 2*N/3, K.begin()));
int4zip last = thrust::make_zip_iterator(thrust::make_tuple(A.begin() + N/3, A.begin() + 2*N/3, A.end(), K.end()));
r3key_unary_op my_unary_op;
r3key_binary_op my_binary_op;
tpl4int init = my_unary_op(*first);
// init = thrust::make_tuple((int) 0, (int) 0, (int) 0, (int) 0);
tpl4int result = thrust::transform_reduce(first, last, my_unary_op, init, my_binary_op);
std::cout << "row 0 = " << result.get<0>() << std::endl;
std::cout << "row 1 = " << result.get<1>() << std::endl;
std::cout << "row 2 = " << result.get<2>() << std::endl;
return 0;
}
笔记:
- 这只是使用
host_vector
. 将其扩展为与 . 一起使用device_vector
,或将其模板化以与其他东西一起使用int
应该很简单。
- 为了完整起见,我使用一元仿函数为每一行的总和减少提供一个非零的初始值。您可能希望将初始值更改为零(零的 4 元组)。