Warp 将原语减少为共享相同值的线程

Question

我面临着减少共享同一变量内容的线程中的值的问题。 更具体地说，为了避免对数组进行原子添加操作，我正在评估每个线程的中间结果，然后最后我想将它们的贡献添加到单个（领导者）线程，负责在内存上写入。我从 CUDA 博客中关于扭曲原语的灵感中获得了灵感。

我最终得到了这个，但得到了错误的结果：

    __device__ void warp_aggregate_add(double *g, size_t offset, unsigned id, double dx, double dy, double dz) {
    unsigned mask = __match_any_sync(0xffffffff, id);

    double sum_dx = 0.;
    double sum_dy = 0.;
    double sum_dz = 0.;

    int leader = __ffs(mask) - 1;
    unsigned single_mask = (1u << leader) | (1u << threadIdx.x);
    
    // Calculate the shuffle distance for the current thread
    int shuffle_distance = threadIdx.x - leader;
    sum_dx += __shfl_down_sync(0xffffffff, dx, shuffle_distance);
    sum_dy += __shfl_down_sync(0xffffffff, dy, shuffle_distance);
    sum_dz += __shfl_down_sync(0xffffffff, dz, shuffle_distance);

    __syncwarp();
    // Accumulate the shuffled values
    // leader only writes to memory
    if (threadIdx.x == leader && id != 0xFFFFFFFF) {
        g[offset + (id * 3)]     += sum_dx;
        g[offset + (id * 3) + 1] += sum_dy;
        g[offset + (id * 3) + 2] += sum_dz;
    }
}

其中 0xFFFFFFFF 是为填充元素的 id 设置的，它不得对总和有贡献。

每个块有 32 个线程，因此我们不必担心使用通道索引。 我想知道是否有一个经过充分测试的程序来将值共享给具有相同 id 值的线程。

Answer 1

如果你想添加具有相同ID的线程，你可以这样做：

__device__ float add_same_id(int id, float value) {
    const auto same_id_mask = __match_any_sync(-1u, id);
    //Needs ampere (sm_80) or newer
    const auto sum = __reduce_add_sync(same_id_mask, value);
    return sum;
}

您可以根据需要将其扩展为多个总和值，但请注意，此方法非常慢，因为如果掩码不是

match_any_sync

，则

reduce_add

和

-1u

都运行缓慢。

更快的方法是使用领导者和原子，原子在共享内存上非常快。

__device__ uint32_t lanemask_lt() {
    uint32_t result;
    asm ("mov.u32 %0, %%lanemask_lt;" : "=r"(result));
    return result;
}

__device__ float add_same_id(int id, float value) {
    __shared__ float sum[32];
    const auto laneid = threadIdx.x % 32;
    sum[laneid] = 0.0f;
    const auto same_id_mask = __match_any_sync(-1u, id);
    const auto Leader = __ffs(same_id_mask) - 1;
    const auto LeaderMask = __ballot_sync(-1u, Leader == laneid);
    const auto LanesBefore = __popc(isLeaderMask & lanemask_lt());
    const auto Dest = __shfl_sync(-1u, LanesBefore, Leader);
    atomicAdd(&sum[Dest], value);
    __syncthreads(); //or __syncwarp(); if only 32 threads in block
    return sum[laneid];    
}