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")
|
|||
|
|
}
|