Cache blocking leans on the hardware’s eviction policy to keep the needed data resident. It works well, but the program has little say in which data is retained or evicted from the cache. Shared memory is a software-managed cache: a block stages exactly the tiles it wants and holds them until it is finished processing them.
The textbook tiling gives each block a
BM×BN tile of C and marches
K in BK-deep steps, staging a
BM×BK slab of A and a
BK×BN slab of B in shared memory –
the tile dimensions are three independent parameters – with one output
per thread:
template<int BM, int BN, int BK>__global__ void sgemm_tiled( int M, int N, int K, const float *A, const float *B, float *C ){ __shared__ float As[BM][BK]; __shared__ float Bs[BK][BN]; const int tx = threadIdx.x, ty = threadIdx.y, tid = ty*BN + tx, nThreads = BM*BN; for ( int brow = blockIdx.y*BM; brow < M; brow += gridDim.y*BM ) { for ( int bcol = blockIdx.x*BN; bcol < N; bcol += gridDim.x*BN ) { float acc = 0.f; for ( int k0 = 0; k0 < K; k0 += BK ) { for ( int e = tid; e < BM*BK; e += nThreads ) As[e/BK][e%BK] = A[(brow+e/BK)*K + (k0+e%BK)]; for ( int e = tid; e < BK*BN; e += nThreads ) Bs[e/BN][e%BN] = B[(k0+e/BN)*N + (bcol+e%BN)]; __syncthreads(); #pragma unroll for ( int k = 0; k < BK; k++ ) acc += As[ty][k] * Bs[k][tx]; __syncthreads(); } C[(brow+ty)*N + (bcol+tx)] = acc; } }}sgemm_tiled).With a 32×32×32 tile, this implementation reaches only 825 GFLOP/s – no better than the naive kernel, and far below cache blocking’s 3762. The cache blocking already had captured the reuse that one-output-per-thread tiling can exploit, so shared memory and its barriers add overhead without adding reuse; with only one output per thread, every multiply-add still consumes two fresh shared-memory reads, and the ceiling simply moves from DRAM bandwidth to shared-memory bandwidth.
The remedy is the one from Section 17.3, now applied on top of shared
memory: give each thread an 8×8 microtile so that each shared-memory
read feeds eight multiply-adds, and stage the
BM×BK and BK×BN
slabs from shared memory, where the reuse is explicit rather than left
to the cache.
template<int BM, int BN, int BK, int TM, int TN>__global__ void sgemm_regblock( int M, int N, int K, const float *A, const float *B, float *C ){ const int nThreads = (BM/TM)*(BN/TN); __shared__ float As[BK][BM]; // transposed: As[k][m] __shared__ float Bs[BK][BN]; const int tid = threadIdx.x; const int threadRow = tid / (BN/TN), threadCol = tid % (BN/TN); const int rowA = tid / BK, colA = tid % BK, strideA = nThreads / BK; const int rowB = tid / BN, colB = tid % BN, strideB = nThreads / BN; for ( int bRow = blockIdx.y*BM; bRow < M; bRow += gridDim.y*BM ) { for ( int bCol = blockIdx.x*BN; bCol < N; bCol += gridDim.x*BN ) { float acc[TM][TN] = {}; for ( int k0 = 0; k0 < K; k0 += BK ) { for ( int off = 0; off < BM; off += strideA ) // stage A slab (transposed) As[colA][rowA+off] = A[(bRow+rowA+off)*K + (k0+colA)]; for ( int off = 0; off < BK; off += strideB ) // stage B slab Bs[rowB+off][colB] = B[(k0+rowB+off)*N + (bCol+colB)]; __syncthreads(); #pragma unroll for ( int k = 0; k < BK; k++ ) { float rA[TM], rB[TN]; for ( int i = 0; i < TM; i++ ) rA[i] = As[k][threadRow*TM + i]; for ( int j = 0; j < TN; j++ ) rB[j] = Bs[k][threadCol*TN + j]; for ( int i = 0; i < TM; i++ ) { for ( int j = 0; j < TN; j++ ) acc[i][j] += rA[i]*rB[j]; } } __syncthreads(); } // ... store acc[TM][TN] back to C ... } }}sgemm_regblock).With a 128×64 block tile and 8×8 microtiles, the register-blocked kernel reaches 6172 GFLOP/s – 1.6× the cache-blocked kernel, and about 74% of cuBLAS’s FP32 result. That gain is what shared memory delivers once it is paired with register blocking: the reuse is explicit and conflict-free instead of left to the cache’s discretion, each slab feeds thousands of multiply-adds, and the arithmetic-to-bytes ratio carries the kernel past the ridge of the roofline into compute-bound territory on the FP32 cores.
Figure 17-3. A thread’s TM×TN microtile is an outer product: each value read from shared memory feeds a whole row or column of multiply-adds.