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.