optimize-shared-memory-tiling

Skill: Optimize Shared Memory Tiling

Purpose

Guide the agent through designing and tuning shared memory tiling strategies for CUDA kernels, covering bank conflict analysis and elimination, tile shape selection, double buffering with async copy, occupancy tradeoffs from shared memory allocation, and the decision of when smem tiling is worth the complexity.

Use this when

• A kernel repeatedly reads the same global memory data from multiple threads and would benefit from staging through a shared memory tile (GEMM, convolution, stencil)
• Profiling shows high shared memory bank conflict rates in Nsight Compute (l1tex__data_bank_conflicts_pipe_lsu_mem_shared)
• You are designing the smem layout for a GEMM or attention tiling kernel and need to choose tile dimensions and padding
• You are adding double buffering to overlap global memory loads with computation using cp.async on SM80+
• An existing kernel has smem usage that limits occupancy and needs to be restructured

Do not use this when

• The access pattern is truly random (no spatial reuse) and shared memory staging will not increase the reuse factor
• The data is small enough to fit in L1 cache across all accesses without explicit smem management: compiler-managed L1 may be sufficient
• The kernel is compute-bound and memory latency is not the bottleneck: smem tiling adds complexity without throughput benefit
• The kernel is a simple elementwise operation with no data reuse: smem staging provides no benefit

Inputs the agent should gather first

• Kernel type and access pattern: describe the data reuse structure — which threads reuse which elements, and along which dimension. E.g., in GEMM, each row of threads reuses a row of A and each column of threads reuses a column of B.
• Tile shape context: what are the thread block dimensions (BM, BN, BK) or equivalent? How many threads are in the block?
• Dtype and element size: fp32 (4 bytes), fp16 (2 bytes), int8 (1 byte) — affects how many elements map to a single bank.
• Hardware target: SM architecture. On all current NVIDIA GPUs (Kepler through Hopper): 32 banks, 4-byte bank width (by default; can be configured to 8-byte via cudaDeviceSetSharedMemConfig). SM80+ has up to 164 KB smem per SM in some configurations.
• Current smem usage: how many bytes of smem are currently allocated per block? How does this affect occupancy?
• Whether async copy is applicable: SM80+ with cp.async, SM90 with cp.async.bulk (TMA). Is the kernel targeting those architectures?

Required reasoning process

1. Map each smem access to a bank. Bank index for a 4-byte access is (address_in_smem_bytes / 4) % 32. For an smem array T smem[rows][cols], element smem[r][c] maps to bank (r * cols + c) * sizeof(T) / 4 % 32 (assuming sizeof(T) is a multiple of 4; for fp16, adjust accordingly).

2. Identify conflicting access patterns. A bank conflict occurs when multiple threads in the same warp access different addresses that map to the same bank in the same memory transaction. No conflict: all threads access the same bank (broadcast). No conflict: all threads access different banks. Conflict: k threads access k different addresses in the same bank (k-way conflict, requiring k serial transactions).

3. Compute the conflict multiplier for the specific access pattern.

   • GEMM A-tile (BM x BK, row-major smem, threads access along the K dimension): for warp accessing a row of A, consecutive threads read consecutive elements of the K dimension. If BK is a multiple of 32, then threads 0..31 read banks 0..31 — no conflict. If BK = 32 and sizeof(T) = 4, this is ideal.
   • GEMM B-tile (BK x BN, row-major smem, threads access down the K dimension): threads in a warp that compute different columns of C all need to read the same column of B. This causes a BK-way broadcast conflict for column-major smem access — wait, a broadcast is not a conflict. But if threads read from different rows of B (different K values) for the same column, the bank is (k * BN + n) % 32. For n fixed, varying k with stride BN: conflicts occur when BN is a multiple of 32. Work through the specific access pattern systematically rather than guessing.

4. Apply smem padding to eliminate conflicts. The standard technique: add pad extra columns to the smem array, changing the stride from BK to BK + pad. This shifts the bank mapping for each row, breaking the periodic conflict pattern. For fp32 with 32-bank smem, adding 1 column of padding (4 bytes) shifts consecutive rows by 1 bank, typically eliminating 32-way column-access conflicts. For fp16, adding 2 columns of padding shifts by 1 bank (since 2 fp16 = 4 bytes = 1 bank). Choose the minimum pad that eliminates the conflict.

5. Verify padding after application. Re-derive the bank mapping for the padded layout and verify that no two threads in a warp access the same bank. Walk through at least: warp accessing a row (consecutive K or N access) and warp accessing a column (strided K access). Confirm that the conflict multiplier is 1 after padding.

6. Choose tile dimensions relative to warp size. The innermost loop access should map 32 consecutive threads to 32 different smem banks. This requires the innermost stride of the smem array to result in 32 distinct bank indices across 32 threads. For row-major smem with row stride of BK (or BN): if BK (or BN) >= 32 and the access is along the row direction, and dtype is fp32, consecutive threads access consecutive elements → consecutive banks → no conflict.

7. Design double buffering with cp.async (SM80+). Allocate two smem buffers: smem_A[2][BM][BK+pad] and smem_B[2][BK][BN+pad]. Use buffer index ping = 0, pong = 1 alternating. While computing on smem_*[ping], issue async copies into smem_*[pong] using cp.async:

   cp.async.cg.global.shared::cta [smem_ptr], [global_ptr], 16;
   

   After issuing all copies for the next tile: cp.async.commit_group(). Before using smem_*[pong] for computation: cp.async.wait_group(1) (leaves at most 1 uncommitted group in flight, i.e., allows the previous group to still be pending but waits for all earlier groups). Then __syncthreads() to ensure all threads see the completed stores to smem.

8. Estimate occupancy impact. Look up the smem capacity per SM for the target architecture:

   • SM70 (V100): 96 KB per SM
   • SM80 (A100): 164 KB per SM (with cudaFuncSetAttribute for max smem)
   • SM86 (A10/A30): 100 KB per SM
   • SM90 (H100): 228 KB per SM Active blocks per SM = floor(smem_per_SM / smem_per_block). If this limits active blocks below the warp-level occupancy target, consider reducing tile size or using a different smem configuration. Request maximum smem with cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes).

9. Consider L1 cache as an alternative to explicit smem. If the access pattern has low spatial reuse but high temporal reuse (the same addresses are re-read by different thread blocks), L2 cache may serve those reads without smem tiling. For access patterns where data is reused only within a single thread block execution, explicit smem tiling is essential. For access patterns with cross-block reuse, L2 is the relevant level.

Kernel design rules

• Determine bank conflicts analytically before writing any smem access code. Do not add padding speculatively without understanding the access pattern.
• Padding must be applied to the dimension that causes the conflict — adding padding to the wrong dimension wastes smem and does not fix the conflict.
• For double buffering: cp.async.commit_group() must be called after all cp.async instructions for one tile are issued. cp.async.wait_group(N) waits until at most N groups are pending. Using wait_group(0) before compute is correct but eliminates the overlap benefit; use wait_group(1) to keep one group in flight while computing on the other.
• __syncthreads() must be called after cp.async.wait_group and before reading from the smem buffer that was just filled. The async copy completes the DMA to smem, but __syncthreads() ensures all threads see the completion and prevents earlier-executing threads from reading stale data.
• For sm_80 and later: request increased smem with cudaFuncSetAttribute. Kernels requiring more than 48 KB per block must explicitly opt in. Failure to call this results in a launch failure.
• Tile sizes must be chosen such that BM * BK * sizeof(dtype) * 2 + BN * BK * sizeof(dtype) * 2 (for double-buffered A and B) fits within the requested smem allocation.

Correctness requirements

• cp.async.wait_group placement: the wait must occur before the computation phase that reads from the newly filled smem buffer. Placing it after the computation reads from uninitialized or stale smem.
• __syncthreads after async copy completion: cp.async.wait_group ensures the DMA is complete, but __syncthreads is still required to ensure all threads in the block see the completion before any thread reads. These are two separate requirements.
• Padding array sizing: if smem_A is declared as __shared__ float smem_A[BM][BK + pad], then all index computations into smem_A must use BK + pad as the stride, not BK. Indexing with BK after declaring with BK + pad accesses wrong elements.
• Double buffer index alternation: the ping/pong buffer index must be correctly alternated on each K-loop iteration. An off-by-one in the buffer index means computing on the wrong (possibly still-in-flight) buffer.
• Smem bank conflict analysis for fp16: for fp16 (2-byte elements), two consecutive elements share one bank. The bank index for smem[i] is (i * 2) / 4 % 32 = i / 2 % 32. Two adjacent elements (i and i+1) map to the same bank if i is even. This means warp accesses to consecutive fp16 elements in smem may cause 2-way conflicts unless addressed with half2 vectorized smem acces

---

Content truncated.