422 lines
11 KiB
Go
422 lines
11 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/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// Benchmarks for the dense LU kernel, the sparse factorisations and the
|
||
// sparse rank-one sweep. Every input is built from fixed literals, so
|
||
// the patterns, the pivots and the orderings are the same on every run.
|
||
|
||
// denseLUInput builds an n×n row-major matrix whose rows view one flat
|
||
// buffer, with a pristine copy the benchmark restores from: the
|
||
// factorisation consumes its argument, and rebuilding the matrix inside
|
||
// the timed region would measure the rebuild.
|
||
func denseLUInput(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*5+j*11)%13) - 6
|
||
}
|
||
// Diagonal dominance keeps the elimination on the
|
||
// well-conditioned side, so the timing measures the kernel.
|
||
flat[i*n+i] += float64(2 * 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 restoreDenseRows(rows [][]float64, pristine []float64) {
|
||
off := 0
|
||
for _, row := range rows {
|
||
copy(row, pristine[off:off+len(row)])
|
||
off += len(row)
|
||
}
|
||
}
|
||
|
||
// BenchmarkDenseLUFactor measures the LU elimination alone at the sizes
|
||
// the solvers reach.
|
||
func BenchmarkDenseLUFactor(b *testing.B) {
|
||
for _, n := range []int{128, 256, 512} {
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
rows, pristine := denseLUInput(n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
restoreDenseRows(rows, pristine)
|
||
base.Factor(rows)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// denseLUSolveInput builds the same matrix as a core array, for the
|
||
// end-to-end solve path.
|
||
func denseLUSolveInput(b *testing.B, n int) (*core.Array, *core.Array) {
|
||
b.Helper()
|
||
flat := make([]float64, n*n)
|
||
for i := range n {
|
||
for j := range n {
|
||
flat[i*n+j] = float64((i*5+j*11)%13) - 6
|
||
}
|
||
flat[i*n+i] += float64(2 * n)
|
||
}
|
||
a, err := core.FromFloats(flat, n, n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
rhs := make([]float64, n)
|
||
for i := range rhs {
|
||
rhs[i] = float64(i%7) - 3
|
||
}
|
||
x, err := core.FromFloats(rhs, n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return a, x
|
||
}
|
||
|
||
// BenchmarkDenseLUSolve measures Solve end to end: the copy of the
|
||
// matrix, the factorisation, the permutation and the substitution.
|
||
func BenchmarkDenseLUSolve(b *testing.B) {
|
||
for _, n := range []int{128, 256, 512} {
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
a, x := denseLUSolveInput(b, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := Solve(a, x); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// sparseCOOFrom assembles a COO matrix from the triplets a builder
|
||
// appended, refusing nothing: the builders below emit canonical,
|
||
// duplicate-free triplets.
|
||
func sparseCOOFrom(b *testing.B, n int, idx []int64, vals []float64) *core.SparseCOO {
|
||
b.Helper()
|
||
indices, err := core.FromInts(idx, len(vals), 2)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(vals)}), []int{n, n})
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return coo
|
||
}
|
||
|
||
// gridLaplacian builds the 5-point Laplacian on a w×h grid in row-major
|
||
// order: symmetric positive definite, banded, and the pattern every
|
||
// sparse direct solver is measured on.
|
||
func gridLaplacian(b *testing.B, w, h int) *core.SparseCOO {
|
||
b.Helper()
|
||
n := w * h
|
||
idx := make([]int64, 0, 5*n)
|
||
vals := make([]float64, 0, 5*n)
|
||
add := func(r, c int, v float64) {
|
||
idx = append(idx, int64(r), int64(c))
|
||
vals = append(vals, v)
|
||
}
|
||
at := func(x, y int) int { return y*w + x }
|
||
for y := range h {
|
||
for x := range w {
|
||
add(at(x, y), at(x, y), 4)
|
||
if x+1 < w {
|
||
add(at(x, y), at(x+1, y), -1)
|
||
add(at(x+1, y), at(x, y), -1)
|
||
}
|
||
if y+1 < h {
|
||
add(at(x, y), at(x, y+1), -1)
|
||
add(at(x, y+1), at(x, y), -1)
|
||
}
|
||
}
|
||
}
|
||
return sparseCOOFrom(b, n, idx, vals)
|
||
}
|
||
|
||
// arrowHead builds the arrowhead matrix of order n: a diagonal plus a
|
||
// dense first row and column. The diagonal dominates the first arrow
|
||
// (n+1 against n−1), so the matrix is positive definite, and the
|
||
// ordering has a genuine choice to make on the arrow.
|
||
func arrowHead(b *testing.B, n int) *core.SparseCOO {
|
||
b.Helper()
|
||
idx := make([]int64, 0, 3*n)
|
||
vals := make([]float64, 0, 3*n)
|
||
add := func(r, c int, v float64) {
|
||
idx = append(idx, int64(r), int64(c))
|
||
vals = append(vals, v)
|
||
}
|
||
add(0, 0, float64(n+1))
|
||
for i := 1; i < n; i++ {
|
||
add(i, i, float64(i+2))
|
||
add(0, i, 1)
|
||
add(i, 0, 1)
|
||
}
|
||
return sparseCOOFrom(b, n, idx, vals)
|
||
}
|
||
|
||
// bandedNonsymmetric builds a banded, diagonally dominant matrix with a
|
||
// periodic spike below the diagonal: every tenth column has a
|
||
// subdiagonal entry larger than its diagonal, so the partial pivoting
|
||
// of the LU factorisation swaps rows and its label bookkeeping is
|
||
// exercised rather than measured cold.
|
||
func bandedNonsymmetric(b *testing.B, n int) *core.SparseCOO {
|
||
b.Helper()
|
||
idx := make([]int64, 0, 6*n)
|
||
vals := make([]float64, 0, 6*n)
|
||
add := func(r, c int, v float64) {
|
||
idx = append(idx, int64(r), int64(c))
|
||
vals = append(vals, v)
|
||
}
|
||
for i := range n {
|
||
d := 6.0
|
||
if i%10 == 3 {
|
||
d = 0.5 // the spike below takes this column's pivot
|
||
}
|
||
add(i, i, d)
|
||
if i+1 < n {
|
||
add(i, i+1, 1+0.25*float64(i%3))
|
||
add(i+1, i, 1.5+0.5*float64(i%5))
|
||
}
|
||
if i+3 < n {
|
||
add(i, i+3, 0.5)
|
||
}
|
||
}
|
||
return sparseCOOFrom(b, n, idx, vals)
|
||
}
|
||
|
||
// swapHeavyBanded builds a tridiagonal matrix that pivots on every
|
||
// column: a small diagonal against a large subdiagonal, so the largest
|
||
// entry at or below the diagonal is always the row below. The factor
|
||
// stays banded, so the measurement is the pivot bookkeeping rather
|
||
// than the elimination arithmetic.
|
||
func swapHeavyBanded(b *testing.B, n int) *core.SparseCOO {
|
||
b.Helper()
|
||
idx := make([]int64, 0, 3*n)
|
||
vals := make([]float64, 0, 3*n)
|
||
add := func(r, c int, v float64) {
|
||
idx = append(idx, int64(r), int64(c))
|
||
vals = append(vals, v)
|
||
}
|
||
for i := range n {
|
||
add(i, i, 0.125)
|
||
if i+1 < n {
|
||
add(i, i+1, 1)
|
||
add(i+1, i, 16)
|
||
}
|
||
}
|
||
return sparseCOOFrom(b, n, idx, vals)
|
||
}
|
||
|
||
// BenchmarkSparseLUFactorSwapped measures the elimination of a system
|
||
// that pivots at every column, so the cost of relabelling L's stored
|
||
// rows is part of the number.
|
||
func BenchmarkSparseLUFactorSwapped(b *testing.B) {
|
||
for _, n := range []int{100, 200, 400} {
|
||
coo := swapHeavyBanded(b, n)
|
||
b.Run(fmt.Sprintf("banded-%d", n), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := NewSparseLU(coo); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// denseRHS builds a deterministic right hand side of length n.
|
||
func denseRHS(b *testing.B, n int) *core.Array {
|
||
b.Helper()
|
||
v := make([]float64, n)
|
||
for i := range v {
|
||
v[i] = float64(i%11) - 5 + 0.5*float64(i%3)
|
||
}
|
||
x, err := core.FromFloats(v, n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return x
|
||
}
|
||
|
||
// BenchmarkSparseCholeskyFactor measures the symbolic and numeric
|
||
// elimination of a few hundred unknowns: the Laplacian grid with the
|
||
// natural order, then the same grid through the reverse Cuthill-McKee
|
||
// order, which meets a much smaller factor.
|
||
func BenchmarkSparseCholeskyFactor(b *testing.B) {
|
||
coo := gridLaplacian(b, 19, 19)
|
||
b.Run("grid-natural", func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := NewSparseCholesky(coo, SparseOrderingNatural); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
b.Run("grid-rcm", func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
coo = arrowHead(b, 300)
|
||
b.Run("arrow-natural", func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := NewSparseCholesky(coo, SparseOrderingNatural); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
|
||
// BenchmarkSparseCholeskySolve measures the solve on a factor built
|
||
// once outside the loop: the gather, the two substitutions and the
|
||
// scatter.
|
||
func BenchmarkSparseCholeskySolve(b *testing.B) {
|
||
for _, size := range []struct {
|
||
name string
|
||
w, h int
|
||
ordering SparseOrdering
|
||
}{
|
||
{"grid-natural", 19, 19, SparseOrderingNatural},
|
||
{"grid-rcm", 19, 19, SparseOrderingReverseCuthillMcKee},
|
||
} {
|
||
coo := gridLaplacian(b, size.w, size.h)
|
||
f, err := NewSparseCholesky(coo, size.ordering)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
rhs := denseRHS(b, size.w*size.h)
|
||
b.Run(size.name, func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := f.Solve(rhs); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// BenchmarkSparseLUFactor measures the left-looking elimination of a
|
||
// banded nonsymmetric system with pivoting.
|
||
func BenchmarkSparseLUFactor(b *testing.B) {
|
||
for _, n := range []int{200, 400} {
|
||
coo := bandedNonsymmetric(b, n)
|
||
b.Run(fmt.Sprintf("banded-%d", n), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := NewSparseLU(coo); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// BenchmarkSparseLUSolve measures the forward and backward substitution
|
||
// on a factor built once outside the loop.
|
||
func BenchmarkSparseLUSolve(b *testing.B) {
|
||
const n = 400
|
||
coo := bandedNonsymmetric(b, n)
|
||
f, err := NewSparseLU(coo)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
rhs := denseRHS(b, n)
|
||
b.Run("banded-400", func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := f.Solve(rhs); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
|
||
// sparseCholRankOneInput factors the tridiagonal Laplacian and returns
|
||
// the factor with a two-entry vector whose support sits on the stored
|
||
// pattern, so both the update and the downdate are accepted.
|
||
func sparseCholRankOneInput(b *testing.B, n int) (*SparseCholesky, *core.Array) {
|
||
b.Helper()
|
||
idx := make([]int64, 0, 3*n)
|
||
vals := make([]float64, 0, 3*n)
|
||
add := func(r, c int, v float64) {
|
||
idx = append(idx, int64(r), int64(c))
|
||
vals = append(vals, v)
|
||
}
|
||
for i := range n {
|
||
add(i, i, 4)
|
||
if i+1 < n {
|
||
add(i, i+1, -1)
|
||
add(i+1, i, -1)
|
||
}
|
||
}
|
||
coo := sparseCOOFrom(b, n, idx, vals)
|
||
f, err := NewSparseCholesky(coo, SparseOrderingNatural)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
x := core.New(core.Float, []int{n}...)
|
||
x.RawFloats()[0] = 0.5
|
||
x.RawFloats()[1] = 0.25
|
||
return f, x
|
||
}
|
||
|
||
// BenchmarkSparseCholRankOnePair measures one rank-one update followed
|
||
// by the matching downdate: each iteration leaves the factor where it
|
||
// found it, so the sweep is timed rather than the refactorisation the
|
||
// modification replaces.
|
||
func BenchmarkSparseCholRankOnePair(b *testing.B) {
|
||
for _, n := range []int{150, 300} {
|
||
f, x := sparseCholRankOneInput(b, n)
|
||
b.Run(fmt.Sprintf("tridiag-%d", n), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if err := f.Update(x); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
if err := f.Downdate(x); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// BenchmarkSparseCholUpdateRefusal measures the refusal path: a vector
|
||
// whose support reaches outside the stored pattern, so every call
|
||
// returns the pattern error after the fill check has walked the support.
|
||
// It is the cost the check pays before it can refuse, and the error
|
||
// construction is deliberately part of it.
|
||
func BenchmarkSparseCholUpdateRefusal(b *testing.B) {
|
||
const n = 300
|
||
f, _ := sparseCholRankOneInput(b, n)
|
||
x := core.New(core.Float, []int{n}...)
|
||
x.RawFloats()[0] = 0.5
|
||
x.RawFloats()[n-1] = 0.25
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if err := f.Update(x); err == nil {
|
||
b.Fatal("expected the update to need fill the pattern does not hold")
|
||
}
|
||
}
|
||
}
|