1292 lines
39 KiB
Go
1292 lines
39 KiB
Go
|
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
|
|
// SPDX-License-Identifier: MIT
|
||
|
|
|
||
|
|
package spmd
|
||
|
|
|
||
|
|
import (
|
||
|
|
"encoding/binary"
|
||
|
|
"math"
|
||
|
|
|
||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
|
|
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Op names a reduction's arithmetic.
|
||
|
|
type Op uint8
|
||
|
|
|
||
|
|
const (
|
||
|
|
// Sum adds: shard partials through the balanced tree, rank arrays
|
||
|
|
// in rank order.
|
||
|
|
Sum Op = iota
|
||
|
|
// Min takes the smallest; NaN never wins.
|
||
|
|
Min
|
||
|
|
// Max takes the largest; NaN never wins.
|
||
|
|
Max
|
||
|
|
// Any reports whether any element is true; Bool arrays only.
|
||
|
|
Any
|
||
|
|
// All reports whether every element is true; Bool arrays only.
|
||
|
|
All
|
||
|
|
// Prod multiplies: shard partials through the balanced product
|
||
|
|
// tree, rank arrays in rank order.
|
||
|
|
Prod
|
||
|
|
)
|
||
|
|
|
||
|
|
func (o Op) String() string {
|
||
|
|
switch o {
|
||
|
|
case Sum:
|
||
|
|
return "sum"
|
||
|
|
case Min:
|
||
|
|
return "min"
|
||
|
|
case Max:
|
||
|
|
return "max"
|
||
|
|
case Any:
|
||
|
|
return "any"
|
||
|
|
case All:
|
||
|
|
return "all"
|
||
|
|
case Prod:
|
||
|
|
return "prod"
|
||
|
|
}
|
||
|
|
return "?"
|
||
|
|
}
|
||
|
|
|
||
|
|
// The payload kinds a shard's contribution can travel as. The kind is
|
||
|
|
// part of the wire, so the root combines exactly what the ranks
|
||
|
|
// folded, never a guess from the dtype.
|
||
|
|
const (
|
||
|
|
kindNone byte = 0 // a rank whose piece is empty
|
||
|
|
kindSumF64 byte = 1 // one float64 fold per block
|
||
|
|
kindSumC128 byte = 2 // one complex fold per block
|
||
|
|
kindSumI64 byte = 3 // one exact int64 sum per block
|
||
|
|
kindExtF64 byte = 4 // one float64 extremum and candidate flag per block
|
||
|
|
kindExtI64 byte = 5 // one exact int64 extremum per block
|
||
|
|
kindFlag byte = 6 // one 0/1 flag for the whole piece
|
||
|
|
kindProdF64 byte = 7 // one float64 product per block
|
||
|
|
kindProdF32 byte = 8 // one native float32 product per block
|
||
|
|
kindProdHalf byte = 9 // one exact half product per block, in float64
|
||
|
|
kindProdI64 byte = 10 // one exact int64 product per block
|
||
|
|
kindArg byte = 11 // one extremum candidate: value, index and flags
|
||
|
|
)
|
||
|
|
|
||
|
|
// intPayload is the integer element type a shard folds exactly.
|
||
|
|
type intPayload interface {
|
||
|
|
int64 | int8 | uint8 | int16 | uint16 | int32 | uint32
|
||
|
|
}
|
||
|
|
|
||
|
|
// AllReduceShards answers the reduction of one global array whose
|
||
|
|
// canonical pieces the ranks hold, every rank receiving the answer.
|
||
|
|
// Each rank folds its own blocks with the single block fold, the block
|
||
|
|
// values gather at rank 0, and core's own tree and extremum rules
|
||
|
|
// combine them: the answer is the single-array reduction's exact bits,
|
||
|
|
// whatever the world's size, whatever the order the frames arrive in.
|
||
|
|
// Any and All answer the integer-class scalar, 1 or 0.
|
||
|
|
func (w *World) AllReduceShards(local *core.Array, span Span, op Op) (core.Scalar, error) {
|
||
|
|
ans, err := w.reduceShards(local, span, op, 0)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
return w.broadcastScalar(ans, 0)
|
||
|
|
}
|
||
|
|
|
||
|
|
// ReduceShards is AllReduceShards with the answer on the root alone;
|
||
|
|
// every other rank receives nil.
|
||
|
|
func (w *World) ReduceShards(local *core.Array, span Span, op Op, root int) (*core.Scalar, error) {
|
||
|
|
if err := checkRoot(root, w.size); err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
ans, err := w.reduceShards(local, span, op, root)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return &ans, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// blockValues is one rank's contribution to a sharded reduction: the
|
||
|
|
// index of its first block, how many blocks it covers, the payload
|
||
|
|
// kind, and the values.
|
||
|
|
type blockValues struct {
|
||
|
|
first int
|
||
|
|
kind byte
|
||
|
|
f []float64
|
||
|
|
f32 []float32
|
||
|
|
c []complex128
|
||
|
|
i []int64
|
||
|
|
oks []bool
|
||
|
|
}
|
||
|
|
|
||
|
|
// blocks reports how many partition blocks the piece covers.
|
||
|
|
func (bv blockValues) blocks() int {
|
||
|
|
switch bv.kind {
|
||
|
|
case kindSumF64, kindSumI64:
|
||
|
|
return len(bv.f) + len(bv.i)
|
||
|
|
case kindSumC128:
|
||
|
|
return len(bv.c)
|
||
|
|
case kindExtF64:
|
||
|
|
return len(bv.oks)
|
||
|
|
case kindExtI64:
|
||
|
|
return len(bv.i)
|
||
|
|
case kindProdF64:
|
||
|
|
return len(bv.f)
|
||
|
|
case kindProdF32:
|
||
|
|
return len(bv.f32)
|
||
|
|
case kindProdHalf, kindProdI64:
|
||
|
|
return len(bv.f) + len(bv.i)
|
||
|
|
}
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
|
||
|
|
// spanBlocks names the partition blocks the span covers: [first, last).
|
||
|
|
// A span that the partition gave out always sits on block boundaries.
|
||
|
|
func spanBlocks(span Span) (first, last, parts int) {
|
||
|
|
parts = core.FoldParts(span.Global)
|
||
|
|
first, last = -1, -1
|
||
|
|
for c := 0; c <= parts; c++ {
|
||
|
|
b := core.FoldBoundary(span.Global, c)
|
||
|
|
if b == span.Lo && first < 0 {
|
||
|
|
first = c
|
||
|
|
}
|
||
|
|
if b == span.Hi {
|
||
|
|
last = c
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return first, last, parts
|
||
|
|
}
|
||
|
|
|
||
|
|
// reduceShards runs the sharded reduction; the answer exists at the
|
||
|
|
// root.
|
||
|
|
func (w *World) reduceShards(local *core.Array, span Span, op Op, root int) (core.Scalar, error) {
|
||
|
|
if err := w.status(); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
if err := w.checkSpan(span, local); err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
if span.Global == 0 {
|
||
|
|
// An empty global axis answers without an exchange: the
|
||
|
|
// single-array reduction answers the same constants, so the
|
||
|
|
// contract holds at the degenerate length too. The dtype
|
||
|
|
// validates first, exactly as a non-empty array would.
|
||
|
|
switch op {
|
||
|
|
case Prod:
|
||
|
|
if dt := local.Dtype(); dt == core.Bool || dt == core.Int8 || dt == core.Uint8 || dt == core.Int16 || dt == core.Uint16 || dt == core.Int32 || dt == core.Uint32 || dt == core.Complex {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: prod of dtype %s is not supported; convert with Astype", dt))
|
||
|
|
}
|
||
|
|
case Min, Max:
|
||
|
|
if dt := local.Dtype(); dt == core.Bool || dt == core.Complex {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: %s of dtype %s has no ordering", op, dt))
|
||
|
|
}
|
||
|
|
case Any, All:
|
||
|
|
if dt := local.Dtype(); dt != core.Bool {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: %s needs Bool elements, got %s", op, dt))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
switch op {
|
||
|
|
case Sum:
|
||
|
|
switch local.Dtype() {
|
||
|
|
case core.Complex:
|
||
|
|
return core.ComplexScalar(0), nil
|
||
|
|
case core.Int, core.Int8, core.Uint8, core.Int16, core.Uint16,
|
||
|
|
core.Int32, core.Uint32, core.Bool:
|
||
|
|
return core.IntScalar(0), nil
|
||
|
|
default:
|
||
|
|
return core.FloatScalar(0), nil
|
||
|
|
}
|
||
|
|
case Any:
|
||
|
|
return core.IntScalar(0), nil
|
||
|
|
case All:
|
||
|
|
return core.IntScalar(1), nil
|
||
|
|
case Prod:
|
||
|
|
// The empty product is the multiplicative identity, the
|
||
|
|
// answer the single-array product keeps.
|
||
|
|
if local.Dtype() == core.Int {
|
||
|
|
return core.IntScalar(1), nil
|
||
|
|
}
|
||
|
|
return core.FloatScalar(1), nil
|
||
|
|
default:
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: %s of an empty array has no answer", op))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
bv, err := w.foldBlocks(local, span, op)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
if err := w.sendTo(root, tagShardsValues, encodeBlockValues(bv)); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
return core.Scalar{}, nil
|
||
|
|
}
|
||
|
|
all := make([]blockValues, w.size)
|
||
|
|
all[w.rank] = bv
|
||
|
|
for r := range w.size {
|
||
|
|
if r == root {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
data, err := w.recvFrom(r, tagShardsValues)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
if all[r], err = decodeBlockValues(data); err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return w.combineShards(all, span, op)
|
||
|
|
}
|
||
|
|
|
||
|
|
// foldBlocks computes the rank's own blocks' values with the same
|
||
|
|
// per-block arithmetic the single-array fold uses, dtype by dtype.
|
||
|
|
func (w *World) foldBlocks(local *core.Array, span Span, op Op) (blockValues, error) {
|
||
|
|
first, last, _ := spanBlocks(span)
|
||
|
|
if first < 0 || last < 0 || first > last {
|
||
|
|
return blockValues{}, base.Errf("spmd: rank %d's span [%d, %d) is not on the partition's block boundaries",
|
||
|
|
w.rank, span.Lo, span.Hi)
|
||
|
|
}
|
||
|
|
bv := blockValues{first: first}
|
||
|
|
if first == last {
|
||
|
|
return bv, nil
|
||
|
|
}
|
||
|
|
n := span.Global
|
||
|
|
count := last - first
|
||
|
|
blockF64 := func(c int) []float64 {
|
||
|
|
return sliceBlock(local.RawFloats()[:local.Len()], n, first, c)
|
||
|
|
}
|
||
|
|
switch op {
|
||
|
|
case Sum:
|
||
|
|
switch local.Dtype() {
|
||
|
|
case core.Float:
|
||
|
|
bv.kind = kindSumF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldRange(blockF64(c))
|
||
|
|
}
|
||
|
|
case core.Float32:
|
||
|
|
bv.kind = kindSumF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldRangeF32(sliceBlock(local.RawFloat32s()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Float16:
|
||
|
|
bv.kind = kindSumF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldRangeF16(sliceBlock(local.RawHalves()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Complex:
|
||
|
|
bv.kind = kindSumC128
|
||
|
|
bv.c = make([]complex128, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.c[c-first] = core.FoldRangeC128(sliceBlock(local.RawComplexes()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Int:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawInts()[:local.Len()], n, first, last)
|
||
|
|
case core.Int8:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawInt8s()[:local.Len()], n, first, last)
|
||
|
|
case core.Uint8:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawUint8s()[:local.Len()], n, first, last)
|
||
|
|
case core.Int16:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawInt16s()[:local.Len()], n, first, last)
|
||
|
|
case core.Uint16:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawUint16s()[:local.Len()], n, first, last)
|
||
|
|
case core.Int32:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawInt32s()[:local.Len()], n, first, last)
|
||
|
|
case core.Uint32:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = exactSums(local.RawUint32s()[:local.Len()], n, first, last)
|
||
|
|
case core.Bool:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
src := local.RawBools()[:local.Len()]
|
||
|
|
bv.i = make([]int64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
var s int64
|
||
|
|
for _, v := range sliceBlock(src, n, first, c) {
|
||
|
|
if v {
|
||
|
|
s++
|
||
|
|
}
|
||
|
|
}
|
||
|
|
bv.i[c-first] = s
|
||
|
|
}
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: sum of dtype %s is not a reduction", local.Dtype())
|
||
|
|
}
|
||
|
|
case Min, Max:
|
||
|
|
greater := op == Max
|
||
|
|
switch local.Dtype() {
|
||
|
|
case core.Float:
|
||
|
|
bv.kind = kindExtF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
bv.oks = make([]bool, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first], bv.oks[c-first] = core.ExtremeRange(blockF64(c), greater)
|
||
|
|
}
|
||
|
|
case core.Float32:
|
||
|
|
bv.kind = kindExtF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
bv.oks = make([]bool, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first], bv.oks[c-first] = core.ExtremeRangeF32(sliceBlock(local.RawFloat32s()[:local.Len()], n, first, c), greater)
|
||
|
|
}
|
||
|
|
case core.Float16:
|
||
|
|
bv.kind = kindExtF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
bv.oks = make([]bool, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first], bv.oks[c-first] = core.ExtremeRangeF16(sliceBlock(local.RawHalves()[:local.Len()], n, first, c), greater)
|
||
|
|
}
|
||
|
|
case core.Int:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawInts()[:local.Len()], n, first, last, greater)
|
||
|
|
case core.Int8:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawInt8s()[:local.Len()], n, first, last, greater)
|
||
|
|
case core.Uint8:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawUint8s()[:local.Len()], n, first, last, greater)
|
||
|
|
case core.Int16:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawInt16s()[:local.Len()], n, first, last, greater)
|
||
|
|
case core.Uint16:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawUint16s()[:local.Len()], n, first, last, greater)
|
||
|
|
case core.Int32:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawInt32s()[:local.Len()], n, first, last, greater)
|
||
|
|
case core.Uint32:
|
||
|
|
bv.kind = kindExtI64
|
||
|
|
bv.i = exactExtremes(local.RawUint32s()[:local.Len()], n, first, last, greater)
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: %s of dtype %s has no ordering", op, local.Dtype())
|
||
|
|
}
|
||
|
|
case Prod:
|
||
|
|
switch local.Dtype() {
|
||
|
|
case core.Float:
|
||
|
|
bv.kind = kindProdF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldProd(sliceBlock(local.RawFloats()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Float32:
|
||
|
|
bv.kind = kindProdF32
|
||
|
|
bv.f32 = make([]float32, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f32[c-first] = core.FoldProdF32(sliceBlock(local.RawFloat32s()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Float16:
|
||
|
|
bv.kind = kindProdHalf
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldProdF16(sliceBlock(local.RawHalves()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Int:
|
||
|
|
bv.kind = kindProdI64
|
||
|
|
bv.i = make([]int64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.i[c-first] = core.FoldProdI64(sliceBlock(local.RawInts()[:local.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: prod of dtype %s is not supported; convert with Astype", local.Dtype())
|
||
|
|
}
|
||
|
|
case Any, All:
|
||
|
|
if local.Dtype() != core.Bool {
|
||
|
|
return blockValues{}, base.Errf("spmd: %s needs Bool elements, got %s", op, local.Dtype())
|
||
|
|
}
|
||
|
|
bv.kind = kindFlag
|
||
|
|
flag := op == All
|
||
|
|
for _, v := range local.RawBools()[:local.Len()] {
|
||
|
|
if op == Any && v {
|
||
|
|
flag = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
if op == All && !v {
|
||
|
|
flag = false
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
bv.i = []int64{bToF(flag)}
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: unknown reduction op %s", op)
|
||
|
|
}
|
||
|
|
return bv, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// sliceBlock is block c of the canonical partition of n elements, seen
|
||
|
|
// inside a slab whose first block is first.
|
||
|
|
func sliceBlock[T any](src []T, n, first, c int) []T {
|
||
|
|
baseIdx := core.FoldBoundary(n, first)
|
||
|
|
return src[core.FoldBoundary(n, c)-baseIdx : core.FoldBoundary(n, c+1)-baseIdx]
|
||
|
|
}
|
||
|
|
|
||
|
|
// exactSums folds each of the piece's blocks into an exact int64 sum.
|
||
|
|
func exactSums[T intPayload](src []T, n, first, last int) []int64 {
|
||
|
|
out := make([]int64, last-first)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
var s int64
|
||
|
|
for _, v := range sliceBlock(src, n, first, c) {
|
||
|
|
s += int64(v)
|
||
|
|
}
|
||
|
|
out[c-first] = s
|
||
|
|
}
|
||
|
|
return out
|
||
|
|
}
|
||
|
|
|
||
|
|
// exactExtremes folds each of the piece's blocks into an exact int64
|
||
|
|
// extremum: the values compare exactly, so the combine needs no flags.
|
||
|
|
func exactExtremes[T intPayload](src []T, n, first, last int, greater bool) []int64 {
|
||
|
|
out := make([]int64, last-first)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
blk := sliceBlock(src, n, first, c)
|
||
|
|
best := int64(blk[0])
|
||
|
|
for _, v := range blk[1:] {
|
||
|
|
if w := int64(v); (greater && w > best) || (!greater && w < best) {
|
||
|
|
best = w
|
||
|
|
}
|
||
|
|
}
|
||
|
|
out[c-first] = best
|
||
|
|
}
|
||
|
|
return out
|
||
|
|
}
|
||
|
|
|
||
|
|
// payloadKind names the kind every non-empty piece carries; the pieces
|
||
|
|
// agree because the program runs one op on one dtype, and the root, whose
|
||
|
|
// own piece may be empty, reads it from the first rank that folded
|
||
|
|
// anything.
|
||
|
|
func payloadKind(all []blockValues) (byte, error) {
|
||
|
|
kind := kindNone
|
||
|
|
for _, bv := range all {
|
||
|
|
if bv.kind == kindNone {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if kind == kindNone {
|
||
|
|
kind = bv.kind
|
||
|
|
} else if bv.kind != kind {
|
||
|
|
return 0, base.Errf("spmd: the shards disagree on the payload kind: %d against %d", bv.kind, kind)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return kind, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// combineShards lays the ranks' pieces over the whole partition and
|
||
|
|
// combines them with the same functions the single-array reduction
|
||
|
|
// combines its chunks with.
|
||
|
|
func (w *World) combineShards(all []blockValues, span Span, op Op) (core.Scalar, error) {
|
||
|
|
first, last, parts := spanBlocks(span)
|
||
|
|
if first < 0 || last < 0 {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: the root's span is not on the partition's boundaries")
|
||
|
|
}
|
||
|
|
kind, kindErr := payloadKind(all)
|
||
|
|
if kindErr != nil {
|
||
|
|
return core.Scalar{}, kindErr
|
||
|
|
}
|
||
|
|
if op != Any && op != All {
|
||
|
|
covered := 0
|
||
|
|
for _, bv := range all {
|
||
|
|
covered += bv.blocks()
|
||
|
|
}
|
||
|
|
if covered != parts {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: the shards' %d blocks do not cover the partition's %d", covered, parts)
|
||
|
|
}
|
||
|
|
// The root joins in rank order, which is the partition's
|
||
|
|
// order only when every piece starts at its own rank's first
|
||
|
|
// block; a piece naming any other start would scramble the
|
||
|
|
// tree's addends into a wrong number returned as success.
|
||
|
|
for r, bv := range all {
|
||
|
|
if bv.kind == kindNone || bv.kind == kindFlag {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if bv.first != r*parts/w.size {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: rank %d's piece starts at block %d, the partition puts it at %d", r, bv.first, r*parts/w.size)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
switch op {
|
||
|
|
case Sum:
|
||
|
|
switch kind {
|
||
|
|
case kindSumF64:
|
||
|
|
vals := make([]float64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.f...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(core.TreeSum(vals)), nil
|
||
|
|
case kindSumC128:
|
||
|
|
vals := make([]complex128, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.c...)
|
||
|
|
}
|
||
|
|
return core.ComplexScalar(core.TreeSum(vals)), nil
|
||
|
|
case kindSumI64:
|
||
|
|
vals := make([]int64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.i...)
|
||
|
|
}
|
||
|
|
return core.IntScalar(core.TreeSum(vals)), nil
|
||
|
|
default:
|
||
|
|
return core.Scalar{}, base.Errf("spmd: sum shards disagree on the payload kind")
|
||
|
|
}
|
||
|
|
case Min, Max:
|
||
|
|
greater := op == Max
|
||
|
|
switch kind {
|
||
|
|
case kindExtF64:
|
||
|
|
vals := make([]float64, 0, parts)
|
||
|
|
oks := make([]bool, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.f...)
|
||
|
|
oks = append(oks, bv.oks...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(core.CombineExtrema(vals, oks, greater)), nil
|
||
|
|
case kindExtI64:
|
||
|
|
var best int64
|
||
|
|
have := false
|
||
|
|
for _, bv := range all {
|
||
|
|
for _, v := range bv.i {
|
||
|
|
if !have || (greater && v > best) || (!greater && v < best) {
|
||
|
|
best, have = v, true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return core.IntScalar(best), nil
|
||
|
|
default:
|
||
|
|
return core.Scalar{}, base.Errf("spmd: %s shards disagree on the payload kind", op)
|
||
|
|
}
|
||
|
|
case Prod:
|
||
|
|
switch kind {
|
||
|
|
case kindProdF64:
|
||
|
|
vals := make([]float64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.f...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(core.TreeProd(vals)), nil
|
||
|
|
case kindProdF32:
|
||
|
|
vals := make([]float32, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.f32...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(float64(core.TreeProd(vals))), nil
|
||
|
|
case kindProdHalf:
|
||
|
|
vals := make([]float64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.f...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(core.TreeProdHalf(vals)), nil
|
||
|
|
case kindProdI64:
|
||
|
|
vals := make([]int64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.i...)
|
||
|
|
}
|
||
|
|
return core.IntScalar(core.TreeProd(vals)), nil
|
||
|
|
default:
|
||
|
|
return core.Scalar{}, base.Errf("spmd: prod shards disagree on the payload kind")
|
||
|
|
}
|
||
|
|
case Any, All:
|
||
|
|
flag := op == All
|
||
|
|
for _, bv := range all {
|
||
|
|
if bv.kind != kindFlag || len(bv.i) == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
v := bv.i[0] != 0
|
||
|
|
if op == Any && v {
|
||
|
|
flag = true
|
||
|
|
}
|
||
|
|
if op == All && !v {
|
||
|
|
flag = false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return core.IntScalar(bToF(flag)), nil
|
||
|
|
}
|
||
|
|
return core.Scalar{}, base.Errf("spmd: unknown reduction op %s", op)
|
||
|
|
}
|
||
|
|
|
||
|
|
// bToF carries a boolean in the int64 slot the integer-class scalar
|
||
|
|
// answers through.
|
||
|
|
func bToF(b bool) int64 {
|
||
|
|
if b {
|
||
|
|
return 1
|
||
|
|
}
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
|
||
|
|
// The shard contribution's wire form: first block index, payload kind,
|
||
|
|
// block count, then the values. Every field is checked against every
|
||
|
|
// other on the way back in.
|
||
|
|
|
||
|
|
func encodeBlockValues(bv blockValues) []byte {
|
||
|
|
buf := binary.LittleEndian.AppendUint64(nil, uint64(bv.first))
|
||
|
|
buf = append(buf, bv.kind)
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, uint64(len(bv.f)))
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, uint64(len(bv.f32)))
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, uint64(len(bv.c)))
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, uint64(len(bv.i)))
|
||
|
|
for _, v := range bv.f {
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, math.Float64bits(v))
|
||
|
|
}
|
||
|
|
for _, v := range bv.c {
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, math.Float64bits(real(v)))
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, math.Float64bits(imag(v)))
|
||
|
|
}
|
||
|
|
for _, v := range bv.i {
|
||
|
|
buf = binary.LittleEndian.AppendUint64(buf, uint64(v))
|
||
|
|
}
|
||
|
|
for _, v := range bv.f32 {
|
||
|
|
buf = binary.LittleEndian.AppendUint32(buf, math.Float32bits(v))
|
||
|
|
}
|
||
|
|
for _, ok := range bv.oks {
|
||
|
|
b := byte(0)
|
||
|
|
if ok {
|
||
|
|
b = 1
|
||
|
|
}
|
||
|
|
buf = append(buf, b)
|
||
|
|
}
|
||
|
|
return buf
|
||
|
|
}
|
||
|
|
|
||
|
|
func decodeBlockValues(data []byte) (blockValues, error) {
|
||
|
|
const head = 8 + 1 + 8*4
|
||
|
|
if len(data) < head {
|
||
|
|
return blockValues{}, base.Errf("spmd: a shard contribution of %d bytes is shorter than its head", len(data))
|
||
|
|
}
|
||
|
|
// Every count is bounded as its unsigned wire value, before any
|
||
|
|
// signed conversion: a top-bit count would come back negative and
|
||
|
|
// sneak past an upper-bound check.
|
||
|
|
firstU := binary.LittleEndian.Uint64(data)
|
||
|
|
if firstU > math.MaxInt {
|
||
|
|
return blockValues{}, base.Errf("spmd: a shard contribution names a first block of %d", firstU)
|
||
|
|
}
|
||
|
|
bv := blockValues{first: int(firstU)}
|
||
|
|
bv.kind = data[8]
|
||
|
|
nf, n32, nc, ni, err := decodeCounts(data[9:])
|
||
|
|
if err != nil {
|
||
|
|
return blockValues{}, err
|
||
|
|
}
|
||
|
|
data = data[head:]
|
||
|
|
want := (nf+ni)*8 + n32*4 + nc*16
|
||
|
|
if bv.kind == kindExtF64 {
|
||
|
|
want += nf
|
||
|
|
}
|
||
|
|
if bv.kind == kindArg {
|
||
|
|
want += 2
|
||
|
|
}
|
||
|
|
if len(data) != want {
|
||
|
|
return blockValues{}, base.Errf("spmd: a shard contribution carries %d value bytes for %d floats, %d float32s, %d complexes and %d integers",
|
||
|
|
len(data), nf, n32, nc, ni)
|
||
|
|
}
|
||
|
|
bv.f = make([]float64, nf)
|
||
|
|
for i := range bv.f {
|
||
|
|
bv.f[i] = math.Float64frombits(binary.LittleEndian.Uint64(data[i*8:]))
|
||
|
|
}
|
||
|
|
data = data[nf*8:]
|
||
|
|
bv.c = make([]complex128, nc)
|
||
|
|
for i := range bv.c {
|
||
|
|
re := math.Float64frombits(binary.LittleEndian.Uint64(data[i*16:]))
|
||
|
|
im := math.Float64frombits(binary.LittleEndian.Uint64(data[i*16+8:]))
|
||
|
|
bv.c[i] = complex(re, im)
|
||
|
|
}
|
||
|
|
data = data[nc*16:]
|
||
|
|
bv.i = make([]int64, ni)
|
||
|
|
for i := range bv.i {
|
||
|
|
bv.i[i] = int64(binary.LittleEndian.Uint64(data[i*8:]))
|
||
|
|
}
|
||
|
|
data = data[ni*8:]
|
||
|
|
bv.f32 = make([]float32, n32)
|
||
|
|
for i := range bv.f32 {
|
||
|
|
bv.f32[i] = math.Float32frombits(binary.LittleEndian.Uint32(data[i*4:]))
|
||
|
|
}
|
||
|
|
data = data[n32*4:]
|
||
|
|
if bv.kind == kindExtF64 {
|
||
|
|
bv.oks = make([]bool, nf)
|
||
|
|
for i := range bv.oks {
|
||
|
|
switch data[i] {
|
||
|
|
case 0:
|
||
|
|
case 1:
|
||
|
|
bv.oks[i] = true
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: a shard's candidate flag %d is neither 0 nor 1", data[i])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if bv.kind == kindArg {
|
||
|
|
bv.oks = make([]bool, 2)
|
||
|
|
for i := range bv.oks {
|
||
|
|
switch data[i] {
|
||
|
|
case 0:
|
||
|
|
case 1:
|
||
|
|
bv.oks[i] = true
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: a shard's candidate flag %d is neither 0 nor 1", data[i])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return bv, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// decodeCounts reads the four block counts as unsigned wire values
|
||
|
|
// and bounds each before any signed conversion: a top-bit count would
|
||
|
|
// come back negative and sneak past an upper-bound check.
|
||
|
|
func decodeCounts(data []byte) (nf, n32, nc, ni int, err error) {
|
||
|
|
raw := [4]uint64{
|
||
|
|
binary.LittleEndian.Uint64(data[0:8]),
|
||
|
|
binary.LittleEndian.Uint64(data[8:16]),
|
||
|
|
binary.LittleEndian.Uint64(data[16:24]),
|
||
|
|
binary.LittleEndian.Uint64(data[24:32]),
|
||
|
|
}
|
||
|
|
for k, u := range raw {
|
||
|
|
if u > maxBlocks {
|
||
|
|
return 0, 0, 0, 0, base.Errf("spmd: a shard contribution names an impossible count %d in slot %d", u, k)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return int(raw[0]), int(raw[1]), int(raw[2]), int(raw[3]), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// maxBlocks bounds the counts a contribution may name before anything
|
||
|
|
// is allocated: the partition itself never exceeds this.
|
||
|
|
const maxBlocks = 1 << 12
|
||
|
|
|
||
|
|
// The scalar answer travels as one fixed frame: a kind byte, then the
|
||
|
|
// int64, float64 and complex slots, each little-endian. Only the slot
|
||
|
|
// the kind names is meaningful; every slot always travels, so the form
|
||
|
|
// has no variable length to get wrong.
|
||
|
|
const (
|
||
|
|
scalarKindInt byte = 1
|
||
|
|
scalarKindFloat byte = 2
|
||
|
|
scalarKindComplex byte = 3
|
||
|
|
)
|
||
|
|
|
||
|
|
const scalarFrameLen = 1 + 8 + 8 + 16
|
||
|
|
|
||
|
|
func encodeScalar(s core.Scalar) []byte {
|
||
|
|
buf := make([]byte, scalarFrameLen)
|
||
|
|
switch {
|
||
|
|
case s.IsComplex():
|
||
|
|
buf[0] = scalarKindComplex
|
||
|
|
binary.LittleEndian.PutUint64(buf[17:], math.Float64bits(real(s.Complex())))
|
||
|
|
binary.LittleEndian.PutUint64(buf[25:], math.Float64bits(imag(s.Complex())))
|
||
|
|
case s.IsFloat():
|
||
|
|
buf[0] = scalarKindFloat
|
||
|
|
binary.LittleEndian.PutUint64(buf[9:], math.Float64bits(s.Float()))
|
||
|
|
default:
|
||
|
|
buf[0] = scalarKindInt
|
||
|
|
binary.LittleEndian.PutUint64(buf[1:], uint64(s.Int()))
|
||
|
|
}
|
||
|
|
return buf
|
||
|
|
}
|
||
|
|
|
||
|
|
func sameShape(a, b []int) bool {
|
||
|
|
if len(a) != len(b) {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
for i := range a {
|
||
|
|
if a[i] != b[i] {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
func decodeScalar(data []byte) (core.Scalar, error) {
|
||
|
|
if len(data) != scalarFrameLen {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: a scalar answer of %d bytes is not the fixed %d", len(data), scalarFrameLen)
|
||
|
|
}
|
||
|
|
switch data[0] {
|
||
|
|
case scalarKindInt:
|
||
|
|
return core.IntScalar(int64(binary.LittleEndian.Uint64(data[1:]))), nil
|
||
|
|
case scalarKindFloat:
|
||
|
|
return core.FloatScalar(math.Float64frombits(binary.LittleEndian.Uint64(data[9:]))), nil
|
||
|
|
case scalarKindComplex:
|
||
|
|
re := math.Float64frombits(binary.LittleEndian.Uint64(data[17:]))
|
||
|
|
im := math.Float64frombits(binary.LittleEndian.Uint64(data[25:]))
|
||
|
|
return core.ComplexScalar(complex(re, im)), nil
|
||
|
|
}
|
||
|
|
return core.Scalar{}, base.Errf("spmd: a scalar answer of kind %d is none of the library's", data[0])
|
||
|
|
}
|
||
|
|
|
||
|
|
// broadcastScalar carries the root's answer to every rank.
|
||
|
|
func (w *World) broadcastScalar(s core.Scalar, root int) (core.Scalar, error) {
|
||
|
|
if w.size == 1 {
|
||
|
|
return s, nil
|
||
|
|
}
|
||
|
|
if w.rank == root {
|
||
|
|
wire := encodeScalar(s)
|
||
|
|
for r := range w.size {
|
||
|
|
if r == root {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if err := w.sendTo(r, tagShardsWhole, wire); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return s, nil
|
||
|
|
}
|
||
|
|
data, err := w.recvFrom(root, tagShardsWhole)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
got, err := decodeScalar(data)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
return got, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Reduce folds the ranks' same-shaped arrays together and answers on
|
||
|
|
// the root alone: element e of the answer is the fold of the ranks'
|
||
|
|
// elements e in rank index order, left to right. The fold's order is
|
||
|
|
// the program's own rank order, never the arrival order, so the same
|
||
|
|
// program and data answer the same bits on one machine and across a
|
||
|
|
// cluster. Every other rank receives nil.
|
||
|
|
func (w *World) Reduce(a *core.Array, op Op, root int) (*core.Array, error) {
|
||
|
|
if err := w.status(); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if err := checkRoot(root, w.size); err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
if err := checkArrayOp(a, op); err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
if w.size == 1 {
|
||
|
|
return a, nil
|
||
|
|
}
|
||
|
|
k := w.size
|
||
|
|
n := a.Len()
|
||
|
|
bounds := func(c int) (int, int) { return c * n / k, (c + 1) * n / k }
|
||
|
|
// Every rank hands its chunk c to the rank that owns c.
|
||
|
|
for c := range k {
|
||
|
|
if c == w.rank {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
lo, hi := bounds(c)
|
||
|
|
wire, err := encodePart(nil, a, []int{hi - lo}, lo, hi-lo)
|
||
|
|
if err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
if err := w.sendTo(c, tagReduce, wire); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// The fold runs as the chunks land: the rank's own chunk seeds the
|
||
|
|
// accumulator, rank 0's chunk replaces it when it arrives, and
|
||
|
|
// every later chunk joins the running accumulator at its rank's
|
||
|
|
// turn. The chain of operands is chunk 0 left to chunk k-1 right,
|
||
|
|
// exactly the rank order it always was, so the bits cannot move,
|
||
|
|
// while folding one chunk overlaps the transfer of the next and
|
||
|
|
// no rank piles the whole exchange up before it computes.
|
||
|
|
lo, hi := bounds(w.rank)
|
||
|
|
ownWire, err := encodePart(nil, a, []int{hi - lo}, lo, hi-lo)
|
||
|
|
if err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
own, err := decodeWire(ownWire)
|
||
|
|
if err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
acc := own
|
||
|
|
for i := range k {
|
||
|
|
chunk := own
|
||
|
|
if i != w.rank {
|
||
|
|
data, rerr := w.recvFrom(i, tagReduce)
|
||
|
|
if rerr != nil {
|
||
|
|
return nil, rerr
|
||
|
|
}
|
||
|
|
chunk, rerr = decodeWire(data)
|
||
|
|
if rerr != nil {
|
||
|
|
return nil, w.fail(rerr)
|
||
|
|
}
|
||
|
|
if chunk.Dtype() != own.Dtype() {
|
||
|
|
return nil, w.fail(base.Errf("spmd: rank %d's %s disagrees with rank %d's %s",
|
||
|
|
i, chunk.Dtype(), w.rank, own.Dtype()))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if i == 0 {
|
||
|
|
acc = chunk // the chain starts at rank 0's chunk
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
var ferr error
|
||
|
|
switch op {
|
||
|
|
case Sum:
|
||
|
|
acc, ferr = core.Add(acc, chunk)
|
||
|
|
case Min:
|
||
|
|
acc, ferr = core.Minimum(acc, chunk)
|
||
|
|
case Max:
|
||
|
|
acc, ferr = core.Maximum(acc, chunk)
|
||
|
|
case Any:
|
||
|
|
acc, ferr = core.Or(acc, chunk)
|
||
|
|
case All:
|
||
|
|
acc, ferr = core.And(acc, chunk)
|
||
|
|
case Prod:
|
||
|
|
acc, ferr = core.Mul(acc, chunk)
|
||
|
|
}
|
||
|
|
if ferr != nil {
|
||
|
|
return nil, w.fail(ferr)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// The folded chunk travels back to the root, which joins the
|
||
|
|
// world's chunks in rank order and reshapes to the input's shape.
|
||
|
|
flat, err := w.gather(acc, root)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return core.Reshape(flat, a.Shape()...)
|
||
|
|
}
|
||
|
|
|
||
|
|
// AllReduce is Reduce with the answer on every rank: one world, one
|
||
|
|
// answer, identical bits everywhere.
|
||
|
|
func (w *World) AllReduce(a *core.Array, op Op) (*core.Array, error) {
|
||
|
|
whole, err := w.Reduce(a, op, 0)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return w.Broadcast(whole, 0)
|
||
|
|
}
|
||
|
|
|
||
|
|
// checkArrayOp refuses the combinations that carry no arithmetic
|
||
|
|
// before any frame moves: Bool has no sum and no ordering, complex has
|
||
|
|
// no ordering, and Any with All are Bool's alone.
|
||
|
|
func checkArrayOp(a *core.Array, op Op) error {
|
||
|
|
dt := a.Dtype()
|
||
|
|
switch op {
|
||
|
|
case Sum:
|
||
|
|
if dt == core.Bool {
|
||
|
|
return base.Errf("spmd: sum of Bool arrays has no arithmetic; Any and All are Bool's reductions")
|
||
|
|
}
|
||
|
|
case Prod:
|
||
|
|
if dt == core.Bool {
|
||
|
|
return base.Errf("spmd: prod of Bool arrays has no arithmetic; Any and All are Bool's reductions")
|
||
|
|
}
|
||
|
|
case Min, Max:
|
||
|
|
if dt == core.Bool || dt == core.Complex {
|
||
|
|
return base.Errf("spmd: %s of dtype %s has no ordering", op, dt)
|
||
|
|
}
|
||
|
|
case Any, All:
|
||
|
|
if dt != core.Bool {
|
||
|
|
return base.Errf("spmd: %s needs Bool elements, got %s", op, dt)
|
||
|
|
}
|
||
|
|
default:
|
||
|
|
return base.Errf("spmd: unknown reduction op %s", op)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// The sharded product, norm and dot extend the reduction families to
|
||
|
|
// the vector operations the iterative solvers live on. All three carry
|
||
|
|
// one-dimensional arrays: the canonical partition cuts the axis, the
|
||
|
|
// shards fold their own blocks with the single block kernels, and the
|
||
|
|
// combine is the same tree the single-array answers use.
|
||
|
|
|
||
|
|
// AllReduceNormShards answers the Lp norm of one global array whose
|
||
|
|
// canonical pieces the ranks hold, every rank receiving the answer.
|
||
|
|
// The power sums fold the canonical blocks with the norm fold's own
|
||
|
|
// per-element arithmetic and combine through the same balanced tree,
|
||
|
|
// so the answer is the single-array norm's exact bits at any world
|
||
|
|
// size. p must be finite and positive; the infinity norm is a maximum
|
||
|
|
// and belongs to Max, not here.
|
||
|
|
func (w *World) AllReduceNormShards(local *core.Array, span Span, p float64) (core.Scalar, error) {
|
||
|
|
ans, err := w.reduceNormShards(local, span, p, 0)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
return w.broadcastScalar(ans, 0)
|
||
|
|
}
|
||
|
|
|
||
|
|
// ReduceNormShards is AllReduceNormShards with the answer on the root
|
||
|
|
// alone; every other rank receives nil.
|
||
|
|
func (w *World) ReduceNormShards(local *core.Array, span Span, p float64, root int) (*core.Scalar, error) {
|
||
|
|
if err := checkRoot(root, w.size); err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
ans, err := w.reduceNormShards(local, span, p, root)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return &ans, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *World) reduceNormShards(local *core.Array, span Span, p float64, root int) (core.Scalar, error) {
|
||
|
|
if err := w.status(); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
if math.IsNaN(p) || p <= 0 || math.IsInf(p, 1) {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: Norm shards need a finite positive p, got %v", p))
|
||
|
|
}
|
||
|
|
if err := w.checkOneDimSpan(local, span); err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
if dt := local.Dtype(); dt != core.Float && dt != core.Float32 && dt != core.Float16 && dt != core.Int {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: norm of dtype %s is not supported; convert with Astype", dt))
|
||
|
|
}
|
||
|
|
if span.Global == 0 {
|
||
|
|
return core.FloatScalar(0), nil
|
||
|
|
}
|
||
|
|
sums, err := w.foldNormBlocks(local, span, p)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
if err := w.sendTo(root, tagShardsValues, encodeBlockValues(blockValues{first: blockFirst(span), kind: kindSumF64, f: sums})); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
return core.Scalar{}, nil
|
||
|
|
}
|
||
|
|
all := make([][]float64, w.size)
|
||
|
|
all[w.rank] = sums
|
||
|
|
for r := range w.size {
|
||
|
|
if r == root {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
data, err := w.recvFrom(r, tagShardsValues)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
bv, err := decodeBlockValues(data)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
if bv.kind != kindSumF64 {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: norm shards disagree on the payload kind"))
|
||
|
|
}
|
||
|
|
all[r] = bv.f
|
||
|
|
}
|
||
|
|
total := make([]float64, 0, core.FoldParts(span.Global))
|
||
|
|
for _, part := range all {
|
||
|
|
total = append(total, part...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(core.NormRoot(core.TreeSum(total), p)), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// foldNormBlocks computes the rank's own blocks' power sums with the
|
||
|
|
// norm fold's per-element arithmetic, dtype by dtype.
|
||
|
|
func (w *World) foldNormBlocks(local *core.Array, span Span, p float64) ([]float64, error) {
|
||
|
|
first, last, _ := spanBlocks(span)
|
||
|
|
if first < 0 || last < 0 || first > last {
|
||
|
|
return nil, base.Errf("spmd: rank %d's span [%d, %d) is not on the partition's block boundaries",
|
||
|
|
w.rank, span.Lo, span.Hi)
|
||
|
|
}
|
||
|
|
n := span.Global
|
||
|
|
count := last - first
|
||
|
|
sums := make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
switch local.Dtype() {
|
||
|
|
case core.Float:
|
||
|
|
sums[c-first] = core.FoldNormPower(sliceBlock(local.RawFloats()[:local.Len()], n, first, c), p)
|
||
|
|
case core.Float32:
|
||
|
|
sums[c-first] = core.FoldNormPowerF32(sliceBlock(local.RawFloat32s()[:local.Len()], n, first, c), p)
|
||
|
|
case core.Float16:
|
||
|
|
sums[c-first] = core.FoldNormPowerF16(sliceBlock(local.RawHalves()[:local.Len()], n, first, c), p)
|
||
|
|
case core.Int:
|
||
|
|
sums[c-first] = core.FoldNormPowerI64(sliceBlock(local.RawInts()[:local.Len()], n, first, c), p)
|
||
|
|
default:
|
||
|
|
return nil, base.Errf("spmd: norm of dtype %s is not supported; convert with Astype", local.Dtype())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return sums, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// blockFirst names the partition block a span starts on.
|
||
|
|
func blockFirst(span Span) int {
|
||
|
|
first, _, _ := spanBlocks(span)
|
||
|
|
return first
|
||
|
|
}
|
||
|
|
|
||
|
|
// AllReduceDotShards answers the dot product of two global arrays the
|
||
|
|
// ranks hold as the same canonical pieces, every rank receiving the
|
||
|
|
// answer. The shards fold their own blocks with the single block dot
|
||
|
|
// kernel and the block values combine through the same balanced tree,
|
||
|
|
// so the answer is the single-array Dot's exact bits at any world
|
||
|
|
// size. The two arrays must carry the same dtype.
|
||
|
|
func (w *World) AllReduceDotShards(x, y *core.Array, span Span) (core.Scalar, error) {
|
||
|
|
ans, err := w.reduceDotShards(x, y, span, 0)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
return w.broadcastScalar(ans, 0)
|
||
|
|
}
|
||
|
|
|
||
|
|
// ReduceDotShards is AllReduceDotShards with the answer on the root
|
||
|
|
// alone; every other rank receives nil.
|
||
|
|
func (w *World) ReduceDotShards(x, y *core.Array, span Span, root int) (*core.Scalar, error) {
|
||
|
|
if err := checkRoot(root, w.size); err != nil {
|
||
|
|
return nil, w.fail(err)
|
||
|
|
}
|
||
|
|
ans, err := w.reduceDotShards(x, y, span, root)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return &ans, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *World) reduceDotShards(x, y *core.Array, span Span, root int) (core.Scalar, error) {
|
||
|
|
if err := w.status(); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
if err := w.checkOneDimSpan(x, span); err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
if y.NDim() != 1 || y.Len() != span.Len() {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: the second shard leads with %d elements against the span's %d", y.Len(), span.Len()))
|
||
|
|
}
|
||
|
|
if x.Dtype() != y.Dtype() {
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: Dot shards carry %s against %s; convert with Astype", x.Dtype(), y.Dtype()))
|
||
|
|
}
|
||
|
|
switch x.Dtype() {
|
||
|
|
case core.Int:
|
||
|
|
case core.Float, core.Float32, core.Float16, core.Complex:
|
||
|
|
default:
|
||
|
|
return core.Scalar{}, w.fail(base.Errf("spmd: dot of dtype %s is not supported; convert with Astype", x.Dtype()))
|
||
|
|
}
|
||
|
|
if span.Global == 0 {
|
||
|
|
if x.Dtype() == core.Int {
|
||
|
|
return core.IntScalar(0), nil
|
||
|
|
}
|
||
|
|
if x.Dtype() == core.Complex {
|
||
|
|
return core.ComplexScalar(0), nil
|
||
|
|
}
|
||
|
|
return core.FloatScalar(0), nil
|
||
|
|
}
|
||
|
|
bv, err := w.foldDotBlocks(x, y, span)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
if w.rank != root {
|
||
|
|
if err := w.sendTo(root, tagShardsValues, encodeBlockValues(bv)); err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
return core.Scalar{}, nil
|
||
|
|
}
|
||
|
|
all := make([]blockValues, w.size)
|
||
|
|
all[w.rank] = bv
|
||
|
|
for r := range w.size {
|
||
|
|
if r == root {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
data, err := w.recvFrom(r, tagShardsValues)
|
||
|
|
if err != nil {
|
||
|
|
return core.Scalar{}, err
|
||
|
|
}
|
||
|
|
if all[r], err = decodeBlockValues(data); err != nil {
|
||
|
|
return core.Scalar{}, w.fail(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return w.combineDot(all, span)
|
||
|
|
}
|
||
|
|
|
||
|
|
// foldDotBlocks computes the rank's own blocks' dot values with the
|
||
|
|
// single block dot kernel, dtype by dtype.
|
||
|
|
func (w *World) foldDotBlocks(x, y *core.Array, span Span) (blockValues, error) {
|
||
|
|
first, last, _ := spanBlocks(span)
|
||
|
|
if first < 0 || last < 0 || first > last {
|
||
|
|
return blockValues{}, base.Errf("spmd: rank %d's span [%d, %d) is not on the partition's block boundaries",
|
||
|
|
w.rank, span.Lo, span.Hi)
|
||
|
|
}
|
||
|
|
n := span.Global
|
||
|
|
count := last - first
|
||
|
|
bv := blockValues{first: first}
|
||
|
|
if count == 0 {
|
||
|
|
return bv, nil
|
||
|
|
}
|
||
|
|
switch x.Dtype() {
|
||
|
|
case core.Float:
|
||
|
|
bv.kind = kindSumF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldDot(
|
||
|
|
sliceBlock(x.RawFloats()[:x.Len()], n, first, c),
|
||
|
|
sliceBlock(y.RawFloats()[:y.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Float32:
|
||
|
|
bv.kind = kindSumF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldDotF32(
|
||
|
|
sliceBlock(x.RawFloat32s()[:x.Len()], n, first, c),
|
||
|
|
sliceBlock(y.RawFloat32s()[:y.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Float16:
|
||
|
|
bv.kind = kindSumF64
|
||
|
|
bv.f = make([]float64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.f[c-first] = core.FoldDotF16(
|
||
|
|
sliceBlock(x.RawHalves()[:x.Len()], n, first, c),
|
||
|
|
sliceBlock(y.RawHalves()[:y.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Complex:
|
||
|
|
bv.kind = kindSumC128
|
||
|
|
bv.c = make([]complex128, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.c[c-first] = core.FoldDotC128(
|
||
|
|
sliceBlock(x.RawComplexes()[:x.Len()], n, first, c),
|
||
|
|
sliceBlock(y.RawComplexes()[:y.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
case core.Int:
|
||
|
|
bv.kind = kindSumI64
|
||
|
|
bv.i = make([]int64, count)
|
||
|
|
for c := first; c < last; c++ {
|
||
|
|
bv.i[c-first] = core.FoldDotI64(
|
||
|
|
sliceBlock(x.RawInts()[:x.Len()], n, first, c),
|
||
|
|
sliceBlock(y.RawInts()[:y.Len()], n, first, c))
|
||
|
|
}
|
||
|
|
default:
|
||
|
|
return blockValues{}, base.Errf("spmd: dot of dtype %s is not supported; convert with Astype", x.Dtype())
|
||
|
|
}
|
||
|
|
return bv, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// combineDot folds the shards' dot values through the balanced tree.
|
||
|
|
func (w *World) combineDot(all []blockValues, span Span) (core.Scalar, error) {
|
||
|
|
first, _, parts := spanBlocks(span)
|
||
|
|
if first < 0 {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: the root's span is not on the partition's boundaries")
|
||
|
|
}
|
||
|
|
covered := 0
|
||
|
|
for r, bv := range all {
|
||
|
|
covered += bv.blocks()
|
||
|
|
if bv.kind == kindNone {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if bv.first != r*parts/w.size {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: rank %d's piece starts at block %d, the partition puts it at %d", r, bv.first, r*parts/w.size)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if covered != parts {
|
||
|
|
return core.Scalar{}, base.Errf("spmd: the shards' %d blocks do not cover the partition's %d", covered, parts)
|
||
|
|
}
|
||
|
|
kind, kindErr := payloadKind(all)
|
||
|
|
if kindErr != nil {
|
||
|
|
return core.Scalar{}, kindErr
|
||
|
|
}
|
||
|
|
switch kind {
|
||
|
|
case kindSumI64:
|
||
|
|
vals := make([]int64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.i...)
|
||
|
|
}
|
||
|
|
return core.IntScalar(core.TreeSum(vals)), nil
|
||
|
|
case kindSumC128:
|
||
|
|
vals := make([]complex128, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.c...)
|
||
|
|
}
|
||
|
|
return core.ComplexScalar(core.TreeSum(vals)), nil
|
||
|
|
case kindSumF64:
|
||
|
|
vals := make([]float64, 0, parts)
|
||
|
|
for _, bv := range all {
|
||
|
|
vals = append(vals, bv.f...)
|
||
|
|
}
|
||
|
|
return core.FloatScalar(core.TreeSum(vals)), nil
|
||
|
|
default:
|
||
|
|
return core.Scalar{}, base.Errf("spmd: dot shards disagree on the payload kind")
|
||
|
|
}
|
||
|
|
}
|