src/tensor/parallel_proofs.bend

bendtorch-core@0.1.0.0 · 227 lines · 19 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
3import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
4import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
5import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
6import bendtorch-core@0.1.0.0/src/storage/clone.bend as Clone
7import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Build
8
9# Proofs that `Build.build` holds, at every output cell, what the
10# one-thread tree computes there from the input state `start`, whatever the
11# split. They need two facts of the operation: a cell hands `start` back
12# unchanged (`Keeps`), and `split` copies `start` into both halves. Each
13# operation then proves only what its cell function computes from `start`.
14
15type Fact<-P: Type> is Data:
16  Fact{-proof: P}
17
18# What cell `position` holds: the F32 its cell function computes from
19# `start`.
20def CellIs(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
21  +p: P, -start: (A & B), +count: Nat, +position: Nat, -observed: Maybe<&2, F32>) -> Type:
22  {observed == Some{Pair.fst(F32, (A & B), cell(p, position, start, Nat.is_lt(position, count)))} : Maybe<&2, F32>}
23
24# Every slot of `out`, the cells `base .. base + 2^depth`, holds what CellIs
25# names for its index.
26def Cells(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
27  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -out: Array<F32>) -> Type:
28  @+path: Model.Path(depth) ->
29    CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(depth, path), base), Model.observe(F32, depth, out, path))
30
31# A cell hands its input state back unchanged.
32def Keeps(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
33  +p: P, -start: (A & B)) -> Type:
34  @+position: Nat -> @inside: Bool -> {Pair.snd(F32, (A & B), cell(p, position, start, inside)) == start : (A & B)}
35
36# A run of the one-thread tree: it hands `start` back and fills its cells.
37def TreeFacts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
38  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -result: Array<F32> & (A & B)) -> Type:
39  {Pair.snd(Array<F32>, (A & B), result) == start : (A & B)} & Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result))
40
41def facts_kept(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
42  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -result: Array<F32> & (A & B),
43  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)>) ->
44  Fact<{Pair.snd(Array<F32>, (A & B), result) == start : (A & B)}>:
45  Fact{Pair.fst({Pair.snd(Array<F32>, (A & B), result) == start : (A & B)},
46    Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result)),
47    {\{Fact: q => q} : Fact<(TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result))> ->
48      TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)}(facts))}
49
50def facts_cells(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
51  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -result: Array<F32> & (A & B),
52  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)>) ->
53  Fact<Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result))>:
54  Fact{Pair.snd({Pair.snd(Array<F32>, (A & B), result) == start : (A & B)},
55    Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result)),
56    {\{Fact: q => q} : Fact<(TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result))> ->
57      TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)}(facts))}
58
59def leaf_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
60  +p: P, -start: (A & B), +count: Nat, +base: Nat, pair: F32 & (A & B),
61  -exact: {pair == cell(p, base, start, Nat.is_lt(base, count)) : F32 & (A & B)},
62  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
63  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 0n, count, base, Build.leaf(~(A & B), pair))>:
64  (value, rest) = pair
65  Fact{(Equal.trans((A & B), rest, Pair.snd(F32, (A & B), cell(p, base, start, Nat.is_lt(base, count))), start,
66      Equal.cong(F32 & (A & B), (A & B), q => Pair.snd(F32, (A & B), q), (value, rest), cell(p, base, start, Nat.is_lt(base, count)), exact),
67      keeps(base)(Nat.is_lt(base, count))),
68    path => Equal.cong(F32 & (A & B), Maybe<&2, F32>, q => Some{Pair.fst(F32, (A & B), q)}, (value, rest),
69      cell(p, base, start, Nat.is_lt(base, count)), exact))}
70
71def position_step(+capacity: Nat, +offset: Nat, +base: Nat) ->
72  {Nat.add(offset, Nat.add(capacity, base)) == Nat.add(Nat.add(capacity, offset), base) : Nat}:
73  Equal.trans(Nat, Nat.add(offset, Nat.add(capacity, base)), Nat.add(capacity, Nat.add(offset, base)), Nat.add(Nat.add(capacity, offset), base),
74    Nats.add_swap(offset, capacity, base), Index.add_assoc(capacity, offset, base))
75
76def node_cell(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
77  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -a: Array<F32>, -b: Array<F32>, +path: Model.Path(1n+inner),
78  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)>,
79  -high: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b)>) ->
80  Fact<CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(1n+inner, path), base), Model.observe(F32, 1n+inner, ANode{a, b}, path))>:
81  Model.Step{+bit, +rest} = path
82  match bit:
83    case False{}:
84      Fact{{\{Fact: q => q} : Fact<(Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a))> ->
85        Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)}(low)(rest)}
86    case True{}:
87      Fact{
88        %position_step(Model.capacity(inner), Model.offset(inner, rest), base) :
89          CellIs(~P, ~A, ~B, ~cell, p, start, count, _, Model.observe(F32, inner, b, rest))
90        {\{Fact: q => q} : Fact<(Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b))> ->
91          Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b)}(high)(rest)}
92
93def node_cells(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
94  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -a: Array<F32>, -b: Array<F32>,
95  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)>,
96  -high: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b)>) ->
97  Fact<Cells(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, ANode{a, b})>:
98  Fact{path => {\{Fact: q => q} :
99    Fact<(CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(1n+inner, path), base), Model.observe(F32, 1n+inner, ANode{a, b}, path)))> ->
100    CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(1n+inner, path), base), Model.observe(F32, 1n+inner, ANode{a, b}, path))}(
101      node_cell(~P, ~A, ~B, ~cell, p, start, inner, count, base, a, b, path, low, high))}
102
103def join_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
104  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -a: Array<F32>, high: Array<F32> & (A & B),
105  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)>,
106  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), high)>) ->
107  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, Build.join(~(A & B), a, high))>:
108  (b, rest) = high
109  Fact{(
110    {\{Fact: q => q} : Fact<({Pair.snd(Array<F32>, (A & B), (b, rest)) == start : (A & B)})> -> {Pair.snd(Array<F32>, (A & B), (b, rest)) == start : (A & B)}}(
111      facts_kept(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), (b, rest), facts)),
112    {\{Fact: q => q} : Fact<(Cells(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, ANode{a, b}))> ->
113      Cells(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, ANode{a, b})}(
114      node_cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a, b, low,
115        facts_cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), (b, rest), facts))))}
116
117def node_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
118  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, rec: A -> B -> Array<F32> & (A & B), low: Array<F32> & (A & B),
119  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, inner, count, base, low)>,
120  -next: @l: A -> @r: B -> @-at: {(l, r) == start : (A & B)} ->
121    Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), rec(l, r))>) ->
122  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, Build.node(~A, ~B, rec, low))>:
123  (a, s) = low
124  (l, r) = s
125  join_facts(~P, ~A, ~B, ~cell, p, start, inner, count, base, a, rec(l, r),
126    facts_cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, (a, (l, r)), facts),
127    next(l)(r)({\{Fact: q => q} : Fact<({Pair.snd(Array<F32>, (A & B), (a, (l, r))) == start : (A & B)})> -> {Pair.snd(Array<F32>, (A & B), (a, (l, r))) == start : (A & B)}}(
128      facts_kept(~P, ~A, ~B, ~cell, p, start, inner, count, base, (a, (l, r)), facts))))
129
130# The depth-1 tree: the leaf of cell `base` next to the leaf of the cell
131# after it, which reads the state the first hands back.
132def pair_join_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
133  +p: P, -start: (A & B), +count: Nat, +base: Nat, x: F32, high: F32 & (A & B),
134  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, 0n, count, base, ALeaf{x})>,
135  -exact: {high == cell(p, Nat.add(1n, base), start, Nat.is_lt(Nat.add(1n, base), count)) : F32 & (A & B)},
136  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
137  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n, count, base, Build.pair_join(~(A & B), x, high))>:
138  (y, rest) = high
139  join_facts(~P, ~A, ~B, ~cell, p, start, 0n, count, base, ALeaf{x}, Build.leaf(~(A & B), (y, rest)), low,
140    leaf_facts(~P, ~A, ~B, ~cell, p, start, count, Build.upper(0n, base), (y, rest), exact, keeps))
141
142def pair_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
143  +p: P, -start: (A & B), +count: Nat, +base: Nat, low: F32 & (A & B),
144  -exact: {low == cell(p, base, start, Nat.is_lt(base, count)) : F32 & (A & B)},
145  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
146  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n, count, base, Build.pair(~P, ~A, ~B, ~cell, p, count, base, low))>:
147  (x, s) = low
148  (l, r) = s
149  pair_join_facts(~P, ~A, ~B, ~cell, p, start, count, base, x, cell(p, Nat.add(1n, base), (l, r), Nat.is_lt(Nat.add(1n, base), count)),
150    facts_cells(~P, ~A, ~B, ~cell, p, start, 0n, count, base, Build.leaf(~(A & B), (x, (l, r))),
151      leaf_facts(~P, ~A, ~B, ~cell, p, start, count, base, (x, (l, r)), exact, keeps)),
152    Equal.cong((A & B), F32 & (A & B), q => cell(p, Nat.add(1n, base), q, Nat.is_lt(Nat.add(1n, base), count)), (l, r), start,
153      {\{Fact: q => q} : Fact<({Pair.snd(Array<F32>, (A & B), Build.leaf(~(A & B), (x, (l, r)))) == start : (A & B)})> ->
154        {Pair.snd(Array<F32>, (A & B), Build.leaf(~(A & B), (x, (l, r)))) == start : (A & B)}}(
155        facts_kept(~P, ~A, ~B, ~cell, p, start, 0n, count, base, Build.leaf(~(A & B), (x, (l, r))),
156          leaf_facts(~P, ~A, ~B, ~cell, p, start, count, base, (x, (l, r)), exact, keeps)))),
157    keeps)
158
159# The one-thread tree hands `start` back and fills every cell as CellIs
160# says.
161def tree_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
162  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, state: (A & B), -at: {state == start : (A & B)},
163  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
164  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, Build.tree(~P, ~A, ~B, ~cell, p, depth, count, base, state))>:
165  match depth:
166    case 0n:
167      leaf_facts(~P, ~A, ~B, ~cell, p, start, count, base, cell(p, base, state, Nat.is_lt(base, count)),
168        Equal.cong((A & B), F32 & (A & B), s => cell(p, base, s, Nat.is_lt(base, count)), state, start, at), keeps)
169    case 1n+inner:
170      match inner:
171        case 0n:
172          pair_facts(~P, ~A, ~B, ~cell, p, start, count, base, cell(p, base, state, Nat.is_lt(base, count)),
173            Equal.cong((A & B), F32 & (A & B), s => cell(p, base, s, Nat.is_lt(base, count)), state, start, at), keeps)
174        case 1n+below:
175          node_facts(~P, ~A, ~B, ~cell, p, start, inner, count, base,
176            l => r => Build.tree(~P, ~A, ~B, ~cell, p, inner, count, Build.upper(inner, base), (l, r)),
177            Build.tree(~P, ~A, ~B, ~cell, p, inner, count, base, state),
178            tree_facts(~P, ~A, ~B, ~cell, p, start, inner, count, base, state, at, keeps),
179            l => r => at => tree_facts(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), (l, r), at, keeps))
180
181def half_equal(-S: Type, -start: S, -state: S, -half: S, -split: S -> S & S, -pick: (S & S) -> S,
182  picked: {half == pick(split(state)) : S}, at: {state == start : S}, copied: {pick(split(start)) == start : S}) ->
183  {half == start : S}:
184  Equal.trans(S, half, pick(split(start)), start,
185    Equal.trans(S, half, pick(split(state)), pick(split(start)), picked, Equal.cong(S, S, s => pick(split(s)), state, start, at)),
186    copied)
187
188def fork_cells(~Q: Data, ~P: Data, ~A: Type, ~B: Type, ~prepare: @+q: Q -> P, ~split: (A & B) -> (A & B) & (A & B),
189  ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
190  +q: Q, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -state: (A & B), pair: (A & B) & (A & B),
191  -exact: {pair == split(state) : (A & B) & (A & B)},
192  -low: A -> B -> Array<F32>, -high: A -> B -> Array<F32>, -at: {state == start : (A & B)},
193  -first: {Pair.fst((A & B), (A & B), split(start)) == start : (A & B)}, -second: {Pair.snd((A & B), (A & B), split(start)) == start : (A & B)},
194  -low_cells: @l: A -> @r: B -> @-at: {(l, r) == start : (A & B)} -> Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, inner, count, base, low(l, r))>,
195  -high_cells: @l: A -> @r: B -> @-at: {(l, r) == start : (A & B)} ->
196    Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, inner, count, Build.upper(inner, base), high(l, r))>) ->
197  Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, 1n+inner, count, base, Build.fork(~A, ~B, pair, low, high))>:
198  (one, two) = pair
199  (l1, r1) = one
200  (l2, r2) = two
201  node_cells(~P, ~A, ~B, ~cell, prepare(q), start, inner, count, base, low(l1, r1), high(l2, r2),
202    low_cells(l1)(r1)(half_equal((A & B), start, state, (l1, r1), split, h => Pair.fst((A & B), (A & B), h),
203      Equal.cong((A & B) & (A & B), (A & B), h => Pair.fst((A & B), (A & B), h), ((l1, r1), (l2, r2)), split(state), exact), at, first)),
204    high_cells(l2)(r2)(half_equal((A & B), start, state, (l2, r2), split, h => Pair.snd((A & B), (A & B), h),
205      Equal.cong((A & B) & (A & B), (A & B), h => Pair.snd((A & B), (A & B), h), ((l1, r1), (l2, r2)), split(state), exact), at, second)))
206
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))