feat(normalization): move dynamic 4PE implementations into solution tree - #113
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c0aa47ec11
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| SCRIPT_DIR = Path(__file__).resolve().parent | ||
| DEFAULT_CMP_DIR = ( | ||
| SCRIPT_DIR.parents[4] | ||
| SCRIPT_DIR.parents[5] |
There was a problem hiding this comment.
Use SCRIPT_DIR.parents[4] here and in the other moved generators/comparators. The move from test/kernel/normalization/... to test/solution/normalization/... did not add a directory level, so parents[5] resolves to benchmark and makes the default tools read and write benchmark/compare/...; the Makefile's CHK_DIR and the RMSNorm runner use benchmark/one-level-arch/compare/.... Consequently, invoking any of these scripts with its documented defaults generates inputs where the ELF cannot read them and looks for outputs where the simulator never writes them.
Useful? React with 👍 / 👎.
| compile_operator "$REPO_ROOT/test/kernel/sort" "sort" | ||
| compile_operator "$REPO_ROOT/test/kernel/deepseek" "deepseek" | ||
| compile_operator "$REPO_ROOT/test/kernel/flashMLA" "flashMLA" | ||
| compile_operator "$REPO_ROOT/test/solution/normalization/rms_norm" "solution/normalization/rms_norm" |
There was a problem hiding this comment.
Keep the existing test/kernel/deepseek invocation when adding the solution cases. That directory and its compile.all script are still present and cover numerous DeepSeek kernels, but this hunk removes their only entry from compile_all.sh; users running the advertised full compilation will therefore stop receiving those ELFs or detecting compilation regressions in them.
Useful? React with 👍 / 👎.
| template <typename TileVec> | ||
| inline void rsqrt_newton(TileVec &out, TileVec &a) { | ||
| auto body = [&](auto &x, auto &t1, auto &t2) { | ||
| TRECIP(x, a); |
There was a problem hiding this comment.
Do not initialize the inverse-square-root Newton iteration with 1/a without range reduction. When a row has mean(x²) + eps below roughly one third, that seed is outside the iteration's convergence basin; for example, a constant input of 0.5 gives a≈0.25, starts at x=4, and converges to -2 rather than +2, flipping the normalized output's sign. Larger denominators can also remain substantially inaccurate after only four steps. The same replacement appears in rms_norm_binary, so both kernels need an actual rsqrt estimate or a properly scaled seed.
Useful? React with 👍 / 👎.
| const uint32_t pe = read_pe_id(); | ||
| kPeBarrier[pe] = phase; | ||
| for (int participant = 0; participant < peNum; ++participant) { | ||
| while (kPeBarrier[participant] < phase) { |
There was a problem hiding this comment.
Make the barrier phase relative to an invocation rather than reusing the fixed values 1, 2, and 3. After the first group_norm_grad call, every slot remains at phase 3; on a second call, the first PE to reach phase 1 sees the other slots already greater than 1 and proceeds immediately into fused-parameter computation before those PEs have produced the new ds/db values. The phase-2 and phase-3 barriers fail similarly, so repeated calls can consume stale workspace data or return before peer writes complete.
Useful? React with 👍 / 👎.
| if (globalA <= 0 || gR <= 0 || tile_a <= 0 || tile_r <= 0 || | ||
| tile_r > tR || tid >= static_cast<uint32_t>(peNum)) { |
There was a problem hiding this comment.
Require tile_r == gR here, or add an R-tile accumulation and output loop. A positive tile_r < gR currently passes validation; for example, tiling {8, 1024, 1, 512} makes each call load and write only the first 512 columns of a row while dividing that partial sum by 1024, leaving the remaining columns untouched. Because the public dynamic entry accepts this tiling rather than reporting it as unsupported, callers can silently receive both an incorrect scale and an incomplete output tensor.
Useful? React with 👍 / 👎.
| const int64_t rid = r > 0 ? rms_bin::GetCacheId(r - 1) : 0; | ||
| gm_f gr(cache + rid * stride, 1, rms_bin::kWsCols); | ||
| TLOAD(sum, gr); |
There was a problem hiding this comment.
Do not assume the number of cache updates r is a power of two for every accepted tiling. For valid values such as gR=9000, powR=8192, and tile_r=4096, the rem/head loops produce three updates: the first two are merged into cache level 1, while the third remains in level 0. GetCacheId(r - 1) then selects only level 0, discarding the first two partial sums and producing a severely incorrect RMS scale. Either reduce all occupied cache levels at the end or reject tilings that do not guarantee a single final cache entry.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Bound cache carries to the allocated level count
When an accepted tiling produces at least 64 reduction updates, cid can exceed the six allocated cache levels. For example, gR=128, powR=64, and tile_r=1 are accepted and produce 64 paired updates; on the last one GetCacheId(63) is 6, so this stores at cache + 6 * stride although valid levels are only 0–5, corrupting memory beyond the documented workspace. Either allocate enough levels dynamically or reject tilings whose update count exceeds the cache capacity.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (peA < tile_a) { | ||
| return; | ||
| } |
There was a problem hiding this comment.
For otherwise accepted tilings where the final PE receives fewer rows than tile_a, this return leaves those rows untouched even though the tail call below is designed to process a smaller block. For example, globalA=10, gR=512, tile_a=2, and tile_r=512 assigns one row to PE3, which returns here and never writes row 9. Allow the short partition to reach the tail path rather than rejecting it.
Useful? React with 👍 / 👎.
| readBinaryFile(CHK_DIR "/input.bin", (uint8_t *)input, | ||
| static_cast<size_t>(g_a) * g_r * sizeof(dtype)); | ||
| input_ready = 1; | ||
| } else { | ||
| while (input_ready == 0) { |
There was a problem hiding this comment.
In an optimized four-PE RES_CHECK run, this hand-written volatile handshake does not prevent the non-volatile input accesses in the inlined kernel from moving across the readiness flag, so worker PEs can begin loading the shared buffers before PE0 has finished publishing them. The repository's res_check_publish_inputs helper places compiler memory barriers before the flag store and after the spin specifically for this purpose; use that helper (and its completion counterpart) in all four new solution drivers instead of the duplicated flag loops.
Useful? React with 👍 / 👎.
2 participants
Summary: move the PR #84 dynamic 4PE normalization kernels and tests into the new solution trees; update includes, data paths, ELF names, documentation and compile_all.sh; retain only dynamic 4PE solution cases; make GroupNormGrad MSE validation scale-aware while retaining atol/rtol checks; configure RMSNorm precision gfrun explicitly for 4PE without instruction tracing. Validation: base 4eb1fa5, LLVM 1ae4ee39, TileOP 804eb035, SuperScalarModel bc7fae0, build -j2. RMSNorm [512,8192], RMSNormBinary [16,16384], GroupNormGrad N32/C16/G8/HxW8192 and GroupNormGrad1D N512/C64/G8/D8 all compile and pass 4PE gfrun (Reach End, R2=0) plus precision comparison. Supersedes the normalization implementation work in #84.