仅使用Warp Shuffle指令实现CUDA矩阵转置失败问题排查

你的代码无法正确完成矩阵转置，核心问题是Warp Shuffle指令的参数逻辑错误，以及对shfl指令的工作原理理解有误：

错误原因 #

1. Shuffle参数逻辑颠倒
   你期望通过val(iy, ix)获取原矩阵第ix行、iy列的元素（即转置后的目标值），这个元素对应负责iy列的线程（线程iy）的col[ix]变量。但你的lambda中：

   auto val = [&tile, &col](int ix, int iy) { return tile.shfl(col[iy], ix); }; 
   

   tile.shfl(col[iy], ix)的含义是：从线程ix中获取它传入的col[iy]值，这拿到的是原矩阵第iy行、ix列的元素（和原数据一致，没有转置），完全不符合你的需求。

2. 对Shuffle指令的误解
   Cooperative Groups的shfl指令工作逻辑是：所有线程同步调用该函数，每个线程传入自己的一个值，函数返回指定源线程（第二个参数）传入的那个值。要拿到线程iy的col[ix]，你需要让所有线程传入自己的col[ix]，然后指定源线程为iy。

修正后的代码 #

方案1：修正Lambda逻辑 #

调整lambda的参数顺序和shuffle的参数，使其正确获取目标线程的对应值：

__global__ void transpose_block(float* mat, size_t sx, size_t sy) 
{
    constexpr size_t size = 8;
    auto tile = cg::tiled_partition<size>(cg::this_thread_block());
    auto ix = tile.thread_rank();

    float col[size];
    for (size_t iy = 0; iy < size; ++iy)
        col[iy] = mat[ix + iy * sx];

    // 修正：从源线程src_thread获取它的col[col_idx]值
    auto val = [&tile, &col](int src_thread, int col_idx) { 
        return tile.shfl(col[col_idx], src_thread); 
    }; 

    for (size_t iy = 0; iy < size; ++iy)
        // 转置后的值 = 原线程iy的col[ix]（原矩阵ix行、iy列）
        mat[ix + iy * sx] = val(iy, ix);
}

方案2：去掉Lambda，直接调用Shuffle #

为了减少逻辑复杂度，也可以直接在循环中调用shfl，避免lambda带来的参数混淆：

__global__ void transpose_block(float* mat, size_t sx, size_t sy) 
{
    constexpr size_t size = 8;
    auto tile = cg::tiled_partition<size>(cg::this_thread_block());
    auto ix = tile.thread_rank();

    float col[size];
    for (size_t iy = 0; iy < size; ++iy)
        col[iy] = mat[ix + iy * sx];

    // 直接获取线程iy的col[ix]，赋值给转置后的对应位置
    for (size_t iy = 0; iy < size; ++iy)
        mat[ix + iy * sx] = tile.shfl(col[ix], iy);
}

验证结果 #

修正后，矩阵转置会正确执行：原矩阵的M[iy][ix]会被替换为M[ix][iy]，输出结果符合预期。

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