Шаблон CUB похож на тяги

Можно ли получить эту функциональность с помощью CUB?

Да, возможно сделать что-то подобное, используя CUB. Большая часть того, что вам нужно, содержится Вот в приведенном ниже фрагменте кода для уменьшения суммы. Кроме того, CUB не выполняет автоматическое копирование количеств обратно в код хоста, поэтому нам нужно это контролировать. Вот одна из возможных реализаций:

$ cat t125.cu
#include <thrust/reduce.h>
#include <thrust/execution_policy.h>
#include <thrust/device_vector.h>
#include <cub/cub.cuh>
#include <iostream>

typedef int mytype;

const int dsize = 10;
const int val  = 1;template <typename T>
T my_cub_reduce(T *begin, T *end){

size_t num_items = end-begin;
T *d_in = begin;
T *d_out, res;
cudaMalloc(&d_out, sizeof(T));
void     *d_temp_storage = NULL;
size_t   temp_storage_bytes = 0;
cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);
// Allocate temporary storage
cudaMalloc(&d_temp_storage, temp_storage_bytes);
// Run sum-reduction
cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items);
cudaMemcpy(&res, d_out, sizeof(T), cudaMemcpyDeviceToHost);
cudaFree(d_out);
cudaFree(d_temp_storage);
return res;
}

template <typename T>
typename thrust::iterator_traits<T>::value_type
my_cub_reduce(T begin, T end){

return my_cub_reduce(thrust::raw_pointer_cast(&(begin[0])), thrust::raw_pointer_cast(&(end[0])));
}

int main(){

mytype *d_data, *h_data;
cudaMalloc(&d_data, dsize*sizeof(mytype));
h_data = (mytype *)malloc(dsize*sizeof(mytype));
for (int i = 0; i < dsize; i++) h_data[i] = val;
cudaMemcpy(d_data, h_data, dsize*sizeof(mytype), cudaMemcpyHostToDevice);
std::cout << "thrust reduce: " << thrust::reduce(thrust::device, d_data, d_data+dsize) << std::endl;
std::cout << "cub reduce:    " << my_cub_reduce(d_data, d_data+dsize) << std::endl;
thrust::device_vector<int> d(5,1);
// using thrust style container iterators and pointers
std::cout << my_cub_reduce(d.begin(), d.end()) << std::endl;
std::cout << my_cub_reduce(thrust::device_pointer_cast(d.data()), thrust::device_pointer_cast(d.data()+d.size())) << std::endl;
}
$ nvcc -arch=sm_61 -o t125 t125.cu
$ ./t125
thrust reduce: 10
cub reduce:    10
5
5
$

РЕДАКТИРОВАТЬ: добавив несколько дополнительных строк кода, мы можем добавить поддержку итераторов и указателей контейнера в стиле тяги. Я обновил код выше, чтобы продемонстрировать это.