CUDA番外 -- thrust简介

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

#include <iostream>
#include <cmath>

#include <thrust/device_vector.h>
#include <thrust/host_vector.h>

thrust::device_vector<int> dev_v1(10);
thrust::device_vector<int> dev_v2(10);
thrust::device_vector<int> dev_v3(10);

// 以0到9进行填充
#include <thrust/sequence.h>
thrust::sequence(dev_v1.begin(), dev_v1.end());

// v2=-v1
#include <thrust/transform.h>
#include <thrust/functional.h>
thrust::transform(
    dev_v1.begin(), dev_v1.end(),
    dev_v2.begin(), thrust::negate<int>()
);

// 以2进行填充
#include <thrust/fill.h>
thrust::fill(dev_v2.begin(), dev_v2.end(), 2);

// v3 = v1 % v2
thrust::transform(
    dev_v1.begin(), dev_v1.end(),
    dev_v2.begin(), dev_v3.begin(),
    thrust::modulus<int>()
);

// 查找与替换
#include <thrust/replace.h>
thrust::replace(dev_v3.begin(), dev_v3.end(), 1, 5);

// 输出
#include <thrust/copy.h>
thrust::copy(
    dev_v3.begin(), dev_v3.end(),
    std::ostream_iterator<int>(std::cout, "\n")
);

// 归约，3个一样
#include <thrust/reduce.h>
int sum = thrust::reduce(dev_v3.begin(), dev_v3.end(), int(0), thrust::plus<int>());
int sum = thrust::reduce(dev_v3.begin(), dev_v3.end(), int(0));
int sum = thrust::reduce(dev_v3.begin(), dev_v3.end());

// 数数
#include <thrust/count.h>
int count = thrust::count(dev_v3.begin(), dev_v3.end(), 1);  // 1的个数

// 求模(l2范数)，也可以使用lambda表达式，但要注意cuda版本间的区别
#include <thrust/transform_reduce.h>
template <class T>
class square {
public:
    __host__ __device__ T operator() (const T &x) const {
        return x * x;
    }
};
float norm = std::sqrt(thrust::transform_reduce(
    dev_v3.begin(), dev_v3.end(), sqaure<int>(), 0, thrust::plus<int>())
);

// 累积
#include <thrust/scan.h>
thrust::inclusive_scan(dev_v1.begin(), dev_v1.end(), dev_v1.begin());
thrust::exclusive_scan(dev_v1.begin(), dev_v1.end(), dev_v1.begin());

// 排序
#include <thrust/sort.h>
thrust::sort(dev_v1.begin(), dev_v1.end());

const int N = 6;
int  keys[] = { 1,  4,  2,  8,  5,  7};
char values[] = {'a', 'b', 'c', 'd', 'e', 'f'};
thrust::sort_by_key(keys, keys + N, values);
// now keys is 1, 2, 4, 5, 7, 8
// now values is a, c, b, e, f, d

thrust::stable_sort(dev_v1.begin(), dev_v1.end(), thrust::greater<int>());

1
2
3
4
5
6
7
8
9

size_t N = 10;
int *raw_ptr = 0;
cudaMalloc((void **)&raw_ptr, N * sizeof(int));
thrust::device_ptr<int> dev_ptr(raw_ptr);  // 这是特性，不作这样的转换可能出错
thrust::fill(dev_ptr, dev_ptr + N, int(0));

size_t N = 10;
thrust::device_ptr<int> dev_ptr = thrust::device_malloc<int>(N);
int *raw_ptr = thrust::raw_pointer_cast(dev_ptr);

1
2
3
4

size_t N = 10000;
thrust::device_vector<float> dev_A(N * N);
float *dev_A_raw = thrust::raw_pointer_cast(dev_A.data());
cublas(..., dev_A_raw, ...); 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

#include <thrust/device_ptr.h>
#include <thrust/host_vector.h>
#include <thrust/transform.h>
#include <thrust/functional.h>
#include <thrust/iterator/counting_iterator.h>

struct GenerateKernel : public thrust::unary_function<int, int> {
    __device__ int operator()(int i) const {
        return i;
    }
};

int N = 10000;
const thrust::counting_iterator<int> first(0);
const thrust::counting_iterator<int> last(N);

int *raw_ptr = nullptr;
cudaMalloc(&raw_ptr, sizeof(int) * N);
thrust::transform(first, last, raw_ptr, GenerateKernel{}); // error
thrust::host_vector<int> vec(raw_ptr, raw_ptr + N);        // error

1
2
3
4
5

thrust::transform(first, last, thrust::device_ptr<int>(raw_ptr), GenerateKernel{});
thrust::host_vector<int> vec(
    thrust::device_ptr<int>(raw_ptr),
    thrust::device_ptr<int>(raw_ptr) + N
);