Layout Follow Ups

[Quafadas]

vecxt — post-#108 implementation plan (work items 1–6)

Follow-on work unlocked by #108 / PR #109 (Layout extraction).
Each phase is independently mergeable. Do them in order; the dependencies are noted where they exist.

---

0. Context for the implementing agent

Prerequisite: PR #109 must be merged to main first. Everything below assumes Matrix[A] holds a
single val layout: Layout and that Layout.linearIndex(row, col) exists.

Build & verify

./mill vecxt.__.test        # all platforms — JVM, JS, Native
./mill vecxt.jvm.test       # JVM only, faster inner loop
./mill bytecodeAudit.test   # static bytecode gate (CI-enforced)
just format                 # scalafmt — run before every commit

Benchmarks (JMH):

Only benchmark single suites. Currently, running the whole benchmark suite takes a very long time and will likely lock up the session.

./mill benchmark.runJmh -jvmArgs --add-modules=jdk.incubator.vector -prof gc -rf json -rff out.json

Note -prof gc is not in the justfile recipe but is required for gc.alloc.rate.norm.

Cross-platform discipline. vecxt/src/ is shared and must compile on JVM, Scala.js and Native.
Platform-specific kernels live in src-jvm/, src-js/, src-native/, src-js-native/. Any change to a
shared file needs ./mill vecxt.__.test, not just the JVM suite.

Constraints that must not be broken (from #105 and #108):

• No case class on hot types — synthesises productElement: Int => Object, which boxes.
• No sealed trait on hot types — makes field reads invokeinterface, defeating escape analysis and
  blocking scalar replacement of FloatVector/VectorMask temporaries.
• HotSpot budgets the audit enforces: MaxTrivialSize 6, MaxInlineSize 35, FreqInlineSize 325.
• @Thin / @HotPath / @AllocFree annotations are read by the audit. @Thin is invalid on inline def —
  the compiler drops the method, so there is nothing to annotate.
• The invariant raw.size == layout.dataLength holds at every checked factory. Do not add a construction
  path that bypasses it.
• The audit baseline lives in bytecode/baseline.json. If a phase legitimately moves bytecode sizes, run
  ./mill bytecodeAudit.updateBaseline and include the baseline diff in the PR with a one-line
  justification per changed entry. Never update the baseline to silence an unexplained regression.

---

Phase A — Close the Layout test gaps (work item 1)

Why first: Phases B, C and F rewrite ~58 loops and ~19 index expressions. This is the net that makes
those rewrites mechanical instead of nerve-wracking.

vecxt/test/src/layout.test.scala already ships 36 example-based tests from #109 and is good. Do not
rewrite it. Two named "(property)" tests are single-example assertions, not generator-driven — that is the gap.

Add to vecxt/test/src/layout.test.scala:

1. linearIndex injectivity. For a generated set of layouts, assert that
   { linearIndex(i, j) : i <- 0 until rows, j <- 0 until cols } has exactly numel distinct values.
   This is the property that actually catches a swapped rowStride/colStride, and nothing currently tests it.
   Exclude broadcast layouts (a zero stride is deliberately non-injective) — assert those separately.
2. linearIndex in-bounds. Every index falls in [0, dataLength).
3. Submatrix offset composition. Taking a contiguous submatrix of a submatrix equals taking the composed
   submatrix directly — same offset, same strides. Mirror the arithmetic in
   MatrixInstance.submatrix (newOffset = m.offset + newRows.head * m.rowStride + newCols.head * m.colStride).
4. transpose round-trip over generated layouts, not the single hard-coded one at line 118.

Generators: hand-roll a for comprehension over small dimension/stride/offset tuples. Do not add a
scalacheck dependency to vecxt — it currently lives only in the laws module, and pulling it into the
core test module for this is disproportionate.

Also add a JMH benchmark covering view creation, which the suite does not measure at all today:
transpose and submatrix inside a @Benchmark body (everything existing builds matrices in
@Setup(Level.Trial)). Put it in benchmark/src/. This is the one workload where a Layout allocation is
the entire cost rather than a rounding error, so it is the honest place to watch it.

Acceptance: ./mill vecxt.__.test green; new benchmark runs and reports.

---

Phase B — Fix the strided copyToArray bug (not in items 1–6; found during analysis)

Should be already implemented : SKIP

Phase C — Unit-stride axis innermost (work item 2)

The largest O(numel) win in this plan.

58 loops across 8 files follow the shape:

while i < m.rows do
  var j = 0
  while j < m.cols do
    val idx = m.layout.linearIndex(i, j)

Rows outer, cols inner. For the column-major layouts these branches actually serve (rowStride == 1), the
inner loop steps memory by colStride — and the destinations (newArr(i + j * m.rows)) stride the same way.
Both sides are cache-hostile and unvectorisable, on exactly the paths every submatrix view takes.

Distribution:

file	loops
vecxt/src/doublematrix.scala	17
vecxt/src-jvm/floatmatrix.scala	12
vecxt/src-jvm/doublematrix.scala	8
vecxt/src-js/doublematrix.scala	8
vecxt/src-native/doublematrix_native.scala	8
vecxt/src/matrixutil.scala	3
vecxt/src/MatrixInstance.scala	1
vecxt/src-jvm/intmatrix.scala	1

Approach. Add to Layout a query for which axis carries unit stride — something like
def unitStrideAxis: Int returning 0 (rows), 1 (cols) or -1 (neither). Keep it a plain def on the
final class; do not introduce an enum or ADT for this.

Then write one shared traversal helper and route the loops through it. It must be inline with an
inline operation parameter so no closure is allocated per element — follow the existing pattern in
DoubleMatrix.reduceAlongDimension, which already does exactly this.

Order the loops so the unit-stride axis is innermost, for both source and destination. Where source and
destination disagree on which axis is contiguous, prefer the destination — writes are more expensive to
scatter than reads.

Do this incrementally. One file per commit, vecxt/src/doublematrix.scala first (largest, shared,
best-covered by tests). Re-run the suite between files.

Acceptance: suite green on all platforms after each file. Benchmark MnistBenchmark.fwd_02_bias_add_b1,
fwd_05_bias_add_b2 and fwd_06_softmax before and after — these are the three benchmarks that exercise
the per-element path, via mapRowsInPlace → m.row(i) → m((i, j)). Report the numbers in the PR.

Out of scope: cache blocking / tiling, and any change to the hasSimpleContiguousMemoryLayout fast paths
that bypass these loops entirely.

---

Phase D — Offset/length overloads on the array kernels (work item 3)

Unlocks SIMD for views, which currently cannot reach it at all.

All 121 defs in vecxt/src-jvm/doublearrays.scala are whole-array — they are extension methods on
vec: Array[Double] that run 0 until vec.length. There is not one (from, len) variant in the file.

The consequence: hasSimpleContiguousMemoryLayout requires offset == 0 && raw.size == numel, so every
submatrix view falls off the SIMD path entirely. But when rowStride == 1, each column is already a
contiguous run of rows doubles at offset + j * colStride. A view is not unvectorisable; it is cols
separate vectorisable segments. The only thing missing is a kernel that accepts a start and a length.

Scope — Double on JVM only for this phase. Do not touch Float/Int or the JS/Native kernels until the
Double results are measured.

1. Add (arr, from, len) forms of the reductions and elementwise ops that the matrix layer actually calls.
   Start with the ones reached from DoubleMatrix: sum/sumSIMD, multInPlace, *=, +, -, /,
   norm. Follow the existing kernel shape — spd.loopBound(len), vector body, scalar tail.
2. Keep the existing whole-array methods as thin forwarders (f(arr) = f(arr, 0, arr.length)) so no call
   site changes and no source compatibility breaks.
3. Add a Layout helper that yields the contiguous segments: when rowStride == 1, segment j is
   (offset + j * colStride, rows); when colStride == 1, segment i is (offset + i * rowStride, cols);
   otherwise none.
4. Route the strided branches in DoubleMatrix through segment dispatch when a unit-stride axis exists.

Watch: for small matrices the per-segment dispatch overhead can exceed the SIMD gain. Benchmark across
sizes and add a size threshold below which the scalar loop from Phase C is used instead. State the threshold
and the data behind it in the PR.

Acceptance: a submatrix view and its deepCopy produce identical results for every routed operation.
Benchmarks showing the crossover point.

---

Phase E — Widen the elementwise compatibility check (work item 4)

Depends on Phase D. Cheap, but only safe once offset-aware kernels exist.

Single choke point — vecxt/src/dimMatCheck.scala:19:

object sameDenseElementWiseMemoryLayoutCheck:
  def apply[A, B](a: Matrix[A], b: Matrix[B]): Boolean = a.layout.sameElementOrderAs(b.layout)

All 17 call sites route through it, so this is one method change with wide reach. Currently
sameElementOrderAs requires dense on both sides. Two matrices with identical strides, offset and shape
touch an identical index set — the operation is a flat loop regardless of whether either is dense.

Add a clause. Use full Layout equality rather than a looser stride comparison: equals includes
dataLength, which guarantees the two backing arrays are the same length — a precondition the whole-array
kernels silently rely on.

Before Phase D, the safe form also needs numel == dataLength, otherwise the whole-array kernel computes
over out-of-view elements (harmless but wasteful — potentially very wasteful for a small view into a large
buffer). After Phase D, drop that extra condition and dispatch to the offset/length kernels instead,
which is the version actually worth having.

Calibration — do not overclaim this one. It does not help transposed operands. For square m,
m is column-major and m.T is row-major, so the layouts are unequal and reindexing is genuinely required.
m + m.T is Phase C's business, not this one. What this fixes is matched views: subviews of a common
parent, and any identically-strided pair.

Acceptance: a test asserting that two identically-strided non-dense views take the fast path and agree
with the slow path elementwise.

---

Phase F — Close the strided gaps (work item 5)

Depends on Phase C — the shared traversal helper is what makes most of these three lines each.

57 ??? sites remain. They are two different things and only one is in scope:

• Strided gaps — if hasSimpleContiguousMemoryLayout then <fast> else ???. Mechanical; the Phase C
  helper closes them. Examples in vecxt/src/doublematrix.scala: *(n: Double) (:54), unary_- (:256),
  tan/tan! (:401, :405), mean (:409), ** (:413), reduceAlongDimension (:420),
  sumSIMD (:475), norm (:479); plus horzcat/vertcat in matrixutil.scala (:257, :275) and the
  fancy-indexing branch in MatrixInstance.scala (:119).
• Missing features — genuinely unimplemented operations, not layout gaps. kronecker
  (doublematrix.scala:168), and the extension (d: Double) operators -, /, +=, -=, /=
  (:15, :16, :19, :20, :21). Out of scope — leave them. They are design decisions, not oversights.

First deliverable of this phase is the classification itself: walk all 57, tag each as strided-gap or
missing-feature, and post the table before writing code. The line numbers above cover
vecxt/src/doublematrix.scala only; floatmatrix.scala (11), intmatrix.scala (10) and the JS/Native
files have not been classified.

Then close the strided-gap bucket, with a test per closed site comparing against deepCopy-then-operate.

Acceptance: every closed site has a view-based test. The remaining ??? are all missing-feature, listed
in the PR description.

---

Not in this plan - don't implement

Phase G — kind as a layout discriminator (work item 6)

Lowest priority. Do last, and only after Phases C–F have settled.

Layout.kind exists as a Byte, is always Layout.Strided = 0, and nothing branches on it — reserved by
#108 for Phase 3 structured layouts.

The cheap use is as a hint that is only ever an optimisation, never a semantic claim: precompute
DenseColMajor / DenseRowMajor / Strided at construction and collapse the
isDenseColMajor / isDenseRowMajor / hasSimpleContiguousMemoryLayout triple-branch into one switch.
If the tag is wrong you get slow, not wrong.

Do not sell this as a speedup. Those are val field loads and branch prediction eats them. The
justification is bytecode size: it shrinks the fast-path guard prefix on every operation, which is the
budget MaxInlineSize 35 / FreqInlineSize 325 is measured against. Judge it on
./mill bytecodeAudit.test output, not on JMH.

Keep the three boolean vals as forwarders so no call site changes.

Explicitly out of scope: anything resembling #108's Phase 3 — Diagonal, Triangular, Symmetric,
Banded, Identity, broadcast layouts, or BLAS dispatch to dsymm/dtrmm/dgbmv. Those need an
invariant-maintenance story (every operation must decide what kind its output carries, and a wrong kind is
silent corruption) that this phase deliberately does not open.

Acceptance: audit output shows guard prefixes shrinking; no behavioural change; suite green.

• Gather-based SIMD indexing for the both-strides-non-unit case. The Vector API only offers
  fromArray(species, arr, offset, indexMap, mapOffset); gathers are intrinsified on x86 AVX2+ but run at
  near-scalar throughput on most microarchitectures, and you would build an int[] index map per call.
  DoubleMatrix.hadamard already shows the likely-better alternative — materialise via deepCopy, then run
  the contiguous kernel. Benchmark gather-vs-materialise as a throwaway before writing any of it.
• # Plan: extract Matrix[A]'s layout fields into a Layout value #108 Phase 3 structured layouts — see Phase G.
• NDArray unification. But note: the unit-stride-axis and segment-decomposition logic in Phases C and D
  is identical to what NDArray needs. Write it as free functions over (shape, strides, offset) rather
  than methods bound to Layout, and the eventual unification gets it for free. This costs nothing now; the
  alternative is writing it twice.

[Quafadas]

@copilot Can you have a look at phase C & D please?