1 import Base 2 import ../storage/model.bend as Model 3 4 # One parallel build for every operation whose output cells are independent: 5 # elementwise, matmul and reduce_to. An operation gives the build a cell 6 # function: `cell(p, position, state, inside)` returns the F32 of output 7 # cell `position` and hands its input state back, and `split(state)` copies 8 # the state for the two halves of a fork. `p = prepare(q)` holds what every 9 # cell needs, such as an index map; each subtree below the forks prepares 10 # it once. The state is a pair `A & B`, such as two input arrays. The cells 11 # `base .. base + 2^depth` form a balanced tree. The top `splits` levels fork, each half with its own copy of the 12 # state, and each subtree below the forks fills its cells in index order on 13 # one thread. `cell`, `split` and `prepare` are templates, so each operation 14 # compiles its own copy of the build; the cell is still called, not inlined. 15 16 def leaf(~S: Type, pair: F32 & S) -> Array & S: 17 (value, state) = pair 18 (ALeaf{value}, state) 19 20 def join(~S: Type, a: Array, high: Array & S) -> Array & S: 21 (b, state) = high 22 (ANode{a, b}, state) 23 24 # The two components of the state go to `rec` one by one, so a closure 25 # call never boxes the state. 26 def node(~A: Type, ~B: Type, rec: A -> B -> Array & (A & B), low: Array & (A & B)) -> Array & (A & B): 27 (a, state) = low 28 (l, r) = state 29 join(~(A & B), a, rec(l, r)) 30 31 # The first cell of the high half of a node whose halves hold 2^inner cells. 32 def upper(+inner: Nat, +base: Nat) -> Nat: 33 Nat.add(Model.capacity(inner), base) 34 35 # The depth-1 tree, cells `base` and `base + 1`: the second cell reads the 36 # state the first hands back, and the node of the two leaves is built here. 37 # The bottom level of every deeper tree runs this, so no leaf returns 38 # through `leaf` and half of the tree's nodes create no closure. The second 39 # cell is `Nat.add(1n, base)`, which is `upper(0n, base)` but a flat add. 40 def pair_join(~S: Type, x: F32, high: F32 & S) -> Array & S: 41 (y, state) = high 42 (ANode{ALeaf{x}, ALeaf{y}}, state) 43 44 def pair(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B), 45 +p: P, +count: Nat, +base: Nat, low: F32 & (A & B)) -> Array & (A & B): 46 (x, state) = low 47 (l, r) = state 48 pair_join(~(A & B), x, cell(p, Nat.add(1n, base), (l, r), Nat.is_lt(Nat.add(1n, base), count))) 49 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. 52 def 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 & (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)) 64 65 def fork(~A: Type, ~B: Type, pair: (A & B) & (A & B), low: A -> B -> Array, high: A -> B -> Array) -> Array: 66 (first, second) = pair 67 (l1, r1) = first 68 (l2, r2) = second 69 a b = low(l1, r1) high(l2, r2) 70 ANode{a, b} 71 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. 74 def 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: 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, A & B, tree(~P, ~A, ~B, ~cell, prepare(q), depth, count, base, state)) 85 86 # A run of the one-thread tree next to its proof of balance: `result` is 87 # the term the code computes, and its first component is balanced at 88 # `depth`. The proofs only follow the code, so they cost nothing at runtime. 89 type Certified<-S: Type, -depth: Nat, -expected: Array & S> is Type: 90 Certified{result: Array & S, same: {result == expected : Array & S}, 91 balanced: Model.Balanced(F32, depth, Pair.fst(Array, S, result))} 92 93 def leaf_certified(~S: Type, pair: F32 & S) -> Certified: 94 (value, state) = pair 95 Certified{leaf(~S, (value, state)), {==}, Unit{}} 96 97 def node_joined(~A: Type, ~B: Type, +inner: Nat, -rec: A -> B -> Array & (A & B), -low: Array & (A & B), 98 a: Array, -l: A, -r: B, same: {(a, (l, r)) == low : Array & (A & B)}, balanced: Model.Balanced(F32, inner, a), 99 high: Certified<(A & B), inner, rec(l, r)>) -> Certified<(A & B), 1n+inner, node(~A, ~B, rec, low)>: 100 Certified{result, high_same, high_balanced} = high 101 (b, rest) = result 102 Certified{(ANode{a, b}, rest), 103 %same : {(ANode{a, b}, rest) == node(~A, ~B, rec, _) : Array & (A & B)} 104 %high_same : {(ANode{a, b}, rest) == join(~(A & B), a, _) : Array & (A & B)} 105 {==}, 106 (balanced, high_balanced)} 107 108 def node_split(~A: Type, ~B: Type, +inner: Nat, -rec: A -> B -> Array & (A & B), -low: Array & (A & B), 109 result: Array & (A & B), same: {result == low : Array & (A & B)}, 110 balanced: Model.Balanced(F32, inner, Pair.fst(Array, A & B, result)), 111 next: @l: A -> @r: B -> Certified<(A & B), inner, rec(l, r)>) -> Certified<(A & B), 1n+inner, node(~A, ~B, rec, low)>: 112 (a, state) = result 113 (l, r) = state 114 node_joined(~A, ~B, inner, rec, low, a, l, r, same, balanced, next(l, r)) 115 116 def node_certified(~A: Type, ~B: Type, +inner: Nat, -rec: A -> B -> Array & (A & B), -low: Array & (A & B), 117 certified: Certified<(A & B), inner, low>, next: @l: A -> @r: B -> Certified<(A & B), inner, rec(l, r)>) -> 118 Certified<(A & B), 1n+inner, node(~A, ~B, rec, low)>: 119 Certified{result, same, balanced} = certified 120 node_split(~A, ~B, inner, rec, low, result, same, balanced, next) 121 122 def pair_joined(~S: Type, x: F32, high: F32 & S) -> Certified: 123 (y, rest) = high 124 Certified{pair_join(~S, x, (y, rest)), {==}, (Unit{}, Unit{})} 125 126 def pair_certified(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B), 127 +p: P, +count: Nat, +base: Nat, low: F32 & (A & B)) -> Certified<(A & B), 1n, pair(~P, ~A, ~B, ~cell, p, count, base, low)>: 128 (x, state) = low 129 (l, r) = state 130 pair_joined(~(A & B), x, cell(p, Nat.add(1n, base), (l, r), Nat.is_lt(Nat.add(1n, base), count))) 131 132 def tree_certified(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B), 133 +p: P, +depth: Nat, +count: Nat, +base: Nat, state: A & B) -> 134 Certified<(A & B), depth, tree(~P, ~A, ~B, ~cell, p, depth, count, base, state)>: 135 match depth: 136 case 0n: 137 leaf_certified(~(A & B), cell(p, base, state, Nat.is_lt(base, count))) 138 case 1n+inner: 139 match inner: 140 case 0n: 141 pair_certified(~P, ~A, ~B, ~cell, p, count, base, cell(p, base, state, Nat.is_lt(base, count))) 142 case 1n+below: 143 node_certified(~A, ~B, inner, l => r => tree(~P, ~A, ~B, ~cell, p, inner, count, upper(inner, base), (l, r)), 144 tree(~P, ~A, ~B, ~cell, p, inner, count, base, state), 145 tree_certified(~P, ~A, ~B, ~cell, p, inner, count, base, state), 146 l => r => tree_certified(~P, ~A, ~B, ~cell, p, inner, count, upper(inner, base), (l, r))) 147 148 # The one-thread tree is balanced at its depth. 149 def tree_balanced(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> A & B -> Bool -> F32 & (A & B), 150 +p: P, +depth: Nat, +count: Nat, +base: Nat, -state: A & B, 151 certified: Certified<(A & B), depth, tree(~P, ~A, ~B, ~cell, p, depth, count, base, state)>) -> 152 Model.Balanced(F32, depth, Pair.fst(Array, A & B, tree(~P, ~A, ~B, ~cell, p, depth, count, base, state))): 153 Certified{result, same, balanced} = certified 154 %same : Model.Balanced(F32, depth, Pair.fst(Array, A & B, _)) 155 balanced 156 157 def fork_balanced(~A: Type, ~B: Type, +inner: Nat, pair: (A & B) & (A & B), -low: A -> B -> Array, -high: A -> B -> Array, 158 low_valid: @l: A -> @r: B -> Model.Balanced(F32, inner, low(l, r)), 159 high_valid: @l: A -> @r: B -> Model.Balanced(F32, inner, high(l, r))) -> 160 Model.Balanced(F32, 1n+inner, fork(~A, ~B, pair, low, high)): 161 (first, second) = pair 162 (l1, r1) = first 163 (l2, r2) = second 164 (low_valid(l1, r1), high_valid(l2, r2)) 165 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. 168 def 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))