feat(normalization): update shapes and tile GroupNormGrad1D stages - #128
Conversation
|
Updated in d1572cb:\n\n- renamed group_norm_grad_pto.hpp to group_norm_grad.hpp\n- renamed group_norm_grad_1d_pto.hpp to group_norm_grad_1d.hpp\n- updated test includes; no GroupNormGrad kernel/test filename now contains _pto\n- corrected both generator/comparator default paths to benchmark/one-level-arch/compare\n\nRegression after renaming:\n- group_norm_grad 4PE gfrun precision: PASS (dx max_abs=0.001953, dgamma=0, dbeta=0)\n- group_norm_grad_1d 4PE gfrun precision: PASS (dx max_abs=0.0009766, dgamma=0.003906, dbeta=0) |
|
Updated group_norm_grad in commit 21c2c14:\n\n- 4PE dynamic shape: N=2, C=32, G=8, HxW=2024 (D=4)\n- runtime-derived tile_hw: 2024\n- gfrun precision: PASS (dx max_abs=0.0004883, dgamma=0, dbeta=0)\n- ELF SHA256: d3262648287d95280ff292d3454aa19e56cbf3837dc35aa33203e39b4bf9ba47 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 21c2c14f7c
ℹ️ 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".
| constexpr int64_t kMinFloatTile = 512 / sizeof(float); | ||
| if constexpr (kGroupWidth < kMinFloatTile) { |
There was a problem hiding this comment.
Route undersized dtype channel tiles to the fallback
When DType=__half and C_CH / G_GRP is between 128 and 255, this condition selects the Tile implementation because it only checks the float minimum, but that implementation also constructs half-precision channel tiles whose documented 512-byte minimum requires at least 256 elements. Such otherwise-valid custom Makefile shapes therefore still hit the invalid active-size problem that the scalar fallback is intended to avoid. Include the dtype tile minimum in this dispatch; the analogous condition in group_norm_grad/src/group_norm_grad.cpp has the same gap.
Useful? React with 👍 / 👎.
21c2c14 to
aa1f0fe
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 99f78bc375
ℹ️ 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".
| using htile = Tile<Location::Vec, dtype, 32, 256, BLayout::RowMajor, -1, -1>; | ||
| using ftile = Tile<Location::Vec, float, 32, 256, BLayout::RowMajor, -1, -1>; |
There was a problem hiding this comment.
Keep the batched dX working set within the Tile RF budget
For the default D=32, tile_g=8 path, dx_groups allocates four distinct physical [32,256] FP32 tiles (xf, dyf, gf, and out) at 32 KiB each, and all four remain live across the operations at lines 210–213. This exceeds the documented 64 KiB Tile RF budget before accounting for h and the vector tiles, matching the reported gfsim Local TLOAD requires one fitting destination Tile abort; shrink the physical tile dimensions or restructure the computation so the default benchmark can run in the simulator.
Useful? React with 👍 / 👎.
| using ht = Tile<Location::Vec, dtype, 32, 256, BLayout::RowMajor, -1, -1>; | ||
| using ft = Tile<Location::Vec, float, 32, 256, BLayout::RowMajor, -1, -1>; | ||
| using vt = Tile<Location::Vec, float, 32, 1, BLayout::RowMajor, -1, 1>; |
There was a problem hiding this comment.
Reduce the Stage B Tile working set before loading inputs
Whenever gb_tile_d <= 256, Stage B selects this helper, which allocates five physical [32,256] FP32 tiles (dyf, xf, tmp, beta, and grad), totaling 160 KiB while several must remain live throughout the N accumulation loop. Thus even after the dX allocation is corrected, the default Stage B path still cannot fit in the 64 KiB Tile RF and is liable to abort under gfsim; use smaller physical blocks or split the beta and gamma reductions into passes.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3dd27e8a73
ℹ️ 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".
| group_norm_grad_gamma_beta<dtype, PE_NUM>(mean, rstd, tiling_info, workspace, | ||
| dgamma, dbeta); |
There was a problem hiding this comment.
Synchronize before the cross-PE gamma/beta reduction
When PEs progress at different rates, this call can read ds/db workspace entries before the PEs responsible for those (n,g) entries finish the spatial stage. In the default tiling, PE 0 owns the single gamma/beta block (bg=8) and reduces all groups, while spatial work for those groups is distributed across all four PEs; merely calling the stages in the same order does not synchronize them. Add an inter-PE barrier before starting gamma/beta to prevent nondeterministic gradients.
Useful? React with 👍 / 👎.
| val = math.nan if mant else math.inf | ||
| else: | ||
| val = math.ldexp(1.0 + mant / 1024.0, exp - 15) | ||
| val = math.ldexp(1.0 + mant / 512.0, exp - 15) |
There was a problem hiding this comment.
Decode the half mantissa with 10-bit scaling
For every finite normalized half with a nonzero mantissa, dividing by 512 instead of 1024 decodes the value incorrectly (for example, a value immediately below 2 is decoded near 3). Because quantize_f16() feeds these decoded values into input generation, forward statistics, and golden gradients, the generated benchmark distribution is systematically distorted rather than representing the requested random FP16 inputs; use the 10-bit half mantissa denominator, as the comparator and 1D generator do.
Useful? React with 👍 / 👎.
| tile_f x_f(1, active_hw); | ||
| tile_f dy_f(1, active_hw); | ||
| tile_f dx_f(1, active_hw); | ||
| tile_f tmp(1, active_hw); |
There was a problem hiding this comment.
Reduce the GroupNorm dX Tile working set
For the default dh=2024, dc=1 path, these declarations create four physical [1,8192] FP32 Tiles (x_f, dy_f, dx_f, and tmp) of 32 KiB each. At TROWEXPANDMUL(dx_f, dy_f, c1), x_f must remain live for the following operation while dy_f and the distinct destination dx_f are also required, already exceeding the 64 KiB Tile RF budget. Restructure this computation or use smaller physical strips so the default GroupNorm benchmark can execute in the simulator.
Useful? React with 👍 / 👎.
| TF sf(rows, cols), bf(rows, cols), t(rows, cols), ga(rows, cols), | ||
| ba(rows, cols); |
There was a problem hiding this comment.
Split the GroupNorm gamma/beta Tile accumulation
Whenever bd <= 256, including the default D=4 configuration, this helper instantiates five physical [32,256] FP32 Tiles (sf, bf, t, ga, and ba), totaling 160 KiB. The two accumulators remain live across the entire N loop while the input and temporary Tiles are loaded, so the Stage B working set cannot fit in the 64 KiB Tile RF; use smaller physical blocks or separate the gamma and beta reductions.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6fdce762ce
ℹ️ 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".
| @@ -57,7 +57,7 @@ void rms_norm_binary(dtype *x, const int64_t *tiling, dtype *out, | |||
| float *workspace, float eps = 1e-6f) { | |||
| static_assert(peNum == 4, "normalization kernels support only 4PE"); | |||
| constexpr int64_t tA = 1; | |||
| constexpr int64_t tR = 8192; | |||
| constexpr int64_t tR = 512; | |||
There was a problem hiding this comment.
Prevent the smaller tiles from overflowing the cache workspace
With the dynamic test's automatically selected 512-column tiles, a supported build such as G_R=65536 performs 64 cache updates per row. The last update computes GetCacheId(63) == 6, but kMaxLevels and workspace_buf provide only levels 0–5, so RMS_BIN_UPDATE_CACHE writes beyond the caller-owned workspace; the previous 8192-column default required only four updates for this shape. Increase the cache level count and workspace allocation based on the maximum update count, or reject shapes that require level 6.
Useful? React with 👍 / 👎.
| parser.add_argument("--g-r", type=int, default=16384) | ||
| parser.add_argument("--tile-a", type=int, default=1) | ||
| parser.add_argument("--tile-r", type=int, default=8192) | ||
| parser.add_argument("--tile-r", type=int, default=512) |
There was a problem hiding this comment.
Finish non-power-of-two cache reductions
For custom dynamic shapes where the new 512-column default produces a non-power-of-two number of partials, the final cache lookup does not represent the whole row. For example, G_R=10000, pow_r=8192, and the default tile_r=512 generate 17 updates; the kernel then loads level GetCacheId(16) == 0, which contains only the seventeenth partial rather than all 17, yielding an incorrect RMS and output. The former 8192-column default produced two partials for this shape, so either merge all populated cache levels at completion or constrain/pad the partial count to a power of two.
Useful? React with 👍 / 👎.
| @@ -98,15 +96,15 @@ void rms_norm(dtype *x, const int64_t *tiling, dtype *out, float eps = 1e-6f) { | |||
| // Physical capacity (Rows×Cols); Valid comes from tiling (tile_a,tile_r). | |||
| // Size must cover ValidRow×ValidCol; SoftCore should not require Rows≥ValidRow. | |||
| constexpr int64_t tA = 1; | |||
| constexpr int64_t tR = 8192; | |||
| constexpr int64_t tR = 512; | |||
There was a problem hiding this comment.
Preserve the documented automatic tile selection
When a direct caller uses the documented tiling[3] <= 0 convention with a reduction wider than 512 columns, tile_r is still set to the full gR, so the new tile_r > tR guard immediately returns without producing output. For example, an 8192-column row with tiling[3] = 0 worked before this capacity reduction and the newly chunked implementation can process it, but it is now rejected. Default the unspecified tile width to min(gR, tR) rather than gR.
Useful? React with 👍 / 👎.
| tile_f x_f(1, active_d); | ||
| tile_f dy_f(1, active_d); | ||
| tile_f t0(1, active_d); | ||
| tile_f t1(1, active_d); |
There was a problem hiding this comment.
Reduce the scalar-group dX working set
When tile_g == 1—which is mandatory for D > 256 and also occurs for configurations such as G == 1—this fallback instantiates four distinct physical [1,8192] FP32 Tiles (x_f, dy_f, t0, and t1), totaling 128 KiB. All four are live while forming dX, so the newly supported large-D/single-group path cannot fit within the 64 KiB Tile RF even though it passes the tiling validation; process smaller physical strips or reuse fewer FP32 destinations.
Useful? React with 👍 / 👎.
| tile_f dy_f(1, vd); | ||
| tile_f x_f(1, vd); | ||
| tile_f t0(1, vd); | ||
| tile_f acc(1, vd); |
There was a problem hiding this comment.
Reduce the large-D dgamma working set
For gb_tile_d > 256, Stage B selects dgamma_group, whose physical column capacity is 8192 regardless of the active strip width. These declarations allocate four FP32 [1,8192] Tiles (dy_f, x_f, t0, and the persistent acc) totaling 128 KiB, and they are simultaneously live in the accumulation loop, so the advertised large-D split path exceeds the 64 KiB Tile RF. Use smaller physical blocks or split the accumulation into passes.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 97a0afc9cc
ℹ️ 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".
| / "solution_normalization_group_norm_grad_group_norm_grad" | ||
| "_DType__half_N32_C16_G8_HxW8192_PE4" | ||
| "_DType__half_N2_C32_G8_HxW2024_PE4" |
There was a problem hiding this comment.
Include the dynamic testcase name in the compare path
When the default generator is used with the newly selected group_norm_grad_dynamic testcase, it writes fixtures under a directory whose basename contains _group_norm_grad_group_norm_grad_, while the Makefile derives CHK_DIR from a target containing _group_norm_grad_group_norm_grad_dynamic_. Consequently, a res_check=on run cannot find the generated inputs, and the comparator mirrors the same incorrect default path; add _dynamic here and in group_norm_grad_data_compare.py.
Useful? React with 👍 / 👎.
Summary
Update the GroupNorm backward test shapes and implement dynamic 4PE GroupNormGrad1D with Tile-only computation.
Validation of current GroupNormGrad1D implementation
Environment: LLVM 553b08045111eb9f9858fa782031f2e19f615a9b; TileOP b8669ce plus issue #100 fix (local commit 013e5658da9ced382257d401cf221697eec4bc54); SuperScalarModel 6a9b739cb8b60c7d87d9b0d57412921532104331.
All three 4PE gfrun accuracy checks PASS (dx, dgamma, dbeta):
git diff --check: PASS.
Known gfsim failures
The following GroupNormGrad1D results were measured before the contiguous c2/c3 layout change (aa1f0fe); gfsim has not been rerun after that change. Separate performance ELFs were built with res_check=off. Real L2 only: --conf fourpe --pto-v02 true -s tlsu.fake_l2_enable=false. All three runs exited with SIGABRT (134), without valid final cycle counts:
Root causes are not yet established; these failures must not be treated as passing performance validation. This validation covers GroupNormGrad1D, not a fresh GroupNormGrad run.
Latest validation notes