CUDA pipeline asynchronous memory copy from global to shared memory - cuda

Question

I'm currently learning how to write fast CUDA kernels. I implemented a tiled matrix multiplication (block size 32x32) which only does coalesc reads/writes from/to global memory and has no bank conflicts when writing/reading from shared memory (it has ~50% of the speed of the pytorch matrix multiplication implementation). Now I tried to use pipelining (two stages) and copy memory asyncronously from global to shared memory (see here, and 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();
}
}
However, I don't know how to make the asynchronous memory copy work. Problem is, I think, that I don't want to copy 32*32 consecutive floats from global memory, but one tile of the matrix (32 times 32 consecutive float). Also, is it possible to somehow transpose (or e.g. use a permuted shared memory layout) while asynchronously loading to prevent later bank conflicts?

Answer 1

The new Hopper architecture (H100 GPU) has a new hardware feature for this, called the tensor memory accelerator (TMA). Software support will come with CUDA 12 later this year.
As far as I understand it this will allow to asynchronously copy tensor tiles with a single command. But, if it works at all on Ampere and older architectures, it might be quite slow in the same way that the emulated cuda::memcpy_async is quite slow on pre-Ampere GPUs in my experience due to missing hardware support.
Not sure if the transpose you mention will be part of that new API but it might:
TMA significantly reduces addressing overhead and improves efficiency with support for different tensor layouts (1D-5D tensors), different memory access modes, reductions, and other features.
When asynchronicity is not required, CUB provides some useful functionality for "transposing" data in cub::BlockLoad (and cub::BlockStore). A downside with these is that they use shared memory as an intermediary only and write to registers or local memory in the end, so for these kinds of tiled matrix multiplication kernels they probably are not of any help. They might add a feature for only copying to shared memory in the future. Maybe these new containers will even support asynchronicity.