CUDA流水线异步全局到共享内存拷贝及矩阵分块加载问题咨询

一、异步拷贝失败的核心原因 #

你当前的cuda::memcpy_async用法完全错误：每个线程都试图拷贝整行32个float，导致1024个线程重复写入共享内存的同一行，不仅逻辑混乱，还会引发数据竞争。你的需求是每个线程加载矩阵分块中对应自己位置的单个元素，这时候应该用**cuda::load_async**而非块级的memcpy_async——因为矩阵的32×32分块在全局内存中并非连续块（行优先存储下，行与行之间存在 stride），无法用单次块拷贝完成。

二、修复异步加载的正确写法 #

替换原来的memcpy_async为cuda::load_async，每个线程异步加载自己对应的元素到共享内存，完全对齐你原来的同步逻辑：

// 异步加载单个元素到共享内存
cuda::load_async(block, 
                 &s_a[stage][threadIdx.y][threadIdx.x], 
                 &a[blockIdx.x*stages_count*32 + stage*32 + threadIdx.y][step + threadIdx.x], 
                 pipeline);

三、流水线逻辑优化（发挥异步优势） #

你原来的“先生产所有stage，再消费”的循环结构完全浪费了流水线的重叠执行能力。正确的两阶段流水线应该是生产一个stage后，立即消费前一个stage，让数据加载和计算重叠：

for(int step=0; step<a.size(1); step+=32) {
    // 预加载第一个stage
    pipeline.producer_acquire();
    cuda::load_async(block, &s_a[0][threadIdx.y][threadIdx.x], 
                     &a[blockIdx.x*stages_count*32 + 0*32 + threadIdx.y][step + threadIdx.x], pipeline);
    pipeline.producer_commit();

    for(int stage=1; stage<stages_count; stage++) {
        // 消费前一个stage的数据
        pipeline.consumer_wait();
        // --- 这里写使用s_a[stage-1]的计算代码 ---
        pipeline.consumer_release();

        // 生产当前stage的数据
        pipeline.producer_acquire();
        cuda::load_async(block, &s_a[stage][threadIdx.y][threadIdx.x], 
                         &a[blockIdx.x*stages_count*32 + stage*32 + threadIdx.y][step + threadIdx.x], pipeline);
        pipeline.producer_commit();
    }

    // 消费最后一个stage的数据
    pipeline.consumer_wait();
    // --- 这里写使用s_a[stages_count-1]的计算代码 ---
    pipeline.consumer_release();
}

四、异步加载时实现无bank冲突的布局 #

完全可以在异步加载时直接转置共享内存布局，从根源上避免bank冲突，无需额外内存填充：
共享内存默认的[stage][y][x]布局中，同一warp的线程访问同一列（x相同，y不同）时，会命中同一个共享内存bank。转置为[stage][x][y]布局后，加载时直接把全局内存的(y, x)元素放到共享内存的(x, y)位置：

// 转置异步加载，消除后续访存的bank冲突
cuda::load_async(block, 
                 &s_a[stage][threadIdx.x][threadIdx.y], 
                 &a[blockIdx.x*stages_count*32 + stage*32 + threadIdx.y][step + threadIdx.x], 
                 pipeline);

这种转置布局在矩阵乘法中尤为实用——当计算C=A*B时，转置B的共享内存块后，后续的乘积累加可以实现连续无冲突的访存。

如果偏好不转置的方式，也可以给共享内存每行加一个填充元素（比如__shared__ float s_a[stages_count][32][33];），让同一列的元素映射到不同bank，但转置方案的内存效率更高。

完整修正代码示例 #

torch::PackedTensorAccessor32<float,2,torch::RestrictPtrTraits> a; // input to the kernel

constexpr unsigned BLOCK_SIZE = 32;
constexpr unsigned stages_count = 2;
// 转置布局的共享内存，消除bank冲突
__shared__ float s_a[stages_count][BLOCK_SIZE][BLOCK_SIZE];

auto block = cooperative_groups::this_thread_block();

__shared__ cuda::pipeline_shared_state<cuda::thread_scope::thread_scope_block, stages_count> shared_state;
auto pipeline = cuda::make_pipeline(block, &shared_state);

int total_steps = a.size(1) / BLOCK_SIZE;
for(int step=0; step < total_steps; step++) {
    int base_col = step * BLOCK_SIZE;

    // 预加载第一个stage
    pipeline.producer_acquire();
    int global_row = blockIdx.x * stages_count * BLOCK_SIZE + 0 * BLOCK_SIZE + threadIdx.y;
    int global_col = base_col + threadIdx.x;
    cuda::load_async(block, &s_a[0][threadIdx.x][threadIdx.y], &a[global_row][global_col], pipeline);
    pipeline.producer_commit();

    for(int stage=1; stage < stages_count; stage++) {
        // 消费前一个stage
        pipeline.consumer_wait();
        // --- 在这里执行s_a[stage-1]相关的乘积累加计算 ---
        pipeline.consumer_release();

        // 生产当前stage
        pipeline.producer_acquire();
        global_row = blockIdx.x * stages_count * BLOCK_SIZE + stage * BLOCK_SIZE + threadIdx.y;
        global_col = base_col + threadIdx.x;
        cuda::load_async(block, &s_a[stage][threadIdx.x][threadIdx.y], &a[global_row][global_col], pipeline);
        pipeline.producer_commit();
    }

    // 消费最后一个stage
    pipeline.consumer_wait();
    // --- 在这里执行s_a[stages_count-1]相关的乘积累加计算 ---
    pipeline.consumer_release();
}

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