code/tt_diffusion_planner/tt/kernels/README.md

# Custom kernels of diffusion-planner-p150

The functional port maps every op to stock ttnn (PLAN.md section 2.12). The optimization rounds add `ttnn.generic_op`
programs whose sources live here (package data in the repo-root `pyproject.toml`, asserted by a `verify:` line of
`tt-model.yaml`). Each one has a device check against the stock graph it replaces (bit for bit where the math is
the same) and went through the hang protocol of PLAN.md section 4.4 (smallest shape under `TT_METAL_WATCHER=2`,
eager twice, a trace): `logs/diffusion-planner/opt_r2/scripts/ln_check.py`, `kcat_check.py`, `smask_check.py`,
`logs/diffusion-planner/opt_r3/scripts/smsm_check.py`, `fattn_check.py`, `kemit_check.py`, `lntr_check.py`.

Per-core runtime-argument counts are compile-time args (probe P3: the generic_op program hash ignores them).

| files | op | used by | knob | status |
|---|---|---|---|---|
| `ln32_reader.cpp`, `ln32_compute.cpp`, `ln32_writer.cpp` | fp32 LayerNorm over the last dim (the 9-program `layer_norm_fp32` decomposition in one program, the same SFPU LLK calls in the same order), optionally with the residual add `h = x + r (* gate)` in front | `tt/ln_kernel.py`, `tt/layers.py` `LayerNorm` (mixers, decoder) | `LN_KERNEL`, `LN_RESID` | bit-identical to the stock decomposition (OPT round 2) |
| `kcat_reader.cpp`, `kcat_compute.cpp`, `kcat_writer.cpp` | operand of the K-concatenated split matmul (optionally after the previous linear's GELU, `KCAT_ACT`): fp32 `x` -> `[bf16(x) \| bf16(x) \| x - bf16(x) \| 1 \| 0...]` | `tt/kcat_kernel.py`, `tt/layers.py` `SplitLinear` (decoder) | `SPLIT_KCAT=2` | bit-identical to the stock-op operand (`SPLIT_KCAT=1`) (OPT round 2) |
| `smask_reader.cpp`, `smask_compute.cpp`, `smask_writer.cpp` | attention score scale + additive mask: `s * scale + mask` in one pass (the two stock binary_ng programs; with a bf16 mask the stock add runs on the FPU, which truncates the scores to TF32, reproduced with an SFPU AND) | `tt/smask_kernel.py`, `tt/attention.py` (decoder self / cross, fusion) | `ATTN_SMASK` | bit-identical (OPT round 2) |
| `smsm_reader.cpp`, `smsm_compute.cpp`, `smsm_writer.cpp` | attention score scale + mask + softmax over one tile row (head, 32 queries) per pass: the smask SFPU sequence into an L1 row, then the stock `ttnn.softmax(numeric_stable=True)` compute (`kernel_lib` reduce / bcast / exp calls in the stock order: FPU row max, `exp(x - max)`, FPU row sum + precise fp32 reciprocal, bcast multiply); the scale either by a scale tile (`mul_binary_tile`) or an immediate (`mul_unary_tile`, the same fp32 SFPU multiply) | `tt/smsm_kernel.py`, `tt/attention.py` (decoder self / cross, fusion) | `ATTN_SMSM` | bit-identical to smask + `ttnn.softmax` (OPT round 3) |
| `fattn_reader.cpp`, `fattn_compute.cpp`, `fattn_writer.cpp` | the whole fp32 matmul attention per (head, query tile row): `Q K^T` into DEST (`matmul_block`, in1 transposed), the smsm phases (scale by an immediate, TF32 truncation, mask, stock softmax `kernel_lib` calls) on the row in L1, `P V` accumulated in DEST in key order; Q / K / V read in place from the projection outputs (or the hoisted cross K / V heads), the output written as tile (r, h) of the merged heads | `tt/fattn_kernel.py`, `tt/decoder.py` / `tt/encoder.py` (decoder self / cross, fusion) | `ATTN_FUSED` | bit-identical to head split + `Q K^T` + smsm + `P V` + head merge (OPT round 3) |
| `ln32s_reader.cpp`, `ln32s_compute.cpp`, `ln32s_writer.cpp` | the fused fp32 LayerNorm (+ residual) with each tile row spread over Wt cores: member j owns column tile j, the root gathers the tiles (semaphore 0), folds them in order and broadcasts mean / rstd (semaphore 1, monotonic counts) | `tt/ln_kernel.py` `layer_norm_fp32_split`, `LayerNorm` (decoder: 11 tile rows on 88 cores) | `LN_SPLIT` | bit-identical (OPT round 2) |
| `ln32_sfpu.h` | SFPU row sum of one fp32 tile with the sum stored to every column: the stock `sfpu_reduce<SUM, Float32, REDUCE_ROW>` arithmetic (same loads, replayed adds, butterfly) with three extra stores, so the LN statistics come out already broadcast (no RISC-V column fill) | `ln32_compute.cpp`, `ln32s_compute.cpp` | `LN_SFPU_BCAST` | bit-identical (OPT round 2) |

`ln32_sfpu.h` is included by the two LN compute kernels through `compiler_include_paths`; the JIT build cache keys a
kernel on its own source, compile-time args and defines, so after editing the header, also touch the including
`.cpp` files (or clear the kernel cache) before measuring.

`KCAT_EMIT` (OPT round 3): `ln32s_*.cpp` and `fattn_*.cpp` can write their output tile as the split operand
`[x_hi | x_hi | x_lo | 1 | 0..]` of the next K-concatenated decoder linear (the `kcat_compute.cpp` LLK calls on the
packed fp32 tile, the `kcat_writer.cpp` layout) instead of the plain output: bit-identical to output + `kcat_operand`.

`LN_TR` (OPT round 3): `ln32_*.cpp` can read the residual per-entity transposed (`[.., E, W, T]`, transposed back with
`transpose_tile`, the stock `ttnn.transpose` LLK with the fp32 unpack-to-dest: exact) and write the output the same
way (via c_18): the two transposes around the mixer's token-mixing MLP, bit-identical.