CUDA中__shfl_down_sync与__match_any_sync配合求和结果异常

问题背景 #

我编写的CUDA核函数foo目标是计算单warp（32线程）中与id[0]拥有相同id的所有vals元素之和。经排查：

• __match_any_sync获取的mask能正确识别同id线程
• if条件可正确筛选出目标线程
• 但__shfl_down_sync计算出的求和结果始终错误

补充信息：

• 输入：vals和ids为长度32的数组，核函数以foo<<<1,32>>>启动（仅一个warp）
• 对比核函数bar：仅对偶数索引值求和，使用固定mask的__shfl_down_sync能得到正确结果
• 运行环境：WIN11、CUDA-12.2、计算能力8.9（CC89）
• 已查阅Nvidia官方文档，未找到明确原因，猜测是__match_any_sync与__shfl_down_sync间存在同步问题

可复现代码 #

#include <cstdio>

__global__ void foo(const int* vals, const int* ids, int* opt) {
    int tid = threadIdx.x;
    int id = ids[tid];
    int val = vals[tid];

    // 获取与id[0]相同的线程mask
    unsigned int thread_mask = __match_any_sync(0xFFFFFFFF, ids[0]);

    // 仅目标线程参与求和
    if (thread_mask & (1 << tid)) {
        // 用shfl_down_sync进行归约求和
        val += __shfl_down_sync(thread_mask, val, 16, 32);
        val += __shfl_down_sync(thread_mask, val, 8, 32);
        val += __shfl_down_sync(thread_mask, val, 4, 32);
        val += __shfl_down_sync(thread_mask, val, 2, 32);
        val += __shfl_down_sync(thread_mask, val, 1, 32);

        if (tid == 0) {
            *opt = val;
        }
    }
}

__global__ void bar(const int* vals, int* opt) {
    int tid = threadIdx.x;
    int val = vals[tid];

    // 固定mask：仅偶数索引线程
    unsigned int thread_mask = 0x55555555;

    if (thread_mask & (1 << tid)) {
        val += __shfl_down_sync(thread_mask, val, 16, 32);
        val += __shfl_down_sync(thread_mask, val, 8, 32);
        val += __shfl_down_sync(thread_mask, val, 4, 32);
        val += __shfl_down_sync(thread_mask, val, 2, 32);
        val += __shfl_down_sync(thread_mask, val, 1, 32);

        if (tid == 0) {
            *opt = val;
        }
    }
}

int main() {
    int vals[32];
    int ids[32];
    int opt_foo, opt_bar;

    // 初始化数据：前16个线程id为0，后16个为1；vals全为1
    for (int i = 0; i < 32; i++) {
        vals[i] = 1;
        ids[i] = (i < 16) ? 0 : 1;
    }

    int *d_vals, *d_ids, *d_opt_foo, *d_opt_bar;
    cudaMalloc(&d_vals, 32 * sizeof(int));
    cudaMalloc(&d_ids, 32 * sizeof(int));
    cudaMalloc(&d_opt_foo, sizeof(int));
    cudaMalloc(&d_opt_bar, sizeof(int));

    cudaMemcpy(d_vals, vals, 32 * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_ids, ids, 32 * sizeof(int), cudaMemcpyHostToDevice);

    foo<<<1,32>>>(d_vals, d_ids, d_opt_foo);
    bar<<<1,32>>>(d_vals, d_opt_bar);

    cudaMemcpy(&opt_foo, d_opt_foo, sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(&opt_bar, d_opt_bar, sizeof(int), cudaMemcpyDeviceToHost);

    printf("foo结果（预期16）：%d\n", opt_foo); // 错误输出示例：8或其他非16值
    printf("bar结果（预期16）：%d\n", opt_bar); // 正确输出：16

    cudaFree(d_vals);
    cudaFree(d_ids);
    cudaFree(d_opt_foo);
    cudaFree(d_opt_bar);

    return 0;
}

示例错误输出 #

foo结果（预期16）：8
bar结果（预期16）：16

问题原因分析 #

核心问题在于对__shfl_down_sync的mask参数语义理解错误：

1. __shfl_down_sync的mask并非指定“参与数据交换的线程”，而是指定“必须执行该shuffle指令的线程集合”。所有执行该指令的线程必须被包含在mask中，否则未被包含的线程会进入未定义行为，破坏warp内的指令同步。
2. 在foo核函数中，__match_any_sync得到的mask仅包含与id[0]同id的线程，未被包含的线程（后16个id=1的线程）并未执行__shfl_down_sync指令，导致warp内指令流不同步，shuffle操作的数据传递出现错误。
3. 对比bar核函数，虽然mask是固定的偶数线程，但所有线程都执行了__shfl_down_sync指令（即使奇数线程的if条件不满足，它们仍然执行了shuffle指令），mask覆盖了所有执行该指令的线程，因此shuffle操作能正常完成。

修复方案 #

修改foo核函数，将__shfl_down_sync的mask设置为全线程mask（0xFFFFFFFF），同时通过条件判断仅让目标线程参与数据累加：

__global__ void foo(const int* vals, const int* ids, int* opt) {
    int tid = threadIdx.x;
    int id = ids[tid];
    int val = vals[tid];

    unsigned int thread_mask = __match_any_sync(0xFFFFFFFF, ids[0]);
    // 目标线程保留值，非目标线程置0
    val = (thread_mask & (1 << tid)) ? val : 0;

    // 全线程mask执行shuffle，保证所有线程指令同步
    val += __shfl_down_sync(0xFFFFFFFF, val, 16, 32);
    val += __shfl_down_sync(0xFFFFFFFF, val, 8, 32);
    val += __shfl_down_sync(0xFFFFFFFF, val, 4, 32);
    val += __shfl_down_sync(0xFFFFFFFF, val, 2, 32);
    val += __shfl_down_sync(0xFFFFFFFF, val, 1, 32);

    if (tid == 0) {
        *opt = val;
    }
}

内容的提问来源于stack exchange，提问作者SnowSR