432 lines
12 KiB
Go
432 lines
12 KiB
Go
|
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
|
|
// SPDX-License-Identifier: MIT
|
||
|
|
|
||
|
|
package base
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"sync"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
|
|
)
|
||
|
|
|
||
|
|
// The dispatch of Factor is the one place where a crew size, a chunk
|
||
|
|
// boundary and a spawn style can silently change the arithmetic, so the
|
||
|
|
// candidates live here beside a bit-identity test and the measurement
|
||
|
|
// that picks the constants. Every candidate must reproduce factorSerial
|
||
|
|
// byte for byte, and the production Factor must reproduce the
|
||
|
|
// parameterised mirror at the shipped constants.
|
||
|
|
|
||
|
|
func benchFactorRows(n int) (rows [][]float64, pristine []float64) {
|
||
|
|
flat := make([]float64, n*n)
|
||
|
|
for i := range n {
|
||
|
|
for j := range n {
|
||
|
|
flat[i*n+j] = float64((i*7+j*13)%11) - 5
|
||
|
|
}
|
||
|
|
flat[i*n+i] += float64(n)
|
||
|
|
}
|
||
|
|
rows = make([][]float64, n)
|
||
|
|
for i := range n {
|
||
|
|
rows[i] = flat[i*n : (i+1)*n]
|
||
|
|
}
|
||
|
|
pristine = make([]float64, len(flat))
|
||
|
|
copy(pristine, flat)
|
||
|
|
return rows, pristine
|
||
|
|
}
|
||
|
|
|
||
|
|
func resetFactorRows(rows [][]float64, pristine []float64) {
|
||
|
|
off := 0
|
||
|
|
for _, row := range rows {
|
||
|
|
copy(row, pristine[off:off+len(row)])
|
||
|
|
off += len(row)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// factorSerial is Factor with the dispatch removed: the reference every
|
||
|
|
// candidate must match byte for byte.
|
||
|
|
func factorSerial(m [][]float64) ([]int, int) {
|
||
|
|
n := len(m)
|
||
|
|
perm := make([]int, n)
|
||
|
|
for i := range perm {
|
||
|
|
perm[i] = i
|
||
|
|
}
|
||
|
|
parity := 1
|
||
|
|
for k := range n {
|
||
|
|
pivot := k
|
||
|
|
for i := k + 1; i < n; i++ {
|
||
|
|
if absOf(m[i][k]) > absOf(m[pivot][k]) {
|
||
|
|
pivot = i
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if m[pivot][k] == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if pivot != k {
|
||
|
|
m[pivot], m[k] = m[k], m[pivot]
|
||
|
|
perm[pivot], perm[k] = perm[k], perm[pivot]
|
||
|
|
parity = -parity
|
||
|
|
}
|
||
|
|
factorRows(m[k+1:], k, n, m[k])
|
||
|
|
}
|
||
|
|
return perm, parity
|
||
|
|
}
|
||
|
|
|
||
|
|
// factorRetired is the dispatch Factor used before the crew moved onto
|
||
|
|
// preallocated slots: one goroutine per chunk through
|
||
|
|
// sync.WaitGroup.Go, so every chunk allocates a closure. It is the
|
||
|
|
// reference the current dispatch is measured against.
|
||
|
|
func factorRetired(m [][]float64, quantum int) ([]int, int) {
|
||
|
|
n := len(m)
|
||
|
|
perm := make([]int, n)
|
||
|
|
for i := range perm {
|
||
|
|
perm[i] = i
|
||
|
|
}
|
||
|
|
parity := 1
|
||
|
|
for k := range n {
|
||
|
|
pivot := k
|
||
|
|
for i := k + 1; i < n; i++ {
|
||
|
|
if absOf(m[i][k]) > absOf(m[pivot][k]) {
|
||
|
|
pivot = i
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if m[pivot][k] == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if pivot != k {
|
||
|
|
m[pivot], m[k] = m[k], m[pivot]
|
||
|
|
perm[pivot], perm[k] = perm[k], perm[pivot]
|
||
|
|
parity = -parity
|
||
|
|
}
|
||
|
|
pivotRow := m[k]
|
||
|
|
rows := m[k+1:]
|
||
|
|
work := len(rows) * (n - k)
|
||
|
|
w := min(work/quantum+1, engine.WorkersFor(len(rows)))
|
||
|
|
if w < 2 {
|
||
|
|
factorRows(rows, k, n, pivotRow)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
chunk := (len(rows) + w - 1) / w
|
||
|
|
if chunk < factorMinRows {
|
||
|
|
w = max(len(rows)/factorMinRows, 1)
|
||
|
|
chunk = (len(rows) + w - 1) / w
|
||
|
|
}
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
for start := 0; start < len(rows); start += chunk {
|
||
|
|
end := min(start+chunk, len(rows))
|
||
|
|
wg.Go(func() {
|
||
|
|
factorRows(rows[start:end], k, n, pivotRow)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
}
|
||
|
|
return perm, parity
|
||
|
|
}
|
||
|
|
|
||
|
|
// factorJobSlot is one crew member's preallocated work slot, mirroring
|
||
|
|
// factorJob for the float64 candidates.
|
||
|
|
type factorJobSlot struct {
|
||
|
|
rows [][]float64
|
||
|
|
pivotRow []float64
|
||
|
|
k, n int
|
||
|
|
start int
|
||
|
|
end int
|
||
|
|
wg *sync.WaitGroup
|
||
|
|
}
|
||
|
|
|
||
|
|
func factorSlotWorker(job *factorJobSlot) {
|
||
|
|
factorRows(job.rows[job.start:job.end], job.k, job.n, job.pivotRow)
|
||
|
|
job.wg.Done()
|
||
|
|
}
|
||
|
|
|
||
|
|
// factorCrew is the shipped dispatch with both tuning constants
|
||
|
|
// exposed.
|
||
|
|
func factorCrew(m [][]float64, quantum, maxCrew int) ([]int, int) {
|
||
|
|
n := len(m)
|
||
|
|
perm := make([]int, n)
|
||
|
|
for i := range perm {
|
||
|
|
perm[i] = i
|
||
|
|
}
|
||
|
|
parity := 1
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
var jobs []factorJobSlot
|
||
|
|
for k := range n {
|
||
|
|
pivot := k
|
||
|
|
for i := k + 1; i < n; i++ {
|
||
|
|
if absOf(m[i][k]) > absOf(m[pivot][k]) {
|
||
|
|
pivot = i
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if m[pivot][k] == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if pivot != k {
|
||
|
|
m[pivot], m[k] = m[k], m[pivot]
|
||
|
|
perm[pivot], perm[k] = perm[k], perm[pivot]
|
||
|
|
parity = -parity
|
||
|
|
}
|
||
|
|
pivotRow := m[k]
|
||
|
|
rows := m[k+1:]
|
||
|
|
work := len(rows) * (n - k)
|
||
|
|
w := min(work/quantum+1, engine.WorkersFor(len(rows)))
|
||
|
|
w = min(w, maxCrew)
|
||
|
|
if w < 2 {
|
||
|
|
factorRows(rows, k, n, pivotRow)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
chunk := (len(rows) + w - 1) / w
|
||
|
|
if chunk < factorMinRows {
|
||
|
|
w = max(len(rows)/factorMinRows, 1)
|
||
|
|
chunk = (len(rows) + w - 1) / w
|
||
|
|
if w < 2 {
|
||
|
|
factorRows(rows, k, n, pivotRow)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if jobs == nil {
|
||
|
|
jobs = make([]factorJobSlot, maxCrew)
|
||
|
|
}
|
||
|
|
spawned := 0
|
||
|
|
for start := 0; start < len(rows); start += chunk {
|
||
|
|
job := &jobs[spawned]
|
||
|
|
job.rows, job.pivotRow, job.k, job.n = rows, pivotRow, k, n
|
||
|
|
job.start, job.end, job.wg = start, min(start+chunk, len(rows)), &wg
|
||
|
|
spawned++
|
||
|
|
}
|
||
|
|
wg.Add(spawned)
|
||
|
|
for i := range spawned {
|
||
|
|
go factorSlotWorker(&jobs[i])
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
}
|
||
|
|
return perm, parity
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestFactorDispatchBitIdentical proves every crew size, chunk boundary
|
||
|
|
// and spawn style leaves the factor and the permutation byte for byte
|
||
|
|
// as the serial reference, and that the shipped Factor matches the
|
||
|
|
// parameterised mirror the tuning benchmark measures.
|
||
|
|
func TestFactorDispatchBitIdentical(t *testing.T) {
|
||
|
|
for _, n := range []int{1, 2, 5, 16, 64, 129, 260} {
|
||
|
|
base, _ := benchFactorRows(n)
|
||
|
|
want := make([][]float64, n)
|
||
|
|
for i := range n {
|
||
|
|
want[i] = append([]float64(nil), base[i]...)
|
||
|
|
}
|
||
|
|
wp, wpar := factorSerial(want)
|
||
|
|
check := func(name string, got [][]float64, gp []int, gpar int) {
|
||
|
|
t.Helper()
|
||
|
|
if len(gp) != len(wp) {
|
||
|
|
t.Fatalf("%s: permutation length %d, want %d", name, len(gp), len(wp))
|
||
|
|
}
|
||
|
|
for i := range gp {
|
||
|
|
if gp[i] != wp[i] {
|
||
|
|
t.Fatalf("%s: permutation[%d] = %d, want %d", name, i, gp[i], wp[i])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for i := range got {
|
||
|
|
for j := range got[i] {
|
||
|
|
if got[i][j] != want[i][j] {
|
||
|
|
t.Fatalf("%s: factor[%d][%d] = %v, want %v", name, i, j, got[i][j], want[i][j])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if gpar != wpar {
|
||
|
|
t.Fatalf("%s: parity %d, want %d", name, gpar, wpar)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for _, q := range []int{1024, 8192, 65536} {
|
||
|
|
for _, w0 := range []int{1, 2, 4, 8, 16, 32} {
|
||
|
|
rows, _ := benchFactorRows(n)
|
||
|
|
gp, gpar := factorCrew(rows, q, w0)
|
||
|
|
check(fmt.Sprintf("crew q=%d w=%d n=%d", q, w0, n), rows, gp, gpar)
|
||
|
|
}
|
||
|
|
rows, _ := benchFactorRows(n)
|
||
|
|
gp, gpar := factorRetired(rows, q)
|
||
|
|
check(fmt.Sprintf("retired q=%d n=%d", q, n), rows, gp, gpar)
|
||
|
|
}
|
||
|
|
rows, _ := benchFactorRows(n)
|
||
|
|
gp, gpar := Factor(rows)
|
||
|
|
check(fmt.Sprintf("Factor n=%d", n), rows, gp, gpar)
|
||
|
|
mirror, _ := benchFactorRows(n)
|
||
|
|
mirrorGP, mirrorPar := factorCrew(mirror, factorWorkQuantum, factorMaxCrew)
|
||
|
|
check(fmt.Sprintf("mirror n=%d", n), mirror, mirrorGP, mirrorPar)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// BenchmarkForkJoinFloor measures an empty fork-join: the spawn, wake
|
||
|
|
// and join of w goroutines that touch nothing. It is the lower bound on
|
||
|
|
// what one dispatch costs, and the reason a pivot whose update is
|
||
|
|
// smaller than this cannot be worth splitting.
|
||
|
|
func BenchmarkForkJoinFloor(b *testing.B) {
|
||
|
|
for _, w := range []int{1, 2, 4, 8, 16, 32} {
|
||
|
|
b.Run(fmt.Sprintf("w=%d", w), func(b *testing.B) {
|
||
|
|
for b.Loop() {
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
for range w {
|
||
|
|
wg.Go(func() {})
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// BenchmarkFactorStylesAB interleaves the dispatch styles iteration by
|
||
|
|
// iteration on separate matrices, so a drift in the machine's speed
|
||
|
|
// lands on every style equally: the reported ns/op per style is the
|
||
|
|
// paired comparison, and the ns/op of the group as a whole is not
|
||
|
|
// comparable across groups.
|
||
|
|
func BenchmarkFactorStylesAB(b *testing.B) {
|
||
|
|
styles := []struct {
|
||
|
|
name string
|
||
|
|
run func(m [][]float64)
|
||
|
|
}{
|
||
|
|
{"serial", func(m [][]float64) { factorSerial(m) }},
|
||
|
|
{"retired/q=16384", func(m [][]float64) { factorRetired(m, 16384) }},
|
||
|
|
{"slots/q=8192/cap=16", func(m [][]float64) { factorCrew(m, 8192, 16) }},
|
||
|
|
{"slots/q=16384/cap=16", func(m [][]float64) { factorCrew(m, 16384, 16) }},
|
||
|
|
{"slots/q=24576/cap=16", func(m [][]float64) { factorCrew(m, 24576, 16) }},
|
||
|
|
{"slots/q=32768/cap=16", func(m [][]float64) { factorCrew(m, 32768, 16) }},
|
||
|
|
}
|
||
|
|
for _, n := range []int{256, 512} {
|
||
|
|
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
|
|
rows := make([][][]float64, len(styles))
|
||
|
|
pristine := make([][]float64, len(styles))
|
||
|
|
for i := range styles {
|
||
|
|
rows[i], pristine[i] = benchFactorRows(n)
|
||
|
|
}
|
||
|
|
elapsed := make([]time.Duration, len(styles))
|
||
|
|
for b.Loop() {
|
||
|
|
for i, s := range styles {
|
||
|
|
resetFactorRows(rows[i], pristine[i])
|
||
|
|
start := time.Now()
|
||
|
|
s.run(rows[i])
|
||
|
|
elapsed[i] += time.Since(start)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for i, s := range styles {
|
||
|
|
b.ReportMetric(float64(elapsed[i].Nanoseconds())/float64(b.N), "ns/op-"+s.name)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// transposeFlatPlain is the untiled copy the tiled TransposeFlat
|
||
|
|
// replaces: the read runs along the rows, the write strides by the row
|
||
|
|
// length.
|
||
|
|
func transposeFlatPlain(a []float64, m, n int) []float64 {
|
||
|
|
out := make([]float64, m*n)
|
||
|
|
for i := range m {
|
||
|
|
for j := range n {
|
||
|
|
out[j*m+i] = a[i*n+j]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return out
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestTransposeFlatTiledBitIdentical pins the tiled copy against the
|
||
|
|
// untiled one: a transpose is a permutation, so it must match value for
|
||
|
|
// value, and the tile boundary must not lose a corner of a ragged
|
||
|
|
// matrix.
|
||
|
|
func TestTransposeFlatTiledBitIdentical(t *testing.T) {
|
||
|
|
for _, dim := range [][2]int{{1, 1}, {1, 7}, {7, 1}, {16, 16}, {17, 16}, {16, 17}, {33, 5}, {5, 33}, {64, 64}} {
|
||
|
|
m, n := dim[0], dim[1]
|
||
|
|
a := make([]float64, m*n)
|
||
|
|
for i := range a {
|
||
|
|
a[i] = float64(i%13) - 6
|
||
|
|
}
|
||
|
|
got := TransposeFlat(a, m, n)
|
||
|
|
want := transposeFlatPlain(a, m, n)
|
||
|
|
if len(got) != len(want) {
|
||
|
|
t.Fatalf("%dx%d: length %d, want %d", m, n, len(got), len(want))
|
||
|
|
}
|
||
|
|
for i := range got {
|
||
|
|
if got[i] != want[i] {
|
||
|
|
t.Fatalf("%dx%d: entry %d = %v, want %v", m, n, i, got[i], want[i])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// BenchmarkTransposeFlatAB interleaves the tiled copy with the untiled
|
||
|
|
// one on separate buffers, so a drift in the machine's speed lands on
|
||
|
|
// both: the reported ns/op per style is the paired comparison.
|
||
|
|
func BenchmarkTransposeFlatAB(b *testing.B) {
|
||
|
|
for _, n := range []int{64, 256} {
|
||
|
|
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
|
|
flat := make([]float64, n*n)
|
||
|
|
for i := range flat {
|
||
|
|
flat[i] = float64(i%17) - 8
|
||
|
|
}
|
||
|
|
var tiled, plain time.Duration
|
||
|
|
for b.Loop() {
|
||
|
|
start := time.Now()
|
||
|
|
out := TransposeFlat(flat, n, n)
|
||
|
|
tiled += time.Since(start)
|
||
|
|
if len(out) != len(flat) {
|
||
|
|
b.Fatal("transpose lost entries")
|
||
|
|
}
|
||
|
|
start = time.Now()
|
||
|
|
out = transposeFlatPlain(flat, n, n)
|
||
|
|
plain += time.Since(start)
|
||
|
|
if len(out) != len(flat) {
|
||
|
|
b.Fatal("transpose lost entries")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
b.ReportMetric(float64(tiled.Nanoseconds())/float64(b.N), "ns/op-tiled")
|
||
|
|
b.ReportMetric(float64(plain.Nanoseconds())/float64(b.N), "ns/op-plain")
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// BenchmarkTransposeFlat measures the tiled transpose alone.
|
||
|
|
func BenchmarkTransposeFlat(b *testing.B) {
|
||
|
|
for _, n := range []int{64, 256} {
|
||
|
|
flat := make([]float64, n*n)
|
||
|
|
for i := range flat {
|
||
|
|
flat[i] = float64(i%17) - 8
|
||
|
|
}
|
||
|
|
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
|
|
b.ReportAllocs()
|
||
|
|
for b.Loop() {
|
||
|
|
out := TransposeFlat(flat, n, n)
|
||
|
|
if len(out) != len(flat) {
|
||
|
|
b.Fatal("transpose lost entries")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// BenchmarkSolveSystemColumns measures the column loop of SolveSystem
|
||
|
|
// with one column and with several: the several-column path permutes
|
||
|
|
// each column through a bitmap the crew reuses, which is what the
|
||
|
|
// one-column path does too.
|
||
|
|
func BenchmarkSolveSystemColumns(b *testing.B) {
|
||
|
|
const n = 256
|
||
|
|
for _, cols := range []int{1, 8} {
|
||
|
|
b.Run(fmt.Sprintf("cols=%d", cols), func(b *testing.B) {
|
||
|
|
rhs := make([][]float64, cols)
|
||
|
|
for c := range rhs {
|
||
|
|
rhs[c] = make([]float64, n)
|
||
|
|
for i := range rhs[c] {
|
||
|
|
rhs[c][i] = float64((i+c)%9) - 4
|
||
|
|
}
|
||
|
|
}
|
||
|
|
b.ReportAllocs()
|
||
|
|
for b.Loop() {
|
||
|
|
m, _ := benchFactorRows(n)
|
||
|
|
work := make([][]float64, cols)
|
||
|
|
for c := range work {
|
||
|
|
work[c] = append([]float64(nil), rhs[c]...)
|
||
|
|
}
|
||
|
|
if _, err := SolveSystem("Solve", m, work); err != nil {
|
||
|
|
b.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|