CUDA pipeline asynchronous memory copy from global to shared memory

chuber 0
Viewed 61

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?

gpu-shared-memory torch cuda
1 Answers

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.

paleonix0
Related