src/tensor/matmul_parallel_proof.bend

bendtorch-core@0.1.0.0 · 287 lines · 20 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
3import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
4import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
5import bendtorch-core@0.1.0.0/src/tensor/matrix_index_proof.bend as MatrixIndex
6import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
7import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
8import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
9import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic
10
11# Proofs that the parallel matmul build writes, at every cell of its output
12# tree, the value the cell's dot names: for two balanced arrays A and B, a
13# cell below `count` holds `DotFold.dot_fold` of the words A and B hold at
14# the cell's row and column, in increasing k from +0.0, and a cell above it
15# holds +0.0. Every read keeps the array it reads and every clone is two
16# copies of it, so the whole build reads A and B.
17
18# The factor pairs the terms `cursor .. cursor + fuel` of cell (row, column)
19# read: left[row * inner + k] and right[k * columns + column].
20def LogicalArr(+fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
21  +row: Nat, +column: Nat, +cursor: Nat, values: List<&2, DotFold.Factors>) -> Data:
22  match fuel values:
23    case 0n Nil{}:
24      Unit
25    case 1n+tail DotFold.Factors{lv, rv} <> rest:
26      ArrStep<tail, dl, dr, A, B, inner, columns, row, column, cursor, lv, rv, rest>
27    case 0n DotFold.Factors{lv, rv} <> rest:
28      Empty
29    case 1n+tail Nil{}:
30      Empty
31
32type ArrStep<-fuel: Nat, -dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -inner: Nat, -columns: Nat,
33  -row: Nat, -column: Nat, -cursor: Nat, -lv: F32, -rv: F32, -values: List<&2, DotFold.Factors>> is Data:
34  ArrStep{-left_observed: {Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))) == Some{lv} : Maybe<&2, F32>},
35    -right_observed: {Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))) == Some{rv} : Maybe<&2, F32>},
36    -remaining: LogicalArr(fuel, dl, dr, A, B, inner, columns, row, column, 1n+cursor, values)}
37
38# A run of `cell_dot`: its total folds the factor pairs it read, and both
39# arrays come back unchanged.
40type DotRun<-fuel: Nat, -dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -inner: Nat, -columns: Nat,
41  -row: Nat, -column: Nat, -cursor: Nat, -accumulator: F32, -result: F32 & (Array<F32> & Array<F32>)> is Data:
42  DotRun{-values: List<&2, DotFold.Factors>, -logical: LogicalArr(fuel, dl, dr, A, B, inner, columns, row, column, cursor, values),
43    -fold: {Pair.fst(F32, Array<F32> & Array<F32>, result) == DotFold.dot_fold(values, accumulator) : F32},
44    -left_kept: {Pair.fst(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, result)) == A : Array<F32>},
45    -right_kept: {Pair.snd(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, result)) == B : Array<F32>}}
46
47def dot_extend(+rest: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
48  +row: Nat, +column: Nat, +cursor: Nat, -accumulator: F32, -x: F32, -y: F32,
49  -ex: {Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))) == Some{x} : Maybe<&2, F32>},
50  -ey: {Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))) == Some{y} : Maybe<&2, F32>},
51  -result: F32 & (Array<F32> & Array<F32>),
52  -run: Parallel.Fact<DotRun<rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, F32.add(accumulator, F32.mul(x, y)), result>>) ->
53  Parallel.Fact<DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, result>>:
54  Parallel.Fact{{\{DotRun: values => logical => fold => left_kept => right_kept =>
55    DotRun{DotFold.Factors{x, y} <> values, ArrStep{ex, ey, logical}, fold, left_kept, right_kept}} :
56      DotRun<rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, F32.add(accumulator, F32.mul(x, y)), result> ->
57      DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, result>}(
58    {\{Parallel.Fact: p => p} : Parallel.Fact<DotRun<rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, F32.add(accumulator, F32.mul(x, y)), result>> ->
59      DotRun<rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, F32.add(accumulator, F32.mul(x, y)), result>}(run))}
60
61# The last term of a run: the reads hold the words of A and B, and both
62# arrays come back.
63def run_last(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
64  +row: Nat, +column: Nat, +cursor: Nat, -accumulator: F32,
65  -la: Array<F32>, -ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>}, -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B),
66  pl: Array<F32> & F32, -exact_l: {pl == Array.get(F32, la, Matrix.left_index(inner, row, cursor)) : Array<F32> & F32},
67  pr: Array<F32> & F32, -exact_r: {pr == Array.get(F32, ra, Matrix.right_index(columns, column, cursor)) : Array<F32> & F32}) ->
68  Parallel.Fact<DotRun<1n, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, Matrix.product_run(0n, 1n, columns, Nat.add(Nat.mul(row, inner), cursor), Nat.add(Nat.mul(cursor, columns), column), accumulator, pl, pr)>>:
69  (kl, x) = pl
70  (kr, y) = pr
71  -ex = Equal.sym(Maybe<&2, F32>, Some{x}, Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))),
72    {\{Parallel.Fact: p => p} : Parallel.Fact<({Some{Refinement.second(F32, (kl, x))} == Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))) : Maybe<&2, F32>})> ->
73      {Some{Refinement.second(F32, (kl, x))} == Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))) : Maybe<&2, F32>}}(
74      Parallel.read_value(dl, la, A, Matrix.left_index(inner, row, cursor), (kl, x), exact_l, el, left_valid)))
75  -ey = Equal.sym(Maybe<&2, F32>, Some{y}, Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))),
76    {\{Parallel.Fact: p => p} : Parallel.Fact<({Some{Refinement.second(F32, (kr, y))} == Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))) : Maybe<&2, F32>})> ->
77      {Some{Refinement.second(F32, (kr, y))} == Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))) : Maybe<&2, F32>}}(
78      Parallel.read_value(dr, ra, B, Matrix.right_index(columns, column, cursor), (kr, y), exact_r, er, right_valid)))
79  -kl_kept = {\{Parallel.Fact: p => p} : Parallel.Fact<({Refinement.first(F32, (kl, x)) == A : Array<F32>})> -> {Refinement.first(F32, (kl, x)) == A : Array<F32>}}(
80    Parallel.read_keeps(dl, la, A, Matrix.left_index(inner, row, cursor), (kl, x), exact_l, el, left_valid))
81  -kr_kept = {\{Parallel.Fact: p => p} : Parallel.Fact<({Refinement.first(F32, (kr, y)) == B : Array<F32>})> -> {Refinement.first(F32, (kr, y)) == B : Array<F32>}}(
82    Parallel.read_keeps(dr, ra, B, Matrix.right_index(columns, column, cursor), (kr, y), exact_r, er, right_valid))
83  dot_extend(0n, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, x, y, ex, ey, (F32.add(accumulator, F32.mul(x, y)), (kl, kr)),
84    Parallel.Fact{DotRun{Nil{}, Unit{}, {==}, kl_kept, kr_kept}})
85
86# The offsets of the next term are the stepped ones.
87def run_stepped(+rest: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
88  +row: Nat, +column: Nat, +cursor: Nat, -value: F32, -l: Array<F32>, -r: Array<F32>,
89  run: Parallel.Fact<DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, value, Matrix.product_run(rest, 1n, columns, Nat.add(Nat.mul(row, inner), 1n+cursor), Nat.add(Nat.mul(1n+cursor, columns), column), value, Array.get(F32, l, Matrix.left_index(inner, row, 1n+cursor)), Array.get(F32, r, Matrix.right_index(columns, column, 1n+cursor)))>>) ->
90  Parallel.Fact<DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, value, Matrix.product_run(rest, 1n, columns, Nat.add(1n, Nat.add(Nat.mul(row, inner), cursor)), Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column)), value, Array.get(F32, l, U32.from_nat(Nat.add(1n, Nat.add(Nat.mul(row, inner), cursor)))), Array.get(F32, r, U32.from_nat(Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column)))))>>:
91  %Equal.sym(Nat, Nat.add(1n, Nat.add(Nat.mul(row, inner), cursor)), Nat.add(Nat.mul(row, inner), 1n+cursor), MatrixIndex.stepped_left(inner, row, cursor)) :
92    Parallel.Fact<DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, value, Matrix.product_run(rest, 1n, columns, _, Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column)), value, Array.get(F32, l, U32.from_nat(_)), Array.get(F32, r, U32.from_nat(Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column)))))>>
93  %Equal.sym(Nat, Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column)), Nat.add(Nat.mul(1n+cursor, columns), column), MatrixIndex.stepped_right(columns, column, cursor)) :
94    Parallel.Fact<DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, value, Matrix.product_run(rest, 1n, columns, Nat.add(Nat.mul(row, inner), 1n+cursor), _, value, Array.get(F32, l, U32.from_nat(Nat.add(Nat.mul(row, inner), 1n+cursor))), Array.get(F32, r, U32.from_nat(_)))>>
95  run
96
97# A term with more after it: its reads hold the words of A and B, and the
98# run goes on from the next term's reads of the arrays they hand back.
99def run_next(+rest: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
100  +row: Nat, +column: Nat, +cursor: Nat, -accumulator: F32,
101  -la: Array<F32>, -ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>}, -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B),
102  pl: Array<F32> & F32, -exact_l: {pl == Array.get(F32, la, Matrix.left_index(inner, row, cursor)) : Array<F32> & F32},
103  pr: Array<F32> & F32, -exact_r: {pr == Array.get(F32, ra, Matrix.right_index(columns, column, cursor)) : Array<F32> & F32},
104  -next: @value: F32 -> @-l: Array<F32> -> @-r: Array<F32> -> @-el: {l == A : Array<F32>} -> @-er: {r == B : Array<F32>} ->
105    Parallel.Fact<DotRun<1n+rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, value, Matrix.product_run(rest, 1n, columns, Nat.add(Nat.mul(row, inner), 1n+cursor), Nat.add(Nat.mul(1n+cursor, columns), column), value, Array.get(F32, l, Matrix.left_index(inner, row, 1n+cursor)), Array.get(F32, r, Matrix.right_index(columns, column, 1n+cursor)))>>) ->
106  Parallel.Fact<DotRun<2n+rest, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, Matrix.product_run(1n+rest, 1n, columns, Nat.add(Nat.mul(row, inner), cursor), Nat.add(Nat.mul(cursor, columns), column), accumulator, pl, pr)>>:
107  (kl, x) = pl
108  (kr, y) = pr
109  -ex = Equal.sym(Maybe<&2, F32>, Some{x}, Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))),
110    {\{Parallel.Fact: p => p} : Parallel.Fact<({Some{Refinement.second(F32, (kl, x))} == Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))) : Maybe<&2, F32>})> ->
111      {Some{Refinement.second(F32, (kl, x))} == Model.observe(F32, dl, A, Refinement.route(dl, Matrix.left_index(inner, row, cursor))) : Maybe<&2, F32>}}(
112      Parallel.read_value(dl, la, A, Matrix.left_index(inner, row, cursor), (kl, x), exact_l, el, left_valid)))
113  -ey = Equal.sym(Maybe<&2, F32>, Some{y}, Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))),
114    {\{Parallel.Fact: p => p} : Parallel.Fact<({Some{Refinement.second(F32, (kr, y))} == Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))) : Maybe<&2, F32>})> ->
115      {Some{Refinement.second(F32, (kr, y))} == Model.observe(F32, dr, B, Refinement.route(dr, Matrix.right_index(columns, column, cursor))) : Maybe<&2, F32>}}(
116      Parallel.read_value(dr, ra, B, Matrix.right_index(columns, column, cursor), (kr, y), exact_r, er, right_valid)))
117  -kl_kept = {\{Parallel.Fact: p => p} : Parallel.Fact<({Refinement.first(F32, (kl, x)) == A : Array<F32>})> -> {Refinement.first(F32, (kl, x)) == A : Array<F32>}}(
118    Parallel.read_keeps(dl, la, A, Matrix.left_index(inner, row, cursor), (kl, x), exact_l, el, left_valid))
119  -kr_kept = {\{Parallel.Fact: p => p} : Parallel.Fact<({Refinement.first(F32, (kr, y)) == B : Array<F32>})> -> {Refinement.first(F32, (kr, y)) == B : Array<F32>}}(
120    Parallel.read_keeps(dr, ra, B, Matrix.right_index(columns, column, cursor), (kr, y), exact_r, er, right_valid))
121  dot_extend(1n+rest, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, x, y, ex, ey,
122    Matrix.product_run(rest, 1n, columns, Nat.add(1n, Nat.add(Nat.mul(row, inner), cursor)), Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column)), F32.add(accumulator, F32.mul(x, y)), Array.get(F32, kl, U32.from_nat(Nat.add(1n, Nat.add(Nat.mul(row, inner), cursor)))), Array.get(F32, kr, U32.from_nat(Nat.add(columns, Nat.add(Nat.mul(cursor, columns), column))))),
123    run_stepped(rest, dl, dr, A, B, inner, columns, row, column, cursor, F32.add(accumulator, F32.mul(x, y)), kl, kr, next(F32.add(accumulator, F32.mul(x, y)))(kl)(kr)(kl_kept)(kr_kept)))
124
125# A run of fuel + 1 terms from cursor folds the factor pairs it read, and
126# both arrays come back.
127def run_eq(+fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
128  +row: Nat, +column: Nat, +cursor: Nat, -accumulator: F32,
129  -la: Array<F32>, -ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>}, -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B),
130  pl: Array<F32> & F32, -exact_l: {pl == Array.get(F32, la, Matrix.left_index(inner, row, cursor)) : Array<F32> & F32},
131  pr: Array<F32> & F32, -exact_r: {pr == Array.get(F32, ra, Matrix.right_index(columns, column, cursor)) : Array<F32> & F32}) ->
132  Parallel.Fact<DotRun<1n+fuel, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, Matrix.product_run(fuel, 1n, columns, Nat.add(Nat.mul(row, inner), cursor), Nat.add(Nat.mul(cursor, columns), column), accumulator, pl, pr)>>:
133  match fuel:
134    case 0n:
135      run_last(dl, dr, A, B, inner, columns, row, column, cursor, accumulator, la, ra, el, er, left_valid, right_valid, pl, exact_l, pr, exact_r)
136    case 1n+rest:
137      run_next(rest, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, la, ra, el, er, left_valid, right_valid, pl, exact_l, pr, exact_r,
138        value => l => r => el => er => run_eq(rest, dl, dr, A, B, inner, columns, row, column, 1n+cursor, value, l, r, el, er, left_valid, right_valid,
139          Array.get(F32, l, Matrix.left_index(inner, row, 1n+cursor)), {==}, Array.get(F32, r, Matrix.right_index(columns, column, 1n+cursor)), {==}))
140
141def dot_run(fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
142  +row: Nat, +column: Nat, +cursor: Nat, -accumulator: F32, la: Array<F32>, ra: Array<F32>,
143  -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>},
144  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
145  Parallel.Fact<DotRun<fuel, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, Matrix.cell_dot(fuel, inner, columns, row, column, cursor, accumulator, la, ra)>>:
146  match fuel:
147    case 0n:
148      Parallel.Fact{DotRun{Nil{}, Unit{}, {==}, el, er}}
149    case 1n+rest:
150      run_eq(rest, dl, dr, A, B, inner, columns, row, column, cursor, accumulator, la, ra, el, er, left_valid, right_valid,
151        Array.get(F32, la, Matrix.left_index(inner, row, cursor)), {==}, Array.get(F32, ra, Matrix.right_index(columns, column, cursor)), {==})
152
153# What output cell `position` holds: +0.0 at or above `count`; below it, a
154# word that folds the factor pairs of row `position / columns` and column
155# `position % columns`, in increasing k from +0.0.
156type CellAt<-dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -inner: Nat, -columns: Nat, -position: Nat,
157  -value: Maybe<&2, F32>> is Data:
158  CellAt{-word: F32, -observed: {value == Some{word} : Maybe<&2, F32>}, -values: List<&2, DotFold.Factors>,
159    -logical: LogicalArr(inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, values),
160    -fold: {word == DotFold.dot_fold(values, 0.0) : F32}}
161
162def CellSpec(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat, +position: Nat, inside: Bool,
163  -value: Maybe<&2, F32>) -> Type:
164  match inside:
165    case False{}:
166      {value == Some{0.0} : Maybe<&2, F32>}
167    case True{}:
168      CellAt<dl, dr, A, B, inner, columns, position, value>
169
170def CellIs(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat, +count: Nat, +position: Nat,
171  -value: Maybe<&2, F32>) -> Type:
172  CellSpec(dl, dr, A, B, inner, columns, position, Nat.is_lt(position, count), value)
173
174# Every slot of `out`, at index `base + offset`, holds what CellIs names.
175def Cells(+depth: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat, +count: Nat, +base: Nat,
176  -out: Array<F32>) -> Type:
177  @+path: Model.Path(depth) ->
178    CellIs(dl, dr, A, B, inner, columns, count, Nat.add(Model.offset(depth, path), base), Model.observe(F32, depth, out, path))
179
180# The matmul cell function of `Parallel.build`, from the input arrays.
181def cell_of(+inner: Nat, +columns: Nat, +position: Nat, A: Array<F32>, B: Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
182  Matrix.cell_at(Matrix.dims(Matrix.Dims{inner, columns}), position, (A, B), inside)
183
184def eta(-X: Type, -Y: Type, pair: X & Y) -> {pair == (Pair.fst(X, Y, pair), Pair.snd(X, Y, pair)) : X & Y}:
185  (x, y) = pair
186  {==}
187
188# Every cell hands both input arrays back.
189def cell_keeps(+dl: Nat, +dr: Nat, A: Array<F32>, B: Array<F32>, +inner: Nat, +columns: Nat, +position: Nat, inside: Bool,
190  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
191  Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, inside)) == (A, B) : Array<F32> & Array<F32>}>:
192  match inside:
193    case False{}:
194      Parallel.Fact{{==}}
195    case True{}:
196      Parallel.Fact{{\{DotRun: values => logical => fold => left_kept => right_kept =>
197        Equal.trans(Array<F32> & Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{})),
198          (Pair.fst(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{}))),
199            Pair.snd(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{})))),
200          (A, B),
201          eta(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{}))),
202          Parallel.pair_equal(Pair.fst(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{}))),
203            Pair.snd(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{}))),
204            A, B, left_kept, right_kept))} :
205        DotRun<inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0,
206          cell_of(inner, columns, position, A, B, True{})> ->
207        {Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{})) == (A, B) : Array<F32> & Array<F32>}}(
208        {\{Parallel.Fact: p => p} :
209          Parallel.Fact<DotRun<inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0,
210            cell_of(inner, columns, position, A, B, True{})>> ->
211          DotRun<inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0,
212            cell_of(inner, columns, position, A, B, True{})>}(
213          dot_run(inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0, A, B,
214            {==}, {==}, left_valid, right_valid)))}
215
216# What the cell function computes is what CellSpec names.
217def cell_spec(+dl: Nat, +dr: Nat, A: Array<F32>, B: Array<F32>, +inner: Nat, +columns: Nat, +position: Nat, inside: Bool,
218  -observed: Maybe<&2, F32>, -computed: {observed == Some{Pair.fst(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, inside))} : Maybe<&2, F32>},
219  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
220  Parallel.Fact<CellSpec(dl, dr, A, B, inner, columns, position, inside, observed)>:
221  match inside:
222    case False{}:
223      Parallel.Fact{computed}
224    case True{}:
225      Parallel.Fact{{\{DotRun: values => logical => fold => left_kept => right_kept =>
226        CellAt{Pair.fst(F32, Array<F32> & Array<F32>, cell_of(inner, columns, position, A, B, True{})), computed, values, logical, fold}} :
227        DotRun<inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0,
228          cell_of(inner, columns, position, A, B, True{})> ->
229        CellAt<dl, dr, A, B, inner, columns, position, observed>}(
230        {\{Parallel.Fact: p => p} :
231          Parallel.Fact<DotRun<inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0,
232            cell_of(inner, columns, position, A, B, True{})>> ->
233          DotRun<inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0,
234            cell_of(inner, columns, position, A, B, True{})>}(
235          dot_run(inner, dl, dr, A, B, inner, columns, Matrix.row_of(columns, position), Matrix.column_of(columns, position), 0n, 0.0, A, B,
236            {==}, {==}, left_valid, right_valid)))}
237
238# Whatever the split, every cell of the parallel build holds what CellIs
239# names: the generic law of `Parallel.build`, with the matmul cell.
240def build_cells(+depth: Nat, +splits: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
241  +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>},
242  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
243  Parallel.Fact<Cells(depth, dl, dr, A, B, inner, columns, count, base, Matrix.build(depth, splits, inner, columns, count, base, la, ra))>:
244  -split = {\{Parallel.Fact: q => q} : Parallel.Fact<({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>} &
245    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>})> ->
246    {Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>} &
247    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}}(
248    Parallel.split_equal(dl, dr, A, B, left_valid, right_valid))
249  -cells = {\{Generic.Fact: q => q} :
250    Generic.Fact<Generic.Cells(~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.cell_at, Matrix.dims(Matrix.Dims{inner, columns}), (A, B), depth, count, base,
251      Matrix.build(depth, splits, inner, columns, count, base, la, ra))> ->
252    Generic.Cells(~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.cell_at, Matrix.dims(Matrix.Dims{inner, columns}), (A, B), depth, count, base,
253      Matrix.build(depth, splits, inner, columns, count, base, la, ra))}(
254    Generic.build_cells(~Matrix.Dims, ~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.dims, ~Elementwise.split, ~Matrix.cell_at,
255      Matrix.Dims{inner, columns}, (A, B), depth, splits, count, base, (la, ra), Parallel.pair_equal(la, ra, A, B, el, er),
256      at => inside => {\{Parallel.Fact: q => q} :
257        Parallel.Fact<({Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, at, A, B, inside)) == (A, B) : Array<F32> & Array<F32>})> ->
258        {Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, at, A, B, inside)) == (A, B) : Array<F32> & Array<F32>}}(
259        cell_keeps(dl, dr, A, B, inner, columns, at, inside, left_valid, right_valid)),
260      Pair.fst({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>},
261        {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}, split),
262      Pair.snd({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>},
263        {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}, split)))
264  Parallel.Fact{path => {\{Parallel.Fact: q => q} :
265    Parallel.Fact<CellSpec(dl, dr, A, B, inner, columns, Nat.add(Model.offset(depth, path), base), Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
266      Model.observe(F32, depth, Matrix.build(depth, splits, inner, columns, count, base, la, ra), path))> ->
267    CellSpec(dl, dr, A, B, inner, columns, Nat.add(Model.offset(depth, path), base), Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
268      Model.observe(F32, depth, Matrix.build(depth, splits, inner, columns, count, base, la, ra), path))}(
269    cell_spec(dl, dr, A, B, inner, columns, Nat.add(Model.offset(depth, path), base), Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
270      Model.observe(F32, depth, Matrix.build(depth, splits, inner, columns, count, base, la, ra), path), cells(path), left_valid, right_valid))}
271
272# A cell named at an equal index, and a CellIs below the count is a CellAt.
273def cell_shift(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat, +count: Nat,
274  -position: Nat, -other: Nat, -value: Maybe<&2, F32>, -same: {position == other : Nat},
275  -at: CellIs(dl, dr, A, B, inner, columns, count, position, value)) ->
276  Parallel.Fact<CellIs(dl, dr, A, B, inner, columns, count, other, value)>:
277  Parallel.Fact{
278    %same : CellIs(dl, dr, A, B, inner, columns, count, _, value)
279    at}
280
281def cell_below(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat, +count: Nat,
282  -position: Nat, -value: Maybe<&2, F32>, -below: {Nat.is_lt(position, count) == True{} : Bool},
283  -at: CellIs(dl, dr, A, B, inner, columns, count, position, value)) ->
284  Parallel.Fact<CellAt<dl, dr, A, B, inner, columns, position, value>>:
285  Parallel.Fact{
286    %below : CellSpec(dl, dr, A, B, inner, columns, position, _, value)
287    at}