src/tensor/parallel_LAWS.bend

bendtorch-core@0.1.0.0 · 79 lines · 4 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
3import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
4
5# The laws of the one parallel build that elementwise, matmul and reduce_to
6# share. An operation hands the build a cell function, a split of its input
7# state and a `prepare` that turns the build's parameters `q` into what
8# every cell reads; these laws hold for every such choice, so each operation
9# proves only what its cell function computes.
10
11type Holds<-P: Type> is Data:
12  Holds{-proof: P}
13
14# For every split, the parallel build is a balanced tree of its depth.
15law parallel_build_balanced:
16  for ~Q: Data
17  for ~P: Data
18  for ~A: Type
19  for ~B: Type
20  for ~prepare: @+q: Q -> P
21  for ~split: A & B -> (A & B) & (A & B)
22  for ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B)
23  for +q: Q
24  for +depth: Nat
25  for +splits: Nat
26  for +count: Nat
27  for +base: Nat
28  for state: (A & B)
29  Model.Balanced(F32, depth, Parallel.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, splits, count, base, state))
30
31# If each cell hands its input state back and the split copies the state
32# into both halves, then for every split each slot of the build, cell
33# base + offset, holds the F32 the cell function computes for it from the
34# input state, the padding past `count` included.
35law parallel_build_cells:
36  for ~Q: Data
37  for ~P: Data
38  for ~A: Type
39  for ~B: Type
40  for ~prepare: @+q: Q -> P
41  for ~split: A & B -> (A & B) & (A & B)
42  for ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B)
43  for +q: Q
44  for +depth: Nat
45  for +splits: Nat
46  for +count: Nat
47  for +base: Nat
48  for state: (A & B)
49  for keeps: @+position: Nat -> @inside: Bool -> {Pair.snd(F32, (A & B), cell(prepare(q), position, state, inside)) == state : (A & B)}
50  for first: {Pair.fst((A & B), (A & B), split(state)) == state : (A & B)}
51  for second: {Pair.snd((A & B), (A & B), split(state)) == state : (A & B)}
52  Holds<@+path: Model.Path(depth) ->
53    {Model.observe(F32, depth, Parallel.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, splits, count, base, state), path) ==
54      Some{Pair.fst(F32, (A & B), cell(prepare(q), Nat.add(Model.offset(depth, path), base), state,
55        Nat.is_lt(Nat.add(Model.offset(depth, path), base), count)))} : Maybe<&2, F32>}>
56
57# Under the same two facts, the parallel build and the one-thread build
58# (no split) hold the same F32 in every slot, so no output word depends on
59# the split or on the thread count.
60law parallel_build_sequential:
61  for ~Q: Data
62  for ~P: Data
63  for ~A: Type
64  for ~B: Type
65  for ~prepare: @+q: Q -> P
66  for ~split: A & B -> (A & B) & (A & B)
67  for ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B)
68  for +q: Q
69  for +depth: Nat
70  for +splits: Nat
71  for +count: Nat
72  for +base: Nat
73  for state: (A & B)
74  for keeps: @+position: Nat -> @inside: Bool -> {Pair.snd(F32, (A & B), cell(prepare(q), position, state, inside)) == state : (A & B)}
75  for first: {Pair.fst((A & B), (A & B), split(state)) == state : (A & B)}
76  for second: {Pair.snd((A & B), (A & B), split(state)) == state : (A & B)}
77  Holds<@+path: Model.Path(depth) ->
78    {Model.observe(F32, depth, Parallel.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, splits, count, base, state), path) ==
79      Model.observe(F32, depth, Parallel.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, 0n, count, base, state), path) : Maybe<&2, F32>}>