WASM PR 4: fold the repeated fused-op arms into macros - #2566
Open
czoli1976 wants to merge 7 commits into
Open
Conversation
…nel and madd macros wasm.rs was a single 3159-line file, while every other linalg backend is a module root plus a subdirectory. Move the three parts that carry no f32 MMM kernel code into wasm/madd.rs, wasm/mmm_i32.rs and wasm/act.rs, and re-export them from the root so every existing path still resolves. Pure relocation: the emitted kernel bodies are unchanged under both +simd128 and +simd128,+relaxed-simd. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… wasm/ module tree Splits the six f32 kernels by accumulator layout: wasm/mmm_f32_gemm.rs holds the two that pack columns into lanes (4x4, 8x8), wasm/mmm_f32_gemv.rs the four that pack rows and fix nr=1 (4x1, 8x1, 16x1, 32x1) — the same split the mmm_f32 and mmv_f32 dispatch callbacks make. The dispatch assertions move to wasm/dispatch_tests.rs. Pure relocation: the emitted kernel bodies are unchanged under both +simd128 and +simd128,+relaxed-simd. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… target Three timing modules lived in src behind #[cfg(test)] + #[ignore], which no other backend does. benches/wasm.rs already supersedes two of them — its own header records that looping four kernels back-to-back biased the in-src version — so those are dropped, and the activation bench moves across. The numerical-consistency check between 16x1 and 32x1 is an assertion rather than a timer, so it stays in src alongside the dispatch tests. The four copies of the kernel-lookup helper in the bench collapse to one; the timing loops keep their own warmup and repetition counts, which differ on purpose. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… against relaxed-simd The backend picks its multiply-add form, its int8 packing and its sigmoid/tanh kernels at compile time on +relaxed-simd, so the single +simd128 CI leg left half of it untested; add a second leg. AddRowColProducts and AddMatMul with k=1 compute the same product and so must agree on whether that multiply-add is fused, which nothing checked — the existing coverage uses operands whose products are exact, where fused and separate agree trivially. Also record why wasm_f32_4x4 sits at TargetOptimized and is unreachable through dispatch, and assert that it stays that way. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… macros Each of the four GEMV kernels spelled the same 24 FusedKerSpec arms out once per accumulator, so widening a kernel meant copying every arm again. The arms that apply one operation uniformly across the accumulators now expand from wasm/fuse.rs, which emits the same statements in the same order; the emitted kernel bodies are unchanged under both +simd128 and +simd128,+relaxed-simd. AddMatMul, AddRowColProducts, AddUnicast and Store stay written out: the first two must agree on whether the multiply-add is fused, and the last two walk the C tile by stride differently in each kernel. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… macros Same treatment as the GEMV kernels, with the two primitives this layout needs: here a per-row operand is splatted across accumulators and a per-column one is loaded as a vector, the opposite of the GEMV kernels, and the 8x8 tile holds each row as a low/high pair that takes one splat between them. The emitted kernel bodies are unchanged under both +simd128 and +simd128,+relaxed-simd. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
WASM_RELAXED_SIMD.md said all six f32 kernels switch to fused multiply-add under +relaxed-simd; only 32x1 and 8x8 do, and the 1.40-1.55x figure applies only to those. The int8 kernel switches too, on instruction rather than rounding, which the document did not mention. The header on wasm_i32_4x4 claimed bit-identity with generic_i32_4x4 without qualification, which holds only without +relaxed-simd — the relaxed dot takes an i7 second operand. The module header still named wasm32-wasi, renamed to wasm32-wasip1. The module gate moves from target_family = "wasm" to target_arch = "wasm32", matching what the kernels actually require (std::arch::wasm32). This only narrows wasm64, which cannot compile these kernels anyway. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #2565. This is the point of the series.
Each of the six f32 kernels spelled the same 24
FusedKerSpecarms out once per accumulator, so 8x8 needed 535 lines to say what 4x1 says in 86, and widening a kernel meant copying every arm again. Of 2342 kernel lines, 2005 were those arms and only 337 were theAddMatMulloops that actually differ.The arms that apply one operation uniformly across the accumulators now expand from
wasm/fuse.rs. Each macro emits the same statements in the same order, so this is a source-level change only. Per kernel the 24 arms become 24 lines.The two layouts need different primitives, and which one an arm reaches for is visible on the arm's own line rather than hidden behind a flag: in the GEMV kernels a per-row operand is a vector load and a per-column one is a splat, and in the GEMM kernels it is the other way round. 8x8 holds each row as a low/high pair that takes one splat between them.
Five arms stay written out per kernel.
AddMatMulandAddRowColProductsboth carry the fused-versus-separate multiply-add choice, and keeping them adjacent and explicit is what makes that choice checkable by eye — #2565 added the test that enforces it.AddUnicastandStorewalk the C tile by stride differently in each kernel.LoadTilediffers too: 4x1 uses an unalignedv128_loadwhere the others dereference, so a shared macro would have changed the alignment immediate for no gain.The last commit corrects documentation that had drifted:
WASM_RELAXED_SIMD.mdsaid all six f32 kernels switch to fused multiply-add under+relaxed-simd, when only 32x1 and 8x8 do and the quoted 1.40–1.55x applies only to those; the header onwasm_i32_4x4claimed bit-identity withgeneric_i32_4x4without qualification, which holds only without+relaxed-simd, since the relaxed dot takes an i7 second operand; and the module gate moves fromtarget_family = "wasm"totarget_arch = "wasm32", matching what the kernels require.Evidence
Same gate as #2563, and it is what makes this reviewable: the emitted body of all eight kernels is byte-identical, under
+simd128, under+simd128,+relaxed-simd, and in the shippedreleaseprofile withlto = true. So the 1153 deleted lines can be read as deleted repetition rather than as a rewrite of six hand-tuned SIMD kernels.Benchmarks are a guard here rather than the claim, since the instructions are unchanged. Best-of-two runs of
benches/wasm.rsacross 24 shapes: median +0.5%, range −1.4% to +2.1%.Tests: 2195 under
+simd128, 2208 under+simd128,+relaxed-simd, withtract-coreandtest-unit-coregreen in both. No new clippy lints.wasm_i32_4x4is deliberately left alone — it already loops over a scalar[[i32; 4]; 4]at about 7 lines an arm, so it never had this problem, and collapsing it would mean a second macro family for scalar accumulators.Across the four PRs
wasm.rsgoes from 3159 lines in one file to 2196 across eight.🍍