Nested multiply-accumulate loops lower to scalar code; targeted `@fmacs` vectorization gives 4.5–7.7× on `samples/spatial/blas/gemv.sptl`

[ycmath]

Summary

Compute blocks of the form

for i16 k, i16 l in [0:K, 0:K] {
    z[k] = z[k] + A_flat[k*K + l] * x[l]
}

(as written in samples/spatial/blas/gemv.sptl, phase 4) currently lower to scalar
nested for (@range(...)) loops. The inner product therefore runs one element per
instruction, and its cost grows as K² while a hand-written CSL kernel using
@fmacs over strided DSDs stays nearly flat. We measured the crossover directly:
at K=16 the generated gemv becomes ~2× slower than a hand-written reference
(gemv-collectives_2d from Cerebras' sdk-examples), and ~4.6× slower at K=32.

We wrote a small, conservative lowering that turns this exact pattern into the
hand-written idiom and would like to contribute it, if there is interest.

Why a compiler change (and not a source-level rewrite)

We first tried to express the vectorizable form at the language level and found
both routes closed:

• await map i16 k in [0:K:1] { z[k] = z[k] + A_flat[(k*K) + l] * x[l] } is
  rejected by _dsd_from_array ("Expected one index variable in array access") —
  the second loop variable l cannot appear in the access.
• A slice assignment z[0:K:1] = z[0:K:1] + A_flat[l:K*K:K] * x[l] is rejected
  by the surface grammar (UnexpectedCharacters: No terminal matches ':'),
  although ArraySlice supports strides at the IR level.

So users cannot opt into the fast form from .sptl today.

Proposed lowering

In emit_for, detect the shape

for (k, l) in [0:Kk, 0:Kl]:  Z[k] = Z[k] + A[k*Ck + l*Cl + C0] * X[f(l)]

and emit the same idiom hand-written kernels use (cf. gemv-collectives_2d):

const __vecmac_dst = @get_dsd(mem1d_dsd, .{ .tensor_access = |i|{Kk} -> Z[i] });
const __vecmac_src = @get_dsd(mem1d_dsd, .{ .tensor_access = |i|{Kk} -> A[i * Ck] });
for (@range(i16, 0, Kl, 1)) |l| {
    const __vecmac_a = @increment_dsd_offset(__vecmac_src, l * Cl, f32);
    @fmacs(__vecmac_dst, __vecmac_dst, __vecmac_a, X[f(l)]);
}

Safety/conservatism:

• fires only on the exact affine MAC shape above, f32 only, ranges starting at 0
  with step 1; anything else falls back to the existing scalar loops;
• aliasing guard: not applied when Z is also read as A or X (the DSD op
  reorders reads relative to the sequential scalar loop);
• respects --disable-dsd, plus a new --disable-mac-vectorization flag wired
  the same way as the existing disable flags;
• @increment_dsd_offset and @fmacs are both available in the SDK line this
  repo targets: @fmacs is already emitted by dsd_ops.py, and
  @increment_dsd_offset is used by Cerebras' own csl-examples at tag
  v1.4.0 (benchmarks/cholesky/pe.csl, lines 113–119).

Measurements

Setup: WSE-2 simfabric (cycle counts are deterministic; every run verified
numerically against a NumPy reference), 4×4 PE grid, gemv.sptl with
PX=PY=4, fp32. Baseline is the unmodified scalar lowering; "hand-written" is
Cerebras' gemv-collectives_2d benchmark instrumented with the same
<time> timestamps the SPADA runtime uses. Kernel-window cycles
(max(stop) − min(start) across PEs):

K	scalar lowering	vectorized	hand-written	speedup (vec vs scalar)
4	490	579	3,127	0.85×
8	1,521	809	3,378	1.88×
16	8,357	1,851	4,090	4.51×
32	28,955	3,744	6,274	7.73×

The scalar lowering crosses above the hand-written kernel at K=16 and is 4.6×
slower by K=32; the vectorized lowering stays below it at every size we
measured. At K=4 the DSD setup constant makes vectorization slightly slower
(0.85×) — if that matters, a size threshold could gate the transform, but we
left the behavior unconditional and are flagging it here instead.

Caveats we want to be upfront about: measurements were taken on the SDK 2.10
simulator (with a small local compatibility layer for 2.10's CSL dialect; the
patch itself is written against current main and does not include that
layer). On unmodified main we verified that the transform fires and the
emitted code retains the SDK 1.4 dialect, and that the full
tests/spatial_ir + tests/stencil_ir suites pass (323 passed / 12 skipped),
plus 6 new unit tests covering the positive case, both aliasing cases, a
non-affine index, a mismatched accumulator, and the disable flag. We have not
run an end-to-end SDK 1.4 simulation.

Offer

The change is ~180 lines of compiler code plus ~110 lines of tests
(syntax/csl/statements.py matcher + emitter, flag plumbing in
lowering/spatial_ir_to_csl.py and cli/compiler.py,
tests/spatial_ir/test_mac_vectorization.py). Happy to open a PR if this is
welcome, and to adjust scope (e.g., f16 via @fmach, a size threshold, or
folding the flag into an existing one) to your preference.

[tbennun]

@ycmath This sounds like a great idea, thank you so much for reporting!

I think that this is the kind of load-bearing change that should be generalized further.

1. Why not support arbitrary levels of nested loops (e.g., PE-local gemm)? It would just consist of further loops with increment_dsd(-old_stride + new_stride) at every step.
2. What about supporting a different order (for (k, l))?
3. Why is the change limited to f32? Our DSD detection works generally across types.
4. (optional) could this approach of affine expressions be generalized to support multi-level nested loops of any DSD operation? I can imagine that Z[k*C1] = Z[k*C1] + X[k*C2] can very easily generate DSDs that look like |k|{K} -> { Z[C1*k] } and a single @fadds.

It should be an easy lift if the infrastructure already exists. I also don't think this requires a different flag over --disable-dsd, as this generalizes DSD detection to ND loops.

[ycmath]

Thanks for the quick and encouraging response — and yes, generalizing this beyond the exact (k,l)-nested f32 MAC shape is the right direction. Concretely, here's how the asks map onto what exists today:

(2) loop order and (3) type generality — will do, small lifts. Operand classification already handles either syntactic order; only the loop-header roles are hardcoded (~a few dozen lines).
For types, your DSD_ASSIGNMENT_MAPPING already carries the full @fmach/@fmachs/@fmacs (and int) dispatch, so the f32 gate can become a dtype dispatch mirroring FMADSDOp._as_csl.
One thing I'd like to verify against the SDK first: whether @increment_dsd_offset(..., T)'s type parameter generalizes cleanly to f16/i16 the way it does for f32.

(5) flag fold-in — agreed. MAC vectorization is already AND-gated behind DISABLE_DSD, so folding removes nothing from --disable-dsd users.
I'll fold it and drop the separate flag. Minor note: one negative test currently isolates MAC vectorization from other DSD emission; it will be rewritten against --disable-dsd, losing a little granularity.

(1) arbitrary nesting depth — one question before I write depth>2 code. You sketched increment_dsd(-old_stride + new_stride) per level, i.e. a relative adjustment.
The current patch uses @increment_dsd_offset with an absolute offset, and I don't see a relative-increment builtin used elsewhere in the tree.
Is there an SDK builtin you had in mind for the delta form, or would you expect per-level absolute recomputation (recomputing the base offset at each outer-loop step)?
Happy to implement either; the delta form is cleaner if the SDK supports it.

(4) affine → any DSD op — worth doing, but it's a real lift, not a small one.
The affine-expression walker is net-new code from this patch; the existing get_dsd_op path fires on ForeachStatement/MapStatement over whole arrays and doesn't walk affine indices inside scalar for loops.
Generalizing means lifting the affine decomposition + DSD construction into a shared utility and extending the aliasing guard from the 2-array MAC special case to N-ary ops.
I'd propose staging it: a first PR with (2)+(3)+(5) on top of the current vectorization (test matrix: order × dtype × aliasing negatives), then a second PR for (1)+(4) once the increment question above is settled.
Does that staging work for you?