perf(parallel): persistent helper pool for parallelFor; finer reduction partials

README.md

@@ -122,6 +122,16 @@ arch is therefore a durable foundation. The numeric loops live in
 on amd64, arm64, ppc64le and loong64 today, behind the same API and held to the
 same tests; riscv64 and s390x are next.
 
+## Goroutines
+
+Operations on large arrays are spread over `GOMAXPROCS` goroutines. After the
+first one, the package keeps up to `GOMAXPROCS-1` helper goroutines alive. When
+an operation ends, they poll for the next one for 200 µs and then block, the way
+OpenBLAS's threads do, so operations issued back to back do not pay to wake
+threads. Waking threads per call made a 2^20-element dot slower on 8 POWER9 cores
+than on 4. The helpers are never stopped: a goroutine-leak check (goleak) that
+runs after ndarray operations will list them.
+
 ## Untrusted input
 
 Shapes, indices and data can come from someone else (a file, a request), so

docs/perf.md

@@ -557,8 +557,41 @@ in scheduling instead: with the caller working, 2.1 of 4 CPUs were busy
 against 3.4 for the original. Its goroutines wait in the run queue of the
 caller's processor until another one is woken to steal them. With the caller
 blocking, the 4-worker case recovers, but the 8-worker gains go with it. The
-code is on the `experiment/parallel-for-scheduling` branch. Not yet solved:
-a scheduler that keeps both.
+code is on the `experiment/parallel-for-scheduling` branch.
+
+**Solved in v0.6.0 with a persistent helper pool.** The cost being paid was
+waking threads, so the fix is to stop putting them to sleep between back-to-back
+operations, as OpenBLAS and OpenMP do. Up to `GOMAXPROCS-1` helper goroutines
+are started once. After an operation they poll for the next one for 200 µs,
+then park. The work keeps owned runs with stealing (row two of the table
+above). Reductions now use four partials per worker: with one each, a helper
+that saw the job late had its only block taken by the caller, who then summed
+two in a row (Sum of 4 Mi on 4 cores: 2.8 of 4 CPUs busy). On POWER9 bare metal,
+each case in a **fresh process**, 6 interleaved rounds, against v0.5.2
+(median):
+
+| | 2 workers | 4 workers | 8 workers | 8 workers, before → after |
+|---|---|---|---|---|
+| Dot 2^20 | 3.29× | 2.41× | **10.8×** | 625 → 58 µs |
+| MatVec 1024² | 2.01× | 1.65× | 5.78× | 323 → 56 µs |
+| SumAxis(1) | 2.06× | 1.80× | 7.09× | 310 → 44 µs |
+| Sum 256 Ki | 2.33× | 3.23× | 5.60× | 61 → 11 µs |
+| Max 256 Ki | 1.75× | 1.34× | 6.78× | 757 → 112 µs |
+| AddInto 256 Ki | 2.22× | 1.91× | 5.57× | 102 → 18 µs |
+| Exp 256 Ki | 1.18× | 1.49× | 1.97× | 1013 → 513 µs |
+| Sum 4 Mi | 1.33× | 1.17× | 1.41× | 431 → 305 µs |
+| Exp 4 Mi | 1.00× | 1.10× (p25 0.84×) | 1.29× | |
+| BroadcastAdd | 0.95× | 1.01× | 1.67× | |
+| Chain (Mul, Add, Sqrt) 4 Mi | | 0.91× | 1.04× | |
+| MatMul 256², 512² (own driver) | | 1.00× | 0.96–0.99× | |
+
+A first combined run, with every benchmark in one process, also showed serial
+operations speeding up (Sqrt of 1 Ki, 1.5–2.2×). Run alone, those are equal.
+That gain belonged to whatever state the earlier benchmarks had left, not to
+this change, so only fresh-process figures are reported here. The cost: after
+an operation, the helpers spend up to 200 µs of CPU polling. They are never
+stopped (see the README, *Goroutines*). Not yet measured on amd64 or on Apple
+M, whose hosts were unreachable or loaded at the time.
 
 ## Log and Log10 (2026-10-04)
 

internal/kernels/gemm.go

@@ -467,8 +467,8 @@ func VecMatP(dst, v, a []float64, k, n int) {
 }
 
 // Dot1DP returns the inner product of two equal-length contiguous vectors,
-// fanned across cores above ParThreshold: each worker dot-products its chunk with
-// the unrolled dotRange and the partials are summed. The chunked four-way
+// fanned across cores above ParThreshold: each chunk is dot-producted with the
+// SIMD kernel and the partials are summed. The chunked four-way
 // grouping reassociates the sum (a few ULP, like SumP / numpy pairwise); callers
 // needing the exact left-to-right value use the scalar Dot1D.
 func Dot1DP(a, b []float64) float64 {
@@ -477,9 +477,10 @@ func Dot1DP(a, b []float64) float64 {
 		return dotSIMD(a, b)
 	}
 	w := numWorkers(n)
-	partials := make([]float64, w)
-	chunk := (n + w - 1) / w
-	parallelFor(w, w, func(lo, hi int) {
+	parts := min(4*w, max(w, n/reduceGrain)) // see mapReduceP
+	partials := make([]float64, parts)
+	chunk := (n + parts - 1) / parts
+	parallelFor(parts, w, func(lo, hi int) {
 		for idx := lo; idx < hi; idx++ {
 			s := idx * chunk
 			if s >= n {

internal/kernels/parallel.go

@@ -3,6 +3,7 @@ package kernels
 import (
 	"runtime"
 	"sync"
+	"sync/atomic"
 )
 
 // Multicore kernels.
@@ -34,21 +35,68 @@ func numWorkers(n int) int {
 	return w
 }
 
-// parallelFor splits [0,n) into w contiguous, near-equal blocks and runs body
-// on each block in its own goroutine, waiting for all to finish. body must be
-// safe to call concurrently on disjoint [lo,hi) ranges.
+// parallelFor runs body over [0,n) split into contiguous blocks, on up to w
+// goroutines counting the caller, and returns once every block is done. body
+// must be safe to call concurrently on disjoint [lo,hi) ranges.
+//
+// The goroutines are the caller and the pool's helpers (pool.go). Each owns a
+// contiguous run of blocks, its static share, and works through it first, then
+// takes what is left of the others' runs, so a helper that wakes late holds
+// nobody back. Block boundaries depend only on n and w, so a result is the same
+// from one run to the next. When another call holds the pool, this one forks
+// its own goroutines instead.
 func parallelFor(n, w int, body func(lo, hi int)) {
-	if w <= 1 {
+	if w <= 1 || n <= 1 {
 		body(0, n)
 		return
 	}
+	w = min(w, n)
+	if !pool.busy.TryLock() {
+		forkJoin(n, w, body)
+		return
+	}
+	defer pool.busy.Unlock()
+	ensureHelpers(w - 1)
+	blocks := min(n, 4*w)
+	q, r := n/blocks, n%blocks // block i starts at i*q + min(i, r): no i*n overflow
+	start := func(i int) int { return i*q + min(i, r) }
+	own := func(g int) int { return g * blocks / w } // owner g: [own(g), own(g+1))
+	taken := make([]paddedCounter, w)
+	var done atomic.Int64
+	run := func(g int) {
+		for k := 0; k < w; k++ {
+			o := (g + k) % w
+			lo, hi := own(o), own(o+1)
+			for {
+				i := lo + int(taken[o].Add(1)) - 1
+				if i >= hi {
+					break
+				}
+				body(start(i), start(i+1))
+				done.Add(1)
+			}
+		}
+	}
+	publish(run, w)
+	run(0)
+	spinUntil(func() bool { return done.Load() == int64(blocks) })
+}
+
+// paddedCounter is an atomic counter alone on its cache line, so owners
+// claiming their own blocks do not contend for one line.
+type paddedCounter struct {
+	atomic.Int64
+	_ [56]byte
+}
+
+// forkJoin splits [0,n) into w contiguous, near-equal blocks, runs body on each
+// in its own goroutine and waits for all: parallelFor's fallback when the pool
+// is serving another call.
+func forkJoin(n, w int, body func(lo, hi int)) {
 	chunk := (n + w - 1) / w
 	var wg sync.WaitGroup
 	for lo := 0; lo < n; lo += chunk {
-		hi := lo + chunk
-		if hi > n {
-			hi = n
-		}
+		hi := min(lo+chunk, n)
 		wg.Add(1)
 		go func(lo, hi int) {
 			defer wg.Done()
@@ -170,16 +218,21 @@ func mapReduceP(a []float64, red func([]float64) float64, combine func([]float64
 	if n < ParThreshold {
 		return red(a)
 	}
+	// Four partials per worker, each of at least reduceGrain elements: with
+	// one per worker, a helper that saw the job late had its only block taken
+	// by the caller, which then summed two in a row (Sum of 4 Mi on 4 POWER9
+	// cores: 2.8 of 4 CPUs busy, 0.76x).
 	w := numWorkers(n)
-	chunk := (n + w - 1) / w
+	parts := min(4*w, max(w, n/reduceGrain))
+	chunk := (n + parts - 1) / parts
 	// active = number of workers whose start s = idx*chunk is < n; the rest own
 	// nothing and are dropped (their start would overrun n and panic a[s:e], and
 	// there is no bit-identical reduction identity to feed combine for Max/Min).
-	// Since chunk = ceil(n/w) >= n/w, active = ceil(n/chunk) is always <= w, so
-	// dropping trailing workers never loses real data.
+	// Since chunk = ceil(n/parts), active = ceil(n/chunk) is always <= parts,
+	// so dropping trailing partials never loses real data.
 	active := (n + chunk - 1) / chunk
 	partials := make([]float64, active)
-	parallelFor(active, active, func(lo, hi int) {
+	parallelFor(active, w, func(lo, hi int) {
 		for idx := lo; idx < hi; idx++ {
 			s := idx * chunk
 			e := s + chunk
@@ -249,6 +302,9 @@ func RunAxisP(k axisKernel, dst, src []float64, outer, axisLen, inner int) {
 	}
 }
 
+// reduceGrain is the fewest elements a reduction partial covers (64 KiB).
+const reduceGrain = 8192
+
 // axisBand is the narrowest column band RunAxisP gives a worker (512 doubles,
 // 4 KiB per row): narrower bands make each row run too short to stream.
 const axisBand = 512

internal/kernels/pool.go

@@ -0,0 +1,96 @@
+package kernels
+
+import (
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// A persistent pool of helper goroutines for parallelFor.
+//
+// Forking goroutines per call wakes parked threads one after another, and a
+// call waited for the slowest wake-up: on a POWER9, a 2^20-element dot ran
+// slower on 8 cores (553 µs) than on 4 (212 µs), with 2.4 of 8 CPUs busy. The
+// helpers here are started once and, after a job, keep polling for the next one
+// for spinWindow before they park, as OpenBLAS's and OpenMP's threads do, so
+// operations issued back to back find them awake. A parked helper is woken when
+// a job is published; until it arrives, the others, the caller included, take
+// its blocks.
+
+// spinWindow is how long, in nanoseconds, an idle helper keeps polling for the
+// next job before it parks. Atomic because the helpers read it while tests set it.
+var spinWindow atomic.Int64
+
+func init() { spinWindow.Store(int64(200 * time.Microsecond)) }
+
+type poolJob struct {
+	run func(g int) // g in [1, w): the helper's owner index
+	w   int
+}
+
+type poolHelper struct {
+	parked atomic.Bool
+	wake   chan struct{}
+}
+
+var pool struct {
+	busy    sync.Mutex // one job at a time; a concurrent caller forks instead
+	grow    sync.Mutex
+	helpers []*poolHelper
+	gen     atomic.Uint64
+	job     atomic.Pointer[poolJob]
+}
+
+// ensureHelpers starts helpers until there are at least k.
+func ensureHelpers(k int) {
+	pool.grow.Lock()
+	defer pool.grow.Unlock()
+	for len(pool.helpers) < k {
+		h := &poolHelper{wake: make(chan struct{}, 1)}
+		i := len(pool.helpers)
+		pool.helpers = append(pool.helpers, h)
+		go helperLoop(i, h, pool.gen.Load())
+	}
+}
+
+// helperLoop runs jobs as helper i: owner index i+1 of each job that has room
+// for it.
+func helperLoop(i int, h *poolHelper, seen uint64) {
+	for {
+		start := time.Now()
+		for spins := 0; pool.gen.Load() == seen; spins++ {
+			if spins&255 != 255 {
+				continue
+			}
+			if time.Since(start) < time.Duration(spinWindow.Load()) {
+				runtime.Gosched()
+				continue
+			}
+			// Park. Re-check after announcing it: a job published in
+			// between either saw the flag and sends a wake-up (wait for it,
+			// so no stale one is left), or this helper takes the flag back
+			// and goes straight to the job.
+			h.parked.Store(true)
+			if pool.gen.Load() == seen || !h.parked.CompareAndSwap(true, false) {
+				<-h.wake
+			}
+			break
+		}
+		seen = pool.gen.Load()
+		if j := pool.job.Load(); i+1 < j.w { // a job is published before gen moves
+			j.run(i + 1)
+		}
+	}
+}
+
+// publish hands run to the first w-1 helpers and wakes the parked ones.
+func publish(run func(g int), w int) {
+	pool.job.Store(&poolJob{run: run, w: w})
+	pool.gen.Add(1)
+	for _, h := range pool.helpers[:w-1] {
+		if h.parked.CompareAndSwap(true, false) {
+			h.wake <- struct{}{}
+		}
+	}
+}