perf(gemm): start workers once per product, synchronise on work done

README.md

@@ -86,8 +86,9 @@ is wrong on subnormals
   beat NumPy at every size measured.
 - **At parity:** `SumAxis(0)`, `Inner` against 16-thread OpenBLAS, `Log` on
   one core.
-- **Slower:** `MatMul` against 16-thread OpenBLAS (0.55× at 1024², 0.3× at
-  256² and below); `Dot`/mat·vec against NumPy's threaded BLAS (0.16×); `Exp`
+- **Slower:** `MatMul` against 16-thread OpenBLAS at 512² and up (0.42× at
+  512², 0.51× at 1024², measured in one run with OpenBLAS; at 256² ours is
+  1.44× faster); `Dot`/mat·vec against NumPy's threaded BLAS (0.16×); `Exp`
   below 256 Ki elements (0.82×); the *allocating* forms of elementwise ops up
   to 256 Ki elements, `Concatenate`/`Stack` and slice copies outside a
   `Workspace` (0.2–0.5×: Go's allocator and garbage collector, which

docs/perf.md

@@ -411,9 +411,13 @@ below 256²). Where the parallel loss is, measured on the Zen 3:
   each ran at 41 GFLOP/s (656 in total), while the parallel 1024² GEMM ran its
   kernel at 15–17 GFLOP/s per core. Per-core efficiency already drops from 40
   to ~30 GFLOP/s at 2 cores.
-- It is not scheduling overhead: 83% of the parallel profile is the
-  micro-kernel itself, so the kernel runs slower when cores share work. Why is
-  **not established**. The shared packed B panel is the suspect (on this part
+- ~~It is not scheduling overhead: 83% of the parallel profile is the
+  micro-kernel itself, so the kernel runs slower when cores share work.~~
+  **Wrong, corrected the same day:** a profile of cycles samples only the cores
+  that run, never the ones that wait. `perf stat` shows 7.1 of 16 CPUs busy
+  during a 1024² product: the loss was mostly waiting, and the fix is in
+  [Parallel GEMM: start once, wait on work](#parallel-gemm-start-once-wait-on-work-2026-10-04).
+  Why the kernel itself also runs slower in parallel is still **not established**. The shared packed B panel is the suspect (on this part
   the last-level cache is split between core complexes, so the panel is partly
   remote), but the drop already shows at 2 cores of one complex, which that
   explanation does not cover.
@@ -426,9 +430,8 @@ below 256²). Where the parallel loss is, measured on the Zen 3:
   rows·cols·k multiply-adds, and a strided pack costs far more per element than
   an FMA, so tiles under ~256×256 lose.
 
-The next step is to find where the kernel stalls at 2 cores (hardware
-counters: cache misses per FMA, serial against parallel) before choosing a
-lever; then an AVX-512 kernel for the hosts that have it.
+The hardware counters were read the same day; see the next-but-one
+section. An AVX-512 kernel for the hosts that have it remains open.
 
 Correctness was checked on real hardware for five of the six 64-bit targets:
 amd64 (Zen 3, FMA path), arm64 (cfarm185), ppc64le (cfarm120), riscv64 (cfarm94)
@@ -637,6 +640,86 @@ parallel 4 Mi sum is bimodal on this host (1.6 ms or 7 ms in either version),
 so only the minima are compared. No NumPy or BLAS is installable there (the
 host has no route to package mirrors), so there is no external reference.
 
+## Parallel GEMM: start once, wait on work (2026-10-04)
+
+**Diagnosis, from hardware counters** (`perf stat`, Zen 3 guest cfarm421,
+steady state, 1024², 3 s per configuration):
+
+| workers | CPUs busy | IPC | instructions |
+|---|---|---|---|
+| 1 | 1.00 | 3.35 | 37.5 G |
+| 2 | 1.91 | 2.00 | 37.6 G |
+| 8 | 5.58 | 2.16 | 172 G (192 products) |
+| 16 | 7.08 | 2.28 | 193 G (226 products) |
+
+The instruction count does not grow, so nobody busy-waits. L1 misses do not
+grow either. What changes is that at 16 workers the CPUs are busy only 44% of
+the time. The driver forked and joined goroutines twice per (jc, pc) block of
+the product, once to pack B and once for the row bands, 16 times in a 1024²
+product. A fork/join wakes parked threads one after another. A micro-benchmark
+on the same guest (16 goroutines, 700 µs of work each) measures its cost at
+**26 µs with 2 workers, 165 µs with 8, 330 µs with 16**. 16 × 330 µs is ~40%
+of a 12.5 ms product.
+
+Two hypotheses were tested and **refuted** first:
+
+- SMT siblings. The guest has 128 vCPUs, as many as the host has threads.
+  Two serial GEMMs pinned on vCPU 16 and any other vCPU slow each other by
+  0–10%, never by half.
+- A block of A larger than L2. With MC = 256, the A block is 512 KiB, all of
+  Zen 3's L2, against BLIS's MC = 72. Sweeping MC over {72, 120, 144, 192, 256}
+  changed nothing beyond the noise at 2 workers, and gave +17% at 16 workers
+  for 1024² but −8% for 512².
+
+**The change.** The workers start once per product. They synchronise on
+**work done**, never on workers arrived: a block's bands start once all of its
+B panels are packed and all bands of the previous block are done, and the
+packing is shared in claims of four NR-wide panels. A first version used a
+barrier that waited for every worker, and it put the wake-up chain straight
+back on the critical path for small products: 96³ on 8 POWER9 cores ran at
+0.57×. B is double-buffered, so block i+1 is packed while block i is still
+being multiplied. Each C tile is still accumulated in pc order, so the result
+is bit-identical to the serial product. The pack buffers moved to two pools,
+one for A and one for B: a pair per worker had held 16 MiB of unused B buffers
+at 16 workers, which the collector empties and the next call zeroes again
+(`memclr` page faults, 20% of the cycles in a KVM spinlock in one profile).
+
+**Measured.** Every figure is the median of interleaved runs against the
+previous driver.
+
+| | 2 workers | 8 workers | 16 workers |
+|---|---|---|---|
+| Zen 3 guest, 1024² | 1.06× | 1.33× | 1.39× |
+| Zen 3 guest, 512² | 1.20× | 1.20× | 1.33× |
+| Zen 3 guest, 256² | 0.92× | 1.19× | 1.24× |
+| Zen 3 guest, 128² | **0.89×** | 1.50× | 1.16× |
+| Zen 3 guest, (1024×256)·(256×1024) | **0.82×** | 1.07× | 0.89× (p25 1.00×) |
+| POWER9 bare metal, 512² | 1.34× | 1.37× | |
+| POWER9 bare metal, 256² | 1.61× | 1.78× | |
+| POWER9 bare metal, (1024×256)·(256×1024) | 1.13× | 1.18× | |
+
+The Zen 3 runs were 12 rounds and the POWER9 runs 6, except 24 for 128² and
+256² at 8 workers. Those two are bimodal on that machine, in both drivers, and
+6 rounds had shown 256² at 0.72×. POWER9 used one thread per core (CPUs 0, 4,
+8, …).
+
+**Still behind.** On the Zen 3 guest, two workers lose 8–18% on mid-size
+products, where a fork/join costs only 26 µs. Bare-metal POWER9 gains at two
+workers, so this is not a property of the scheme, and no second code path was
+added for it. Against OpenBLAS with 16 threads, measured three ways in one run:
+
+| n | OpenBLAS-16 | previous driver | new driver |
+|---|---|---|---|
+| 256 | 708 µs | 1.06× | **1.44×** |
+| 512 | 789 µs | 0.29× | 0.42× |
+| 1024 | 4.60 ms | 0.43× | 0.51× |
+
+(× = OpenBLAS time / ours, so higher is faster). On this guest, absolute times
+drift through the day: the previous driver measured 0.54–0.66× of OpenBLAS in
+the morning. Only ratios taken in the same run are comparable. Not measured on
+Apple M-series, whose machine was shared with other jobs at the time. The bands
+stay dynamic, so its slower efficiency cores are handled as before.
+
 ## SIMD coverage
 
 - **amd64 (SSE2)** ships hand-vectorized `sum` (4-accumulator `ADDPD`), `sqrt`

internal/kernels/barrier_test.go

@@ -0,0 +1,41 @@
+package kernels
+
+import (
+	"sync/atomic"
+	"testing"
+)
+
+// TestSpinUntil waits on a condition another goroutine makes true, both
+// spinning and yielding from the first poll.
+func TestSpinUntil(t *testing.T) {
+	for _, iters := range []int{spinIters, 0} {
+		saved := spinIters
+		spinIters = iters
+		var flag atomic.Bool
+		go func() {
+			for i := 0; i < 1000; i++ {
+				_ = i * i
+			}
+			flag.Store(true)
+		}()
+		spinUntil(flag.Load)
+		spinIters = saved
+	}
+}
+
+// TestGemmParallelYielding runs the parallel GEMM with every wait yielding at
+// once, as on an oversubscribed machine, and with more workers than CPUs.
+func TestGemmParallelYielding(t *testing.T) {
+	saved := spinIters
+	spinIters = 0
+	defer func() { spinIters = saved }()
+	withMaxProcs(8, func() {
+		withThresholds(1<<14, 1, func() {
+			for _, s := range []struct{ m, k, n int }{{50, 70, 90}, {13, 300, 7}, {97, 31, 600}} {
+				a, b := intMat(s.m, s.k, 3), intMat(s.k, s.n, 4)
+				assertGemm(t, s.m, s.k, s.n, a, b)
+				withBlocks(MR, 8, 2*NR, func() { assertGemm(t, s.m, s.k, s.n, a, b) })
+			}
+		})
+	})
+}

internal/kernels/gemm.go

@@ -62,9 +62,11 @@ type Mat struct {
 // rowMajor views a contiguous row-major matrix with ld elements per row.
 func rowMajor(data []float64, ld int) Mat { return Mat{Data: data, RS: ld, CS: 1} }
 
-// packBuf is one worker's reusable pair of pack buffers: paBuf for the MC*KC A
-// panel, pbBuf for the KC*NC B panel. Pooling them keeps the per-call GEMM
-// allocation-free above the goroutine launch, which matters at small n.
+// packBuf is a pair of pack buffers: pa for the MC*KC A panel, pb for the KC*NC
+// B panel. The two come from separate pools, because the parallel GEMM needs
+// one pa per worker but only two pb in all: drawing a pair per worker held
+// 16 MiB of unused B buffers at 16 workers, which the garbage collector empties
+// out of the pool and the next call zeroes again.
 type packBuf struct {
 	pa []float64
 	pb []float64
@@ -83,27 +85,29 @@ func roundUp(x, m int) int { return (x + m - 1) / m * m }
 func paCap() int { return roundUp(blockMC, MR) * blockKC }
 func pbCap() int { return blockKC * roundUp(blockNC, NR) }
 
-var packPool = sync.Pool{New: func() any {
-	return &packBuf{
-		pa: make([]float64, paCap()),
-		pb: make([]float64, pbCap()),
-	}
-}}
+var paPool, pbPool sync.Pool
 
-// getPackBuf returns a pooled buffer pair sized for the current block params. If
-// the pooled buffers are too small for (possibly test-tuned) larger blocks it
-// reallocates them, so a test that raises blockMC/KC/NC still gets valid scratch.
-func getPackBuf() *packBuf {
-	b := packPool.Get().(*packBuf)
-	if cap(b.pa) < paCap() {
-		b.pa = make([]float64, paCap())
-	}
-	if cap(b.pb) < pbCap() {
-		b.pb = make([]float64, pbCap())
+// getPA and getPB return a pooled pack buffer sized for the current block
+// parameters, allocating one when the pool is empty or its buffer is too small
+// (a test may raise blockMC/KC/NC). putPA and putPB return it.
+func getPA() []float64  { return getScratch(&paPool, paCap()) }
+func getPB() []float64  { return getScratch(&pbPool, pbCap()) }
+func putPA(b []float64) { paPool.Put(&b) }
+func putPB(b []float64) { pbPool.Put(&b) }
+
+func getScratch(p *sync.Pool, n int) []float64 {
+	if b, ok := p.Get().(*[]float64); ok && cap(*b) >= n {
+		return (*b)[:n]
 	}
-	b.pa = b.pa[:paCap()]
-	b.pb = b.pb[:pbCap()]
-	return b
+	return make([]float64, n)
+}
+
+// getPackBuf returns a pooled pair, for the single-worker paths.
+func getPackBuf() *packBuf { return &packBuf{pa: getPA(), pb: getPB()} }
+
+func putPackBuf(b *packBuf) {
+	putPA(b.pa)
+	putPB(b.pb)
 }
 
 // packGemmRows computes the row band [r0,r1) of dst = a(m x k) * b(k x n) with
@@ -190,30 +194,6 @@ func packB(b Mat, pc, kc, jc, nc int, pb []float64) {
 	}
 }
 
-// packBParMin is the B-block size, in elements, from which packBP splits the
-// packing across workers (64 Ki doubles = 512 KiB). A var so tests can pin it.
-var packBParMin = 1 << 16
-
-// packBP is packB with the NR-wide panels split across w goroutines. MatMulP
-// packs the shared B block once per (jc, pc) round, and done by one goroutine
-// it is a serial step while every other worker waits at the round's start:
-// a 1 MiB strided copy, eight times in a 1024^2 product.
-//
-// Below packBParMin elements the goroutine launch costs more than the copy
-// (measured on M4: a 128x128 block packed in parallel made 96^3..128^2
-// products 8-15% slower), so a small block is packed by the caller alone.
-func packBP(b Mat, pc, kc, jc, nc int, pb []float64, w int) {
-	if kc*nc < packBParMin {
-		packB(b, pc, kc, jc, nc, pb)
-		return
-	}
-	panels := (nc + NR - 1) / NR
-	parallelFor(panels, min(w, panels), func(lo, hi int) {
-		j0 := lo * NR
-		packB(b, pc, kc, jc+j0, min(hi*NR, nc)-j0, pb[j0*kc:])
-	})
-}
-
 // packA copies A[ic:ic+mc][pc:pc+kc] into MR-tall row panels: for each panel ir
 // the layout is pa[ir*kc + p*MR + r] = A(ic+ir+r, pc+p), with the r >= mr rows
 // of an edge panel zero-filled so the micro-kernel reads a full MR.
@@ -312,7 +292,7 @@ func MatMulStridedP(dst []float64, a, b Mat, m, k, n int) {
 	if m*n < GemmThreshold {
 		buf := getPackBuf()
 		packGemmRows(0, m, dst, a, b, k, n, buf)
-		packPool.Put(buf)
+		putPackBuf(buf)
 		return
 	}
 	w := runtime.GOMAXPROCS(0)
@@ -339,41 +319,80 @@ func MatMulStridedP(dst []float64, a, b Mat, m, k, n int) {
 		w = nBands // never spawn more goroutines than there is work for
 	}
 
-	// One shared B-pack buffer (packed once per jc/pc block) plus one private
-	// A-pack buffer per worker. The shared pb is sized like a normal pooled B
-	// buffer; getPackBuf grows it for any test-raised block sizes.
-	shared := getPackBuf()
-	pb := shared.pb
-	defer packPool.Put(shared)
+	// The workers are started once per call and stay for every (jc, pc) block,
+	// instead of being forked and joined per block. A fork/join wakes the
+	// parked worker threads one after another; measured on a 16-vCPU Zen 3
+	// guest it costs 26 µs at 2 workers, 165 µs at 8 and 330 µs at 16, against
+	// ~0.8 ms of work per block of a 1024² product, and there were two per
+	// block (pack B, then the bands): ~40% of the wall time.
+	//
+	// The workers synchronise on work done, never on workers arrived: a block's
+	// bands start once all of its B panels are packed and all bands of the
+	// previous block are done. A worker still waking up holds nobody back, so
+	// the wake-up chain stays off the critical path; a barrier that waited for
+	// every worker put it right back on it for small products (96³ on 8 POWER9
+	// cores: 0.57x). Claimed work never waits, so the scheme cannot deadlock.
+	//
+	// pb is double-buffered, so block i can be packed while block i-1 is being
+	// multiplied. Tiles of C are still accumulated in pc order, since block i's
+	// bands wait for block i-1's.
+	pbs := [2][]float64{getPB(), getPB()}
+	defer putPB(pbs[0])
+	defer putPB(pbs[1])
 
+	type block struct{ jc, nc, pc, kc, chunks int }
+	const chunk = 4 // B panels packed per claim
+	var blocks []block
 	for jc := 0; jc < n; jc += blockNC { // L3 column block of B
 		nc := min(blockNC, n-jc)
+		panels := (nc + NR - 1) / NR
 		for pc := 0; pc < k; pc += blockKC { // L2 contraction block
-			kc := min(blockKC, k-pc)
-			packBP(b, pc, kc, jc, nc, pb, w) // pack B once, by all workers
+			blocks = append(blocks, block{jc, nc, pc, min(blockKC, k-pc), (panels + chunk - 1) / chunk})
+		}
+	}
+	type progress struct {
+		packNext, packDone, bandNext, bandDone atomic.Int64
+		_                                      [32]byte // keep blocks' counters off one cache line
+	}
+	prog := make([]progress, len(blocks))
+	bandsDone := func(i int) bool { return i < 0 || prog[i].bandDone.Load() == int64(nBands) }
 
-			var next atomic.Int64 // shared band cursor
-			var wg sync.WaitGroup
-			for g := 0; g < w; g++ {
-				wg.Add(1)
-				go func() {
-					defer wg.Done()
-					ab := getPackBuf()
-					for {
-						bi := int(next.Add(1)) - 1
-						if bi >= nBands {
-							break
-						}
-						r0 := bi * bandRows
-						r1 := min(r0+bandRows, m)
-						gemmBand(r0, r1, jc, nc, pc, kc, dst, a, n, ab.pa, pb)
-					}
-					packPool.Put(ab)
-				}()
+	work := func() {
+		pa := getPA()
+		defer putPA(pa)
+		for i, bl := range blocks {
+			// pbs[i%2] is free: its last reader was block i-2, and to get here
+			// this worker passed block i-1's wait, which requires every band of
+			// block i-2 done.
+			pr, pb := &prog[i], pbs[i%2]
+			for {
+				c := int(pr.packNext.Add(1)) - 1
+				if c >= bl.chunks {
+					break
+				}
+				j0, j1 := c*chunk*NR, min((c+1)*chunk*NR, bl.nc)
+				packB(b, bl.pc, bl.kc, bl.jc+j0, j1-j0, pb[j0*bl.kc:])
+				pr.packDone.Add(1)
+			}
+			spinUntil(func() bool { return pr.packDone.Load() == int64(bl.chunks) && bandsDone(i-1) })
+			for {
+				bi := int(pr.bandNext.Add(1)) - 1
+				if bi >= nBands {
+					break
+				}
+				r0 := bi * bandRows
+				gemmBand(r0, min(r0+bandRows, m), bl.jc, bl.nc, bl.pc, bl.kc, dst, a, n, pa, pb)
+				pr.bandDone.Add(1)
 			}
-			wg.Wait()
 		}
 	}
+	var wg sync.WaitGroup
+	for g := 1; g < w; g++ {
+		wg.Add(1)
+		go func() { defer wg.Done(); work() }()
+	}
+	work() // the caller is worker 0
+	wg.Wait()
 }
 
 // MatMul computes dst = a(m x k) * b(k x n) serially with the packed GEMM (the
@@ -383,7 +402,7 @@ func MatMulStridedP(dst []float64, a, b Mat, m, k, n int) {
 func MatMul(dst, a, b []float64, m, k, n int) {
 	buf := getPackBuf()
 	packGemmRows(0, m, dst, rowMajor(a, k), rowMajor(b, n), k, n, buf)
-	packPool.Put(buf)
+	putPackBuf(buf)
 }
 
 // dotRange returns sum(a[i]*b[i]) over equal-length slices, using four