def build_cells

bendtorch-core@0.1.0.0 · src/tensor/parallel_proofs.bend:207

import Base
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Build

207# Whatever the split, every cell of the parallel build holds what the
208# one-thread tree computes there from `start`.
209def build_cells(~Q: Data, ~P: Data, ~A: Type, ~B: Type, ~prepare: @+q: Q -> P, ~split: (A & B) -> (A & B) & (A & B), ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
210  +q: Q, -start: (A & B), +depth: Nat, +splits: Nat, +count: Nat, +base: Nat, state: (A & B), -at: {state == start : (A & B)},
211  -keeps: Keeps(~P, ~A, ~B, ~cell, prepare(q), start),
212  -first: {Pair.fst((A & B), (A & B), split(start)) == start : (A & B)}, -second: {Pair.snd((A & B), (A & B), split(start)) == start : (A & B)}) ->
213  Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, depth, count, base, Build.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, splits, count, base, state))>:
214  match depth splits:
215    case 1n+inner 1n+rest:
216      fork_cells(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, start, inner, count, base, state, split(state), {==},
217        l => r => Build.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, base, (l, r)),
218        l => r => Build.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, Build.upper(inner, base), (l, r)),
219        at, first, second,
220        l => r => at => build_cells(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, start, inner, rest, count, base, (l, r), at, keeps, first, second),
221        l => r => at => build_cells(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, start, inner, rest, count, Build.upper(inner, base), (l, r), at, keeps, first, second))
222    case 0n _:
223      facts_cells(~P, ~A, ~B, ~cell, prepare(q), start, 0n, count, base, Build.tree(~P, ~A, ~B, ~cell, prepare(q), 0n, count, base, state),
224        tree_facts(~P, ~A, ~B, ~cell, prepare(q), start, 0n, count, base, state, at, keeps))
225    case 1n+inner 0n:
226      facts_cells(~P, ~A, ~B, ~cell, prepare(q), start, 1n+inner, count, base, Build.tree(~P, ~A, ~B, ~cell, prepare(q), 1n+inner, count, base, state),
227        tree_facts(~P, ~A, ~B, ~cell, prepare(q), start, 1n+inner, count, base, state, at, keeps))