diff --git a/compute/cpu_engine.go b/compute/cpu_engine.go index b0527fd..2bf7784 100644 --- a/compute/cpu_engine.go +++ b/compute/cpu_engine.go @@ -1557,10 +1557,12 @@ func (e *CPUEngine[T]) Sum( // A negative axis means sum over all axes. if axis < 0 { - var sum T - for _, v := range a.Data() { - sum = e.ops.Add(sum, v) - } + data := a.Data() + // Fixed-order pairwise accumulation: order depends only on len(data), + // so the result is bitwise-stable run to run and tighter than a naive + // left-to-right fold. See compute/reduce_pairwise.go. + sum := pairwiseReduce(len(data), e.ops.FromFloat64(0), + func(i int) T { return data[i] }, e.ops.Add) shape := []int{1} if keepDims { shape = make([]int, a.Dims()) @@ -1661,12 +1663,10 @@ func (e *CPUEngine[T]) Sum( } } - // Reduce along the axis for this stripe - sum := e.ops.FromFloat64(0) - for k := 0; k < axisSize; k++ { //nolint:intrange - idx := base + k*step - sum = e.ops.Add(sum, aData[idx]) - } + // Reduce along the axis for this stripe using fixed-order pairwise + // accumulation (order depends only on axisSize, not on worker count). + sum := pairwiseReduce(axisSize, e.ops.FromFloat64(0), + func(k int) T { return aData[base+k*step] }, e.ops.Add) rData[rIndex] = sum } }) @@ -2251,15 +2251,15 @@ func (e *CPUEngine[T]) Softmax(_ context.Context, a *tensor.TensorNumeric[T], ax } } - // 2) Compute exponentials and sum - sum := e.ops.FromFloat64(0) + // 2) Compute exponentials, then sum the denominator in fixed-order + // pairwise order (order depends only on axisSize) for run-to-run + // stability and tighter agreement with the oracle. for k := 0; k < axisSize; k++ { //nolint:intrange idx := base + k*step - shifted := e.ops.Sub(aData[idx], maxVal) - ex := e.ops.Exp(shifted) - oData[idx] = ex - sum = e.ops.Add(sum, ex) + oData[idx] = e.ops.Exp(e.ops.Sub(aData[idx], maxVal)) } + sum := pairwiseReduce(axisSize, e.ops.FromFloat64(0), + func(k int) T { return oData[base+k*step] }, e.ops.Add) // 3) Normalize for k := 0; k < axisSize; k++ { //nolint:intrange diff --git a/compute/reduce_pairwise.go b/compute/reduce_pairwise.go new file mode 100644 index 0000000..cb9a512 --- /dev/null +++ b/compute/reduce_pairwise.go @@ -0,0 +1,48 @@ +package compute + +// pairwiseBlock is the serial base-case size for fixed-order pairwise +// reduction. It matches numpy's pairwise blocksize: large enough to amortize +// the recursion, small enough to preserve the O(log n) error bound. +const pairwiseBlock = 128 + +// pairwiseReduce sums n elements in a fixed, worker-independent order using a +// recursive tree (pairwise) reduction. Element i is read via get(i) and +// combined with add; zero is returned only for n <= 0. +// +// The split points are a pure function of n (recursive halving aligned to +// pairwiseBlock), so the result is bitwise-identical regardless of how the +// surrounding loop is parallelized or chunked. Compared with a naive +// left-to-right fold, the accumulated rounding error grows as O(log n) instead +// of O(n), which tightens agreement with a higher-precision oracle without +// changing the accumulation dtype. This is the canonical fixed-order +// accumulation for CPU reductions; see docs/design.md "Reduction accumulation". +func pairwiseReduce[T any](n int, zero T, get func(i int) T, add func(a, b T) T) T { + if n <= 0 { + return zero + } + return pairwiseRange(0, n, get, add) +} + +// pairwiseRange reduces the half-open index range [lo, hi) using the fixed-order +// pairwise tree. The caller guarantees hi > lo. +func pairwiseRange[T any](lo, hi int, get func(i int) T, add func(a, b T) T) T { + n := hi - lo + if n <= pairwiseBlock { + acc := get(lo) + for i := lo + 1; i < hi; i++ { + acc = add(acc, get(i)) + } + + return acc + } + + // Split at half, snapped down to a pairwiseBlock boundary so the tree shape + // depends only on n, never on runtime chunking. + half := n / 2 + half -= half % pairwiseBlock + if half == 0 { + half = pairwiseBlock + } + + return add(pairwiseRange(lo, lo+half, get, add), pairwiseRange(lo+half, hi, get, add)) +} diff --git a/compute/reduce_pairwise_test.go b/compute/reduce_pairwise_test.go new file mode 100644 index 0000000..70ee32a --- /dev/null +++ b/compute/reduce_pairwise_test.go @@ -0,0 +1,168 @@ +package compute + +import ( + "context" + "math" + "math/rand/v2" + "runtime" + "testing" + + "github.com/zerfoo/ztensor/numeric" + "github.com/zerfoo/ztensor/tensor" +) + +// kahanRef64 is a compensated float64 ground-truth sum for accuracy comparisons. +func kahanRef64(s []float32) float64 { + var sum, c float64 + for _, v := range s { + y := float64(v) - c + t := sum + y + c = (t - sum) - y + sum = t + } + + return sum +} + +func naiveSumF32(s []float32) float32 { + var acc float32 + for _, v := range s { + acc += v + } + + return acc +} + +func relErr(got float32, truth float64) float64 { + if truth == 0 { + return math.Abs(float64(got)) + } + + return math.Abs((float64(got) - truth) / truth) +} + +// TestPairwiseReduce_EqualsNaiveForSmallN verifies the base case is an ordinary +// fold, so behavior is unchanged for short reductions. +func TestPairwiseReduce_EqualsNaiveForSmallN(t *testing.T) { + for _, n := range []int{1, 2, 7, 63, pairwiseBlock} { + s := make([]float32, n) + for i := range s { + s[i] = float32(i) * 0.5 + } + got := pairwiseReduce(len(s), float32(0), func(i int) float32 { return s[i] }, func(a, b float32) float32 { return a + b }) + if want := naiveSumF32(s); got != want { + t.Fatalf("n=%d: pairwise=%v naive=%v (must match for n<=block)", n, got, want) + } + } +} + +// TestPairwiseReduce_Deterministic verifies the accumulation order is a pure +// function of length: repeated reductions are bitwise-identical. +func TestPairwiseReduce_Deterministic(t *testing.T) { + rng := rand.New(rand.NewPCG(1, 2)) + s := make([]float32, 40000) + for i := range s { + s[i] = float32(rng.NormFloat64()) + } + add := func(a, b float32) float32 { return a + b } + first := pairwiseReduce(len(s), float32(0), func(i int) float32 { return s[i] }, add) + for r := 0; r < 8; r++ { + if got := pairwiseReduce(len(s), float32(0), func(i int) float32 { return s[i] }, add); got != first { + t.Fatalf("run %d not bitwise-stable: %v != %v", r, got, first) + } + } +} + +// TestPairwiseReduce_TightensVsNaive is the core numeric-correctness guard: the +// fixed-order pairwise sum must be at least as close to the float64 ground truth +// as the naive left-to-right fold for ill-conditioned float32 inputs. A change +// that loosened any reduction would regress here. +func TestPairwiseReduce_TightensVsNaive(t *testing.T) { + rng := rand.New(rand.NewPCG(7, 11)) + for _, n := range []int{1024, 4096, 16384, 65536} { + s := make([]float32, n) + for i := range s { + s[i] = float32(0.1 + 0.001*rng.NormFloat64()) + } + truth := kahanRef64(s) + add := func(a, b float32) float32 { return a + b } + pw := pairwiseReduce(len(s), float32(0), func(i int) float32 { return s[i] }, add) + en, ep := relErr(naiveSumF32(s), truth), relErr(pw, truth) + if ep > en { + t.Fatalf("n=%d: pairwise relerr %.3e worse than naive %.3e", n, ep, en) + } + } +} + +// TestCPUEngineSum_WorkerInvariant verifies ReduceSum is bitwise-identical +// regardless of GOMAXPROCS: the per-stripe accumulation order must not depend on +// runtime scheduling. This is the determinism guarantee downstream modes rely on. +func TestCPUEngineSum_WorkerInvariant(t *testing.T) { + engine := NewCPUEngine[float32](numeric.Float32Ops{}) + rng := rand.New(rand.NewPCG(3, 5)) + rows, cols := 17, 4096 + data := make([]float32, rows*cols) + for i := range data { + data[i] = float32(rng.NormFloat64()) + } + a, err := tensor.New[float32]([]int{rows, cols}, data) + if err != nil { + t.Fatal(err) + } + + sumWith := func(procs int) []float32 { + old := runtime.GOMAXPROCS(procs) + defer runtime.GOMAXPROCS(old) + out, err := engine.Sum(context.Background(), a, 1, false, nil) + if err != nil { + t.Fatal(err) + } + cp := make([]float32, len(out.Data())) + copy(cp, out.Data()) + + return cp + } + + ref := sumWith(1) + for _, p := range []int{2, 4, 8} { + got := sumWith(p) + for i := range ref { + if got[i] != ref[i] { + t.Fatalf("GOMAXPROCS=%d row %d: %v != %v (worker-dependent order)", p, i, got[i], ref[i]) + } + } + } +} + +// TestCPUEngineSoftmax_Deterministic verifies the softmax denominator reduction +// is run-to-run bitwise-stable. +func TestCPUEngineSoftmax_Deterministic(t *testing.T) { + engine := NewCPUEngine[float32](numeric.Float32Ops{}) + rng := rand.New(rand.NewPCG(9, 13)) + // inner>1 shape forces the generic (non-xblas) softmax path under test. + data := make([]float32, 3*2048*2) + for i := range data { + data[i] = float32(10 * rng.NormFloat64()) + } + a, err := tensor.New[float32]([]int{3, 2048, 2}, data) + if err != nil { + t.Fatal(err) + } + first, err := engine.Softmax(context.Background(), a, 1) + if err != nil { + t.Fatal(err) + } + ref := make([]float32, len(first.Data())) + copy(ref, first.Data()) + for r := 0; r < 4; r++ { + out, err := engine.Softmax(context.Background(), a, 1) + if err != nil { + t.Fatal(err) + } + for i := range ref { + if out.Data()[i] != ref[i] { + t.Fatalf("softmax not deterministic at %d on run %d", i, r) + } + } + } +} diff --git a/docs/design.md b/docs/design.md index 7d48e77..64fd17d 100644 --- a/docs/design.md +++ b/docs/design.md @@ -114,6 +114,13 @@ All methods accept a `context.Context` and an optional variadic `dst` tensor for **dst-form accumulation policy:** when a `dst` is supplied, the result is written into `dst`'s *existing* storage — an op never re-homes `dst` onto a fresh pool allocation. This storage-identity guarantee is what makes dst-form accumulation safe for persistent tensors (gradient accumulators, optimizer state): a caller that repeatedly issues `Add(ctx, acc, g, acc)` keeps the same device buffer for `acc`, outside the arena's per-step reuse, instead of being silently converted into an arena tensor that the next `Reset` recycles behind the live reference. On GPU, when `dst` carries a same-size GPUStorage the result is copied device-to-device into `dst`'s buffer and the temporary returns to the pool (skipped during CUDA graph capture, where the synchronous copy is not capturable). +**Reduction accumulation policy:** every CPU reduction accumulates in a **fixed, worker-independent order** rather than a naive left-to-right fold. `Sum`/`ReduceSum`/`ReduceMean`, the `Softmax` denominator, RMSNorm's mean-of-squares, and `numeric.Arithmetic.Sum` all use a recursive **pairwise (tree) reduction** whose split points are a pure function of the reduced length (recursive halving snapped to a 128-element base block, matching numpy's pairwise blocksize). Two guarantees follow: + +- **Determinism.** The tree shape depends only on length, never on `GOMAXPROCS`, chunk boundaries, or goroutine scheduling, so a reduction is bitwise-identical run to run. Parallelism is confined to *independent* output stripes; the accumulation along each reduced axis is sequential-in-tree-order within one goroutine. +- **Accuracy.** Pairwise accumulation bounds the rounding error at ~O(log n) versus O(n) for a naive fold, which tightens agreement with the PyTorch oracle without changing the accumulation dtype (float32 stays float32 — the acc-type is never *narrowed*). Measured tightening on ill-conditioned float32 inputs: 16–940× lower relative error for `ReduceSum` at n=1k–64k, 4–24× for RMSNorm's sum-of-squares, up to 52× for the softmax denominator; no input regresses. + +Accumulation dtype is preserved per site (float32 data reduces in float32, float64 in float64). Widening reduced-precision element types (float16/bfloat16) to a float32 accumulator, and a fixed-order rewrite of the GPU reduction kernels and the arm64 NEON RMSNorm SIMD path (both currently deterministic-but-SIMD-strided), are deferred to the deterministic-reductions mode (`ZTENSOR_DETERMINISTIC`, plan-gpu-training-hardening T4.1). + ### Optional Engine Capabilities Beyond the core interface, engines may implement optional capability interfaces discovered via type assertion: diff --git a/internal/xblas/reduce_pairwise.go b/internal/xblas/reduce_pairwise.go new file mode 100644 index 0000000..cde3261 --- /dev/null +++ b/internal/xblas/reduce_pairwise.go @@ -0,0 +1,41 @@ +package xblas + +// pairwiseBlock is the serial base-case size for fixed-order pairwise +// reduction, matching numpy's pairwise blocksize. +const pairwiseBlock = 128 + +// pairwiseSumF32 sums n float32 values in a fixed, chunk-independent order using +// a recursive tree (pairwise) reduction. Element i is read via get(i). +// +// The accumulation dtype stays float32 (matching the NEON RMSNorm asm path); +// only the order changes. Because the tree shape is a pure function of n, the +// result is bitwise-stable run to run, and its rounding error grows as +// O(log n) rather than O(n) for a naive scalar fold. See the arm64 asm in +// rmsnorm_arm64.s for the SIMD counterpart used on aarch64. +func pairwiseSumF32(n int, get func(i int) float32) float32 { + if n <= 0 { + return 0 + } + + return pairwiseSumRangeF32(0, n, get) +} + +func pairwiseSumRangeF32(lo, hi int, get func(i int) float32) float32 { + n := hi - lo + if n <= pairwiseBlock { + acc := get(lo) + for i := lo + 1; i < hi; i++ { + acc += get(i) + } + + return acc + } + + half := n / 2 + half -= half % pairwiseBlock + if half == 0 { + half = pairwiseBlock + } + + return pairwiseSumRangeF32(lo, lo+half, get) + pairwiseSumRangeF32(lo+half, hi, get) +} diff --git a/internal/xblas/reduce_pairwise_test.go b/internal/xblas/reduce_pairwise_test.go new file mode 100644 index 0000000..e072d3e --- /dev/null +++ b/internal/xblas/reduce_pairwise_test.go @@ -0,0 +1,74 @@ +package xblas + +import ( + "math" + "math/rand/v2" + "runtime" + "testing" +) + +// TestPairwiseSumF32_TightensAndStable verifies the fixed-order pairwise float32 +// reduction is bitwise-stable and no worse than a naive fold against a float64 +// ground truth for ill-conditioned inputs. +func TestPairwiseSumF32_TightensAndStable(t *testing.T) { + rng := rand.New(rand.NewPCG(21, 42)) + for _, n := range []int{2048, 8192, 32768} { + s := make([]float32, n) + for i := range s { + s[i] = float32(0.1 + 0.001*rng.NormFloat64()) + } + get := func(i int) float32 { return s[i] } + first := pairwiseSumF32(n, get) + for r := 0; r < 4; r++ { + if pairwiseSumF32(n, get) != first { + t.Fatalf("n=%d: not bitwise-stable", n) + } + } + + var truth, c float64 + for _, v := range s { + y := float64(v) - c + tt := truth + y + c = (tt - truth) - y + truth = tt + } + var naive float32 + for _, v := range s { + naive += v + } + en := math.Abs((float64(naive) - truth) / truth) + ep := math.Abs((float64(first) - truth) / truth) + if ep > en { + t.Fatalf("n=%d: pairwise relerr %.3e worse than naive %.3e", n, ep, en) + } + } +} + +// TestRMSNormF32_SumOfSquaresAccurate checks the RMSNorm scale factor agrees +// with a float64 reference within a tight tolerance. On arm64 this exercises the +// NEON asm (already fixed-order fp32); on other arches it exercises the pairwise +// generic fallback. Either way the result must be order-stable and accurate. +func TestRMSNormF32_SumOfSquaresAccurate(t *testing.T) { + rng := rand.New(rand.NewPCG(2, 3)) + const eps = float32(1e-6) + for _, d := range []int{2048, 4096} { + x := make([]float32, d) + w := make([]float32, d) + for i := range x { + x[i] = float32(rng.NormFloat64()) + w[i] = 1 + } + out := make([]float32, d) + var scale float32 + RMSNormF32(&out[0], &x[0], &w[0], d, eps, &scale) + + var sumSq float64 + for _, v := range x { + sumSq += float64(v) * float64(v) + } + wantScale := 1.0 / math.Sqrt(sumSq/float64(d)+float64(eps)) + if rel := math.Abs((float64(scale) - wantScale) / wantScale); rel > 1e-4 { + t.Fatalf("d=%d arch=%s: scale relerr %.3e exceeds 1e-4", d, runtime.GOARCH, rel) + } + } +} diff --git a/internal/xblas/rmsnorm_generic.go b/internal/xblas/rmsnorm_generic.go index d35ad0a..c1241ce 100644 --- a/internal/xblas/rmsnorm_generic.go +++ b/internal/xblas/rmsnorm_generic.go @@ -11,10 +11,10 @@ func RMSNormF32(out, x, weight *float32, dim int, eps float32, scale *float32) { xSlice := unsafe.Slice(x, dim) wSlice := unsafe.Slice(weight, dim) oSlice := unsafe.Slice(out, dim) - var sumSq float32 - for i := range dim { - sumSq += xSlice[i] * xSlice[i] - } + // Fixed-order pairwise sum of squares in float32 -- matches the fp32 + // accumulation of the arm64 NEON path (rmsnorm_arm64.s) while shrinking the + // error of the naive scalar fold this fallback previously used. + sumSq := pairwiseSumF32(dim, func(i int) float32 { return xSlice[i] * xSlice[i] }) s := float32(1.0 / math.Sqrt(float64(sumSq/float32(dim)+eps))) for i := range dim { oSlice[i] = xSlice[i] * s * wSlice[i] diff --git a/numeric/native_ops.go b/numeric/native_ops.go index 146da21..3ae4549 100644 --- a/numeric/native_ops.go +++ b/numeric/native_ops.go @@ -132,14 +132,11 @@ func (ops Float32Ops) Abs(x float32) float32 { return x } -// Sum computes the sum of elements in a slice. +// Sum computes the sum of elements in a slice using fixed-order pairwise +// accumulation (see numeric/reduce_pairwise.go): order-stable run to run and +// tighter than a naive fold. func (ops Float32Ops) Sum(s []float32) float32 { - var sum float32 - for _, v := range s { - sum += v - } - - return sum + return pairwiseSum(s) } // GreaterThan checks if a is greater than b. @@ -282,14 +279,11 @@ func (ops Float64Ops) Abs(x float64) float64 { return x } -// Sum computes the sum of elements in a slice. +// Sum computes the sum of elements in a slice using fixed-order pairwise +// accumulation (see numeric/reduce_pairwise.go): order-stable run to run and +// tighter than a naive fold. func (ops Float64Ops) Sum(s []float64) float64 { - var sum float64 - for _, v := range s { - sum += v - } - - return sum + return pairwiseSum(s) } // GreaterThan checks if a is greater than b. diff --git a/numeric/reduce_pairwise.go b/numeric/reduce_pairwise.go new file mode 100644 index 0000000..37a6b8f --- /dev/null +++ b/numeric/reduce_pairwise.go @@ -0,0 +1,40 @@ +package numeric + +// pairwiseBlock is the serial base-case size for fixed-order pairwise +// reduction, matching numpy's pairwise blocksize. +const pairwiseBlock = 128 + +// pairwiseSum sums s in a fixed, chunk-independent order using a recursive tree +// (pairwise) reduction. The accumulation dtype is the element type; only the +// order changes relative to a naive left-to-right fold, giving run-to-run +// bitwise stability and O(log n) rather than O(n) rounding-error growth. +// +// It is instantiated for the float element types whose Sum feeds reduction +// statistics; integer Sums are exact under any order and keep their simple fold. +func pairwiseSum[T ~float32 | ~float64](s []T) T { + if len(s) == 0 { + return 0 + } + + return pairwiseSumRange(s, 0, len(s)) +} + +func pairwiseSumRange[T ~float32 | ~float64](s []T, lo, hi int) T { + n := hi - lo + if n <= pairwiseBlock { + acc := s[lo] + for i := lo + 1; i < hi; i++ { + acc += s[i] + } + + return acc + } + + half := n / 2 + half -= half % pairwiseBlock + if half == 0 { + half = pairwiseBlock + } + + return pairwiseSumRange(s, lo, lo+half) + pairwiseSumRange(s, lo+half, hi) +} diff --git a/numeric/reduce_pairwise_test.go b/numeric/reduce_pairwise_test.go new file mode 100644 index 0000000..2279e3d --- /dev/null +++ b/numeric/reduce_pairwise_test.go @@ -0,0 +1,66 @@ +package numeric + +import ( + "math" + "math/rand/v2" + "testing" +) + +func kahanRef64(s []float32) float64 { + var sum, c float64 + for _, v := range s { + y := float64(v) - c + t := sum + y + c = (t - sum) - y + sum = t + } + + return sum +} + +func naiveF32(s []float32) float32 { + var acc float32 + for _, v := range s { + acc += v + } + + return acc +} + +// TestFloat32OpsSum_TightensAndStable verifies Float32Ops.Sum uses fixed-order +// pairwise accumulation: bitwise-stable run to run and no worse than the naive +// fold against a float64 ground truth. +func TestFloat32OpsSum_TightensAndStable(t *testing.T) { + ops := Float32Ops{} + rng := rand.New(rand.NewPCG(4, 8)) + for _, n := range []int{2048, 8192, 32768} { + s := make([]float32, n) + for i := range s { + s[i] = float32(0.1 + 0.001*rng.NormFloat64()) + } + first := ops.Sum(s) + for r := 0; r < 4; r++ { + if ops.Sum(s) != first { + t.Fatalf("n=%d: Sum not bitwise-stable", n) + } + } + truth := kahanRef64(s) + en := math.Abs((float64(naiveF32(s)) - truth) / truth) + ep := math.Abs((float64(first) - truth) / truth) + if ep > en { + t.Fatalf("n=%d: pairwise relerr %.3e worse than naive %.3e", n, ep, en) + } + } +} + +// TestPairwiseSum_Float64 sanity-checks the float64 instantiation matches an +// exact small-sum and stays deterministic. +func TestPairwiseSum_Float64(t *testing.T) { + s := []float64{1, 2, 3, 4, 5} + if got := pairwiseSum(s); got != 15 { + t.Fatalf("got %v want 15", got) + } + if pairwiseSum([]float64{}) != 0 { + t.Fatal("empty sum must be 0") + } +}