513 lines
13 KiB
Go
513 lines
13 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package linalg
|
||
|
||
import (
|
||
"fmt"
|
||
"testing"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
)
|
||
|
||
// Decomposition benchmarks for the shapes the solver benchmarks in
|
||
// perf_bench_test.go do not cover: the tall least-squares system, the
|
||
// pivoted QR sweep, the complex and matrix-function paths, and the SVD
|
||
// shapes whose cost sits in the reconstruction kernels rather than in
|
||
// the bidiagonalisation. Every input is a fixed literal formula, so the
|
||
// timings compare like with like across runs.
|
||
|
||
// benchDecompFloats fills an m×n row-major matrix with a deterministic
|
||
// literal pattern: a diagonally weighted band plus a bounded pseudo-random
|
||
// ripple.
|
||
func benchDecompFloats(m, n int, diagonal float64) []float64 {
|
||
v := make([]float64, m*n)
|
||
for i := range m {
|
||
for j := range n {
|
||
v[i*n+j] = 0.5*float64((i*11+j*7)%17) - 4 + 0.25*float64((i*j*5)%13)
|
||
}
|
||
v[i*n+i%n] += diagonal
|
||
}
|
||
return v
|
||
}
|
||
|
||
// benchDecompSPD builds a symmetric strictly diagonally dominant
|
||
// matrix, which is positive definite however the ripple lands.
|
||
func benchDecompSPD(n int) []float64 {
|
||
v := make([]float64, n*n)
|
||
for i := range n {
|
||
for j := range n {
|
||
v[i*n+j] = float64((i*5+j*3)%9) - 4
|
||
}
|
||
}
|
||
for i := range n {
|
||
for j := i + 1; j < n; j++ {
|
||
avg := (v[i*n+j] + v[j*n+i]) / 2
|
||
v[i*n+j], v[j*n+i] = avg, avg
|
||
}
|
||
row := 0.0
|
||
for j := range n {
|
||
if j != i {
|
||
row += v[i*n+j]
|
||
if v[i*n+j] < 0 {
|
||
row += 2 * -v[i*n+j]
|
||
}
|
||
}
|
||
}
|
||
v[i*n+i] += row + float64(n)
|
||
}
|
||
return v
|
||
}
|
||
|
||
// benchDecompHermitian builds a Hermitian matrix with a real diagonal
|
||
// whose entries are well separated, so the Jacobi sweep deflates
|
||
// quickly.
|
||
func benchDecompHermitian(n int) []complex128 {
|
||
v := make([]complex128, n*n)
|
||
for i := range n {
|
||
v[i*n+i] = complex(4*float64(i)+8, 0)
|
||
for j := i + 1; j < n; j++ {
|
||
re := 0.5*float64((i*3+j*5)%7) - 1.5
|
||
im := 0.25*float64((i+j)%5) - 0.5
|
||
v[i*n+j] = complex(re, im)
|
||
v[j*n+i] = complex(re, -im)
|
||
}
|
||
}
|
||
return v
|
||
}
|
||
|
||
func benchDecompArray(b *testing.B, v []float64, m, n int) *core.Array {
|
||
b.Helper()
|
||
a, err := core.FromFloats(v, m, n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return a
|
||
}
|
||
|
||
func benchDecompComplexArray(b *testing.B, v []complex128, m, n int) *core.Array {
|
||
b.Helper()
|
||
a, err := core.FromComplexes(v, m, n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return a
|
||
}
|
||
|
||
// BenchmarkCholesky512 measures the blocked Cholesky sweep above the
|
||
// size where the solver benchmark's 256 stops.
|
||
func BenchmarkCholesky512(b *testing.B) {
|
||
a := benchDecompArray(b, benchDecompSPD(512), 512, 512)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := Cholesky(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkLeastSquares512x32 measures the tall least-squares route,
|
||
// where the QR factor's orthogonal accumulation dominates.
|
||
func BenchmarkLeastSquares512x32(b *testing.B) {
|
||
const m, n = 512, 32
|
||
a := benchDecompArray(b, benchDecompFloats(m, n, 3), m, n)
|
||
rhs, err := core.FromFloats(benchDecompFloats(m, 1, 1), m)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := LeastSquares(a, rhs); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkRRQR192x48 measures the pivoted QR sweep, whose cost is the
|
||
// per-step column-norm scan and the reflector application.
|
||
func BenchmarkRRQR192x48(b *testing.B) {
|
||
const m, n = 192, 48
|
||
a := benchDecompArray(b, benchDecompFloats(m, n, 3), m, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, _, _, err := RRQR(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkRRQR128x32 measures the pivoted QR sweep at the size where a
|
||
// single reflector's work is small enough that crew sizing decides it.
|
||
func BenchmarkRRQR128x32(b *testing.B) {
|
||
const m, n = 128, 32
|
||
a := benchDecompArray(b, benchDecompFloats(m, n, 3), m, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, _, _, err := RRQR(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSVDReconstruct1024x16 measures the SVD shapes whose cost sits
|
||
// in the reconstruction kernels: the full QR factor is far larger than
|
||
// the bidiagonalisation.
|
||
func BenchmarkSVDReconstruct1024x16(b *testing.B) {
|
||
const m, n = 1024, 16
|
||
a := benchDecompArray(b, benchDecompFloats(m, n, 2), m, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, _, err := SVD(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkEigenComplex8 measures the complex Hermitian Jacobi sweep at
|
||
// the largest size its convergence floor admits: the sweep's off-mass
|
||
// threshold sits below the rounding floor the rotations leave behind
|
||
// above 8, so a bigger input errors out before it measures anything.
|
||
func BenchmarkEigenComplex8(b *testing.B) {
|
||
const n = 8
|
||
a := benchDecompComplexArray(b, benchDecompHermitian(n), n, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, err := EigenComplex(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSVDComplex96x48 measures the direct complex bidiagonalisation
|
||
// and its Golub-Reinsch sweep.
|
||
func BenchmarkSVDComplex96x48(b *testing.B) {
|
||
const m, n = 96, 48
|
||
flat := benchDecompFloats(m, n, 2)
|
||
v := make([]complex128, m*n)
|
||
for i := range m * n {
|
||
v[i] = complex(flat[i], 0.5*float64((i*3)%7)-1.5)
|
||
}
|
||
a := benchDecompComplexArray(b, v, m, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, _, err := SVDComplex(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkEigenGeneral32 measures the Hessenberg reduction, the shifted
|
||
// QR sweep and the eigenvector back-substitution on a general matrix.
|
||
func BenchmarkEigenGeneral32(b *testing.B) {
|
||
const n = 32
|
||
v := benchDecompFloats(n, n, 4)
|
||
a := benchDecompArray(b, v, n, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, err := EigenGeneral(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSchurComplex48 measures the complex Schur decomposition the
|
||
// matrix functions are built on.
|
||
func BenchmarkSchurComplex48(b *testing.B) {
|
||
const n = 48
|
||
flat := benchDecompFloats(n, n, 3)
|
||
v := make([]complex128, n*n)
|
||
for i := range n * n {
|
||
v[i] = complex(flat[i], 0.25*float64((i*5)%9)-1)
|
||
}
|
||
a := benchDecompComplexArray(b, v, n, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, err := SchurComplex(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkMatrixExp32 measures the Padé scaling-and-squaring kernel.
|
||
func BenchmarkMatrixExp32(b *testing.B) {
|
||
const n = 32
|
||
v := benchDecompFloats(n, n, 0)
|
||
v[0] = -0.5 // keep the 1-norm inside the unscaled degrees
|
||
a := benchDecompArray(b, v, n, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := MatrixExp(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkMatrixSqrt64 measures the symmetric eigen route of a matrix
|
||
// function, whose cost is the eigendecomposition plus two triple
|
||
// products.
|
||
func BenchmarkMatrixSqrt64(b *testing.B) {
|
||
const n = 64
|
||
a := benchDecompArray(b, benchDecompSPD(n), n, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := MatrixSqrt(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkMatrixLogSchur32 measures the Schur-Parlett logarithm, the
|
||
// square-root walk and the Mercator series included.
|
||
func BenchmarkMatrixLogSchur32(b *testing.B) {
|
||
const n = 32
|
||
flat := benchDecompSPD(n)
|
||
v := benchDecompComplexSPD(flat, n)
|
||
a := benchDecompComplexArray(b, v, n, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := MatrixLog(a); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// benchDecompComplexSPD lifts a symmetric positive definite real matrix
|
||
// to complex with a small positive imaginary part on the strict upper
|
||
// triangle and its conjugate below, keeping the spectrum off the
|
||
// non-positive real axis.
|
||
func benchDecompComplexSPD(flat []float64, n int) []complex128 {
|
||
v := make([]complex128, n*n)
|
||
for i := range n {
|
||
for j := range n {
|
||
v[i*n+j] = complex(flat[i*n+j], 0)
|
||
}
|
||
}
|
||
for i := range n {
|
||
for j := i + 1; j < n; j++ {
|
||
im := 0.05 * float64((i+j)%4+1)
|
||
v[i*n+j] = complex(flat[i*n+j], im)
|
||
v[j*n+i] = complex(flat[j*n+i], -im)
|
||
}
|
||
}
|
||
return v
|
||
}
|
||
|
||
// BenchmarkTikhonov256x64 measures the SVD solve route, whose cost is a
|
||
// full SVD of the system matrix plus the Uᵀb and V·d products.
|
||
func BenchmarkTikhonov256x64(b *testing.B) {
|
||
const m, n = 256, 64
|
||
a := benchDecompArray(b, benchDecompFloats(m, n, 2), m, n)
|
||
rhs, err := core.FromFloats(benchDecompFloats(m, 1, 1), m)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SolveTikhonov(a, rhs, 0.5); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestDecompWorkerCountBitIdentity pins the parallel splits: every
|
||
// kernel this file benchmarks must answer bit for bit the same whether
|
||
// the crew is one goroutine or the machine's full width. A split that
|
||
// moved an addend or a boundary shows up here before the oracle's
|
||
// smaller pinned shapes see it.
|
||
func TestDecompWorkerCountBitIdentity(t *testing.T) {
|
||
prev := engine.SetNumWorkers(1)
|
||
defer engine.SetNumWorkers(prev)
|
||
|
||
const (
|
||
cholN = 320
|
||
lsM = 384
|
||
lsN = 24
|
||
rrM = 160
|
||
rrN = 40
|
||
svdM = 384
|
||
svdN = 24
|
||
eigN = 192
|
||
cplxN = 8
|
||
)
|
||
spd := benchDecompSPD(cholN)
|
||
tall := benchDecompFloats(lsM, lsN, 3)
|
||
rr := benchDecompFloats(rrM, rrN, 3)
|
||
wide := benchDecompFloats(svdM, svdN, 2)
|
||
herm := benchDecompHermitian(cplxN)
|
||
square := benchDecompSPD(eigN)
|
||
complexSquare := make([]complex128, cplxN*cplxN)
|
||
flatComplex := benchDecompFloats(cplxN, cplxN, 3)
|
||
for i := range cplxN * cplxN {
|
||
complexSquare[i] = complex(flatComplex[i], 0.25*float64((i*5)%9)-1)
|
||
}
|
||
|
||
spdA := mustFromFloats(t, spd, cholN, cholN)
|
||
tallA := mustFromFloats(t, tall, lsM, lsN)
|
||
rrA := mustFromFloats(t, rr, rrM, rrN)
|
||
wideA := mustFromFloats(t, wide, svdM, svdN)
|
||
hermA := mustFromComplexes(t, herm, cplxN, cplxN)
|
||
squareA := mustFromFloats(t, square, eigN, eigN)
|
||
complexA := mustFromComplexes(t, complexSquare, cplxN, cplxN)
|
||
rhs := mustFromFloats(t, benchDecompFloats(lsM, 1, 1), lsM)
|
||
rrRHS := mustFromFloats(t, benchDecompFloats(rrM, 1, 1), rrM)
|
||
// A tall system wide enough that the blocked Qᵀb dispatches a crew,
|
||
// so the split is compared against the serial order too.
|
||
const lsWideN = 48
|
||
wideLS := mustFromFloats(t, benchDecompFloats(lsM, lsWideN, 3), lsM, lsWideN)
|
||
|
||
type snap struct {
|
||
name string
|
||
run func() []*core.Array
|
||
}
|
||
cases := []snap{
|
||
{"Cholesky", func() []*core.Array {
|
||
l, err := Cholesky(spdA)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{l}
|
||
}},
|
||
{"LeastSquares", func() []*core.Array {
|
||
x, err := LeastSquares(tallA, rhs)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{x}
|
||
}},
|
||
{"LeastSquaresBlockedQtb", func() []*core.Array {
|
||
x, err := LeastSquares(wideLS, rhs)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{x}
|
||
}},
|
||
{"RRQR", func() []*core.Array {
|
||
q, r, perm, rank, err := RRQR(rrA)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
pf := make([]float64, len(perm)+1)
|
||
for i, p := range perm {
|
||
pf[i] = float64(p)
|
||
}
|
||
pf[len(perm)] = float64(rank)
|
||
return []*core.Array{q, r, mustFromFloats(t, pf, len(pf), 1)}
|
||
}},
|
||
{"SolveRRQR", func() []*core.Array {
|
||
x, err := SolveRRQR(rrA, rrRHS)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{x}
|
||
}},
|
||
{"SVD", func() []*core.Array {
|
||
u, s, vt, err := SVD(wideA)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{u, s, vt}
|
||
}},
|
||
{"Eigen", func() []*core.Array {
|
||
v, q, err := Eigen(squareA)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{v, q}
|
||
}},
|
||
{"EigenComplex", func() []*core.Array {
|
||
v, q, err := EigenComplex(hermA)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{v, q}
|
||
}},
|
||
{"SVDComplex", func() []*core.Array {
|
||
u, s, vh, err := SVDComplex(complexA)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return []*core.Array{u, s, vh}
|
||
}},
|
||
}
|
||
|
||
serial := make([][]*core.Array, len(cases))
|
||
for i, c := range cases {
|
||
serial[i] = c.run()
|
||
}
|
||
engine.SetNumWorkers(0) // the machine's full width
|
||
for i, c := range cases {
|
||
got := c.run()
|
||
for k := range got {
|
||
if got[k].Len() != serial[i][k].Len() {
|
||
t.Fatalf("%s: result %d length %d under the full crew, %d serial",
|
||
c.name, k, got[k].Len(), serial[i][k].Len())
|
||
}
|
||
if !rawBitsEqual(got[k], serial[i][k]) {
|
||
t.Fatalf("%s: result %d differs bitwise between the serial and parallel crew", c.name, k)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// rawBitsEqual compares two arrays' payloads bit for bit, complex
|
||
// payloads included.
|
||
func rawBitsEqual(a, b *core.Array) bool {
|
||
if a.Dtype() == core.Complex || b.Dtype() == core.Complex {
|
||
ac, bc := a.RawComplexes(), b.RawComplexes()
|
||
if len(ac) != len(bc) {
|
||
return false
|
||
}
|
||
for i := range ac {
|
||
if ac[i] != bc[i] {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
af, bf := a.RawFloats(), b.RawFloats()
|
||
if len(af) != len(bf) {
|
||
return false
|
||
}
|
||
for i := range af {
|
||
if af[i] != bf[i] {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// BenchmarkLeastSquaresTall measures the solve at the shapes the
|
||
// reflector route exists for: many more rows than columns, where forming
|
||
// Q would dominate everything else.
|
||
func BenchmarkLeastSquaresTall(b *testing.B) {
|
||
for _, c := range []struct{ m, n int }{{512, 32}, {2048, 64}, {8192, 16}} {
|
||
a := make([]float64, c.m*c.n)
|
||
s := uint64(20260920)
|
||
for i := range a {
|
||
s = s*6364136223846793005 + 1442695040888963407
|
||
a[i] = float64((s>>40)%9+1) * 0.5
|
||
}
|
||
bm := make([]float64, c.m)
|
||
for i := range bm {
|
||
bm[i] = float64(i%11) - 5
|
||
}
|
||
am, err := core.FromFloats(a, c.m, c.n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
bv, err := core.FromFloats(bm, c.m, 1)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
b.Run(fmt.Sprintf("%dx%d", c.m, c.n), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := LeastSquares(am, bv); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|