tensor/internal/core/bench_narrow_test.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import "testing"

// The narrow element types' kernels: element-wise arithmetic, the
// comparisons, the scalar maps, the full-array reductions and Where,
// each against the int64 or float64 kernel of the same element count,
// so a narrow walk is judged against the widest loop it could have
// widened into. Every payload holds one million elements, the shape the
// BenchmarkAdd1M family reports at.

// narrowBenchLen is the element count every payload in this file fills.
const narrowBenchLen = 1 << 20

// benchNarrowInts fills an n-element signed integer slice with a
// deterministic mix of positive and negative values.
func benchNarrowInts[T ~int8 | ~int16 | ~int32](n int) []T {
	src := make([]T, n)
	for i := range src {
		src[i] = T(i%251) - 125
	}
	return src
}

// benchNarrowUints fills an n-element unsigned integer slice with a
// deterministic mix of small values.
func benchNarrowUints[T ~uint8 | ~uint16 | ~uint32](n int) []T {
	src := make([]T, n)
	for i := range src {
		src[i] = T(i % 251)
	}
	return src
}

// benchNarrowPair builds two one-million-element arrays of dtype dt
// with the same deterministic fill, the pair the binary kernels walk.
func benchNarrowPair(dt Dtype) (*Array, *Array) {
	must := func(a *Array, err error) *Array {
		if err != nil {
			panic(err)
		}
		return a
	}
	fillInt := func() []int64 {
		v := make([]int64, narrowBenchLen)
		for i := range v {
			v[i] = int64(i%251) - 125
		}
		return v
	}
	fillFloat := func() []float64 {
		v := make([]float64, narrowBenchLen)
		for i := range v {
			v[i] = float64(i%251) - 125
		}
		return v
	}
	switch dt {
	case Int8:
		return must(Int8sFromArray(benchNarrowInts[int8](narrowBenchLen), narrowBenchLen)),
			must(Int8sFromArray(benchNarrowInts[int8](narrowBenchLen), narrowBenchLen))
	case Uint8:
		return must(Uint8sFromArray(benchNarrowUints[uint8](narrowBenchLen), narrowBenchLen)),
			must(Uint8sFromArray(benchNarrowUints[uint8](narrowBenchLen), narrowBenchLen))
	case Int16:
		return must(Int16sFromArray(benchNarrowInts[int16](narrowBenchLen), narrowBenchLen)),
			must(Int16sFromArray(benchNarrowInts[int16](narrowBenchLen), narrowBenchLen))
	case Uint16:
		return must(Uint16sFromArray(benchNarrowUints[uint16](narrowBenchLen), narrowBenchLen)),
			must(Uint16sFromArray(benchNarrowUints[uint16](narrowBenchLen), narrowBenchLen))
	case Int32:
		return must(Int32sFromArray(benchNarrowInts[int32](narrowBenchLen), narrowBenchLen)),
			must(Int32sFromArray(benchNarrowInts[int32](narrowBenchLen), narrowBenchLen))
	case Uint32:
		return must(Uint32sFromArray(benchNarrowUints[uint32](narrowBenchLen), narrowBenchLen)),
			must(Uint32sFromArray(benchNarrowUints[uint32](narrowBenchLen), narrowBenchLen))
	case Int:
		return must(FromInts(fillInt(), narrowBenchLen)), must(FromInts(fillInt(), narrowBenchLen))
	default: // Float
		return must(FromFloats(fillFloat(), narrowBenchLen)), must(FromFloats(fillFloat(), narrowBenchLen))
	}
}

func benchNarrowAdd(b *testing.B, dt Dtype) {
	b.Helper()
	a, c := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Add(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowAddInt8(b *testing.B)    { benchNarrowAdd(b, Int8) }
func BenchmarkNarrowAddUint8(b *testing.B)   { benchNarrowAdd(b, Uint8) }
func BenchmarkNarrowAddInt16(b *testing.B)   { benchNarrowAdd(b, Int16) }
func BenchmarkNarrowAddUint16(b *testing.B)  { benchNarrowAdd(b, Uint16) }
func BenchmarkNarrowAddInt32(b *testing.B)   { benchNarrowAdd(b, Int32) }
func BenchmarkNarrowAddUint32(b *testing.B)  { benchNarrowAdd(b, Uint32) }
func BenchmarkNarrowAddInt64(b *testing.B)   { benchNarrowAdd(b, Int) }
func BenchmarkNarrowAddFloat64(b *testing.B) { benchNarrowAdd(b, Float) }

func benchNarrowMul(b *testing.B, dt Dtype) {
	b.Helper()
	a, c := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Mul(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowMulInt8(b *testing.B)   { benchNarrowMul(b, Int8) }
func BenchmarkNarrowMulUint16(b *testing.B) { benchNarrowMul(b, Uint16) }
func BenchmarkNarrowMulInt32(b *testing.B)  { benchNarrowMul(b, Int32) }
func BenchmarkNarrowMulInt64(b *testing.B)  { benchNarrowMul(b, Int) }

func benchNarrowLt(b *testing.B, dt Dtype) {
	b.Helper()
	a, c := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Lt(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowLtInt8(b *testing.B)    { benchNarrowLt(b, Int8) }
func BenchmarkNarrowLtUint8(b *testing.B)   { benchNarrowLt(b, Uint8) }
func BenchmarkNarrowLtInt16(b *testing.B)   { benchNarrowLt(b, Int16) }
func BenchmarkNarrowLtUint32(b *testing.B)  { benchNarrowLt(b, Uint32) }
func BenchmarkNarrowLtInt64(b *testing.B)   { benchNarrowLt(b, Int) }
func BenchmarkNarrowLtFloat64(b *testing.B) { benchNarrowLt(b, Float) }

func benchNarrowEqI(b *testing.B, dt Dtype) {
	b.Helper()
	a, _ := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := EqI(a, 17); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowEqIInt8(b *testing.B)   { benchNarrowEqI(b, Int8) }
func BenchmarkNarrowEqIUint16(b *testing.B) { benchNarrowEqI(b, Uint16) }
func BenchmarkNarrowEqIInt64(b *testing.B)  { benchNarrowEqI(b, Int) }

func benchNarrowLtF(b *testing.B, dt Dtype) {
	b.Helper()
	a, _ := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LtF(a, 100.5); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowLtFInt32(b *testing.B) { benchNarrowLtF(b, Int32) }
func BenchmarkNarrowLtFInt64(b *testing.B) { benchNarrowLtF(b, Int) }

func benchNarrowSum(b *testing.B, dt Dtype) {
	b.Helper()
	a, _ := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		Sum(a)
	}
}

func BenchmarkNarrowSumInt8(b *testing.B)    { benchNarrowSum(b, Int8) }
func BenchmarkNarrowSumUint8(b *testing.B)   { benchNarrowSum(b, Uint8) }
func BenchmarkNarrowSumInt16(b *testing.B)   { benchNarrowSum(b, Int16) }
func BenchmarkNarrowSumUint16(b *testing.B)  { benchNarrowSum(b, Uint16) }
func BenchmarkNarrowSumInt32(b *testing.B)   { benchNarrowSum(b, Int32) }
func BenchmarkNarrowSumUint32(b *testing.B)  { benchNarrowSum(b, Uint32) }
func BenchmarkNarrowSumInt64(b *testing.B)   { benchNarrowSum(b, Int) }
func BenchmarkNarrowSumFloat64(b *testing.B) { benchNarrowSum(b, Float) }
func BenchmarkNarrowSumBool(b *testing.B) {
	a, _ := FromBools(make([]bool, narrowBenchLen), narrowBenchLen)
	bs := a.RawBools()
	for i := range bs {
		bs[i] = i%3 != 0
	}
	b.ReportAllocs()
	for b.Loop() {
		Sum(a)
	}
}

func benchNarrowMin(b *testing.B, dt Dtype) {
	b.Helper()
	a, _ := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Min(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowMinInt8(b *testing.B)   { benchNarrowMin(b, Int8) }
func BenchmarkNarrowMinUint32(b *testing.B) { benchNarrowMin(b, Uint32) }
func BenchmarkNarrowMinInt64(b *testing.B)  { benchNarrowMin(b, Int) }

func benchNarrowDot(b *testing.B, dt Dtype) {
	b.Helper()
	a, c := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Dot(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowDotInt8(b *testing.B)   { benchNarrowDot(b, Int8) }
func BenchmarkNarrowDotUint16(b *testing.B) { benchNarrowDot(b, Uint16) }
func BenchmarkNarrowDotInt32(b *testing.B)  { benchNarrowDot(b, Int32) }
func BenchmarkNarrowDotInt64(b *testing.B)  { benchNarrowDot(b, Int) }

func benchNarrowAddF(b *testing.B, dt Dtype) {
	b.Helper()
	a, _ := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		AddF(a, 0.5)
	}
}

func BenchmarkNarrowAddFInt8(b *testing.B)  { benchNarrowAddF(b, Int8) }
func BenchmarkNarrowAddFInt32(b *testing.B) { benchNarrowAddF(b, Int32) }
func BenchmarkNarrowAddFInt64(b *testing.B) { benchNarrowAddF(b, Int) }

func benchNarrowMulI(b *testing.B, dt Dtype) {
	b.Helper()
	a, _ := benchNarrowPair(dt)
	b.ReportAllocs()
	for b.Loop() {
		MulI(a, 3)
	}
}

func BenchmarkNarrowMulIInt8(b *testing.B)   { benchNarrowMulI(b, Int8) }
func BenchmarkNarrowMulIUint32(b *testing.B) { benchNarrowMulI(b, Uint32) }
func BenchmarkNarrowMulIInt64(b *testing.B)  { benchNarrowMulI(b, Int) }

func benchNarrowWhere(b *testing.B, dt Dtype) {
	b.Helper()
	a, c := benchNarrowPair(dt)
	cond, _ := FromInts(make([]int64, narrowBenchLen), narrowBenchLen)
	ci := cond.RawInts()
	for i := range ci {
		if i%2 == 0 {
			ci[i] = 1
		}
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Where(cond, a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowWhereInt8(b *testing.B)   { benchNarrowWhere(b, Int8) }
func BenchmarkNarrowWhereUint16(b *testing.B) { benchNarrowWhere(b, Uint16) }
func BenchmarkNarrowWhereInt64(b *testing.B)  { benchNarrowWhere(b, Int) }

func BenchmarkNarrowAndBool(b *testing.B) {
	a, _ := FromBools(make([]bool, narrowBenchLen), narrowBenchLen)
	c, _ := FromBools(make([]bool, narrowBenchLen), narrowBenchLen)
	ab, cb := a.RawBools(), c.RawBools()
	for i := range ab {
		ab[i] = i%3 != 0
		cb[i] = i%5 != 0
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := And(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowNotBool(b *testing.B) {
	a, _ := FromBools(make([]bool, narrowBenchLen), narrowBenchLen)
	ab := a.RawBools()
	for i := range ab {
		ab[i] = i%3 != 0
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Not(a); err != nil {
			b.Fatal(err)
		}
	}
}

// The mixed-width pairs: a narrow operand against a different narrow
// width and against Int, both promoted walks a single dense payload
// read from each side.
func BenchmarkNarrowMixedAddI8I16(b *testing.B) {
	a, _ := benchNarrowPair(Int8)
	c, _ := benchNarrowPair(Int16)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Add(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowMixedAddU8Int(b *testing.B) {
	a, _ := benchNarrowPair(Uint8)
	c, _ := benchNarrowPair(Int)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Add(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkNarrowMixedLtI8U16(b *testing.B) {
	a, _ := benchNarrowPair(Int8)
	c, _ := benchNarrowPair(Uint16)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Lt(a, c); err != nil {
			b.Fatal(err)
		}
	}
}