Vectorize nested multiply-accumulate loops into @fmac* DSD operations - #72
Open
ycmath wants to merge 1 commit into
Open
Vectorize nested multiply-accumulate loops into @fmac* DSD operations#72ycmath wants to merge 1 commit into
ycmath wants to merge 1 commit into
Conversation
Detect the two-variable nested MAC pattern Z[k] = Z[k] + A[affine(k, l)] * X[f(l)] in emit_for and lower it to a strided base DSD over A plus a per-l @increment_dsd_offset and the dtype-matched @FMac* builtin, instead of scalar loops. On samples/spatial/blas/gemv.sptl this gives 4.5-7.7x (issue spcl#69). Per the discussion in spcl#69: - Loop-header order is irrelevant: the variable indexing the accumulator takes the k role, the other is the reduction variable (ask 2). - Dtypes dispatch like FMADSDOp._as_csl: @fmach (f16), @fmachs (f16 multiply / f32 accumulate), @fmacs (f32) (ask 3). - No separate flag: the vectorizer is gated behind --disable-dsd like all other DSD detection (ask 5). Conservative by construction: falls back to scalar loops on aliasing of Z with A or X, non-affine indices, unsupported dtype combinations, non-unit steps, or a nonzero k start. Arbitrary nesting depth and generalization to other DSD ops (asks 1 and 4) are deferred to a follow-up PR pending the increment_dsd semantics question. Co-Authored-By: Claude Fable 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.
Implements the targeted MAC-loop vectorization from #69, incorporating the staging agreed there: this PR delivers the base vectorization together with asks (2), (3), and (5); asks (1) and (4) are deferred to a follow-up PR pending the
increment_dsdsemantics question raised in the issue thread.What this does
emit_fornow recognizes the two-variable nested multiply-accumulate patternand lowers it to a strided base DSD over
Aplus a per-l@increment_dsd_offsetand one@fmac*per column, instead of scalar loops — the same idiom as the handwritten CSL GEMV. Onsamples/spatial/blas/gemv.sptlthis measured 4.5–7.7× (details in #69).Zwith coefficient 1 takes thekrole, the other is the reduction variable — sofor (l, k)vectorizes identically tofor (k, l). Per-element accumulation order overlis preserved bitwise in both orders.FMADSDOp._as_csl: all-f16 →@fmach, f32 accumulate with f16 scalar →@fmachs, all-f32 →@fmacs; the@increment_dsd_offsetelement type followsA. Unsupported combinations fall back to scalar loops.--disable-dsdlike all other DSD detection.Conservative by construction — falls back to scalar code on: aliasing of
ZwithAorX(DSD ops reorder reads relative to the sequential loop), non-affine indices, non-unit steps, nonzerokstart, or any dtype combination outside the table above.Tests
tests/spatial_ir/test_mac_vectorization.py: 11 tests covering the positive matrix (both loop orders × f32/f16/mixed dtypes) and negative cases (aliasing in both orders, non-affine index, accumulator mismatch, unsupported dtypes,--disable-dsd). Full suite: 386 passed, 12 skipped on this branch.@increment_dsd_offsetis available from SDK 1.x (Cerebras' own csl-examples v1.4.0 cholesky uses it), so generated code stays compatible with the SDK 1.4 pinned in CI; the f16 form of the type parameter was additionally verified to compile with cslc.🤖 Generated with Claude Code