feat(x86): AVX-512 row-reduction kernel (sum/max/mean, f32) — validated on AVX-512 hardware - #133
Conversation
…ed on-device
First artifact of the "real optimizing-compiler backend tier" work: an AVX-512
vectorized row-wise reduction over the last axis for the S-series reduction
family (reduce_sum / mean / max) — the CPU analog of the AMX/AVX-512 GEMM lane,
so reductions get a genuine vectorized CPU kernel instead of only the numpy
reference.
- tessera_x86_avx512_reduce_f32: 16-lane __m512 accumulate + _mm512_reduce_{add,
max}_ps horizontal reduce, scalar column tail. Scalar reference alongside for
validation.
- test_reduce: AVX-512 vs scalar-reference vs an independent hand-computed
expectation, across sum/max/mean × shapes (16-aligned, %16 tail, sub-16
scalar-only, wide, degenerate). Validated ON-DEVICE on this AVX-512 host
(avx512f/bw/dq/vl) — ALL PASSED (15/15).
Proves the CPU half of the per-op multi-device completeness path (CPU AVX +
ROCm). Follow-ons to fully complete the reduce_sum op: Graph IR → x86 lowering
wiring + manifest x86 slot, the ROCm reduction kernel (gfx1151), and the
dashboard move. NOTE: the full tessera_x86_backend library has pre-existing
gcc build breakage in unrelated kernels (amx_runtime/amx_gemm/vnni) — this
kernel builds + runs standalone; that breakage is a separate cleanup.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
![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: fb2f7f9408
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| vacc = _mm512_max_ps(vacc, _mm512_loadu_ps(row + c)); | ||
| float acc = _mm512_reduce_max_ps(vacc); | ||
| for (; c < cols; ++c) acc = row[c] > acc ? row[c] : acc; |
There was a problem hiding this comment.
When a reduced row contains a NaN, this max path ignores it (for example a row like [0, NaN, 1, ...] returns the finite maximum) because both _mm512_max_ps and the scalar tail use ordered max/comparison. The existing reduction reference path maps max to NumPy (np.amax in _apple_gpu_dispatch_reduce), which propagates NaNs, so wiring reduce_max to this kernel will silently diverge from the current CPU/reference behavior for NaN-containing tensors; handle NaNs explicitly in the vector, tail, and reference paths.
Useful? React with 👍 / 👎.
…PR#133 review + CI) Two fixes: - Review (P2 — NaN semantics): reduce_max ignored NaN (plain MAXPS / ordered `>` return the finite max), diverging from the reference `np.amax` which PROPAGATES NaN. Now track NaN explicitly — vector body via _mm512_cmp_ps_mask(v,v, _CMP_UNORD_Q), scalar tail + reference via v!=v — and force a quiet NaN out when any row element is NaN. (sum/mean already propagate via `+`.) Test gains NaN-in-vector-body + NaN-in-scalar-tail cases for all three kinds; 19/19 pass on-device. - CI (drift): the new `tessera_x86_avx512_reduce_f32` C-ABI symbol is auto-inventoried by the runtime_abi dashboard, so the on-disk CSV was stale. Regenerated runtime_abi (+ fleet); 17 dashboards back in sync, 62 drift-gate tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Both addressed in NaN semantics (P2): you're right — CI (drift): the new |
…letes the reduction family's ROCm device (#134) The ROCm analog of the x86 AVX-512 reduction lane (PR #133), continuing the per-op multi-device completeness path. A tessera_rocm.reduce directive + generate-rocm-reduce-kernel pass emit a row-reduction kernel: one workgroup per row, lanes stride the last axis and tree-reduce through LDS (identity-seeded combine = +/max/min; mean divides by K). f32 reduce regardless of storage. Runtime lane rocm_reduce_compiled folds an arbitrary reduced `axis` to a [outer, inner] last-axis reduction by transposing the reduced axes to the end (matching _apple_gpu_dispatch_reduce); keepdims supported. Handles tessera.sum/mean/max/min (+ amax/amin) by op name; f16/bf16/f32. Full wiring: ODS op + pass (registered + CMake) + executor + execution_matrix row + KNOWN_EXECUTORS + 6 _ROCM_COMPILED entries + fixtures. Validated on gfx1151 vs numpy across dtype × shape × axis incl. rank-3, reduce-all, keepdims (test_rocm_reduce_compiled.py) + a GPU-free codegen gate. 185 gate tests pass; mypy + ruff clean; 17 dashboards in sync. With the AVX-512 CPU kernel (#133), the S2 reduction family now has a real optimized kernel + on-device test on BOTH devices we have hardware for. Co-authored-by: gstoner <angstroms01@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 participant
First artifact of the real optimizing-compiler backend tier — proving the per-op, per-device completeness path on the CPU side.
The S-series reduction family (
reduce_sum/mean/max) previously had only the numpy reference on CPU (today only matmul/gemm have a native AVX/AMX kernel). This adds a genuine AVX-512 vectorized row-reduction kernel:tessera_x86_avx512_reduce_f32— 16-lane__m512accumulate +_mm512_reduce_{add,max}_pshorizontal reduce, scalar column tail; with a scalar reference alongside.test_reduce— AVX-512 vs scalar-reference vs an independent hand-computed expectation, across sum/max/mean × shapes (16-aligned,%16tail, sub-16 scalar-only, wide 1024, degenerate 1-col).Validated on-device on this AVX-512 host (
avx512f/bw/dq/vl): ALL PASSED (15/15).This is the CPU half of the per-op multi-device path (CPU AVX + ROCm gfx1151) we'll repeat on the NVIDIA and M1 Max boxes. Follow-ons to fully complete the
reduce_sumop: Graph IR → x86 lowering wiring + the manifestx86slot, the ROCm reduction kernel, and the honest dashboard move.Note: the full
tessera_x86_backendlibrary has pre-existing gcc build breakage in unrelated kernels (amx_runtimecpuidscope,amx_gemmforward-decl,vnniintrinsic type) — developed on clang/Mac. This kernel builds + runs standalone; that breakage is a separate cleanup.🤖 Generated with Claude Code