tensor/internal/core/axis_bench_test.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)
		}
	}
}