feat: initial release

Assisted-by: GLM 5.3 Flash

internal/core/bench_probes_test.go

@@ -0,0 +1,744 @@
+// 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())
+					}
+				})
+			}
+		})
+	}
+}