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 access.

## Performance requirements
- Bank conflict elimination should target 0-conflict or at worst 2-way conflict for all warp access patterns in the inner loop. Nsight Compute's `l1tex__data_bank_conflicts_pipe_lsu_mem_shared` counter shows actual conflicts; aim for near-zero.
- Double buffering with `cp.async` should produce measurable overlap between global memory loads and compute for sufficiently large tiles. Profile with and without double buffering to confirm the benefit. For small tiles (BK < 8), the overlap may be negligible.
- Smem occupancy target: design tiles so that at least 2 blocks are active per SM (for latency hiding). For SM80 with 164 KB smem: 2 active blocks → each block can use up to 82 KB.
- Vectorized `cp.async` (16-byte transfers, `cp.async.cg.shared.global ... 16`) is more efficient than scalar or 4-byte async copies. Use 16-byte async copies when the tile row size is a multiple of 16 bytes and the source pointer is 16-byte aligned.

## Output format
The final response must include:
1. **Bank conflict analysis**: derive the bank mapping for each smem access in the inner loop. Show the calculation explicitly.
2. **Conflict resolution**: specify the padding required (which dimension, how many elements). Verify the bank mapping after padding.
3. **Smem layout specification**: exact array declarations with dimensions and padding. ASCII diagram of the layout if helpful.
4. **Tile size recommendation**: BM, BN, BK (or equivalent), with justification based on occupancy and register budget.
5. **Double buffering implementation** (if SM80+): smem buffer declarations, `cp.async` call pattern, fence placement, sync placement, buffer index management.
6. **Occupancy estimate**: smem per block, active blocks per SM, theoretical occupancy percentage.
7. **Before/after comparison**: show the smem access code before and after the optimization, with the specific change highlighted.

## Common failure modes
- **Padding the wrong dimension**: adding padding columns to a row-major smem array when the conflict is caused by row-access (column strided) patterns. The conflict is along the column direction, and padding adds extra bytes per row — which shifts bank indices for column accesses. This may or may not fix the conflict. Derive the bank mapping with and without the padding to verify.
- **Wrong async copy fence placement**: `cp.async.wait_group(1)` is called before all copies for the current tile have been committed (`cp.async.commit_group()` was not called). The wait completes but the last batch of copies is still in flight. Fix: always call `commit_group` immediately after all `cp.async` instructions for a tile are issued, before the `wait_group` for the preceding tile.
- **Missing `__syncthreads` after wait_group**: the DMA completes (wait_group returns) but threads read smem before the completion propagates. Fix: always follow `wait_group` with `__syncthreads()`.
- **Exceeding per-SM smem without setting max**: requesting more than 48 KB per block without `cudaFuncSetAttribute` causes a silent launch failure (the kernel does not run, or runs with the old smem limit and silently overwrites memory). Fix: always call `cudaFuncSetAttribute` for kernels using > 48 KB smem.
- **Double buffer off-by-one**: the computation phase reads from buffer `ping` while the async load is writing to buffer `ping` (not `pong`). This happens when the buffer index is not correctly alternated. Fix: trace the buffer index manually through two loop iterations to verify correctness.
- **fp16 smem bank conflicts**: declaring `__shared__ half smem[M][N]` and accessing with consecutive thread indices causes 2-way bank conflicts (two adjacent fp16 elements share a bank). Fix: use `half2` loads from smem, or declare as `__shared__ __half2 smem[M][N/2]` and access in pairs.

## Review checklist
- [ ] Has the bank index been computed analytically for each smem access in the inner loop, not just assumed?
- [ ] Is the padding applied to the correct dimension (the one causing the stride-32 bank conflict pattern)?
- [ ] Are all smem index computations using the padded stride (BK + pad) consistently?
- [ ] For double buffering: is `cp.async.commit_group()` called after all `cp.async` instructions for a tile?
- [ ] Is `cp.async.wait_group(N)` called with the correct N (1 for double buffering, not 0 which eliminates overlap)?
- [ ] Is `__syncthreads()` called after `wait_group` and before reading the newly filled buffer?
- [ ] Is the double buffer `ping/pong` index correctly alternated on every loop iteration?
- [ ] Is `cudaFuncSetAttribute` called before launch for any kernel using > 48 KB of smem?
- [ ] Does the total smem per block (including padding bytes and double buffer factor) fit within the requested allocation?
- [ ] Has Nsight Compute been used to verify that bank conflicts are eliminated (not just estimated)?
- [ ] Is the occupancy impact of the smem allocation acceptable (at least 2 active blocks per SM)?