关于 CUDA 实现一个归一化的调研

事实上，我们在日常的学习过程中经常会遇到归一化的问题，归一化看起来很简单，最朴素的版本是：

__global__ void normalize_vector(float *d_vector, int n) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < n) {
        // 每个线程计算向量的平方和
        float sum = 0.0f;
        for (int i = 0; i < n; i++) {
            sum += d_vector[i] * d_vector[i];
        }

        // 求平方根（向量的模）
        float magnitude = sqrt(sum);

        // 每个线程除以模（单位化向量）
        if (magnitude != 0) {
            d_vector[idx] /= magnitude;
        }
    }
}

然而这样显然是一个很烂的实现，每个线程都计算了平方和，显然平方和才是一个大规模的计算。不过这里要考虑这里的 n 的大小，如果 n 很小也许是划算的。

比较朴素的一个算法是

template<int N_ITEMS_PER_THREAD=8>
__device__ __forceinline__ void normalize_vector_256(float* reg_data) {
    // 假设当前 Warp 有 32 个线程，每个线程持有 reg_data[0..7]

    // 1. Thread-local 平方和
    float local_sum_sq = 0.0f;
    #pragma unroll
    for (int i = 0; i < N_ITEMS_PER_THREAD; ++i) {
        local_sum_sq += reg_data[i] * reg_data[i];
    }

    // 2. Warp-level 规约 (无需 Shared Memory)
    // 这是一个标准的 Butterfly Reduction 或者 XOR Reduction
    #pragma unroll
    for (int offset = 16; offset > 0; offset /= 2) {
        local_sum_sq += __shfl_xor_sync(0xffffffff, local_sum_sq, offset);
    }
    // 此时所有线程的 local_sum_sq 都是全局平方和

    // 3. 计算缩放因子 (利用 rsqrt 指令加速)
    float scale = __frsqrt_rn(local_sum_sq);

    // 4. 应用归一化
    #pragma unroll
    for (int i = 0; i < N_ITEMS_PER_THREAD; ++i) {
        reg_data[i] *= scale;
    }
}

这样子的话，我们会遇到一个除 0 的问题，在数值上存在不稳定性。

所以就可以用 Blue's Method 的办法。

Blue 的思路是维护两个尺度，一个 scale 直接用最大绝对值，一个 $\frac{x_i}{scale}$ 的范数，显然我们先求范数的时候就不会出现上溢的可能。