feat: initial release

Assisted-by: GLM 5.3 Flash

internal/core/einsum_bench_test.go

@@ -0,0 +1,98 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package core
+
+import "testing"
+
+// Einsum benchmarks guard both the fast paths and the general engine.
+// The batched shapes are the training-loop workloads; the general
+// engine shapes are the ones the dispatch table declines.
+
+func benchVals(b *testing.B, seed, n int) []float64 {
+	b.Helper()
+	v := make([]float64, n)
+	for i := range v {
+		v[i] = float64(i%17)*float64(seed%5) + float64(i%7) - 6
+	}
+	return v
+}
+
+func benchMat(b *testing.B, seed int, shape ...int) *Array {
+	b.Helper()
+	n := 1
+	for _, d := range shape {
+		n *= d
+	}
+	a, err := FromFloats(benchVals(b, seed, n), shape...)
+	if err != nil {
+		b.Fatal(err)
+	}
+	return a
+}
+
+// BenchmarkEinsumMatMulPath exercises the dispatch table's MatMul fast
+// path as the control.
+func BenchmarkEinsumMatMulPath(b *testing.B) {
+	a := benchMat(b, 1, 64, 64)
+	c := benchMat(b, 2, 64, 64)
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Einsum("ij,jk->ik", a, c); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkEinsumBatchedMatMul is the batched product the table has no
+// fast path for: it falls to the general engine.
+func BenchmarkEinsumBatchedMatMul(b *testing.B) {
+	a := benchMat(b, 3, 16, 64, 64)
+	c := benchMat(b, 4, 16, 64, 64)
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Einsum("bij,bjk->bik", a, c); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkEinsumBatchedSmall keeps the general engine's fixed costs
+// visible at a size where they dominate.
+func BenchmarkEinsumBatchedSmall(b *testing.B) {
+	a := benchMat(b, 5, 4, 32, 32)
+	c := benchMat(b, 6, 4, 32, 32)
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Einsum("bij,bjk->bik", a, c); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkEinsumEllipsisMatMul is the ellipsis spelling of the batched
+// product.
+func BenchmarkEinsumEllipsisMatMul(b *testing.B) {
+	a := benchMat(b, 7, 8, 48, 48)
+	c := benchMat(b, 8, 48, 48)
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Einsum("...ij,jk->...ik", a, c); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// BenchmarkEinsumBilinear sums a shared label across three operands,
+// the pattern attention and tensordot shapes reduce to.
+func BenchmarkEinsumBilinear(b *testing.B) {
+	a := benchMat(b, 9, 16, 32)
+	c := benchMat(b, 10, 32, 32)
+	d := benchMat(b, 11, 32, 16)
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Einsum("ik,kj,jl->il", a, c, d); err != nil {
+			b.Fatal(err)
+		}
+	}
+}