基本的 CUDA 加载和扭曲转置

Sim*_*mon 0 cuda

我想使用 CUDA 9.0 的 shuffle 操作实现基本的阻塞加载和扭曲转置。我知道 cub 和 trove 实现,但我仅限于使用 nvrtc 进行编译,并且标准头文件包含使这些库难以满足。我不是在寻找任何花哨的东西,只是对维度为 2 的幂的数据进行一些整数、浮点和双洗牌。

可视化经线大小为 8 的示例,我想从:

             correlation
             0    1    2    3

lane 0       0    8   16   24
lane 1       1    9   17   25
lane 2       2   10   18   26
lane 3       3   11   19   27
lane 4       4   12   20   28
lane 5       5   13   21   29
lane 6       6   14   22   30 
lane 7       7   15   23   31 
Run Code Online (Sandbox Code Playgroud)

到这个结构:

             correlation
             0    1    2    3

lane 0       0    1    2    3
lane 1       8    9   10   11
lane 2       16  17   18   19
lane 3       24  25   26   27 
lane 4       4    5    6    7
lane 5       12  13   14   15
lane 6       20  21   22   23
lane 7       28  29   30   31 
Run Code Online (Sandbox Code Playgroud)

我觉得这应该很简单,但我无法弄清楚我做错了什么。我认为基本的换位循环应该如下所示:

int loads[ncorrs];
int values[ncorrs];
int lane_id = threadIdx.x & (warp_size - 1);
// 0 0 0 0 4 4 4 4 8 8 8 8 ....
int base_idx = lane_id & (warp_size - ncorrs);
// 0 1 2 3 0 1 2 3 0 1 2 3
int src_corr = lane_id & (ncorrs - 1);

for(int corr=0; corr < ncorrs; ++corr)
{
    int src_lane = base_idx + corr;
    values[corr] = __shfl_sync(mask, loads[src_corr],
                                 src_lane, warp_size);
}
Run Code Online (Sandbox Code Playgroud)

因此,鉴于上面的示例数据,如果我们在第 5 道,我预计应该发生以下索引:

base_idx == 4;
src_corr == 1;

corr == [0, 1, 2, 3]
src_lane == [4, 5, 6, 7]
values == [12, 13, 14 15]
Run Code Online (Sandbox Code Playgroud)

但是发生了以下情况(33 个来自数据的后面):

             correlation
             0    1    2    3

lane 0       0    0    0    0
lane 1       4    4    4    4
lane 2       12  12   12   12
lane 3       16  16   16   16
lane 4       20  20   20   20
lane 5       24  24   24   24
lane 6       28  28   28   28 
lane 7       33  33   33   33 
Run Code Online (Sandbox Code Playgroud)

我做错了什么?经纱尺寸为 32 的完整实现:

#include <cstdlib>
#include <cstdio>

#include "cuda.h"

#define ncorr 4
#define warp_size 32

template <int ncorrs>
__global__ void kernel(
    int * input,
    int * output,
    int N)
{
    // This should provide 0 0 0 0 4 4 4 4 8 8 8 8 ...
    #define base_idx(lane_id) (lane_id & (warp_size - ncorrs))
    // This should provide 0 1 2 3 0 1 2 3 0 1 2 3
    #define corr_idx(lane_id) (lane_id & (ncorrs - 1))


    int n = blockIdx.x*blockDim.x + threadIdx.x;
    int lane_id = threadIdx.x & (warp_size - 1);

    if(n >= N)
        { return; }

    // Input correlation handled by this thread
    int src_corr = corr_idx(lane_id);
    int mask = __activemask();

    if(threadIdx.x == 0)
        { printf("mask %d\n", mask); }

    int loads[ncorrs];
    int values[ncorrs];

    #pragma unroll (ncorrs)
    for(int corr=0; corr < ncorrs; ++corr)
        { loads[corr] = input[n + corr*N]; }

    __syncthreads();

    printf("[%d, %d] %d %d %d %d\n",
           lane_id, base_idx(lane_id),
           loads[0], loads[1],
           loads[2], loads[3]);

    #pragma unroll (ncorrs)
    for(int corr=0; corr < ncorrs; ++corr)
    {
        int src_lane = base_idx(lane_id) + corr;
        values[corr] = __shfl_sync(mask, loads[src_corr],
                                     src_lane, warp_size);
    }

    printf("[%d, %d] %d %d %d %d\n",
           lane_id, base_idx(lane_id),
           values[0], values[1],
           values[2], values[3]);


    #pragma unroll (ncorrs)
    for(int corr=0; corr < ncorrs; ++corr)
        { output[n + corr*N] = values[corr]; }
}

void print_data(int * data, int N)
{
    for(int n=0; n < N; ++n)
    {
        printf("% -3d: ", n);
        for(int c=0; c < ncorr; ++c)
        {
            printf("%d ", data[n*ncorr + c]);
        }
        printf("\n");
    }
}

int main(void)
{
    int * host_input;
    int * host_output;

    int * device_input;
    int * device_output;
    int N = 32;

    host_input = (int *) malloc(sizeof(int)*N*ncorr);
    host_output = (int *) malloc(sizeof(int)*N*ncorr);

    printf("malloc done\n");

    cudaMalloc((void **) &device_input, sizeof(int)*N*ncorr);
    cudaMalloc((void **) &device_output, sizeof(int)*N*ncorr);

    printf("cudaMalloc done\n");

    for(int i=0; i < N*ncorr; ++i)
        { host_input[i] = i; }

    print_data(host_input, N);

    dim3 block(256, 1, 1);
    dim3 grid((block.x + N - 1) / N, 1, 1);

    cudaMemcpy(device_input, host_input,
               sizeof(int)*N*ncorr, cudaMemcpyHostToDevice);

    printf("memcpy done\n");

    kernel<4> <<<grid, block>>> (device_input, device_output, N);

    cudaMemcpy(host_output, device_output,
               sizeof(int)*N*ncorr, cudaMemcpyDeviceToHost);

    print_data(host_output, N);

    cudaFree(device_input);
    cudaFree(device_output);

    free(host_input);
    free(host_output);
}
Run Code Online (Sandbox Code Playgroud)

编辑 1:澄清了视觉示例的扭曲大小为 8,而完整代码的扭曲大小为 32

Rob*_*lla 5

我做错了什么?

TL;DR:简而言之,您将相同的输入值传输到多个输出值。这是一个示例,在这行代码中:

    values[corr] = __shfl_sync(mask, loads[src_corr],
                                 src_lane, warp_size);
Run Code Online (Sandbox Code Playgroud)

表示的数量loads[src_corr]循环不变的。因此,您将该值传输到 4 个经线通道(通过 4 次循环迭代),这意味着该值占用了 4 个输出值(这正是您的打印输出数据显示的内容)。这不适合转置。

从更长远的角度来看,从您的代码中使用另一个示例:

我不确定我是否能读懂你的想法,但你可能对经纱洗牌操作感到困惑。可能您已经假设目标车道可以从源车道loads[]阵列中选择所需的值。不是这种情况。目标通道只能选择源通道提供的任何值。让我们来看看你的循环:

// This should provide 0 0 0 0 4 4 4 4 8 8 8 8 ...
#define base_idx(lane_id) (lane_id & (warp_size - ncorrs))
// This should provide 0 1 2 3 0 1 2 3 0 1 2 3
#define corr_idx(lane_id) (lane_id & (ncorrs - 1))


int n = blockIdx.x*blockDim.x + threadIdx.x;
int lane_id = threadIdx.x & (warp_size - 1);

...

// Input correlation handled by this thread
int src_corr = corr_idx(lane_id);
int mask = __activemask();

...

int loads[ncorrs];
int values[ncorrs];

...

#pragma unroll (ncorrs)
for(int corr=0; corr < ncorrs; ++corr)
{
    int src_lane = base_idx(lane_id) + corr;
    values[corr] = __shfl_sync(mask, loads[src_corr], src_lane, warp_size);
}
Run Code Online (Sandbox Code Playgroud)

在上述循环的第一次通过时,src_lanefor warp 通道 0、1、2 和 3 都将是 0。从上面摘录的代码中可以明显看出这一点,或者如果您不确定,则将其打印出来。这意味着经线 0-3 将请求经线 0 提供的任何值。经线 0 提供的值是loads[src_corr],但src_corr这里的解释是它对经线 0 的任何值。因此,只有一个值将分配到经线 0-3。对于转置,这不可能是正确的;输出中的 4 个位置没有显示输入值。

为了解决这个问题,我们需要修改这两个的计算src_lanesrc_corr。我们还需要修改每个经线通道的存储位置(索引),在每次循环时(我称之为新变量dest。)我们可以认为src_lane是定义我的线程将接收的目标值。我们可以认为src_corr是定义我将在该循环迭代中将哪些值发布到其他线程。destvalues[]我将存储当前接收到的值在我的数组中的位置。我们可以通过仔细研究输入值 inloads[]与期望输出位置之间的关系来推断出必要的模式values[],考虑到源和目标的适当经线通道。在循环的第一次通过时,我们需要这种模式:

warp lane: 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 ... 
src_lane:  0  8 16 24  1  9 17 25  2 10 18 26  3 11 19 27  4 ... (where my data comes from)
src_corr:  0  0  0  0  0  0  0  0  1  1  1  1  1  1  1  1  2 ... (which value I am transmitting)
dest:      0  1  2  3  0  1  2  3  0  1  2  3  0  1  2  3  0 ... (where I store the received value)
Run Code Online (Sandbox Code Playgroud)

在循环的第二遍,我们需要这种模式:

warp lane: 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 ... 
src_lane:  8 16 24  0  9 17 25  1 10 18 26  2 11 19 27  3 19 ... (where my data comes from)
src_corr:  3  3  3  3  3  3  3  3  0  0  0  0  0  0  0  0  1 ... (which value I am transmitting)
dest:      1  2  3  0  1  2  3  0  1  2  3  0  1  2  3  0  1 ... (where I store the received value)
Run Code Online (Sandbox Code Playgroud)

对循环的第 3 次和第 4 次进行相应的更改。如果我们在 shuffle 循环的代码中实现了这些模式,它可能看起来像这样:

$ cat t352.cu
#include <cstdlib>
#include <cstdio>

#include <assert.h>
#define ncorr 4
#define warp_size 32

template <int ncorrs>
__global__ void kernel(
    int * input,
    int * output,
    int N)
{
    // This should provide 0 0 0 0 4 4 4 4 8 8 8 8 ...
    #define base_idx(lane_id) (lane_id & (warp_size - ncorrs))
    // This should provide 0 1 2 3 0 1 2 3 0 1 2 3
    #define corr_idx(lane_id) (lane_id & (ncorrs - 1))


    int n = blockIdx.x*blockDim.x + threadIdx.x;
    int lane_id = threadIdx.x & (warp_size - 1);

    if(n >= N)
        { return; }

    // Input correlation handled by this thread
    int mask = __activemask();

    if(threadIdx.x == 0)
        { printf("mask %d\n", mask); }

    int loads[ncorrs];
    int values[ncorrs];

    #pragma unroll (ncorrs)
    for(int corr=0; corr < ncorrs; ++corr)
        { loads[corr] = input[n + corr*N]; }

    __syncthreads();

    printf("[%d, %d] %d %d %d %d\n",
           lane_id, base_idx(lane_id),
           loads[0], loads[1],
           loads[2], loads[3]);
    #pragma unroll (ncorrs)
    for(int corr=0; corr < ncorrs; ++corr)
    {
        int src_lane = ((lane_id+corr)%ncorrs)*(warp_size/ncorrs) + (lane_id/ncorrs);
        int src_corr = ((ncorrs-corr)+(lane_id/(warp_size/ncorrs)))%ncorrs;
        int dest = (lane_id+corr)%ncorrs;
        values[dest] = __shfl_sync(mask, loads[src_corr],
                                     src_lane, warp_size);
    }

    printf("[%d, %d] %d %d %d %d\n",
           lane_id, base_idx(lane_id),
           values[0], values[1],
           values[2], values[3]);


    #pragma unroll (ncorrs)
    for(int corr=0; corr < ncorrs; ++corr)
        { output[n + corr*N] = values[corr]; }
}

void print_data(int * data, int N)
{
    for(int n=0; n < N; ++n)
    {
        printf("% -3d: ", n);
        for(int c=0; c < ncorr; ++c)
        {
            printf("%d ", data[n*ncorr + c]);
        }
        printf("\n");
    }
}

int main(void)
{
    int * host_input;
    int * host_output;

    int * device_input;
    int * device_output;
    int N = 32;

    host_input = (int *) malloc(sizeof(int)*N*ncorr);
    host_output = (int *) malloc(sizeof(int)*N*ncorr);

    printf("malloc done\n");

    cudaMalloc((void **) &device_input, sizeof(int)*N*ncorr);
    cudaMalloc((void **) &device_output, sizeof(int)*N*ncorr);

    printf("cudaMalloc done\n");

    for(int i=0; i < N*ncorr; ++i)
        { host_input[i] = i; }

    print_data(host_input, N);

    dim3 block(256, 1, 1);
    dim3 grid((block.x + N - 1) / N, 1, 1);

    cudaMemcpy(device_input, host_input,
               sizeof(int)*N*ncorr, cudaMemcpyHostToDevice);

    printf("memcpy done\n");

    kernel<4> <<<grid, block>>> (device_input, device_output, N);

    cudaMemcpy(host_output, device_output,
               sizeof(int)*N*ncorr, cudaMemcpyDeviceToHost);

    print_data(host_output, N);
    cudaFree(device_input);
    cudaFree(device_output);

    free(host_input);
    free(host_output);
}
$ nvcc -o t352 t352.cu
$ cuda-memcheck ./t352
========= CUDA-MEMCHECK
malloc done
cudaMalloc done
 0 : 0 1 2 3
 1 : 4 5 6 7
 2 : 8 9 10 11
 3 : 12 13 14 15
 4 : 16 17 18 19
 5 : 20 21 22 23
 6 : 24 25 26 27
 7 : 28 29 30 31
 8 : 32 33 34 35
 9 : 36 37 38 39
 10: 40 41 42 43
 11: 44 45 46 47
 12: 48 49 50 51
 13: 52 53 54 55
 14: 56 57 58 59
 15: 60 61 62 63
 16: 64 65 66 67
 17: 68 69 70 71
 18: 72 73 74 75
 19: 76 77 78 79
 20: 80 81 82 83
 21: 84 85 86 87
 22: 88 89 90 91
 23: 92 93 94 95
 24: 96 97 98 99
 25: 100 101 102 103
 26: 104 105 106 107
 27: 108 109 110 111
 28: 112 113 114 115
 29: 116 117 118 119
 30: 120 121 122 123
 31: 124 125 126 127
memcpy done
mask -1
[0, 0] 0 32 64 96
[1, 0] 1 33 65 97
[2, 0] 2 34 66 98
[3, 0] 3 35 67 99
[4, 4] 4 36 68 100
[5, 4] 5 37 69 101
[6, 4] 6 38 70 102
[7, 4] 7 39 71 103
[8, 8] 8 40 72 104
[9, 8] 9 41 73 105
[10, 8] 10 42 74 106
[11, 8] 11 43 75 107
[12, 12] 12 44 76 108
[13, 12] 13 45 77 109
[14, 12] 14 46 78 110
[15, 12] 15 47 79 111
[16, 16] 16 48 80 112
[17, 16] 17 49 81 113
[18, 16] 18 50 82 114
[19, 16] 19 51 83 115
[20, 20] 20 52 84 116
[21, 20] 21 53 85 117
[22, 20] 22 54 86 118
[23, 20] 23 55 87 119
[24, 24] 24 56 88 120
[25, 24] 25 57 89 121
[26, 24] 26 58 90 122
[27, 24] 27 59 91 123
[28, 28] 28 60 92 124
[29, 28] 29 61 93 125
[30, 28] 30 62 94 126
[31, 28] 31 63 95 127
[0, 0] 0 8 16 24
[1, 0] 32 40 48 56
[2, 0] 64 72 80 88
[3, 0] 96 104 112 120
[4, 4] 1 9 17 25
[5, 4] 33 41 49 57
[6, 4] 65 73 81 89
[7, 4] 97 105 113 121
[8, 8] 2 10 18 26
[9, 8] 34 42 50 58
[10, 8] 66 74 82 90
[11, 8] 98 106 114 122
[12, 12] 3 11 19 27
[13, 12] 35 43 51 59
[14, 12] 67 75 83 91
[15, 12] 99 107 115 123
[16, 16] 4 12 20 28
[17, 16] 36 44 52 60
[18, 16] 68 76 84 92
[19, 16] 100 108 116 124
[20, 20] 5 13 21 29
[21, 20] 37 45 53 61
[22, 20] 69 77 85 93
[23, 20] 101 109 117 125
[24, 24] 6 14 22 30
[25, 24] 38 46 54 62
[26, 24] 70 78 86 94
[27, 24] 102 110 118 126
[28, 28] 7 15 23 31
[29, 28] 39 47 55 63
[30, 28] 71 79 87 95
[31, 28] 103 111 119 127
 0 : 0 32 64 96
 1 : 1 33 65 97
 2 : 2 34 66 98
 3 : 3 35 67 99
 4 : 4 36 68 100
 5 : 5 37 69 101
 6 : 6 38 70 102
 7 : 7 39 71 103
 8 : 8 40 72 104
 9 : 9 41 73 105
 10: 10 42 74 106
 11: 11 43 75 107
 12: 12 44 76 108
 13: 13 45 77 109
 14: 14 46 78 110
 15: 15 47 79 111
 16: 16 48 80 112
 17: 17 49 81 113
 18: 18 50 82 114
 19: 19 51 83 115
 20: 20 52 84 116
 21: 21 53 85 117
 22: 22 54 86 118
 23: 23 55 87 119
 24: 24 56 88 120
 25: 25 57 89 121
 26: 26 58 90 122
 27: 27 59 91 123
 28: 28 60 92 124
 29: 29 61 93 125
 30: 30 62 94 126
 31: 31 63 95 127
========= ERROR SUMMARY: 0 errors
$
Run Code Online (Sandbox Code Playgroud)
  1. 我相信上面的代码相当清楚地展示了 32x4 -> 4x32 转置。我认为它与您提供的代码“最接近”。它不会执行您在图表中描绘的一组 4x8 转置。

  2. 我承认的计算src_corrsrc_lane以及dest尚未完全优化。但是它们会生成正确的索引。我假设您可以弄清楚如何从您已有的模式中最佳地生成那些。

  3. 我认为上面的代码完全有可能在其他维度上有错误。除了 32x4 的情况外,我没有在任何东西上尝试过。尽管如此,我认为我已经指出了您的代码的根本错误,并展示了获得正确索引的途径。

  4. 正方形矩阵转置到32x32可以在使用经电平来进行更简单的方法