为何CUDA矩阵转置循环展开未加速反而变慢？如何优化？

问题背景 #

学习CUDA时发现，循环展开在数组求和核函数中能有效加速，但实现矩阵转置时，循环展开版本反而比原版本更慢。以下是相关代码与性能数据：

原矩阵转置核函数 #

__global__ void transform_matrix_read_col(
    int* mat_a , int* mat_b , size_t row_num , size_t col_num
){
    int ix = threadIdx.x + blockDim.x * blockIdx.x;
    int iy = threadIdx.y + blockDim.y * blockIdx.y;
    int row_idx = iy*col_num + ix;
    int col_idx = ix*row_num + iy;
    if(ix < col_num && iy < row_num){
        mat_b[row_idx] = mat_a[col_idx];
    }
}

循环展开后的核函数 #

__global__ void transform_matrix_read_col_unrool(
    int* mat_a , int* mat_b , size_t row_num , size_t col_num
){
    int ix = threadIdx.x +(blockDim.x * blockIdx.x * 4);
    int iy = threadIdx.y + blockDim.y * blockIdx.y;
    int row_idx = iy*col_num + ix;
    int col_idx = ix*row_num + iy;
    if(ix < col_num && iy < row_num){
        mat_b[row_idx] = mat_a[col_idx];
        mat_b[row_idx + blockDim.x*1] = mat_a[col_idx + row_num*blockDim.x*1];
        mat_b[row_idx + blockDim.x*2] = mat_a[col_idx + row_num*blockDim.x*2];
        mat_b[row_idx + blockDim.x*3] = mat_a[col_idx + row_num*blockDim.x*3];
    }
}

主函数调用代码 #

size_t width = 128 , height = 128,
    array_size = width*height,array_bytes = array_size * sizeof(int);
    int* matrix_data = nullptr,*output_data = nullptr;
    cudaMallocHost(&matrix_data, array_bytes);
    cudaMallocHost(&output_data, array_bytes);
    util::init_array_int(matrix_data,array_size);// 随机生成整数数组

    int* matrix_data_dev = nullptr,* output_matrix_dev = nullptr;
    cudaMalloc(&matrix_data_dev, array_bytes);
    cudaMemcpy(matrix_data_dev, matrix_data, array_bytes, cudaMemcpyHostToDevice);
    cudaMalloc(&output_matrix_dev, array_bytes);
    
    dim3 block(32,16);
    dim3 grid((width-1)/block.x+1,(height-1)/block.y+1);
    dim3 gridUnrool4((width-1)/(block.x*4)+1,(height-1)/block.y +1);

    transform_matrix_read_col<<<grid,block>>>(matrix_data_dev, output_matrix_dev, height, width);
    cudaDeviceSynchronize();

    transform_matrix_read_col_unrool<<<gridUnrool4,block>>>(matrix_data_dev, output_matrix_dev, height, width);
    cudaDeviceSynchronize();

性能统计（RTX 3090 + nsys） #

CUDA Kernel Statistics:

 Time(%)  Total Time (ns)  Instances  Average   Minimum  Maximum                                     Name                                     
 -------  ---------------  ---------  --------  -------  -------  ---------------------------------------------------------------------------

     6.3            3,456          1   3,456.0    3,456    3,456  transform_matrix_read_col_unrool(int*, int*, unsigned long, unsigned long) 
     5.2            2,880          1   2,880.0    2,880    2,880  transform_matrix_read_col(int*, int*, unsigned long, unsigned long)        

循环展开失效的核心原因 #

1. 全局内存访问模式彻底恶化
   原核函数中，mat_a的访问是列方向，本身属于非连续访问（CUDA全局内存要求连续地址才能合并访问，带宽利用率最高）。循环展开后，每个线程访问的mat_a地址间隔为row_num * blockDim.x * N（N为展开步长），这个间隔远大于内存事务的粒度，完全破坏了合并访问的可能性，导致全局内存带宽利用率骤降，这是性能变慢的主要原因。

2. SM线程利用率下降
   原grid的x维度为(128-1)/32 +1 =4，展开后grid的x维度变为(128-1)/(32*4)+1=1，线程块总数从4*(128/16)=32减少到1*8=8。RTX3090的SM数量为82，过少的线程块无法充分填满所有SM，导致硬件资源闲置。

3. 边界处理隐患（非本次性能问题原因，但需注意）
   展开后的if判断仅检查了ix < col_num，但后续三次访问的row_idx + blockDim.x*N可能超出数组边界，本次测试中width=128刚好匹配32*4，但换其他尺寸会导致越界错误。

矩阵转置的优化方案 #

1. 共享内存分块转置（经典最优方案） #

利用共享内存的高带宽和低延迟，将全局内存的非连续访问转化为连续访问，转置操作在共享内存中完成，再连续写入全局内存。同时添加额外列避免共享内存bank冲突：

__global__ void transform_matrix_shared(int* mat_a, int* mat_b, size_t row_num, size_t col_num) {
    __shared__ int tile[32][33]; // 额外一列消除bank冲突

    int x = threadIdx.x + blockDim.x * blockIdx.x;
    int y = threadIdx.y + blockDim.y * blockIdx.y;

    // 连续加载全局内存到共享内存
    if (x < col_num && y < row_num) {
        tile[threadIdx.y][threadIdx.x] = mat_a[y * col_num + x];
    }
    __syncthreads(); // 等待所有线程完成加载

    // 转置后连续写入全局内存
    x = threadIdx.y + blockDim.y * blockIdx.y;
    y = threadIdx.x + blockDim.x * blockIdx.x;
    if (x < row_num && y < col_num) {
        mat_b[y * row_num + x] = tile[threadIdx.x][threadIdx.y];
    }
}

调用时使用dim3 block(32,32)的线程块维度，grid维度计算为((col_num-1)/32+1, (row_num-1)/32+1)。

2. 正确的循环展开方式 #

如果要使用循环展开，需保证每个线程处理连续的内存地址，比如让每个线程处理同一行的4个连续元素：

__global__ void transform_matrix_unroll_correct(int* mat_a, int* mat_b, size_t row_num, size_t col_num) {
    int x = threadIdx.x *4 + blockDim.x * blockIdx.x;
    int y = threadIdx.y + blockDim.y * blockIdx.y;
    if (y < row_num) {
        for(int i=0; i<4; i++){
            if(x+i < col_num){
                mat_b[y*col_num + x+i] = mat_a[(x+i)*row_num + y];
            }
        }
    }
}

这种展开方式下，mat_a的访问虽然还是列方向，但每个线程内的4次访问是连续的列元素，能部分利用全局内存的合并访问特性。

3. 使用高度优化的库函数 #

直接调用CUDA官方库实现转置，比如cuBLAS的cublasSgeam（支持矩阵转置）或Thrust库的thrust::transpose，这些函数经过NVIDIA工程师的深度优化，能充分利用硬件特性。

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