A memory-bound Tensor Core kernel calls for asynchronous copy. The
cp.async mechanism (Ampere and later, through
<cuda_pipeline.h>) streams data from global memory
directly into shared memory without staging it through registers, and –
crucially – without blocking the requesting threads. That makes
double buffering possible: while the Tensor Cores grind through
the current k-step’s slabs, the next step’s slabs are already
streaming into a second shared-memory buffer.
The overlap of asynchronous memcpy with Tensor Core processing echoes
the age-old CUDA optimization of Section 6.3, where CPU/GPU overlap is
used to maximize performance of pageable memcpy: for host-to-device
memcpy (for example), the copy ping-pongs between two staging buffers,
with appropriate synchronization, reading from one staging buffer with
the GPU while the other is written by the CPU. The only bubbles in the
execution pipeline happen when the GPU can read faster than the CPU can
write, or vice versa. The differences with cp.async, of
course, are twofold: the participating memories are device memory and
shared memory instead of device memory and a page-locked staging buffer;
and the copies are requested by a running kernel instead of
CPU-arbitrated DMA traffic.
int buf = 0;stage( buf, 0 ); // prime the first slabfor ( int k0 = 0; k0 < K; k0 += BK ) { __pipeline_wait_prior( 0 ); // this step's slab has landed __syncthreads(); if ( k0 + BK < K ) stage( buf ^ 1, k0 + BK ); // prefetch next -- overlaps the MMAs below WMMA_COMPUTE( As[buf], Bs[buf] ) __syncthreads(); buf ^= 1;}cp.async
(sgemm_wmma_db).where stage() issues one 16-byte
__pipeline_memcpy_async per thread per slab and commits the
group. The prefetch of the next slab is launched before the
current step’s tensor-core work, so the copy and the compute overlap. On
the 128×64 kernel, double-buffering lifts throughput from 8718 to
9987 GFLOP/s – a 14.5% gain, and the best hand-rolled
result in the chapter.
Figure 17-5. Double-buffering: each MMA waits for its buffer’s cp.async to land, while the next buffer streams in behind the current MMAs.
The FP32 register-blocked kernel of Section 17.4 records a similar performance improvement from the same treatment: its 128×64 tile does little enough arithmetic per slab that its load, too, sits on the critical path waiting to be hidden. Asynchronous copy overlaps the global load with computation, so it helps whenever a kernel has real arithmetic to hide the transfer behind and the load is otherwise exposed. Both GEMM kernels qualify; bandwidth-bound reduction (Chapter 12) and scan (Chapter 13) do not, because for those workloads, streaming the data is the work – there is not enough computation to hide the transfer. The rule is the one that already explained the scan result:
Asynchronous data movement improves performance when there is computation to hide the memory load behind, and the load is otherwise on the critical path.
Tensor Cores push GEMM furthest across that line: the faster the multiply-accumulates retire, the larger the fraction of the runtime that would otherwise be spent waiting on memory, and the more the overlap is worth.