def tree

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

import Base

50# The cells `base .. base + 2^depth` as a balanced tree, built in index order
51# on one thread. A cell at or past `count` is padding.
52def tree(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B),
53  +p: P, +depth: Nat, +count: Nat, +base: Nat, state: A & B) -> Array<F32> & (A & B):
54  match depth:
55    case 0n:
56      leaf(~(A & B), cell(p, base, state, Nat.is_lt(base, count)))
57    case 1n+inner:
58      match inner:
59        case 0n:
60          pair(~P, ~A, ~B, ~cell, p, count, base, cell(p, base, state, Nat.is_lt(base, count)))
61        case 1n+below:
62          node(~A, ~B, l => r => tree(~P, ~A, ~B, ~cell, p, inner, count, upper(inner, base), (l, r)),
63            tree(~P, ~A, ~B, ~cell, p, inner, count, base, state))