tensor/internal/core/bench_probes_test.go

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

package core

import (
	"cmp"
	"math"
	"slices"
	"testing"
	"time"

	"sourcedock.dev/petrbalvin/tensor/internal/engine"
)

// Paired probes for two dispatch decisions: the Argwhere packing
// strategy and the sort radix crew cap and digit width. Each probe
// interleaves every style on the same deterministic
// fixture in one process, so a drift in the machine's speed lands on
// all styles equally; the reported ns/op per style is the paired
// comparison, and the ns/op of a group as a whole is not comparable
// across groups.

// probeSink accumulates checksums so no measured walk is elided.
var probeSink int64

// --- Argwhere variants -------------------------------------------------

// probeArgwhereBaseline is the append-growth Argwhere the merge
// replaced: per-chunk slices grown by append, concatenated in chunk
// order, then copied once more through FromInts.
func probeArgwhereBaseline(a *Array) *Array {
	if a.dt == Complex {
		return nil
	}
	n := a.Len()
	parts := make([][]int64, 1)
	chunk := n
	if n >= copyMinPerWorker {
		w := workersFor(n)
		chunk = (n + w - 1) / w
		parts = make([][]int64, (n+chunk-1)/chunk)
	}
	parallelMin(len(parts), 1, func(s, e int) {
		for c := s; c < e; c++ {
			start := c * chunk
			parts[c] = probeArgwhereRun(a, start, min(start+chunk, n))
		}
	})
	total := 0
	for _, p := range parts {
		total += len(p)
	}
	rows := make([]int64, 0, total)
	for _, p := range parts {
		rows = append(rows, p...)
	}
	nnz := len(rows) / a.NDim()
	out, _ := FromInts(rows, nnz, a.NDim())
	return out
}

// probeArgwhereRun is the baseline per-chunk collector: the odometer
// seeded from the flat start, coordinates appended as they are met.
func probeArgwhereRun(a *Array, start, end int) []int64 {
	ndim := a.NDim()
	coord := make([]int, ndim)
	if start > 0 {
		rest := start
		for d := ndim - 1; d >= 0; d-- {
			coord[d] = rest % a.shape[d]
			rest /= a.shape[d]
		}
	}
	var rows []int64
	appendCoord := func() {
		for d := range ndim {
			rows = append(rows, int64(coord[d]))
		}
	}
	if !a.isContiguous() {
		for i := start; i < end; i++ {
			if !isZero(a, i) {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
		return rows
	}
	switch a.dt {
	case Int:
		p := a.ints
		for i := start; i < end; i++ {
			if p[i] != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	case Float16:
		p := a.halves
		for i := start; i < end; i++ {
			if p[i]&0x7FFF != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	case Float32:
		p := a.floats32
		for i := start; i < end; i++ {
			if p[i] != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	default:
		p := a.floats
		for i := start; i < end; i++ {
			if p[i] != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	}
	return rows
}

// probeArgwhereCountFill is the count-then-fill variant: an
// odometer-free count pass, then one fill pass writing exact disjoint
// slots of a single allocation.
func probeArgwhereCountFill(a *Array) *Array {
	if a.dt == Complex {
		return nil
	}
	n := a.Len()
	ndim := a.NDim()
	chunk := n
	chunks := 1
	if n >= copyMinPerWorker {
		w := workersFor(n)
		chunk = (n + w - 1) / w
		chunks = (n + chunk - 1) / chunk
	}
	counts := make([]int, chunks)
	parallelMin(chunks, 1, func(s, e int) {
		for c := s; c < e; c++ {
			start := c * chunk
			counts[c] = probeArgwhereCount(a, start, min(start+chunk, n))
		}
	})
	total := 0
	for _, cnt := range counts {
		total += cnt
	}
	out := &Array{shape: []int{total, ndim}, dt: Int}
	out.alloc(total * ndim)
	offsets := make([]int, chunks)
	off := 0
	for c, cnt := range counts {
		offsets[c] = off
		off += cnt * ndim
	}
	parallelMin(chunks, 1, func(s, e int) {
		for c := s; c < e; c++ {
			start := c * chunk
			probeArgwhereFill(a, out.ints[offsets[c]:], start, min(start+chunk, n))
		}
	})
	return out
}

func probeArgwhereCount(a *Array, start, end int) int {
	count := 0
	if !a.isContiguous() {
		for i := start; i < end; i++ {
			if !isZero(a, i) {
				count++
			}
		}
		return count
	}
	switch a.dt {
	case Int:
		p := a.ints
		for i := start; i < end; i++ {
			if p[i] != 0 {
				count++
			}
		}
	case Float16:
		p := a.halves
		for i := start; i < end; i++ {
			if p[i]&0x7FFF != 0 {
				count++
			}
		}
	case Float32:
		p := a.floats32
		for i := start; i < end; i++ {
			if p[i] != 0 {
				count++
			}
		}
	default:
		p := a.floats
		for i := start; i < end; i++ {
			if p[i] != 0 {
				count++
			}
		}
	}
	return count
}

func probeArgwhereFill(a *Array, dst []int64, start, end int) {
	ndim := a.NDim()
	coord := make([]int, ndim)
	if start > 0 {
		rest := start
		for d := ndim - 1; d >= 0; d-- {
			coord[d] = rest % a.shape[d]
			rest /= a.shape[d]
		}
	}
	pos := 0
	appendCoord := func() {
		for d := range ndim {
			dst[pos] = int64(coord[d])
			pos++
		}
	}
	if !a.isContiguous() {
		for i := start; i < end; i++ {
			if !isZero(a, i) {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
		return
	}
	switch a.dt {
	case Int:
		p := a.ints
		for i := start; i < end; i++ {
			if p[i] != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	case Float16:
		p := a.halves
		for i := start; i < end; i++ {
			if p[i]&0x7FFF != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	case Float32:
		p := a.floats32
		for i := start; i < end; i++ {
			if p[i] != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	default:
		p := a.floats
		for i := start; i < end; i++ {
			if p[i] != 0 {
				appendCoord()
			}
			advanceOdometer(coord, a.shape)
		}
	}
}

// probeArgwhereMerge is the single-pass merge with a configurable first
// buffer capacity per chunk; bufMin 0 is pure append growth.
func probeArgwhereMerge(a *Array, bufMin int) *Array {
	if a.dt == Complex {
		return nil
	}
	n := a.Len()
	ndim := a.NDim()
	chunk := n
	chunks := 1
	if n >= copyMinPerWorker {
		w := workersFor(n)
		chunk = (n + w - 1) / w
		chunks = (n + chunk - 1) / chunk
	}
	parts := make([][]int64, chunks)
	parallelMin(chunks, 1, func(s, e int) {
		for c := s; c < e; c++ {
			start := c * chunk
			end := min(start+chunk, n)
			buf := make([]int64, 0, min(bufMin, (end-start)*ndim))
			parts[c] = argwhereAppend(a, buf, start, end)
		}
	})
	total := 0
	for _, p := range parts {
		total += len(p)
	}
	out := &Array{shape: []int{total / ndim, ndim}, dt: Int}
	out.alloc(total)
	off := 0
	for _, p := range parts {
		copy(out.ints[off:off+len(p)], p)
		off += len(p)
	}
	return out
}

// argwhereStyles lists the Argwhere variants both probes walk.
func argwhereStyles() []struct {
	name string
	run  func(*Array) *Array
} {
	return []struct {
		name string
		run  func(*Array) *Array
	}{
		{"production", func(a *Array) *Array { out, _ := Argwhere(a); return out }},
		{"baseline-growth", probeArgwhereBaseline},
		{"count-then-fill", probeArgwhereCountFill},
		{"merge-pure", func(a *Array) *Array { return probeArgwhereMerge(a, 0) }},
	}
}

// BenchmarkArgwhereVariantsAB interleaves the Argwhere packing variants
// on the 4096x4096 sparse fixture of BenchmarkArgwhere4M.
func BenchmarkArgwhereVariantsAB(b *testing.B) {
	a := benchSparseInts(4096*4096, 64)
	styles := argwhereStyles()
	elapsed := make([]time.Duration, len(styles))
	for b.Loop() {
		for i, s := range styles {
			start := time.Now()
			out := s.run(a)
			elapsed[i] += time.Since(start)
			probeSink += int64(out.Len())
			if out.Len() > 0 {
				probeSink += out.RawInts()[0] + out.RawInts()[out.Len()-1]
			}
		}
	}
	for i, s := range styles {
		b.ReportMetric(float64(elapsed[i].Nanoseconds())/float64(b.N), "ns/op-"+s.name)
	}
}

// BenchmarkArgwhereVariantAllocs reports each variant's deterministic
// allocation count on the same fixture, one sub-benchmark per style.
func BenchmarkArgwhereVariantAllocs(b *testing.B) {
	a := benchSparseInts(4096*4096, 64)
	for _, s := range argwhereStyles() {
		b.Run(s.name, func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				out := s.run(a)
				probeSink += int64(out.Len())
			}
		})
	}
}

// --- Sort radix configurations ----------------------------------------

// probeDigitSpec returns the digit shifts, bucket count and exact digit
// mask of a probe radix over the key range diff at the given width.
func probeDigitSpec(diff uint64, bits int) (shifts []uint, buckets int, digits uint32) {
	if bits == 8 {
		shifts = make([]uint, 8)
		for j := range shifts {
			shifts[j] = uint(8 * j)
			if diff>>shifts[j]&0xFF != 0 {
				digits |= 1 << j
			}
		}
		return shifts, 256, digits
	}
	shifts = append([]uint(nil), radix12Shifts[:]...)
	return shifts, radixBuckets, digitMask12Of(diff)
}

func probeWorkerCap(n, wcap int) int {
	if wcap == 0 {
		return engine.WorkersFor(n)
	}
	return min(engine.WorkersFor(n), wcap)
}

// probeArgSortRadixUint64 mirrors argSortRadixUint64 with the parallel
// crew cap and digit width exposed; the serial paths are production's
// own functions, so only the parallel core varies.
func probeArgSortRadixUint64(keys []uint64, idx []int, diff uint64, wcap, bits int) {
	n := len(idx)
	if n < radixSeqMin {
		slices.SortStableFunc(idx, func(x, y int) int {
			return cmp.Compare(keys[x], keys[y])
		})
		return
	}
	if diff == 0 {
		return
	}
	if n < radixParMin {
		argSortRadixSerial(keys, idx, digitMaskOf(diff))
		return
	}
	buf := make([]int, n)
	src, dst := idx, buf
	shifts, buckets, digits := probeDigitSpec(diff, bits)
	w := probeWorkerCap(n, wcap)
	chunk := (n + w - 1) / w
	if chunk < radixParMin || chunk >= n {
		w = 1
		chunk = n
	}
	hist := make([]int, w*buckets)
	passes := 0
	for j, shift := range shifts {
		if digits&(1<<j) == 0 {
			continue
		}
		clear(hist)
		forEachRadixChunk(n, chunk, func(row, start, end int) {
			base := row * buckets
			for _, i := range src[start:end] {
				hist[base+(int(keys[i]>>shift)&(buckets-1))]++
			}
		})
		sum := 0
		for bb := range buckets {
			for c := range w {
				cnt := hist[c*buckets+bb]
				hist[c*buckets+bb] = sum
				sum += cnt
			}
		}
		forEachRadixChunk(n, chunk, func(row, start, end int) {
			base := row * buckets
			for _, i := range src[start:end] {
				d := base + (int(keys[i]>>shift) & (buckets - 1))
				dst[hist[d]] = i
				hist[d]++
			}
		})
		src, dst = dst, src
		passes++
	}
	if passes%2 == 1 {
		copy(idx, buf)
	}
}

// probeRadixSortUint64 mirrors radixSortUint64 with the same exposures.
func probeRadixSortUint64(vals []uint64, diff uint64, wcap, bits int) {
	n := len(vals)
	if n < radixSeqMin {
		slices.Sort(vals)
		return
	}
	if diff == 0 {
		return
	}
	if n < radixParMin {
		probeRadixSortSerial(vals, diff)
		return
	}
	buf := make([]uint64, n)
	src, dst := vals, buf
	shifts, buckets, digits := probeDigitSpec(diff, bits)
	w := probeWorkerCap(n, wcap)
	chunk := (n + w - 1) / w
	if chunk < radixParMin || chunk >= n {
		w = 1
		chunk = n
	}
	hist := make([]int, w*buckets)
	passes := 0
	for j, shift := range shifts {
		if digits&(1<<j) == 0 {
			continue
		}
		clear(hist)
		forEachRadixChunk(n, chunk, func(row, start, end int) {
			base := row * buckets
			for _, v := range src[start:end] {
				hist[base+(int(v>>shift)&(buckets-1))]++
			}
		})
		sum := 0
		for bb := range buckets {
			for c := range w {
				cnt := hist[c*buckets+bb]
				hist[c*buckets+bb] = sum
				sum += cnt
			}
		}
		forEachRadixChunk(n, chunk, func(row, start, end int) {
			base := row * buckets
			for _, v := range src[start:end] {
				d := base + (int(v>>shift) & (buckets - 1))
				dst[hist[d]] = v
				hist[d]++
			}
		})
		src, dst = dst, src
		passes++
	}
	if passes%2 == 1 {
		copy(vals, buf)
	}
}

// probeRadixSortSerial is the byte-digit serial counting sort for the
// probe's small-n path.
func probeRadixSortSerial(vals []uint64, diff uint64) {
	n := len(vals)
	buf := make([]uint64, n)
	src, dst := vals, buf
	mask := digitMaskOf(diff)
	passes := 0
	var count [256]int
	for b := range 8 {
		if mask&(1<<b) == 0 {
			continue
		}
		shift := uint(8 * b)
		clear(count[:])
		for _, v := range src {
			count[v>>shift&0xFF]++
		}
		sum := 0
		for bb, c := range count {
			count[bb], sum = sum, sum+c
		}
		for _, v := range src {
			d := v >> shift & 0xFF
			dst[count[d]] = v
			count[d]++
		}
		src, dst = dst, src
		passes++
	}
	if passes%2 == 1 {
		copy(vals, buf)
	}
}

// probeSortInt mirrors Sort's int path with the probe radix.
func probeSortInt(a *Array, wcap, bits int) *Array {
	ints, _, _, _, _ := a.cloneData()
	keys := make([]uint64, len(ints))
	hi, lo := uint64(0), ^uint64(0)
	for i, v := range ints {
		k := uint64(v) ^ (1 << 63)
		keys[i] = k
		hi |= k
		lo &= k
	}
	probeRadixSortUint64(keys, hi^lo, wcap, bits)
	for i, k := range keys {
		ints[i] = int64(k ^ (1 << 63))
	}
	return &Array{shape: a.Shape(), dt: Int, ints: ints}
}

// probeSortFloatsRadix mirrors sortFloatsRadix with the probe radix.
func probeSortFloatsRadix(floats []float64, wcap, bits int) []float64 {
	keys := make([]uint64, len(floats))
	finite := floats[:0]
	m := 0
	hi, lo := uint64(0), ^uint64(0)
	for _, v := range floats {
		if v != v {
			continue
		}
		k := floatSortKey(v)
		finite = append(finite, v)
		keys[m] = k
		m++
		hi |= k
		lo &= k
	}
	keys = keys[:m]
	probeRadixSortUint64(keys, hi^lo, wcap, bits)
	for i, k := range keys {
		finite[i] = floatKeyToFloat(k)
	}
	for range len(floats) - m {
		finite = append(finite, math.NaN())
	}
	return finite
}

// probeSortFloat mirrors Sort's float path with the probe radix.
func probeSortFloat(a *Array, wcap, bits int) *Array {
	_, _, _, floats, _ := a.cloneData()
	return &Array{shape: a.Shape(), dt: Float, floats: probeSortFloatsRadix(floats, wcap, bits)}
}

// probeArgSort mirrors ArgSort's int and float paths with the probe
// radix; the fixture dtypes of the named benchmarks are int and float,
// which are the paths this mirror carries.
func probeArgSort(a *Array, wcap, bits int) *Array {
	n := a.Len()
	idx := make([]int, n)
	m := a.materialise()
	if a.dt == Int {
		ints := m.ints
		keys := make([]uint64, n)
		hi, lo := uint64(0), ^uint64(0)
		for i := range n {
			k := uint64(ints[i]) ^ (1 << 63)
			keys[i] = k
			idx[i] = i
			hi |= k
			lo &= k
		}
		probeArgSortRadixUint64(keys, idx, hi^lo, wcap, bits)
	} else {
		vals := m.floats
		keys := make([]uint64, n)
		kept := make([]int, 0, n)
		var nans []int
		hi, lo := uint64(0), ^uint64(0)
		for i := range n {
			v := vals[i]
			if v != v {
				nans = append(nans, i)
				continue
			}
			k := floatSortKey(v)
			keys[i] = k
			kept = append(kept, i)
			hi |= k
			lo &= k
		}
		probeArgSortRadixUint64(keys, kept, hi^lo, wcap, bits)
		copy(idx, kept)
		copy(idx[len(kept):], nans)
	}
	out := make([]int64, len(idx))
	parallelMin(len(out), copyMinPerWorker, func(s, e int) {
		for i := s; i < e; i++ {
			out[i] = int64(idx[i])
		}
	})
	return &Array{shape: []int{len(out)}, dt: Int, ints: out}
}

// probeFloatSeed7 rebuilds the seed-7 float fixture of BenchmarkSort1M
// and BenchmarkArgSort1M.
func probeFloatSeed7() *Array {
	a, _ := FromFloats(make([]float64, 1<<20), 1<<20)
	g := NewGenerator(7)
	src, _ := Floats(g, 1<<20)
	copy(a.RawFloats(), src.RawFloats())
	return a
}

// probeSortConfigs are the crew-cap and digit-width combinations the
// interleaved probe walks; wcap 0 means the machine's full worker count.
func probeSortConfigs() []struct {
	name string
	wcap int
	bits int
} {
	return []struct {
		name string
		wcap int
		bits int
	}{
		{"cap=8/12b", 8, 12},
		{"cap=16/12b", 16, 12},
		{"cap=uncapped/12b", 0, 12},
		{"cap=8/8b", 8, 8},
		{"cap=16/8b", 16, 8},
		{"cap=uncapped/8b", 0, 8},
	}
}

// BenchmarkSortRadixConfigAB interleaves the six radix configurations
// on the fixtures of the five named sort benchmarks.
func BenchmarkSortRadixConfigAB(b *testing.B) {
	fixtures := []struct {
		name string
		mk   func() *Array
		run  func(*Array, int, int) *Array
	}{
		{"SortInt1M", func() *Array { return benchIntsShape(1<<10, 1<<20) }, probeSortInt},
		{"Sort1M", probeFloatSeed7, probeSortFloat},
		{"SortFloat1M", func() *Array { return benchFloatsShape(1 << 20) }, probeSortFloat},
		{"ArgSortInt1M", func() *Array { return benchIntsShape(1<<10, 1<<20) }, probeArgSort},
		{"ArgSort1M", probeFloatSeed7, probeArgSort},
	}
	cfgs := probeSortConfigs()
	for _, f := range fixtures {
		b.Run(f.name, func(b *testing.B) {
			a := f.mk()
			elapsed := make([]time.Duration, len(cfgs))
			for b.Loop() {
				for i, c := range cfgs {
					start := time.Now()
					out := f.run(a, c.wcap, c.bits)
					elapsed[i] += time.Since(start)
					probeSink += int64(out.Len())
				}
			}
			for i, c := range cfgs {
				b.ReportMetric(float64(elapsed[i].Nanoseconds())/float64(b.N), "ns/op-"+c.name)
			}
		})
	}
}

// BenchmarkSortRadixConfigAllocs reports each configuration's
// deterministic allocation count on the SortFloat1M and SortInt1M
// fixtures, one sub-benchmark per configuration.
func BenchmarkSortRadixConfigAllocs(b *testing.B) {
	for _, fx := range []struct {
		name string
		mk   func() *Array
		run  func(*Array, int, int) *Array
	}{
		{"SortInt1M", func() *Array { return benchIntsShape(1<<10, 1<<20) }, probeSortInt},
		{"SortFloat1M", func() *Array { return benchFloatsShape(1 << 20) }, probeSortFloat},
	} {
		b.Run(fx.name, func(b *testing.B) {
			a := fx.mk()
			for _, c := range probeSortConfigs() {
				b.Run(c.name, func(b *testing.B) {
					b.ReportAllocs()
					for b.Loop() {
						out := fx.run(a, c.wcap, c.bits)
						probeSink += int64(out.Len())
					}
				})
			}
		})
	}
}