编程模型和语言层

实现一个MUSA环境中运行的并行累加操作的示例。这段代码展示了在 GPU 上进行大规模数据计算的并行处理能力，利用了线程块、warp 内同步和原子操作等技术，适用于需要高效执行的累加操作等场景。

  • 核函数（sum）
    • 该核函数通过线程块（block）和线程（thread）并行处理数组的部分元素，执行累加操作。
    • 使用 warp 内部的规约操作加快累加速度，并通过原子加法确保不同线程的结果正确写入共享的输出变量。
  • 主程序
    • 在主函数中，首先定义了数组的大小、线程块和网格的配置。
    • 使用 float 作为数据类型和 long 作为索引类型（类型模板支持灵活替换其他数据类型）。
    • 主机端分配和初始化输入数组 h_a，并将其复制到设备端 d_a，然后为结果变量 h_b 分配内存。
    • 核函数在 GPU 上并行启动，计算完成后将结果从设备端复制回主机并打印。
  • 并行计算
    • 并行累加：多个线程块并行处理数组中的不同部分，每个线程累加自己负责的元素。
    • warp 规约：在同一 warp 内使用 __shfl_down_sync 操作进行规约加速。
    • 原子加法：使用 atomicAdd 确保多个线程同时更新输出变量时不产生冲突。

示例代码：

#include "musa_runtime.h"

__global__ void VectorAdd(int* a, int* b) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    a[idx] = a[idx] + b[idx];
}

int main() {
    const size_t numElements = 4096;
    const size_t sizeBytes = numElements * sizeof(int);

    musaStream_t stream;
    musaStreamCreate(&stream);

    int *hA = nullptr, *hB = nullptr;
    int *dA = nullptr, *dB = nullptr;
    hA = reinterpret_cast<int*>(malloc(sizeBytes));
    hB = reinterpret_cast<int*>(malloc(sizeBytes));
    musaMalloc(&dA, sizeBytes);
    musaMalloc(&dB, sizeBytes);

    for (int i = 0; i < numElements; ++i) {
        hA[i] = i;
        hB[i] = 2 * i;
    }

    musaMemcpy(dA, hA, sizeBytes, musaMemcpyHostToDevice);
    musaMemcpy(dB, hB, sizeBytes, musaMemcpyHostToDevice);

    int threadsPerBlock = 1024;
    int blocksPerGrid = (numElements + threadsPerBlock - 1) / threadsPerBlock;

    VectorAdd<<<blocksPerGrid, threadsPerBlock>>>(dA, dB);
    musaError_t err = musaGetLastError();

    musaMemcpy(hA, dA, sizeBytes, musaMemcpyDeviceToHost);

    musaStreamDestroy(stream);
    free(hA);
    free(hB);
    musaFree(dA);
    musaFree(dB);

    return 0;
}

结果输出：

Sum result: 1e+06