kernels: SumP on several cores again, bit-identical to floats.SumPairwise

README.md

@@ -183,9 +183,11 @@ which this module imports: on amd64 (AVX2) and arm64 (NEON) the elementwise
 `Exp`, `Log` and the correctly rounded float32 `Exp`/`Log` (the AVX2 and NEON
 kernels were written here and moved down). **`Sum` is NumPy's pairwise sum**
 (`floats.SumPairwise`): bit-identical to `np.sum` of a contiguous 1-D float64
-array, on every architecture, at every size. It replaces lane-parallel sums whose
-last bits differed from NumPy's and from one architecture to the next, and it is
-serial (a multicore split would regroup NumPy's order). The other architectures
+array, on every architecture, at every size and whatever the core count. It
+replaces lane-parallel sums whose last bits differed from NumPy's and from one
+architecture to the next. Large sums still run on several cores: the split cuts
+NumPy's pairwise tree where NumPy cuts it and adds the subtrees back up the same
+tree, so the bits do not depend on the number of workers. The other architectures
 keep their own elementwise kernels here until floats has them.
 
 ## Goroutines

internal/kernels/parallel.go

@@ -266,14 +266,75 @@ func mapReduceP(a []float64, red func([]float64) float64, combine func([]float64
 	return combine(partials)
 }
 
-// SumP returns the sum of all elements of a in NumPy's order: floats.SumPairwise,
-// bit-identical to np.sum of a contiguous 1-D float64 array. It is serial at every
-// size: NumPy's pairwise order is a fixed recursion, and summing chunks on several
-// cores and folding the partials would regroup it (the result would then depend on
-// the core count). An order-preserving parallel split (chunks cut where NumPy's
-// recursion cuts) is possible and not done here. Callers that need the exact
-// sequential left-to-right value use Sum.
-func SumP(a []float64) float64 { return sumSIMD(a) }
+// SumP returns the sum of all elements of a in NumPy's order: bit-identical to
+// floats.SumPairwise (np.sum of a contiguous 1-D float64 array), at every size and
+// whatever the core count. NumPy's pairwise sum is a fixed binary tree (a block of
+// at most 128 elements is the leaf; a longer one splits at n/2 rounded down to a
+// multiple of 8, and the two halves are added), so the parallel form cuts the same
+// tree: it descends the top levels exactly as NumPy does, until there are about
+// four subtrees per worker of at least reduceGrain elements, sums each subtree with
+// floats.SumPairwise on its own core, and adds the results back up the same tree.
+// The only difference from the serial recursion is the sign of a subtree of only
+// -0 values (floats.SumPairwise starts from +0), which the final +0 + result of
+// NumPy's reduction erases. Callers that need the exact sequential left-to-right
+// value use Sum.
+func SumP(a []float64) float64 {
+	n := len(a)
+	if n < ParThreshold {
+		return sumSIMD(a)
+	}
+	w := numWorkers(n)
+	depth := 0
+	for 1<<depth < 4*w && n>>(depth+1) >= reduceGrain {
+		depth++
+	}
+	if depth == 0 {
+		return sumSIMD(a)
+	}
+	var cuts []int // leaf boundaries, in tree order: leaf i is a[cuts[i]:cuts[i+1]]
+	cuts = append(cuts, 0)
+	var leaves func(lo, n, depth int)
+	leaves = func(lo, n, depth int) {
+		if depth == 0 || n <= pairwiseBlock {
+			cuts = append(cuts, lo+n)
+			return
+		}
+		n2 := pairwiseSplit(n)
+		leaves(lo, n2, depth-1)
+		leaves(lo+n2, n-n2, depth-1)
+	}
+	leaves(0, n, depth)
+	vals := make([]float64, len(cuts)-1)
+	parallelFor(len(vals), w, func(lo, hi int) {
+		for i := lo; i < hi; i++ {
+			vals[i] = sumSIMD(a[cuts[i]:cuts[i+1]])
+		}
+	})
+	next := 0
+	var fold func(n, depth int) float64
+	fold = func(n, depth int) float64 {
+		if depth == 0 || n <= pairwiseBlock {
+			v := vals[next]
+			next++
+			return v
+		}
+		n2 := pairwiseSplit(n)
+		l := fold(n2, depth-1)
+		return l + fold(n-n2, depth-1)
+	}
+	return fold(n, depth)
+}
+
+// pairwiseBlock is the longest block NumPy's pairwise sum adds with eight
+// accumulators (PW_BLOCKSIZE); a longer one splits at pairwiseSplit.
+const pairwiseBlock = 128
+
+// pairwiseSplit is where DOUBLE_pairwise_sum cuts a block of n > 128 elements: n/2
+// rounded down to a multiple of 8.
+func pairwiseSplit(n int) int {
+	n2 := n / 2
+	return n2 - n2%8
+}
 
 // reduceP folds per-chunk results of the associative scalar reducer red with
 // red itself, giving a result identical to the serial kernel.

internal/kernels/sum_numpy_test.go

@@ -35,3 +35,35 @@ func TestSumMatchesNumPy(t *testing.T) {
 		}
 	}
 }
+
+// TestSumPParallelIsBitIdentical: SumP cuts NumPy's pairwise tree and sums the
+// subtrees on several cores; the result is floats.SumPairwise's bits for every
+// length and core count (a subtree of only -0 included: the final +0 erases it).
+func TestSumPParallelIsBitIdentical(t *testing.T) {
+	negZero := math.Copysign(0, -1)
+	for _, procs := range []int{1, 2, 3, 4, 7, 16} {
+		withMaxProcs(procs, func() {
+			withThresholds(64, 1<<14, func() {
+				for _, n := range []int{64, 129, 1000, 4097, 100000, 1<<20 + 13, 3<<20 - 7} {
+					a := randVec(n, int64(n))
+					for i := range a {
+						a[i] *= 1e8 // wide dynamic range: any regrouping would show
+						if i%97 == 0 {
+							a[i] = 1e-9
+						}
+					}
+					want := sumSIMD(a)
+					if got := SumP(a); math.Float64bits(got) != math.Float64bits(want) {
+						t.Fatalf("procs=%d n=%d: SumP %v (%#x), SumPairwise %v (%#x)", procs, n, got, math.Float64bits(got), want, math.Float64bits(want))
+					}
+					for i := range a {
+						a[i] = negZero
+					}
+					if got := SumP(a); math.Float64bits(got) != math.Float64bits(sumSIMD(a)) {
+						t.Fatalf("procs=%d n=%d all -0: SumP %v", procs, n, got)
+					}
+				}
+			})
+		})
+	}
+}