133 lines
3.7 KiB
Go
133 lines
3.7 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package core
|
||
|
||
import "testing"
|
||
|
||
// Along-dimension axis reductions: the softmax, normalisation and loss
|
||
// primitives. Run with -bench before and after any change to axis.go.
|
||
// The 2-D shapes line up with BenchmarkNorm2D and BenchmarkProd2D; the
|
||
// 3-D 32×17×9 shape exercises every dim position against a
|
||
// non-power-of-two stride.
|
||
|
||
// benchFloat2D builds an n×n float64 array with a deterministic mix of
|
||
// small values.
|
||
func benchFloat2D(b *testing.B, n int) *Array {
|
||
b.Helper()
|
||
a, _ := FromFloats(make([]float64, n*n), n, n)
|
||
for i := range a.RawFloats() {
|
||
a.RawFloats()[i] = float64(i%9) - 4
|
||
}
|
||
return a
|
||
}
|
||
|
||
// benchFloat3D builds a 32×17×9 float64 array with a deterministic mix
|
||
// of small values.
|
||
func benchFloat3D(b *testing.B) *Array {
|
||
b.Helper()
|
||
a, _ := FromFloats(make([]float64, 32*17*9), 32, 17, 9)
|
||
for i := range a.RawFloats() {
|
||
a.RawFloats()[i] = float64(i%11) - 5
|
||
}
|
||
return a
|
||
}
|
||
|
||
func benchAxis1(b *testing.B, f func(*Array, int) (*Array, error)) {
|
||
b.Helper()
|
||
a := benchFloat2D(b, 512)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := f(a, 1); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
func BenchmarkSumAxis2D(b *testing.B) { benchAxis1(b, SumAxis) }
|
||
func BenchmarkMeanAxis2D(b *testing.B) { benchAxis1(b, MeanAxis) }
|
||
func BenchmarkMinAxis2D(b *testing.B) { benchAxis1(b, MinAxis) }
|
||
func BenchmarkMaxAxis2D(b *testing.B) { benchAxis1(b, MaxAxis) }
|
||
func BenchmarkArgMaxAxis2D(b *testing.B) { benchArgAxis1(b, false) }
|
||
func BenchmarkArgMinAxis2D(b *testing.B) { benchArgAxis1(b, true) }
|
||
|
||
func benchArgAxis1(b *testing.B, wantMin bool) {
|
||
b.Helper()
|
||
a := benchFloat2D(b, 512)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
var err error
|
||
if wantMin {
|
||
_, err = ArgMinAxis(a, 1)
|
||
} else {
|
||
_, err = ArgMaxAxis(a, 1)
|
||
}
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// The 3-D cases walk every dim position: dim 2 is the contiguous
|
||
// trailing run, dim 0 the coarsest split, dim 1 the interior stride.
|
||
func BenchmarkSumAxis3DDim0(b *testing.B) { benchAxis3D(b, SumAxis, 0) }
|
||
func BenchmarkSumAxis3DDim1(b *testing.B) { benchAxis3D(b, SumAxis, 1) }
|
||
func BenchmarkSumAxis3DDim2(b *testing.B) { benchAxis3D(b, SumAxis, 2) }
|
||
func BenchmarkMinAxis3DDim0(b *testing.B) { benchAxis3D(b, MinAxis, 0) }
|
||
func BenchmarkMinAxis3DDim1(b *testing.B) { benchAxis3D(b, MinAxis, 1) }
|
||
func BenchmarkMinAxis3DDim2(b *testing.B) { benchAxis3D(b, MinAxis, 2) }
|
||
|
||
func benchAxis3D(b *testing.B, f func(*Array, int) (*Array, error), dim int) {
|
||
b.Helper()
|
||
a := benchFloat3D(b)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := f(a, dim); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSumAxis2DInt covers the int payload of SumAxis; the int path
|
||
// folds straight into the accumulator with no float scratch.
|
||
func BenchmarkSumAxis2DInt(b *testing.B) {
|
||
a, _ := FromInts(make([]int64, 512*512), 512, 512)
|
||
for i := range a.RawInts() {
|
||
a.RawInts()[i] = int64(i%9) - 4
|
||
}
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SumAxis(a, 1); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkSumAxis2DFloat32 covers the float32 scratch path: folds go
|
||
// through a float64 scratch and round once at the end.
|
||
func BenchmarkSumAxis2DFloat32(b *testing.B) {
|
||
af, _ := FromFloats(make([]float64, 512*512), 512, 512)
|
||
for i := range af.RawFloats() {
|
||
af.RawFloats()[i] = float64(i%9) - 4
|
||
}
|
||
a, _ := Astype(af, Float32)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, err := SumAxis(a, 1); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// BenchmarkTopK2D covers the repeated-argmax selection along the
|
||
// trailing dimension, the loss and beam-search hot path.
|
||
func BenchmarkTopK2D(b *testing.B) {
|
||
a := benchFloat2D(b, 512)
|
||
b.ReportAllocs()
|
||
for b.Loop() {
|
||
if _, _, err := TopK(a, 16, 1); err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
}
|
||
}
|