src/tensor/parallel.bend

bendtorch-core@0.1.0.0 · 185 lines · 20 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/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
16def leaf(~S: Type, pair: F32 & S) -> Array<F32> & S:
17  (value, state) = pair
18  (ALeaf{value}, state)
19
20def join(~S: Type, a: Array<F32>, high: Array<F32> & S) -> Array<F32> & 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.
26def node(~A: Type, ~B: Type, rec: A -> B -> Array<F32> & (A & B), low: Array<F32> & (A & B)) -> Array<F32> & (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.
32def 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.
40def pair_join(~S: Type, x: F32, high: F32 & S) -> Array<F32> & S:
41  (y, state) = high
42  (ANode{ALeaf{x}, ALeaf{y}}, state)
43
44def 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<F32> & (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.
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))
64
65def fork(~A: Type, ~B: Type, pair: (A & B) & (A & B), low: A -> B -> Array<F32>, high: A -> B -> Array<F32>) -> Array<F32>:
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.
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))
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.
89type Certified<-S: Type, -depth: Nat, -expected: Array<F32> & S> is Type:
90  Certified{result: Array<F32> & S, same: {result == expected : Array<F32> & S},
91    balanced: Model.Balanced(F32, depth, Pair.fst(Array<F32>, S, result))}
92
93def leaf_certified(~S: Type, pair: F32 & S) -> Certified<S, 0n, leaf(~S, pair)>:
94  (value, state) = pair
95  Certified{leaf(~S, (value, state)), {==}, Unit{}}
96
97def node_joined(~A: Type, ~B: Type, +inner: Nat, -rec: A -> B -> Array<F32> & (A & B), -low: Array<F32> & (A & B),
98  a: Array<F32>, -l: A, -r: B, same: {(a, (l, r)) == low : Array<F32> & (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<F32> & (A & B)}
104    %high_same : {(ANode{a, b}, rest) == join(~(A & B), a, _) : Array<F32> & (A & B)}
105    {==},
106    (balanced, high_balanced)}
107
108def node_split(~A: Type, ~B: Type, +inner: Nat, -rec: A -> B -> Array<F32> & (A & B), -low: Array<F32> & (A & B),
109  result: Array<F32> & (A & B), same: {result == low : Array<F32> & (A & B)},
110  balanced: Model.Balanced(F32, inner, Pair.fst(Array<F32>, 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
116def node_certified(~A: Type, ~B: Type, +inner: Nat, -rec: A -> B -> Array<F32> & (A & B), -low: Array<F32> & (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
122def pair_joined(~S: Type, x: F32, high: F32 & S) -> Certified<S, 1n, pair_join(~S, x, high)>:
123  (y, rest) = high
124  Certified{pair_join(~S, x, (y, rest)), {==}, (Unit{}, Unit{})}
125
126def 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
132def 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.
149def 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<F32>, 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<F32>, A & B, _))
155  balanced
156
157def fork_balanced(~A: Type, ~B: Type, +inner: Nat, pair: (A & B) & (A & B), -low: A -> B -> Array<F32>, -high: A -> B -> Array<F32>,
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.
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))