CUDA管道异步内存从全局复制到共享内存

Question

我目前正在学习如何编写快速 CUDA 内核。我实现了一个平铺矩阵乘法（块大小 32x32），它仅从/向全局内存进行合并读取/写入，并且在从共享内存写入/读取时没有存储体冲突（它的速度约为 pytorch 矩阵乘法实现的 50%） ）。现在我尝试使用管道（两个阶段）并将内存从全局异步复制到共享内存（请参阅here和here）。

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

constexpr unsigned stages_count = 2;
__shared__ float s_a[stages_count][32][32];

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);

for(int step=0; step<a.size(1); step+=32) {
   for(int stage=0; stage<stages_count; stage++) {
        pipeline.producer_acquire();

        // what i would like to do (this works but is not asynchronous)
        s_a[stage][threadIdx.y][threadIdx.x] = a[blockIdx.x*stages_count*32 + stage*32 + threadIdx.y][step + threadIdx.x];

        // this does not work
        cuda::memcpy_async(block,
                           &s_a[stage][threadIdx.y][0],
                           &a[blockIdx.x*stages_count*32 + stage*32 + threadIdx.y][step],
                           sizeof(float) * 32,
                           pipeline);

        pipeline.producer_commit();
    }
    for(int stage=0; stage<stages_count; stage++) {
        pipeline.consumer_wait();
        // use shared memory
        pipeline.consumer_release();
    }
}

但是，我不知道如何使异步内存复制工作。我认为问题是，我不想从全局内存中复制 32*32 连续浮点，而是复制矩阵的一个图块（32 乘以 32 连续浮点）。另外，是否可以在异步加载时以某种方式转置（或例如使用置换的共享内存布局）以防止以后的存储体冲突？

Answer 1

新的Hopper架构（H100 GPU）为此提供了一个新的硬件功能，称为张量内存加速器（TMA）。今年晚些时候，CUDA 12 将提供软件支持。

据我了解，这将允许使用单个命令异步复制张量图块。但是，如果它在 Ampere 和较旧的架构上工作，它可能会非常慢，就像根据我的经验，由于缺少硬件支持，模拟的

cuda::memcpy_async

在 Ampere 之前的 GPU 上相当慢。

不确定您提到的转置是否会成为新 API 的一部分，但可能：

TMA 通过支持不同的张量布局（1D-5D 张量）、不同的内存访问模式、缩减和其他功能，显着减少寻址开销并提高效率。

当不需要时，CUB提供了一些有用的功能，用于在

cub::BlockLoad

（和

cub::BlockStore

）中“转置”数据。这些的缺点是它们仅使用共享内存作为中介并最终写入寄存器或本地内存，因此对于这些类型的平铺矩阵乘法内核它们可能没有任何帮助。他们“可能会添加一个功能”，以便将来仅复制到共享内存。也许这些新容器甚至会支持异步性。

Answer 2