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