4. 第一个Jetson CUDA算子

向量加法是CUDA编程的“Hello World”，但它同时覆盖了CUDA算子开发最核心的四个环节：主机与设备内存管理、核函数编写与线程索引映射、编译运行、错误检查与性能初测。把这四件事做扎实，后面写卷积、矩阵乘、归一化等算子时，只是计算逻辑变复杂，工程骨架不变。

一、问题定义与数据约定

给定两个长度为 N 的单精度浮点数组 A 和 B，计算 C[i] = A[i] + B[i]。在Jetson上，CPU与GPU共享同一块物理内存（统一内存架构），但CUDA编程模型仍区分主机端（host）与设备端（device）指针。为便于观察，本节取 N = 1<<20（约104.8万），即 1,048,576 个元素，每个数组占 4 MiB。约定：所有数组按行主序一维展开，索引 i 从 0 开始，线程块大小 blockDim = 256，网格大小 gridDim = (N + 255) / 256 = 4096。

二、核函数编写与线程索引映射

核函数用 __global__ 修饰，返回类型必须为 void。每个线程通过 blockIdx.x、blockDim.x、threadIdx.x 计算全局索引，并用边界判断防止越界访问——这是所有CUDA算子的必备防御。

__global__ void vectorAdd(const float* __restrict__ A,
                          const float* __restrict__ B,
                          float* __restrict__ C, int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        C[i] = A[i] + B[i];
    }
}

这里 __restrict__ 告诉编译器三个指针不重叠，便于生成更优的加载指令。索引映射关系是：全局线程号 = 块号 × 块内线程数 + 块内线程号。若 N 不能被 blockDim 整除，最后一个块中会有部分线程因 i < n 不成立而空转，这是正常且必要的。

三、主机端代码：内存分配、拷贝与启动

主机端流程为：分配主机内存 → 分配设备内存 → 拷贝输入到设备 → 启动核函数 → 拷贝结果回主机 → 释放内存。Jetson支持统一内存（cudaMallocManaged），但为清晰展示数据流，这里先用显式拷贝方式。

int main() {
    const int N = 1 << 20;
    size_t bytes = N * sizeof(float);

    float *h_A = (float*)malloc(bytes);
    float *h_B = (float*)malloc(bytes);
    float *h_C = (float*)malloc(bytes);
    for (int i = 0; i < N; ++i) { h_A[i] = 1.0f; h_B[i] = 2.0f; }

    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, bytes);
    cudaMalloc(&d_B, bytes);
    cudaMalloc(&d_C, bytes);

    cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice);

    int blockSize = 256;
    int gridSize  = (N + blockSize - 1) / blockSize;
    vectorAdd<<<gridSize, blockSize>>>(d_A, d_B, d_C, N);

    cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost);

    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    free(h_A); free(h_B); free(h_C);
    return 0;
}

注意 cudaMalloc 的第一个参数是二级指针，因为需要修改指针本身的值。核函数启动语法 <<<gridSize, blockSize>>> 是异步的，主机不会等待其完成，因此后续的 cudaMemcpy 默认会与同一流中的核函数同步，保证结果正确。

四、编译与运行

在Jetson上使用nvcc编译。以Orin Nano（算力8.7）为例，指定架构可减少编译时间并避免警告：

nvcc -O2 -arch=sm_87 vector_add.cu -o vector_add
./vector_add

若不确定设备算力，可用 nvcc -arch=native（CUDA 11.5+）或先运行 deviceQuery 查看。编译时若出现“no kernel image is available”错误，说明 -arch 与设备不匹配，需调整。

五、错误检查：从“能跑”到“可信”

CUDA错误分两类：启动错误（如参数非法、内存不足）和异步执行错误（如越界访问）。核函数启动后必须检查：

cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
    printf("Kernel launch failed: %s\n", cudaGetErrorString(err));
}
err = cudaDeviceSynchronize();
if (err != cudaSuccess) {
    printf("Kernel execution failed: %s\n", cudaGetErrorString(err));
}

此外，结果验证不能只看“程序没崩”。对向量加法，可用CPU串行计算参考值并逐元素比较，容差取 1e-5：

bool ok = true;
for (int i = 0; i < N; ++i) {
    if (fabsf(h_C[i] - (h_A[i] + h_B[i])) > 1e-5f) { ok = false; break; }
}
printf("Result %s\n", ok ? "PASS" : "FAIL");

六、性能初测：用事件计时

用CUDA事件测量核函数执行时间，比 clock() 更准确。注意：只计时核函数，不含内存拷贝，才能反映计算本身的开销。

cudaEvent_t start, stop;
cudaEventCreate(&start); cudaEventCreate(&stop);
cudaEventRecord(start);
vectorAdd<<<gridSize, blockSize>>>(d_A, d_B, d_C, N);
cudaEventRecord(stop);
cudaEventSynchronize(stop);
float ms = 0;
cudaEventElapsedTime(&ms, start, stop);
printf("Kernel time: %.3f ms\n", ms);

以Orin Nano为例，1M元素向量加法核函数耗时通常在几十微秒量级（示例假设约 0.05 ms），有效带宽 = 3 × 4 MiB / 时间。若结果远高于预期，优先检查：是否误将内存拷贝计入、blockSize是否过小、是否发生越界导致重试。性能初测的意义不是追求极限，而是建立“改动—观察—判断”的闭环：改blockSize、改gridSize、加__restrict__，看时间如何变化。