Thrust/CUDA 多次复制一个数组并结合另一个数组的值

Thrust/CUDA replicate an array multiple times combined with the values of another array

假设我有两个数组

A = {1, 2, 3}

B = {10,20,30,40,50}

我想生成一个大小为

的新数组

sizeof(A) * sizeof(B)

我想复制 B sizeof(A) 次,并且在每次重复时 i,结果数组应该添加 A[i]。所以结果会是这样的

{11,21,31,41,51,12,22,32,42,52,13,23,33,43,53}

这个任务可以解释为一个二维问题,其中输出数组可以被视为维度 sizeof(A) 乘以 sizeof(B) 的矩阵。这样,我们就可以使用 2D CUDA 索引来实现所需的功能。此 2D 实现的示例 CUDA C++ 代码如下所示:

#include <iostream>
#include <cuda_runtime.h>
#include <cassert>

using namespace std;

__global__ void kernel_replicate(int* a, int* b, int* c, int alen, int blen, int clen)
{
    const int ai = blockIdx.x * blockDim.x + threadIdx.x;
    const int bi = blockIdx.y * blockDim.y + threadIdx.y;

    if(ai<alen && bi<blen)
    {
        const int ci = ai * blen + bi;
        c[ci] = a[ai] + b[bi];
    }
}


void replicate_device(int* a, int* b, int* c, int alen, int blen, int clen)
{
    dim3 block(16,16);
    dim3 grid;
    grid.x = (alen + block.x - 1) / block.x;
    grid.y = (blen + block.y - 1) / block.y;

    kernel_replicate<<<grid, block>>>(a,b,c,alen,blen,clen);

    assert(cudaSuccess == cudaDeviceSynchronize());
}


void replicate(int* a, int* b, int* c, int alen, int blen, int clen)
{
    int *ad, *bd, *cd;

    size_t abytes = alen * sizeof(int);
    size_t bbytes = blen * sizeof(int);
    size_t cbytes = clen * sizeof(int);

    cudaMalloc(&ad, abytes);
    cudaMalloc(&bd, bbytes);
    cudaMalloc(&cd, cbytes);

    cudaMemcpy(ad,a, abytes, cudaMemcpyHostToDevice);
    cudaMemcpy(bd,b, bbytes, cudaMemcpyHostToDevice);

    replicate_device(ad,bd,cd, alen,blen,clen);

    cudaMemcpy(c,cd, cbytes, cudaMemcpyDeviceToHost);


    cudaFree(ad);
    cudaFree(bd);
    cudaFree(cd);
}


int main()
{
    const int alen = 3;
    const int blen = 5;
    const int clen = alen * blen;

    int A[alen] = {1,2,3};
    int B[blen] = {10,20,30,40,50};
    int C[clen] = {0};

    replicate(A,B,C,alen, blen, clen);

    for(int i=0; i<alen; i++)
    {
        cout<<A[i]<<" ";
    }
    cout<<endl;

    for(int i=0; i<blen; i++)
    {
        cout<<B[i]<<" ";
    }
    cout<<endl;

    for(int i=0; i<clen; i++)
    {
        cout<<C[i]<<" ";
    }
    cout<<endl;


    return 0;
}