feat: initial release

Assisted-by: GLM 5.3 Flash

spmd/bench_test.go

@@ -0,0 +1,118 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package spmd
+
+import (
+	"strconv"
+	"testing"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// The collective benchmarks time one world of four ranks running the
+// named collective back to back: every rank loops b.N times, so the
+// per-operation number is one round of the collective across the
+// world. In-process links put the channel and the arithmetic on the
+// clock, not the network.
+
+func BenchmarkAllReduceShards(b *testing.B) {
+	const gn = 1 << 20
+	whole := fixtureArray(gn)
+	if err := Launch(4, func(w *World) error {
+		span := mustPartition(b, gn, w.Size(), w.Rank())
+		local, err := dealRows(whole, span)
+		if err != nil {
+			return err
+		}
+		for i := 0; i < b.N; i++ {
+			if _, err := w.AllReduceShards(local, span, Sum); err != nil {
+				return err
+			}
+		}
+		return w.Barrier()
+	}); err != nil {
+		b.Fatal(err)
+	}
+}
+
+func BenchmarkAllReduce(b *testing.B) {
+	const per = 1 << 18
+	a, err := core.FromFloats(fixture(per), per)
+	if err != nil {
+		b.Fatal(err)
+	}
+	if err := Launch(4, func(w *World) error {
+		for i := 0; i < b.N; i++ {
+			if _, err := w.AllReduce(a, Sum); err != nil {
+				return err
+			}
+		}
+		return w.Barrier()
+	}); err != nil {
+		b.Fatal(err)
+	}
+}
+
+// BenchmarkReduce times the same-shaped-arrays fold on one million
+// float64 elements across world sizes 2, 4 and 8 for Sum and Min: the
+// fold runs as the chunks land, each arriving chunk joining the
+// running accumulator at its rank's turn.
+func BenchmarkReduce(b *testing.B) {
+	const n = 1 << 20
+	a, err := core.FromFloats(fixture(n), n)
+	if err != nil {
+		b.Fatal(err)
+	}
+	for _, size := range []int{2, 4, 8} {
+		for _, op := range []Op{Sum, Min} {
+			b.Run(op.String()+"/"+strconv.Itoa(size), func(b *testing.B) {
+				b.ReportAllocs()
+				if err := Launch(size, func(w *World) error {
+					for i := 0; i < b.N; i++ {
+						if _, err := w.Reduce(a, op, 0); err != nil {
+							return err
+						}
+					}
+					return w.Barrier()
+				}); err != nil {
+					b.Fatal(err)
+				}
+			})
+		}
+	}
+}
+
+func BenchmarkBroadcast(b *testing.B) {
+	a := fixtureArray(1 << 18)
+	if err := Launch(4, func(w *World) error {
+		for i := 0; i < b.N; i++ {
+			if _, err := w.Broadcast(a, 0); err != nil {
+				return err
+			}
+		}
+		return w.Barrier()
+	}); err != nil {
+		b.Fatal(err)
+	}
+}
+
+func BenchmarkScatterGather(b *testing.B) {
+	const gn = 1 << 20
+	a := fixtureArray(gn)
+	if err := Launch(4, func(w *World) error {
+		for i := 0; i < b.N; i++ {
+			local, span, err := w.Scatter(a, 0)
+			if err != nil {
+				return err
+			}
+			if _, err := w.Gather(local, 0); err != nil {
+				return err
+			}
+			_ = span
+		}
+		return w.Barrier()
+	}); err != nil {
+		b.Fatal(err)
+	}
+}