cuda算子--矩阵转置

复制代码

#include <cuda_runtime.h>

#define TILE 32
#define BLOCK_ROWS 8   // block 的 y 方向线程数，配合 TILE 减少 padding 浪费

// 用 padding 避免 shared bank conflict
__global__ void transpose_shared_kernel(const float* __restrict__ input,
                                        float* __restrict__ output,
                                        int rows, int cols) {
    __shared__ float tile[TILE][TILE + 1];

    const int tx = threadIdx.x;   // [0, TILE)
    const int ty = threadIdx.y;   // [0, BLOCK_ROWS)

    // ---- 读阶段 ----
    // 该 block 负责 input 的子块：
    //   行 [blockIdx.y*TILE, blockIdx.y*TILE + TILE)
    //   列 [blockIdx.x*TILE, blockIdx.x*TILE + TILE)
    const int col = blockIdx.x * TILE + tx;   // input 列
    const int row = blockIdx.y * TILE + ty;   // input 行起点

    #pragma unroll
    for (int j = 0; j < TILE; j += BLOCK_ROWS) {
        if (col < cols && (row + j) < rows)
            tile[ty + j][tx] = input[(size_t)(row + j) * cols + col];
    }

    __syncthreads();

    // ---- 写阶段 ----
    // 转置后，output 的子块：
    //   行 [blockIdx.x*TILE, blockIdx.x*TILE + TILE)  ← 原来是 input 的列
    //   列 [blockIdx.y*TILE, blockIdx.y*TILE + TILE)  ← 原来是 input 的行
    const int outCol     = blockIdx.y * TILE + tx;   // 输出列 = 原输入行
    const int outRowBase = blockIdx.x * TILE + ty;   // 输出行起点 = 原输入列

    #pragma unroll
    for (int j = 0; j < TILE; j += BLOCK_ROWS) {
        const int outRow = outRowBase + j;           // 关键：j 加到行上
        if (outCol < rows && outRow < cols)
            output[(size_t)outRow * rows + outCol] = tile[tx][ty + j];
    }
}

extern "C" void solve(const float* input, float* output, int rows, int cols) {
    if (rows <= 0 || cols <= 0) return;

    dim3 threadsPerBlock(TILE, BLOCK_ROWS);
    dim3 blocksPerGrid((cols + TILE - 1) / TILE,
                       (rows + TILE - 1) / TILE);

    matrix_transpose_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, rows, cols);
    cudaDeviceSynchronize();
}

关键是

for (int j = 0; j < TILE; j += BLOCK_ROWS) {

const int outRow = outRowBase + j; // 关键：j 加到行上

if (outCol < rows && outRow < cols)

output $(size_t)outRow \* rows + outCol$ = tile $tx$ $ty + j$ ;

}

如果让block都等于0的话，那么base row和col就是ty，tx，那么最后就是output $ty$ $tx$ =tile $tx$ ty]很简单的转置。接下来考虑block x增加，此时base row下移，但他不应该影响与tile的对应关系，因为tile只有一块，所以应该还是和block x等于0时一样，去tile的同样的位置去找。需要注意的是，为什么base row对应的是block x，而不是y，因为这样out才对应input矩阵的block级转置，而thread是对应tile级的转置。