319 lines
8.6 KiB
Go
319 lines
8.6 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"
|
||
)
|
||
|
||
// Benchmarks for the sparse surface: the matrix-vector and matrix-matrix
|
||
// products, the coordinate-to-compressed construction, and the Krylov
|
||
// solvers that drive them. Every input is a fixed literal pattern, so a
|
||
// run is reproducible and two builds compare like with like.
|
||
|
||
// benchSink keeps a kernel's result reachable so the compiler cannot
|
||
// elide the call it came from.
|
||
var (
|
||
benchSinkCSR *SparseCSR
|
||
benchSinkCSC *SparseCSC
|
||
benchSinkArr *core.Array
|
||
)
|
||
|
||
// laplacianTriples returns the flat (row, col, value) triples of the
|
||
// n×n symmetric banded Laplacian of the given half-bandwidth: 2·width
|
||
// on the diagonal and -1 at every offset up to width. It is symmetric
|
||
// positive definite by strict diagonal dominance, which is what the
|
||
// Lanczos and conjugate-gradient benchmarks need, and its non-zeros per
|
||
// row are 2·width+1 at every row.
|
||
func laplacianTriples(n, width int) []float64 {
|
||
tri := make([]float64, 0, n*(2*width+1)*3)
|
||
for i := range n {
|
||
tri = append(tri, float64(i), float64(i), float64(2*width))
|
||
for d := 1; d <= width; d++ {
|
||
j := i + d
|
||
if j >= n {
|
||
break
|
||
}
|
||
tri = append(tri, float64(i), float64(j), -1)
|
||
tri = append(tri, float64(j), float64(i), -1)
|
||
}
|
||
}
|
||
return tri
|
||
}
|
||
|
||
// duplicateTriples returns the triples of a matrix whose coordinates
|
||
// repeat: entry k carries row (k·13) mod n and column (k·37) mod n, so
|
||
// each coordinate appears about len/n times and the construction path
|
||
// has to merge duplicates rather than only sort distinct pairs. The
|
||
// values step along a literal modular sequence, so no merge cancels by
|
||
// accident.
|
||
func duplicateTriples(n, len_ int) []float64 {
|
||
tri := make([]float64, 0, len_*3)
|
||
for k := range len_ {
|
||
row := (k * 13) % n
|
||
col := (k * 37) % n
|
||
v := float64((k*7)%23) - 11
|
||
if v == 0 {
|
||
v = 1.5
|
||
}
|
||
tri = append(tri, float64(row), float64(col), v)
|
||
}
|
||
return tri
|
||
}
|
||
|
||
// benchCOO builds a sparse matrix from flat (row, col, value) triples.
|
||
func benchCOO(b *testing.B, rows, cols int, tri []float64) *core.SparseCOO {
|
||
b.Helper()
|
||
nnz := len(tri) / 3
|
||
idx := make([]int64, 0, nnz*2)
|
||
vals := make([]float64, 0, nnz)
|
||
for i := range nnz {
|
||
idx = append(idx, int64(tri[i*3]), int64(tri[i*3+1]))
|
||
vals = append(vals, tri[i*3+2])
|
||
}
|
||
indices, err := core.FromInts(idx, nnz, 2)
|
||
if err != nil {
|
||
b.Fatalf("FromInts: %v", err)
|
||
}
|
||
values, err := core.FromFloats(vals, nnz)
|
||
if err != nil {
|
||
b.Fatalf("FromFloats: %v", err)
|
||
}
|
||
coo, err := core.NewSparseCOO(indices, values, []int{rows, cols})
|
||
if err != nil {
|
||
b.Fatalf("NewSparseCOO: %v", err)
|
||
}
|
||
return coo
|
||
}
|
||
|
||
// benchVector returns a dense vector of length n whose entries come
|
||
// from a literal modular sequence, deliberately without zeros so a
|
||
// scaled product cannot skip work.
|
||
func benchVector(b *testing.B, n int) *core.Array {
|
||
b.Helper()
|
||
v := make([]float64, n)
|
||
for i := range n {
|
||
v[i] = float64((i*11)%17)/8 - 1
|
||
if v[i] == 0 {
|
||
v[i] = 0.25
|
||
}
|
||
}
|
||
arr, err := core.FromFloats(v, n)
|
||
if err != nil {
|
||
b.Fatalf("FromFloats: %v", err)
|
||
}
|
||
return arr
|
||
}
|
||
|
||
func BenchmarkSparseCSRMatVec(b *testing.B) {
|
||
for _, n := range []int{256, 131072} {
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 1))
|
||
csr, err := CSRFromCOO(coo)
|
||
if err != nil {
|
||
b.Fatalf("CSRFromCOO: %v", err)
|
||
}
|
||
x := benchVector(b, n)
|
||
for b.Loop() {
|
||
benchSinkArr, _ = csr.MatVec(x)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func BenchmarkSparseCSCMatVec(b *testing.B) {
|
||
for _, n := range []int{256, 131072} {
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 1))
|
||
csc, err := CSCFromCOO(coo)
|
||
if err != nil {
|
||
b.Fatalf("CSCFromCOO: %v", err)
|
||
}
|
||
x := benchVector(b, n)
|
||
for b.Loop() {
|
||
benchSinkArr, _ = csc.MatVec(x)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func BenchmarkSparseCSRMatMulSparse(b *testing.B) {
|
||
const n = 256
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 4))
|
||
csr, err := CSRFromCOO(coo)
|
||
if err != nil {
|
||
b.Fatalf("CSRFromCOO: %v", err)
|
||
}
|
||
for b.Loop() {
|
||
benchSinkCSR, _ = csr.MatMulSparse(csr)
|
||
}
|
||
})
|
||
}
|
||
|
||
func BenchmarkSparseCSRMatMulDense(b *testing.B) {
|
||
const n = 512
|
||
const cols = 4
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 2))
|
||
csr, err := CSRFromCOO(coo)
|
||
if err != nil {
|
||
b.Fatalf("CSRFromCOO: %v", err)
|
||
}
|
||
v := make([]float64, n*cols)
|
||
for i := range v {
|
||
v[i] = float64((i*11)%17)/8 - 1
|
||
}
|
||
x, err := core.FromFloats(v, n, cols)
|
||
if err != nil {
|
||
b.Fatalf("FromFloats: %v", err)
|
||
}
|
||
for b.Loop() {
|
||
benchSinkArr, _ = csr.MatMulDense(x)
|
||
}
|
||
})
|
||
}
|
||
|
||
func BenchmarkCSRFromCOODuplicates(b *testing.B) {
|
||
const n = 1500
|
||
tri := duplicateTriples(n, 6000)
|
||
b.Run("nnz=6000", func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, tri)
|
||
for b.Loop() {
|
||
benchSinkCSR, _ = CSRFromCOO(coo)
|
||
}
|
||
})
|
||
}
|
||
|
||
func BenchmarkCSCFromCOODuplicates(b *testing.B) {
|
||
const n = 1500
|
||
tri := duplicateTriples(n, 6000)
|
||
b.Run("nnz=6000", func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, tri)
|
||
for b.Loop() {
|
||
benchSinkCSC, _ = CSCFromCOO(coo)
|
||
}
|
||
})
|
||
}
|
||
|
||
// BenchmarkSpEigenBanded measures the real symmetric Krylov solve end
|
||
// to end: the coordinate construction, the banded matrix-vector
|
||
// products and the full reorthogonalisation of every Lanczos step.
|
||
func BenchmarkSpEigenBanded(b *testing.B) {
|
||
const n = 800
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 1))
|
||
for b.Loop() {
|
||
benchSinkArr, _, _ = SpEigen(coo, 6, core.NewGenerator(7))
|
||
}
|
||
})
|
||
}
|
||
|
||
// BenchmarkSpEigenGeneralBanded measures the Arnoldi solve on the same
|
||
// matrix: one explicit orthogonalisation per basis column instead of
|
||
// the three-term recurrence.
|
||
func BenchmarkSpEigenGeneralBanded(b *testing.B) {
|
||
const n = 400
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 1))
|
||
for b.Loop() {
|
||
benchSinkArr, _, _ = SpEigenGeneral(coo, 4, core.NewGenerator(7))
|
||
}
|
||
})
|
||
}
|
||
|
||
func BenchmarkSpEigenComplexBanded(b *testing.B) {
|
||
const n = 400
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
idx := make([]int64, 0, n*6)
|
||
vals := make([]complex128, 0, n*3)
|
||
for i := range n {
|
||
idx = append(idx, int64(i), int64(i))
|
||
vals = append(vals, 2+0i)
|
||
if i+1 < n {
|
||
idx = append(idx, int64(i), int64(i+1))
|
||
vals = append(vals, 1i)
|
||
idx = append(idx, int64(i+1), int64(i))
|
||
vals = append(vals, -1i)
|
||
}
|
||
}
|
||
indices, err := core.FromInts(idx, len(vals), 2)
|
||
if err != nil {
|
||
b.Fatalf("FromInts: %v", err)
|
||
}
|
||
values, err := core.FromComplexes(vals, len(vals))
|
||
if err != nil {
|
||
b.Fatalf("FromComplexes: %v", err)
|
||
}
|
||
coo, err := core.NewSparseCOO(indices, values, []int{n, n})
|
||
if err != nil {
|
||
b.Fatalf("NewSparseCOO: %v", err)
|
||
}
|
||
for b.Loop() {
|
||
benchSinkArr, _, _ = SpEigenComplex(coo, 4, core.NewGenerator(7))
|
||
}
|
||
})
|
||
}
|
||
|
||
// BenchmarkSpExpApplyBanded measures the Krylov projection of the
|
||
// matrix exponential action, which shares the Lanczos kernel with the
|
||
// symmetric eigensolver.
|
||
func BenchmarkSpExpApplyBanded(b *testing.B) {
|
||
const n = 400
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
coo := benchCOO(b, n, n, laplacianTriples(n, 1))
|
||
v := benchVector(b, n)
|
||
for b.Loop() {
|
||
benchSinkArr, _ = SpExpApply(coo, v, 0)
|
||
}
|
||
})
|
||
}
|
||
|
||
// BenchmarkSpSolveComplexCGBanded measures the conjugate-gradient solve
|
||
// whose every step is a complex sparse matrix-vector product with a
|
||
// Jacobi preconditioner.
|
||
func BenchmarkSpSolveComplexCGBanded(b *testing.B) {
|
||
const n = 256
|
||
b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
|
||
idx := make([]int64, 0, n*6)
|
||
vals := make([]complex128, 0, n*3)
|
||
for i := range n {
|
||
idx = append(idx, int64(i), int64(i))
|
||
vals = append(vals, 4+0i)
|
||
if i+1 < n {
|
||
idx = append(idx, int64(i), int64(i+1))
|
||
vals = append(vals, 1i)
|
||
idx = append(idx, int64(i+1), int64(i))
|
||
vals = append(vals, -1i)
|
||
}
|
||
}
|
||
indices, err := core.FromInts(idx, len(vals), 2)
|
||
if err != nil {
|
||
b.Fatalf("FromInts: %v", err)
|
||
}
|
||
values, err := core.FromComplexes(vals, len(vals))
|
||
if err != nil {
|
||
b.Fatalf("FromComplexes: %v", err)
|
||
}
|
||
coo, err := core.NewSparseCOO(indices, values, []int{n, n})
|
||
if err != nil {
|
||
b.Fatalf("NewSparseCOO: %v", err)
|
||
}
|
||
ones := make([]complex128, n)
|
||
for i := range ones {
|
||
ones[i] = 1
|
||
}
|
||
rhs, err := core.FromComplexes(ones, n)
|
||
if err != nil {
|
||
b.Fatalf("FromComplexes: %v", err)
|
||
}
|
||
for b.Loop() {
|
||
benchSinkArr, _ = SpSolveComplexCG(coo, rhs, 1e-12, 400)
|
||
}
|
||
})
|
||
}
|