def build

bendtorch-core@0.1.0.0 · src/tensor/parallel.bend:72

import Base

72# The parallel build. Every cell gets the F32 the one-thread tree computes
73# there, so no output word depends on `splits` or on the thread count.
74def build(~Q: Data, ~P: Data, ~A: Type, ~B: Type, ~prepare: @+q: Q -> P, ~split: A & B -> (A & B) & (A & B),
75  ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B),
76  +q: Q, +depth: Nat, +splits: Nat, +count: Nat, +base: Nat, state: A & B) -> Array<F32>:
77  match depth splits:
78    case 1n+inner 1n+rest:
79      +high = upper(inner, base)
80      fork(~A, ~B, split(state),
81        l => r => build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, base, (l, r)),
82        l => r => build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, high, (l, r)))
83    case _ _:
84      Pair.fst(Array<F32>, A & B, tree(~P, ~A, ~B, ~cell, prepare(q), depth, count, base, state))