Kernel Fusion - Phase 8 SIMD design

[Quafadas]

Phase 8 — JVM SIMD backend + benchmark (design)

Stacks on copilot/implement-phase-2 (PR #94). Reuses KernelIR / ScalarExpr /
Schedule / FusedRunner unchanged — this phase adds a second JVM executor and
the benchmark that proves (or disproves) the fusion win. Scope of the first cut:
JVM only, F64 only.

Goal

Replace the per-element scalar tree-walk (KernelExecutor, the Phase-8 "first cut")
with a vectorised executor using jdk.incubator.vector, and land a benchmark
comparing fused-SIMD vs fused-scalar vs unfused vecxt. The --add-modules jdk.incubator.vector flag is already wired via vecIncubatorFlag.

Prerequisite folded into this PR: broadcast correctness fix

The current ScalarExpr.Load(buf, numel) → inputs(k)(i % numel) is incorrect for
non-prefix broadcasts (e.g. [1,4] → [3,4] in column-major) and is also SIMD-hostile
(a wrapping gather can't be a contiguous vector load). Fix once, serving both:

• Scalar operand (numel == 1): SIMD splat / scalar broadcast — trivially correct.
• Contiguous operand (numel == outNumel): contiguous vector load.
• General broadcast: pre-materialise the operand to outNumel once, before the
  kernel runs, using the interpreter's correct strided logic. The kernel then only ever
  sees contiguous loads + splats.

Add the missing parity test: param[Double]("v", [1,4]).broadcastTo([3,4]) vs the
interpreter (currently absent — every existing .broadcastTo test broadcasts a scalar).

1. SIMD executor — vectorised tree interpretation

Walk the same ScalarExpr tree, but per vector block rather than per element:

val sp    = DoubleVector.SPECIES_PREFERRED
val bound = sp.loopBound(n)          // n - n % L
var i = 0
while i < bound do
  evalVec(expr, i, sp, inputs).intoArray(out, i)
  i += sp.length()
while i < n do                       // scalar tail (reuse existing evalScalar)
  out(i) = evalScalar(expr, i, inputs)
  i += 1

evalVec mirrors evalScalar but returns a DoubleVector:

ScalarExpr node	evalVec mapping
Load (contiguous)	DoubleVector.fromArray(sp, buf, i)
Load (numel==1)	DoubleVector.broadcast(sp, buf(0))
Lit(v)	DoubleVector.broadcast(sp, v)
SUnary(Neg/Abs)	.neg() / .abs()
SUnary(Sqrt)	.lanewise(SQRT)
SUnary(Sin/Cos/Tan/Exp/Log)	.lanewise(SIN/COS/TAN/EXP/LOG) — SVML on x64
SUnary(Reciprocal)	one.div(v)
SUnary(Not)	v.compare(EQ, 0.0) → zero.blend(one, mask)
SBinary(Add/Sub/Mul/Div)	.add/.sub/.mul/.div
SBinary(Min/Max)	.lanewise(MIN/MAX)
SBinary(Pow)	.lanewise(POW, b) — SVML on x64
SBinary(Lt/Lte/Gt/Gte/Eq/Neq)	a.compare(OP, b) → zero.blend(one, mask)
SBinary(And/Or)	on the 0/1 encoding: .mul / .lanewise(MAX) (or mask ops)
Select(c,x,y)	mask = c.compare(NE, 0.0); evalVec(y).blend(evalVec(x), mask)

Transcendentals are vectorised directly (VectorOperators.SIN/COS/TAN/EXP/LOG/POW
exist and are SVML-backed intrinsics on x64). No scalar fallback needed for them. On
non-x64 JVMs the runtime may fall back internally, but the result stays correct — only
throughput varies by platform.

2. Reductions under SIMD

Keep the running total in a DoubleVector accumulator and combine vertically
(lane-wise) in the hot loop. Do not call reduceLanes — cross-lane (horizontal)
reductions are comparatively slow, so keep the hot loop purely vertical and fold the
accumulator to a scalar with an explicit lane loop at the end, off the hot path.

val L     = sp.length()
val bound = sp.loopBound(n)
var acc   = DoubleVector.zero(sp)         // Product: broadcast(1.0); Min: +inf; Max: -inf
var i     = 0
while i < bound do                        // hot loop — vertical combine only
  acc = acc.add(evalVec(body, i, sp, inputs))   // Product: .mul; Min/Max: .lanewise(MIN/MAX)
  i += L

// final fold to scalar: explicit lane loop, NOT reduceLanes
val lanes = new Array[Double](L)
acc.intoArray(lanes, 0)
var s = lanes(0)
var k = 1
while k < L do { s += lanes(k); k += 1 }  // combine op matches the ReduceOp

while i < n do { s += evalScalar(body, i, inputs); i += 1 }   // scalar tail

• Sum → add (identity 0.0), Product → mul (1.0), Min/Max → lanewise(MIN/MAX)
  (±inf) — applied consistently in the loop, the lane fold, and the tail.
• Optional latency-hiding: use several independent accumulators (acc0..accN)
  unrolled over the loop to break the loop-carried dependency chain, then combine them
  vertically (acc0.add(acc1)…) before the lane fold. Same numerics, higher throughput on
  latency-bound reductions.
• All/Any: accumulate a boolean/mask vector, fold at the end (no short-circuit under
  SIMD — acceptable). ArgMax/ArgMin: stay scalar in cut 1 (index tracking under
  SIMD is fiddly).

3. What stays scalar / deferred in cut 1

• ArgMax / ArgMin.
• Non-F64 dtypes (matches Schedule's current F64-only restriction).
• JS / Native SIMD (Phase 9 — those backends stay scalar while-loops).
• Per-kernel codegen (bytecode / generated loops). Tree-interpretation-per-block is
  the pragmatic cut; codegen is the ceiling if interpretation overhead proves to eat the
  memory-traffic savings on cheap kernels.

4. Correctness / parity strategy

Every SIMD kernel is parity-tested against Interpreter.eval (the Phase-7 oracle), which
uses java.lang.Math. Because SVML transcendentals share Math's ULP spec but are not
bit-identical (and are no longer semi-monotonic):

• Elementwise algebraic kernels (+ - * /, min/max, sqrt, compares, select): each
  output lane runs the identical scalar ops as the interpreter, so results are bit-exact —
  strict 1e-12 relative holds trivially.
• Reductions: the SIMD accumulation order (lane-interleaved + lane fold) differs from
  the interpreter's sequential fold, so results are not bit-identical even for Sum
  (FP addition is non-associative). Use a relative tolerance (~1e-12, looser for very
  large N where accumulated rounding grows), not bit-equality.
• Transcendental kernels: use a combined tolerance |a − b| ≤ atol + rtol·|b| with
  e.g. atol = 1e-12, rtol = 1e-9. atol covers near-zero results (sin(kπ)), rtol
  covers the normal range (real divergence is a few ULP, ~1e-15).
• Keep golden input ranges well-conditioned; don't put pathologically large trig
  arguments in goldens (range-reduction divergence) unless the tolerance is widened there.
• Edge cases to test explicitly: n < L, n not a multiple of L (tail loop),
  single-element and scalar-output reductions, splat + pre-materialised broadcast.

5. Benchmark

Add entries to the existing benchmark/ module (JMH via mill; same incubator flag).
Compare fused-SIMD vs fused-scalar vs unfused vecxt on kernels where fusion can
actually pay off — i.e. multiple ops (a single a+b has nothing to fuse and unfused
vecxt's own vectorised loop will match it):

Kernel	Expression	Stresses
axpy	a*x + y	memory-bound, cheap ops
Horner poly	((a*x+b)*x+c)*x+d	best case — many intermediates removed
dot / sum-sq	sum(x*x)	fused reduction
logistic	1/(1+exp(-x))	vector transcendental (SVML) path

• Sweep N ∈ {1e2, 1e4, 1e6, 1e8} to expose the memory-bound crossover and where
  interpretation overhead stops mattering.
• Report elements/sec (or ns/op). Baselines: the existing vectorAddition / sum
  benchmarks.
• Success = fused-SIMD beats unfused on the multi-op kernels (especially Horner and
  the reduction). Record actual numbers in the PR — if a kernel loses, say so; that's
  the signal for whether codegen is needed.

6. Deliverables checklist

• Broadcast fix (pre-materialise) + non-scalar broadcast parity test
• KernelExecutorSimd (or a mode flag) — vectorised evalVec + scalar tail
• SIMD reductions (Sum/Product/Min/Max/All/Any); ArgMax/ArgMin scalar
• Parity tests vs interpreter (tolerance strategy above; tail-loop edges)
• NDArray.fuse { … } opt-in entry point (or a clean FusedRunner façade for the bench)
• Benchmark module entries + measured numbers in the PR description
• Confirm the incubator flag is applied to the benchmark module's forkArgs

Open questions / risks

• SVML availability: x64 (CI = ubuntu-latest) is accelerated; aarch64 (Apple silicon)
  may fall back internally. Correctness holds; document that benchmark numbers are
  platform-specific.
• Interpretation overhead: per-block dispatch is amortised over L lanes but not free.
  If cheap kernels don't beat unfused, that's the trigger for per-kernel codegen (a later
  phase), not a failure of the IR.
• Transcendental tolerance: confirm rtol = 1e-9 is acceptable for the project's
  correctness bar, or tighten with a higher-accuracy vector math path later.

Sources

• JEP 448: Vector API (Sixth Incubator)
• JDK-8267743: separate library for x86 Intel SVML assembly intrinsics
• VectorOperators (jdk.incubator.vector) API docs