252 lines
6.3 KiB
Go
252 lines
6.3 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package linalg
|
||
|
||
import (
|
||
"testing"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// Iterative solver benchmarks: the Krylov solves and the least-squares
|
||
// recursions, with and without the ILU(0) preconditioner. Every system
|
||
// is built from fixed literals so a run is reproducible, and each one
|
||
// is sized to finish well inside a second, which keeps the timings on
|
||
// the solvers' own per-iteration work rather than on input
|
||
// construction.
|
||
|
||
// benchGridCOO builds the five-point stencil of an nx by ny grid as an
|
||
// n x n COO matrix (n = nx·ny) in natural row ordering. The diagonal is
|
||
// 4 and the vertical coupling −1 on both sides. A nonzero skew makes
|
||
// the horizontal coupling −1∓skew, which breaks the symmetry while
|
||
// keeping the matrix diagonally dominant; skew 0 leaves it symmetric
|
||
// positive-definite.
|
||
func benchGridCOO(b *testing.B, nx, ny int, skew float64) *core.SparseCOO {
|
||
b.Helper()
|
||
n := nx * ny
|
||
idx := make([]int64, 0, 5*n*2)
|
||
vals := make([]float64, 0, 5*n)
|
||
add := func(i, j int, v float64) {
|
||
idx = append(idx, int64(i), int64(j))
|
||
vals = append(vals, v)
|
||
}
|
||
for r := range ny {
|
||
for c := range nx {
|
||
i := r*nx + c
|
||
add(i, i, 4)
|
||
if c+1 < nx {
|
||
add(i, i+1, -1-skew)
|
||
}
|
||
if c > 0 {
|
||
add(i, i-1, -1+skew)
|
||
}
|
||
if r+1 < ny {
|
||
add(i, i+nx, -1)
|
||
}
|
||
if r > 0 {
|
||
add(i, i-nx, -1)
|
||
}
|
||
}
|
||
}
|
||
indices, err := core.FromInts(idx, len(vals), 2)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
values, err := core.FromFloats(vals, len(vals))
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
coo, err := core.NewSparseCOO(indices, values, []int{n, n})
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return coo
|
||
}
|
||
|
||
// benchTallCOO builds the overdetermined (2n − 2) x n matrix whose
|
||
// first n rows are the identity and whose remaining rows are the
|
||
// second-difference stencil over three adjacent columns. It is full
|
||
// rank and banded, so LSQR and LSMR converge in a bounded number of
|
||
// steps.
|
||
func benchTallCOO(b *testing.B, n int) *core.SparseCOO {
|
||
b.Helper()
|
||
m := 2*n - 2
|
||
idx := make([]int64, 0, (4*n)*2)
|
||
vals := make([]float64, 0, 4*n)
|
||
add := func(i, j int, v float64) {
|
||
idx = append(idx, int64(i), int64(j))
|
||
vals = append(vals, v)
|
||
}
|
||
for j := range n {
|
||
add(j, j, 1)
|
||
}
|
||
for r := range n - 2 {
|
||
i := n + r
|
||
add(i, r, -1)
|
||
add(i, r+1, 2)
|
||
add(i, r+2, -1)
|
||
}
|
||
indices, err := core.FromInts(idx, len(vals), 2)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
values, err := core.FromFloats(vals, len(vals))
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
coo, err := core.NewSparseCOO(indices, values, []int{m, n})
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return coo
|
||
}
|
||
|
||
// benchPatternRHS builds the deterministic right-hand side of length n
|
||
// from a fixed integer pattern.
|
||
func benchPatternRHS(b *testing.B, n int) *core.Array {
|
||
b.Helper()
|
||
v := make([]float64, n)
|
||
for i := range v {
|
||
v[i] = float64(i%7) - 3
|
||
}
|
||
a, err := core.FromFloats(v, n)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
return a
|
||
}
|
||
|
||
// BenchmarkGMRESSparse measures a restarted GMRES solve of a
|
||
// nonsymmetric five-point system.
|
||
func BenchmarkGMRESSparse(b *testing.B) {
|
||
const n = 400
|
||
coo := benchGridCOO(b, 20, 20, 0.5)
|
||
c, err := cooToCSR(coo, "BenchmarkGMRESSparse")
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
rowStart, colIdx, vals := c.rowStart, c.colIdx, c.vals
|
||
op := func(v *core.Array) (*core.Array, error) {
|
||
out := core.New(core.Float, n)
|
||
dst := out.RawFloats()
|
||
for i := range n {
|
||
sum := 0.0
|
||
for p := rowStart[i]; p < rowStart[i+1]; p++ {
|
||
sum += vals[p] * v.FloatAt(colIdx[p])
|
||
}
|
||
dst[i] = sum
|
||
}
|
||
return out, nil
|
||
}
|
||
bv := benchPatternRHS(b, n)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := GMRES(op, bv, 30, 200, 1e-10); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSpSolveCG measures a conjugate-gradient solve of a
|
||
// symmetric five-point system under the Jacobi preconditioner.
|
||
func BenchmarkSpSolveCG(b *testing.B) {
|
||
coo := benchGridCOO(b, 20, 20, 0)
|
||
bv := benchPatternRHS(b, 400)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SpSolve(coo, bv, 1e-10, 0); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSpSolveCGILU measures the same system under the ILU(0)
|
||
// preconditioner.
|
||
func BenchmarkSpSolveCGILU(b *testing.B) {
|
||
coo := benchGridCOO(b, 20, 20, 0)
|
||
bv := benchPatternRHS(b, 400)
|
||
ilu, err := NewSparseILU(coo)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SpSolve(coo, bv, 1e-10, 0, ilu); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSpSolveBiCGSTAB measures a BiCGSTAB solve of a nonsymmetric
|
||
// five-point system under the Jacobi preconditioner.
|
||
func BenchmarkSpSolveBiCGSTAB(b *testing.B) {
|
||
coo := benchGridCOO(b, 20, 20, 0.5)
|
||
bv := benchPatternRHS(b, 400)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SpSolveBiCGSTAB(coo, bv, 1e-10, 0); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSpSolveBiCGSTABILU measures the same system under the
|
||
// ILU(0) preconditioner.
|
||
func BenchmarkSpSolveBiCGSTABILU(b *testing.B) {
|
||
coo := benchGridCOO(b, 20, 20, 0.5)
|
||
bv := benchPatternRHS(b, 400)
|
||
ilu, err := NewSparseILU(coo)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SpSolveBiCGSTAB(coo, bv, 1e-10, 0, ilu); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSpLSQR measures the Golub-Kahan least-squares recursion of
|
||
// SpLSQR on an overdetermined banded system.
|
||
func BenchmarkSpLSQR(b *testing.B) {
|
||
const n = 400
|
||
coo := benchTallCOO(b, n)
|
||
bv := benchPatternRHS(b, 2*n-2)
|
||
_, info, err := SpLSQR(coo, bv, 1e-10, 0, 0)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
steps := float64(info.Iterations)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, err := SpLSQR(coo, bv, 1e-10, 0, 0); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
// The system and the tolerance are fixed, so the step count is too:
|
||
// reporting it shows the size of the work the timing covers. It goes
|
||
// after the loop, which clears any metric reported before it.
|
||
b.ReportMetric(steps, "steps")
|
||
}
|
||
|
||
// BenchmarkSpLSMR measures the LSMR recursion on the same system.
|
||
func BenchmarkSpLSMR(b *testing.B) {
|
||
const n = 400
|
||
coo := benchTallCOO(b, n)
|
||
bv := benchPatternRHS(b, 2*n-2)
|
||
_, info, err := SpLSMR(coo, bv, 1e-10, 0, 0)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
steps := float64(info.Iterations)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, err := SpLSMR(coo, bv, 1e-10, 0, 0); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
b.ReportMetric(steps, "steps")
|
||
}
|