def balanced

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

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model

166# The parallel build returns a balanced tree of its depth, whatever the
167# split. It is the certificate the output storage carries, erased at runtime.
168def balanced(~Q: Data, ~P: Data, ~A: Type, ~B: Type, ~prepare: @+q: Q -> P, ~split: A & B -> (A & B) & (A & B),
169  ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B),
170  +q: Q, +depth: Nat, +splits: Nat, +count: Nat, +base: Nat, state: A & B) ->
171  Model.Balanced(F32, depth, build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, splits, count, base, state)):
172  match depth splits:
173    case 1n+inner 1n+rest:
174      +high = upper(inner, base)
175      fork_balanced(~A, ~B, inner, split(state),
176        l => r => build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, base, (l, r)),
177        l => r => build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, high, (l, r)),
178        l => r => balanced(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, base, (l, r)),
179        l => r => balanced(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, high, (l, r)))
180    case 0n _:
181      tree_balanced(~P, ~A, ~B, ~cell, prepare(q), 0n, count, base, state,
182        tree_certified(~P, ~A, ~B, ~cell, prepare(q), 0n, count, base, state))
183    case 1n+inner 0n:
184      tree_balanced(~P, ~A, ~B, ~cell, prepare(q), 1n+inner, count, base, state,
185        tree_certified(~P, ~A, ~B, ~cell, prepare(q), 1n+inner, count, base, state))