让您的索引(偏移)数组为 idx[]。让你的数据数组为 A[],让扫描的结果在 B[] 中。
扫描整个数组 A[],将输出存储在 B[] 中。
对于 idx[i] 处的每个元素,转到 B[] 中的该索引负 1,检索该值,然后使用 idx[i-1] 处的元素来索引 B[] 中的负 1 并减去该值,然后减去A[] 中相同索引 idx[i](不负 1)的结果。
重新扫描 A 到 B。
举个简单的例子:
idx: 0 2 5
0: 1 1 1 1 1 1 1 1
1: 1 2 3 4 5 6 7 8
2: 1 1 -1 1 1 -2 1 1
3: 1 2 1 2 3 1 2 3
在上面的示例中,步骤 2 中的 -1 计算为步骤 1 中索引 (2-1) 处的扫描值减去步骤 1 中索引 (0-1) 处的扫描值(假设为零),然后从原始数据值中减去。步骤 2 中的 -2 计算为步骤 1 中索引 (5-1) 处的扫描值减去步骤 1 中索引 (2-1) 处的扫描值,从原始数据值中减去。
这是一个例子:
$ cat t453.cu
#include <cub/cub.cuh>
#include <iostream>
template <int TPB, int IPT, typename T>
__global__ void k(T *data, int *idx, int n){
// Specialize BlockScan for a 1D block of TPB threads on type T
__shared__ T sdata[TPB*IPT*2];
sdata[threadIdx.x*IPT] = 1;
__syncthreads();
typedef cub::BlockScan<T, TPB> BlockScan;
// Allocate shared memory for BlockScan
__shared__ typename BlockScan::TempStorage temp_storage;
// Obtain a segment of consecutive items that are blocked across threads
int thread_data[IPT];
thread_data[0] = sdata[threadIdx.x*IPT];
// Collectively compute the block-wide exclusive prefix sum
BlockScan(temp_storage).InclusiveSum(thread_data, thread_data);
__syncthreads();
sdata[IPT*(threadIdx.x+TPB)] = thread_data[0];
if ((threadIdx.x < n) && (threadIdx.x > 0)) // assume the first element if idx points to 0
sdata[idx[threadIdx.x]*IPT] -= (sdata[((idx[threadIdx.x]-1)+TPB)*IPT] - ((threadIdx.x == 1)?0:sdata[((idx[threadIdx.x-1]-1)+TPB)*IPT]));
__syncthreads();
thread_data[0] = sdata[threadIdx.x*IPT];
BlockScan(temp_storage).InclusiveSum(thread_data, thread_data);
__syncthreads();
data[threadIdx.x] = thread_data[0];
}
typedef int dtype;
const int nTPB = 256;
int main(){
int h_idx[] = {0, 4, 7, 32, 55, 99, 104, 200};
int n = sizeof(h_idx)/sizeof(h_idx[0]);
std::cout << "n = " << n << std::endl;
int *d_idx;
cudaMalloc(&d_idx, n*sizeof(d_idx[0]));
cudaMemcpy(d_idx, h_idx, n*sizeof(h_idx[0]), cudaMemcpyHostToDevice);
dtype *h_data, *d_data;
h_data = new dtype[nTPB];
cudaMalloc(&d_data, nTPB*sizeof(dtype));
k<nTPB, 1><<<1,nTPB>>>(d_data, d_idx, n);
cudaMemcpy(h_data, d_data, nTPB*sizeof(dtype), cudaMemcpyDeviceToHost);
dtype sum;
int idx = 0;
for (int i = 0; i < nTPB; i++){
if (i == h_idx[idx]) {sum = 0; idx++;}
sum++;
std::cout << "gpu: " << h_data[i] << " cpu: " << sum << std::endl;
}
}
$ nvcc -o t453 t453.cu
$ cuda-memcheck ./t453
========= CUDA-MEMCHECK
n = 8
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 26 cpu: 26
gpu: 27 cpu: 27
gpu: 28 cpu: 28
gpu: 29 cpu: 29
gpu: 30 cpu: 30
gpu: 31 cpu: 31
gpu: 32 cpu: 32
gpu: 33 cpu: 33
gpu: 34 cpu: 34
gpu: 35 cpu: 35
gpu: 36 cpu: 36
gpu: 37 cpu: 37
gpu: 38 cpu: 38
gpu: 39 cpu: 39
gpu: 40 cpu: 40
gpu: 41 cpu: 41
gpu: 42 cpu: 42
gpu: 43 cpu: 43
gpu: 44 cpu: 44
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 26 cpu: 26
gpu: 27 cpu: 27
gpu: 28 cpu: 28
gpu: 29 cpu: 29
gpu: 30 cpu: 30
gpu: 31 cpu: 31
gpu: 32 cpu: 32
gpu: 33 cpu: 33
gpu: 34 cpu: 34
gpu: 35 cpu: 35
gpu: 36 cpu: 36
gpu: 37 cpu: 37
gpu: 38 cpu: 38
gpu: 39 cpu: 39
gpu: 40 cpu: 40
gpu: 41 cpu: 41
gpu: 42 cpu: 42
gpu: 43 cpu: 43
gpu: 44 cpu: 44
gpu: 45 cpu: 45
gpu: 46 cpu: 46
gpu: 47 cpu: 47
gpu: 48 cpu: 48
gpu: 49 cpu: 49
gpu: 50 cpu: 50
gpu: 51 cpu: 51
gpu: 52 cpu: 52
gpu: 53 cpu: 53
gpu: 54 cpu: 54
gpu: 55 cpu: 55
gpu: 56 cpu: 56
gpu: 57 cpu: 57
gpu: 58 cpu: 58
gpu: 59 cpu: 59
gpu: 60 cpu: 60
gpu: 61 cpu: 61
gpu: 62 cpu: 62
gpu: 63 cpu: 63
gpu: 64 cpu: 64
gpu: 65 cpu: 65
gpu: 66 cpu: 66
gpu: 67 cpu: 67
gpu: 68 cpu: 68
gpu: 69 cpu: 69
gpu: 70 cpu: 70
gpu: 71 cpu: 71
gpu: 72 cpu: 72
gpu: 73 cpu: 73
gpu: 74 cpu: 74
gpu: 75 cpu: 75
gpu: 76 cpu: 76
gpu: 77 cpu: 77
gpu: 78 cpu: 78
gpu: 79 cpu: 79
gpu: 80 cpu: 80
gpu: 81 cpu: 81
gpu: 82 cpu: 82
gpu: 83 cpu: 83
gpu: 84 cpu: 84
gpu: 85 cpu: 85
gpu: 86 cpu: 86
gpu: 87 cpu: 87
gpu: 88 cpu: 88
gpu: 89 cpu: 89
gpu: 90 cpu: 90
gpu: 91 cpu: 91
gpu: 92 cpu: 92
gpu: 93 cpu: 93
gpu: 94 cpu: 94
gpu: 95 cpu: 95
gpu: 96 cpu: 96
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 26 cpu: 26
gpu: 27 cpu: 27
gpu: 28 cpu: 28
gpu: 29 cpu: 29
gpu: 30 cpu: 30
gpu: 31 cpu: 31
gpu: 32 cpu: 32
gpu: 33 cpu: 33
gpu: 34 cpu: 34
gpu: 35 cpu: 35
gpu: 36 cpu: 36
gpu: 37 cpu: 37
gpu: 38 cpu: 38
gpu: 39 cpu: 39
gpu: 40 cpu: 40
gpu: 41 cpu: 41
gpu: 42 cpu: 42
gpu: 43 cpu: 43
gpu: 44 cpu: 44
gpu: 45 cpu: 45
gpu: 46 cpu: 46
gpu: 47 cpu: 47
gpu: 48 cpu: 48
gpu: 49 cpu: 49
gpu: 50 cpu: 50
gpu: 51 cpu: 51
gpu: 52 cpu: 52
gpu: 53 cpu: 53
gpu: 54 cpu: 54
gpu: 55 cpu: 55
gpu: 56 cpu: 56
========= ERROR SUMMARY: 0 errors
$
这仍然需要您将数组的“结束”填充到线程块大小。根据您的描述,我假设这应该是可能的,无论如何它对 cub 来说基本上是必要的;cub 期望使用线程块中的每个线程。
对于较大的数组,上述方法可以直接扩展为使用DeviceScan。第 1 步是第一次扫描。第 2 步将是单独的内核启动。第 3 步是第二次扫描。
如果你想让每个线程块对一个段执行扫描,你不需要填充每个段。您只需要填充数组的“结束”,这样最后一次扫描就可以了,甚至这个“填充”操作也可以通过条件加载来完成,而不是实际的填充操作。这是一个例子:
$ cat t455.cu
#include <cub/cub.cuh>
#include <iostream>
template <int TPB, int IPT, typename T>
__global__ void k(T *data, int *idx){
int lidx = threadIdx.x;
// Specialize BlockScan for a 1D block of TPB threads on type T
typedef cub::BlockScan<T, TPB> BlockScan;
// Allocate shared memory for BlockScan
__shared__ typename BlockScan::TempStorage temp_storage;
// Obtain a segment of consecutive items that are blocked across threads
int thread_data[IPT];
thread_data[0] = ((lidx+idx[blockIdx.x])>=idx[blockIdx.x+1])?0:data[lidx+idx[blockIdx.x]];
// Collectively compute the block-wide inclusive prefix sum
BlockScan(temp_storage).InclusiveSum(thread_data, thread_data);
__syncthreads();
if ((lidx+idx[blockIdx.x]) < idx[blockIdx.x+1])
data[lidx+idx[blockIdx.x]] = thread_data[0];
}
typedef int dtype;
const int nTPB = 128; // sized with IPT to handle the largest segment
const int DS = 256;
int main(){
int h_idx[] = {0, 4, 7, 32, 55, 99, 104, 200, 256};
int n = sizeof(h_idx)/sizeof(h_idx[0]);
std::cout << "n = " << n << std::endl;
int *d_idx;
cudaMalloc(&d_idx, n*sizeof(d_idx[0]));
cudaMemcpy(d_idx, h_idx, n*sizeof(h_idx[0]), cudaMemcpyHostToDevice);
dtype *h_data, *d_data;
h_data = new dtype[DS];
for (int i = 0; i < DS; i++) h_data[i] = 1;
cudaMalloc(&d_data, DS*sizeof(dtype));
cudaMemcpy(d_data, h_data, DS*sizeof(h_data[0]), cudaMemcpyHostToDevice);
k<nTPB, 1><<<n-1,nTPB>>>(d_data, d_idx);
cudaMemcpy(h_data, d_data, DS*sizeof(dtype), cudaMemcpyDeviceToHost);
dtype sum;
int idx = 0;
for (int i = 0; i < DS; i++){
if (i == h_idx[idx]) {sum = 0; idx++;}
sum++;
std::cout << "gpu: " << h_data[i] << " cpu: " << sum << std::endl;
}
}
$ nvcc -o t455 t455.cu
$ cuda-memcheck ./t455
========= CUDA-MEMCHECK
n = 9
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 26 cpu: 26
gpu: 27 cpu: 27
gpu: 28 cpu: 28
gpu: 29 cpu: 29
gpu: 30 cpu: 30
gpu: 31 cpu: 31
gpu: 32 cpu: 32
gpu: 33 cpu: 33
gpu: 34 cpu: 34
gpu: 35 cpu: 35
gpu: 36 cpu: 36
gpu: 37 cpu: 37
gpu: 38 cpu: 38
gpu: 39 cpu: 39
gpu: 40 cpu: 40
gpu: 41 cpu: 41
gpu: 42 cpu: 42
gpu: 43 cpu: 43
gpu: 44 cpu: 44
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 26 cpu: 26
gpu: 27 cpu: 27
gpu: 28 cpu: 28
gpu: 29 cpu: 29
gpu: 30 cpu: 30
gpu: 31 cpu: 31
gpu: 32 cpu: 32
gpu: 33 cpu: 33
gpu: 34 cpu: 34
gpu: 35 cpu: 35
gpu: 36 cpu: 36
gpu: 37 cpu: 37
gpu: 38 cpu: 38
gpu: 39 cpu: 39
gpu: 40 cpu: 40
gpu: 41 cpu: 41
gpu: 42 cpu: 42
gpu: 43 cpu: 43
gpu: 44 cpu: 44
gpu: 45 cpu: 45
gpu: 46 cpu: 46
gpu: 47 cpu: 47
gpu: 48 cpu: 48
gpu: 49 cpu: 49
gpu: 50 cpu: 50
gpu: 51 cpu: 51
gpu: 52 cpu: 52
gpu: 53 cpu: 53
gpu: 54 cpu: 54
gpu: 55 cpu: 55
gpu: 56 cpu: 56
gpu: 57 cpu: 57
gpu: 58 cpu: 58
gpu: 59 cpu: 59
gpu: 60 cpu: 60
gpu: 61 cpu: 61
gpu: 62 cpu: 62
gpu: 63 cpu: 63
gpu: 64 cpu: 64
gpu: 65 cpu: 65
gpu: 66 cpu: 66
gpu: 67 cpu: 67
gpu: 68 cpu: 68
gpu: 69 cpu: 69
gpu: 70 cpu: 70
gpu: 71 cpu: 71
gpu: 72 cpu: 72
gpu: 73 cpu: 73
gpu: 74 cpu: 74
gpu: 75 cpu: 75
gpu: 76 cpu: 76
gpu: 77 cpu: 77
gpu: 78 cpu: 78
gpu: 79 cpu: 79
gpu: 80 cpu: 80
gpu: 81 cpu: 81
gpu: 82 cpu: 82
gpu: 83 cpu: 83
gpu: 84 cpu: 84
gpu: 85 cpu: 85
gpu: 86 cpu: 86
gpu: 87 cpu: 87
gpu: 88 cpu: 88
gpu: 89 cpu: 89
gpu: 90 cpu: 90
gpu: 91 cpu: 91
gpu: 92 cpu: 92
gpu: 93 cpu: 93
gpu: 94 cpu: 94
gpu: 95 cpu: 95
gpu: 96 cpu: 96
gpu: 1 cpu: 1
gpu: 2 cpu: 2
gpu: 3 cpu: 3
gpu: 4 cpu: 4
gpu: 5 cpu: 5
gpu: 6 cpu: 6
gpu: 7 cpu: 7
gpu: 8 cpu: 8
gpu: 9 cpu: 9
gpu: 10 cpu: 10
gpu: 11 cpu: 11
gpu: 12 cpu: 12
gpu: 13 cpu: 13
gpu: 14 cpu: 14
gpu: 15 cpu: 15
gpu: 16 cpu: 16
gpu: 17 cpu: 17
gpu: 18 cpu: 18
gpu: 19 cpu: 19
gpu: 20 cpu: 20
gpu: 21 cpu: 21
gpu: 22 cpu: 22
gpu: 23 cpu: 23
gpu: 24 cpu: 24
gpu: 25 cpu: 25
gpu: 26 cpu: 26
gpu: 27 cpu: 27
gpu: 28 cpu: 28
gpu: 29 cpu: 29
gpu: 30 cpu: 30
gpu: 31 cpu: 31
gpu: 32 cpu: 32
gpu: 33 cpu: 33
gpu: 34 cpu: 34
gpu: 35 cpu: 35
gpu: 36 cpu: 36
gpu: 37 cpu: 37
gpu: 38 cpu: 38
gpu: 39 cpu: 39
gpu: 40 cpu: 40
gpu: 41 cpu: 41
gpu: 42 cpu: 42
gpu: 43 cpu: 43
gpu: 44 cpu: 44
gpu: 45 cpu: 45
gpu: 46 cpu: 46
gpu: 47 cpu: 47
gpu: 48 cpu: 48
gpu: 49 cpu: 49
gpu: 50 cpu: 50
gpu: 51 cpu: 51
gpu: 52 cpu: 52
gpu: 53 cpu: 53
gpu: 54 cpu: 54
gpu: 55 cpu: 55
gpu: 56 cpu: 56
========= ERROR SUMMARY: 0 errors
$