[Python知识库]cuda任意长度矢量求和兼容数据分块与小数据量时平铺

本代码实现计算矢量求和时，当矢量长度比较小，一个线程负责计算一个个维度的矢量求和，即数据平铺，当矢量长度较大，采取数据分块，一个线程可能负责多个维度的矢量求和。未使用优化手段，代码中包含对于计算正确性的检验、与CPU计算耗时的对比。

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <stdio.h>
#include <random>
#include "time.h"

#define N 130712

int numBlocks = (N + 512 - 1) / 512 > 512 ? 512 : (N + 512 - 1) / 512;

__global__ void addKernel(int* c, const int* a, const int* b)
{
    int offset = blockDim.x * gridDim.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    while (idx < N)
    {
        c[idx] = a[idx] + b[idx];
        idx += offset;
    }
}

void addCPU(int* c, int* a, int* b) {
    for (int i = 0; i < N; ++i)
        c[i] = a[i] + b[i];
}

int main()
{
    int* a = new int[N];
    int* b = new int[N];
    int* c1 = new int[N];
    int* c2 = new int[N];
    srand(0);
    for (int i = 0; i < N; ++i) {
        a[i] = rand() % 1000;
        b[i] = rand() % 1000;
    }

    clock_t start, end;
    double elapsedTime;
    start = clock();
    addCPU(c1, a, b);
    end = clock();
    elapsedTime = (double)(end - start);
    printf("time to generate CPU:% 5.3f ms\n", elapsedTime);

    int* dev_a, * dev_b, * dev_c;
    cudaMalloc((void**)&dev_c, N * sizeof(int));
    cudaMalloc((void**)&dev_a, N * sizeof(int));
    cudaMalloc((void**)&dev_b, N * sizeof(int));
    cudaMemcpy(dev_a, a, N * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(dev_b, b, N * sizeof(int), cudaMemcpyHostToDevice);

    cudaEvent_t start1, stop1;
    cudaEventCreate(&start1);
    cudaEventCreate(&stop1);
    cudaEventRecord(start1, 0);
    addKernel << <numBlocks, 512 >> > (dev_c, dev_a, dev_b);
    cudaEventRecord(stop1, 0);
    cudaEventSynchronize(stop1);
    float elapsedTime1;
    cudaEventElapsedTime(&elapsedTime1, start1, stop1);
    printf("time to generate GPU:% 5.3f ms\n", elapsedTime1);
    cudaEventDestroy(start1);
    cudaEventDestroy(stop1);

    cudaMemcpy(c2, dev_c, N * sizeof(int), cudaMemcpyDeviceToHost);

    bool flag = true;
    for (int i = 0; i < N; ++i) {
        if (c1[i] != c2[i]) {
            flag = false;
            break;
        }
    }
    if (flag) printf("Consistent!!!\n");
    else printf("Not consistent!!!\n");

    delete[] a;
    delete[] b;
    delete[] c1;
    delete[] c2;
    cudaFree(dev_c);
    cudaFree(dev_a);
    cudaFree(dev_b);

    return 0;
}