src/tensor/contracts/matmul_proof.bend

bendtorch-core@0.1.0.0 · 398 lines · 27 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
4import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
5import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
6import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
7import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
8import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
9import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
10import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
11import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
12import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
13import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
14import bendtorch-core@0.1.0.0/src/tensor/contracts/fold_records.bend as Records
15import bendtorch-core@0.1.0.0/src/tensor/contracts/dot_proof.bend as DotProof
16import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Laws
17import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as Par
18import bendtorch-core@0.1.0.0/src/tensor/read_contract.bend as Reads
19import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
20import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
21import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
22import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
23
24type MatrixFacts<-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32> is Data:
25  MatrixFacts{left_dimensions: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
26    right_dimensions: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
27    output_dimensions: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
28    compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}}
29
30def Facts(-left: Shape.Shape, -right: Shape.Shape, result: Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>) -> Data:
31  match result:
32    case Fail{error}:
33      Unit
34    case Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}}:
35      MatrixFacts<left, right, output, rows, inner, right_inner, columns>
36
37def facts_output(-left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
38  le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
39  re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
40  compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
41  result: Shape.ConstructResult<[rows, columns]>) ->
42  Facts(left, right, Matrix.plan_output(left, right, rows, inner, right_inner, columns, le, re, compatible, result)):
43  match result:
44    case Shape.SError{error}:
45      Unit{}
46    case Shape.SBuilt{shape, preserved}:
47      MatrixFacts{le, re, preserved, compatible}
48
49def facts_compatible(-left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
50  le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
51  re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
52  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))>) ->
53  Facts(left, right, Matrix.plan_compatible(left, right, rows, inner, right_inner, columns, le, re, decision)):
54  match decision:
55    case Shape.BFalse{equal}:
56      Unit{}
57    case Shape.BTrue{equal}:
58      facts_output(left, right, rows, inner, right_inner, columns, le, re,
59        Lifecycle.natural_equal(U32.to_nat(inner), U32.to_nat(right_inner), equal), Shape.construct([rows, columns]))
60
61def facts_right(-left: Shape.Shape, +right: Shape.Shape, +rows: U32, +inner: U32,
62  le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}, dimensions: List<&2, U32>,
63  same: {dimensions == Shape.dimensions(right) : List<&2, U32>}) ->
64  Facts(left, right, Matrix.plan_right(left, right, rows, inner, le, dimensions, same)):
65  match dimensions:
66    case +right_inner <> +columns <> Nil{}:
67      facts_compatible(left, right, rows, inner, right_inner, columns, le,
68        Equal.sym(List<&2, U32>, [right_inner, columns], Shape.dimensions(right), same),
69        Shape.bool_case(Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))))
70    case Nil{}:
71      Unit{}
72    case head <> Nil{}:
73      Unit{}
74    case first <> second <> third <> rest:
75      Unit{}
76
77def facts_left(+left: Shape.Shape, +right: Shape.Shape, dimensions: List<&2, U32>,
78  same: {dimensions == Shape.dimensions(left) : List<&2, U32>}) ->
79  Facts(left, right, Matrix.plan_left(left, right, dimensions, same)):
80  match dimensions:
81    case +rows <> +inner <> Nil{}:
82      facts_right(left, right, rows, inner, Equal.sym(List<&2, U32>, [rows, inner], Shape.dimensions(left), same),
83        Shape.dimensions(right), {==})
84    case Nil{}:
85      Unit{}
86    case head <> Nil{}:
87      Unit{}
88    case first <> second <> third <> rest:
89      Unit{}
90
91def plan_facts(+left: Shape.Shape, +right: Shape.Shape, +plan: Matrix.Plan<left, right>,
92  planned: {Matrix.plan(left, right) == Done{plan} : Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>}) ->
93  Facts(left, right, Done{plan}):
94  %planned : Facts(left, right, _)
95  facts_left(left, right, Shape.dimensions(left), {==})
96
97def lt_from(+a: Nat, +b: Nat, below: {Nat.is_lt(a, b) == True{} : Bool}) -> Shape.LT(a, b):
98  match a b:
99    case 0n 0n:
100      Empty.absurd(Shape.LT(0n, 0n), Arith.clash(below))
101    case 0n 1n+q:
102      Unit{}
103    case 1n+p 0n:
104      Empty.absurd(Shape.LT(1n+p, 0n), Arith.clash(below))
105    case 1n+p 1n+q:
106      lt_from(p, q, below)
107
108def lt_bound(-a: Nat, -bound: Nat, -other: Nat, same: {bound == other : Nat}, below: {Nat.is_lt(a, bound) == True{} : Bool}) ->
109  {Nat.is_lt(a, other) == True{} : Bool}:
110  %same : {Nat.is_lt(a, _) == True{} : Bool}
111  below
112
113def row_of(+columns: Nat, +position: Nat) -> Nat:
114  Nat.div(position, Nat.mul(columns, 1n))
115
116def column_of(+columns: Nat, +position: Nat) -> Nat:
117  Nat.div(Nat.mod(position, Nat.mul(columns, 1n)), 1n)
118
119def div_one(+number: Nat) -> {Nat.div(number, 1n) == number : Nat}:
120  %Equal.trans(Nat, Nat.add(Nat.mul(number, 1n), 0n), Nat.mul(number, 1n), number,
121    Conversion.add_zero(Nat.mul(number, 1n)), Nats.mul_one_right(number)) : {Nat.div(_, 1n) == number : Nat}
122  Equal.cong(Nat & Nat, Nat, pair => Pair.fst(Nat, Nat, pair), Nat.divmod(Nat.add(Nat.mul(number, 1n), 0n), 1n), (number, 0n),
123    Equal.sym(Nat & Nat, (number, 0n), Nat.divmod(Nat.add(Nat.mul(number, 1n), 0n), 1n), Nats.divmod_exact(number, 1n, 0n, Unit{})))
124
125def row_at(+row: Nat, +column: Nat, +columns: Nat, below: {Nat.is_lt(column, columns) == True{} : Bool}) ->
126  {row_of(columns, Nat.add(Nat.mul(row, columns), column)) == row : Nat}:
127  %Equal.sym(Nat, Nat.mul(columns, 1n), columns, Nats.mul_one_right(columns)) :
128    {Nat.div(Nat.add(Nat.mul(row, columns), column), _) == row : Nat}
129  Equal.cong(Nat & Nat, Nat, pair => Pair.fst(Nat, Nat, pair), Nat.divmod(Nat.add(Nat.mul(row, columns), column), columns), (row, column),
130    Equal.sym(Nat & Nat, (row, column), Nat.divmod(Nat.add(Nat.mul(row, columns), column), columns),
131      Nats.divmod_exact(row, columns, column, lt_from(column, columns, below))))
132
133def column_at(+row: Nat, +column: Nat, +columns: Nat, below: {Nat.is_lt(column, columns) == True{} : Bool}) ->
134  {column_of(columns, Nat.add(Nat.mul(row, columns), column)) == column : Nat}:
135  %Equal.sym(Nat, Nat.mul(columns, 1n), columns, Nats.mul_one_right(columns)) :
136    {Nat.div(Nat.mod(Nat.add(Nat.mul(row, columns), column), _), 1n) == column : Nat}
137  %Equal.cong(Nat & Nat, Nat, pair => Pair.snd(Nat, Nat, pair), (row, column), Nat.divmod(Nat.add(Nat.mul(row, columns), column), columns),
138    Nats.divmod_exact(row, columns, column, lt_from(column, columns, below))) : {Nat.div(_, 1n) == column : Nat}
139  div_one(column)
140
141def output_dimensions(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
142  facts: MatrixFacts<left, right, output, rows, inner, right_inner, columns>) -> {Shape.dimensions(output) == [rows, columns] : List<&2, U32>}:
143  MatrixFacts{le, re, oe, compatible} = facts
144  oe
145
146def cell_below(+output: Shape.Shape, +rows: U32, +columns: U32, +oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
147  +row: Nat, +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
148  +column: Nat, +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool}) ->
149  {Nat.is_lt(Nat.add(Nat.mul(row, U32.to_nat(columns)), column), Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}:
150  %Equal.sym(List<&2, U32>, Shape.dimensions(output), [rows, columns], oe) :
151    {Nat.is_lt(Nat.add(Nat.mul(row, U32.to_nat(columns)), column), Shape.logical_numel(_)) == True{} : Bool}
152  %Equal.sym(Nat, Nat.mul(U32.to_nat(columns), 1n), U32.to_nat(columns), Nats.mul_one_right(U32.to_nat(columns))) :
153    {Nat.is_lt(Nat.add(Nat.mul(row, U32.to_nat(columns)), column), Nat.mul(U32.to_nat(rows), _)) == True{} : Bool}
154  DotProof.lt_mul_add(row, U32.to_nat(rows), column, U32.to_nat(columns), row_below, column_below)
155
156def kept_value(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +history: List<&2, StorageLaws.Op>,
157  -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, history)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
158  +index: U32, inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool},
159  same: {Replayed.value(shape, fill, history, index) == Replayed.value(shape, fill, ops, index) : F32}) ->
160  {Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, history), fits}}, index)) ==
161    Some{Replayed.value(shape, fill, ops, index)} : Maybe<&2, F32>}:
162  Equal.trans(Maybe<&2, F32>,
163    Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, history), fits}}, index)),
164    Some{Replayed.value(shape, fill, history, index)}, Some{Replayed.value(shape, fill, ops, index)},
165    PointwiseProof.replayed_point(shape, fill, history, fits, index, inside),
166    Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(shape, fill, history, index), Replayed.value(shape, fill, ops, index), same))
167
168def left_dimensions(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
169  facts: MatrixFacts<left, right, output, rows, inner, right_inner, columns>) -> {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}:
170  MatrixFacts{le, re, oe, compatible} = facts
171  le
172
173def right_dimensions(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
174  facts: MatrixFacts<left, right, output, rows, inner, right_inner, columns>) -> {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}:
175  MatrixFacts{le, re, oe, compatible} = facts
176  re
177
178def inner_compatible(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -rows: U32, -inner: U32, -right_inner: U32, -columns: U32,
179  facts: MatrixFacts<left, right, output, rows, inner, right_inner, columns>) -> {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}:
180  MatrixFacts{le, re, oe, compatible} = facts
181  compatible
182
183# The parallel product over replayed inputs, in transport form: P holds for
184# the product if it holds for both inputs replayed with histories that keep
185# their words, next to the output replayed with the writes of every cell.
186def ParallelFinal(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +inner: U32, +columns: U32,
187  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
188  -s: Storage.Storage, -P: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type) -> Type:
189  @+left_kept: List<&2, StorageLaws.Op> ->
190  @left_values: (@+other: U32 -> {Replayed.value(left, left_fill, left_kept, other) == Replayed.value(left, left_fill, left_ops, other) : F32}) ->
191  @left_words: {Laws.words(left, left_fill, left_kept, Shape.logical_numel(Shape.dimensions(left))) ==
192    Laws.words(left, left_fill, left_ops, Shape.logical_numel(Shape.dimensions(left))) : List<&2, F32>} ->
193  @-left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_kept)) == Shape.logical_numel(Shape.dimensions(left)) : Nat} ->
194  @+right_kept: List<&2, StorageLaws.Op> ->
195  @right_values: (@+other: U32 -> {Replayed.value(right, right_fill, right_kept, other) == Replayed.value(right, right_fill, right_ops, other) : F32}) ->
196  @right_words: {Laws.words(right, right_fill, right_kept, Shape.logical_numel(Shape.dimensions(right))) ==
197    Laws.words(right, right_fill, right_ops, Shape.logical_numel(Shape.dimensions(right))) : List<&2, F32>} ->
198  @-right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_kept)) == Shape.logical_numel(Shape.dimensions(right)) : Nat} ->
199  @-output_fits: {Lifecycle.logical_length(StorageLaws.replay(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), s)) == Shape.logical_numel(Shape.dimensions(output)) : Nat} ->
200    P(Done{(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_fits}},
201      Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_kept), right_fits}},
202      Core.Packed{output, Core.TTensor{StorageLaws.replay(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), s), output_fits}})})
203
204def parallel_lift(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
205  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
206  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
207  -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
208  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
209  +planned: {Matrix.plan(left, right) == Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
210    Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>},
211  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
212  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
213  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
214  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat},
215  -P: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type,
216  h: ParallelFinal(left, right, output, inner, columns, left_fill, left_ops, right_fill, right_ops, StorageLaws.created(Shape.count(output), 0.0), P)) ->
217  P(Matrix.matmul(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}}, Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})):
218  +plan : Matrix.Plan<left, right> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
219  +facts : MatrixFacts<left, right, output, rows, inner, right_inner, columns> = plan_facts(left, right, plan, planned)
220  +le_live = left_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
221  +re_live = right_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
222  +oe_live = output_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
223  +compatible_live = inner_compatible(left, right, output, rows, inner, right_inner, columns, facts)
224  +left_extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(left), [rows, inner], le_live)
225  +right_extents = Equal.trans(List<&2, Nat>, Shape.extents(right), [U32.to_nat(right_inner), U32.to_nat(columns)], [U32.to_nat(inner), U32.to_nat(columns)],
226    Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(right), [right_inner, columns], re_live),
227    Equal.cong(Nat, List<&2, Nat>, n => [n, U32.to_nat(columns)], U32.to_nat(right_inner), U32.to_nat(inner),
228      Equal.sym(Nat, U32.to_nat(inner), U32.to_nat(right_inner), compatible_live)))
229  +output_extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(output), [rows, columns], oe_live)
230  %Equal.sym(Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>, Matrix.plan(left, right), Done{plan}, planned) :
231    P(Matrix.parallel_planned(left, right, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}, _))
232  Replayed.new_lift(output, 0.0,
233    result => P(Matrix.parallel_allocated(left, right, output, U32.to_nat(inner), U32.to_nat(columns), Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}, result)),
234    fo => Owned.created_lift(Shape.count(output), 0.0, Conversion.shape_count_within(output),
235      st => (ParallelFinal(left, right, output, inner, columns, left_fill, left_ops, right_fill, right_ops, st, P) ->
236        @-fo: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(output)) : Nat} ->
237        P(Matrix.parallel_allocated(left, right, output, U32.to_nat(inner), U32.to_nat(columns), Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}, Done{Core.TTensor{st, fo}}))),
238      zero => final => fo =>
239        Par.storage_shift(output, StorageLaws.replay(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output)))),
240          st => k => P(Done{(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}}, Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}, Core.Packed{output, Core.TTensor{st, k}})}),
241          final(left_ops)(other => {==})({==})(left_fits)(right_ops)(other => {==})({==})(right_fits))(fo),
242      d => fits => f => final => fo =>
243        Par.input_take(left, left_fill, left_ops,
244          taken => P(Matrix.parallel_output(output, U32.to_nat(inner), U32.to_nat(columns), d, U32.to_nat(Shape.count(output)), f, fo,
245            taken, Elementwise.take(Core.pack(right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits})))),
246          left_kept => left_values => left_words => ha => dl => fa => lsrc => left_kept_fits =>
247            Par.input_take(right, right_fill, right_ops,
248              taken => P(Matrix.parallel_output(output, U32.to_nat(inner), U32.to_nat(columns), d, U32.to_nat(Shape.count(output)), f, fo,
249                (Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_kept_fits}}, Api.arr(ha, dl, fa)), taken)),
250              right_kept => right_values => right_words => hb => dr => fb => rsrc => right_kept_fits =>
251                Par.storage_shift(output,
252                  StorageLaws.replay(Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), Storage.Owned{d, U32.to_nat(Shape.count(output)), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
253                  Storage.Owned{d, U32.to_nat(Shape.count(output)),
254                    Matrix.build(d, Matrix.product_splits(d, U32.to_nat(inner)), U32.to_nat(inner), U32.to_nat(columns), U32.to_nat(Shape.count(output)), 0n, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)),
255                    Matrix.build_balanced(d, Matrix.product_splits(d, U32.to_nat(inner)), U32.to_nat(inner), U32.to_nat(columns), U32.to_nat(Shape.count(output)), 0n, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)),
256                    f},
257                  Par.output_eq(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb,
258                    lsrc, rsrc, output, U32.to_nat(rows), left_extents, right_extents, output_extents, d, fits, f),
259                  st => k => P(Done{(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_kept_fits}},
260                    Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_kept), right_kept_fits}},
261                    Core.Packed{output, Core.TTensor{st, k}})}),
262                  final(left_kept)(left_values)(left_words)(left_kept_fits)(right_kept)(right_values)(right_words)(right_kept_fits))(fo))(right_fits))(left_fits))(h)(fo))
263
264def parallel_cell(+left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
265  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +output: Shape.Shape,
266  +rows: U32, +inner: U32, +columns: U32, +oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
267  +row: Nat, +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
268  +column: Nat, +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool}) ->
269  {Replayed.value(output, 0.0, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), U32.from_nat(Nat.add(Nat.mul(row, U32.to_nat(columns)), column))) == DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0) : F32}:
270  +position = Nat.add(Nat.mul(row, U32.to_nat(columns)), column)
271  +below = cell_below(output, rows, columns, oe, row, row_below, column, column_below)
272  +exact = PointwiseProof.word_exact(position, output, below)
273  +room : {Nat.is_le(U32.to_nat(Shape.count(output)), Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool} =
274    %Equal.sym(Nat, U32.to_nat(Shape.count(output)), Shape.logical_numel(Shape.dimensions(output)), Reads.count_exact(output)) : {Nat.is_le(_, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}
275    Conversion.le_reflexive(Shape.logical_numel(Shape.dimensions(output)))
276  Equal.trans(F32, Replayed.value(output, 0.0, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), U32.from_nat(position)),
277    Par.cell_word(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(U32.from_nat(position))), DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0),
278    Par.writes_value(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, output, U32.to_nat(Shape.count(output)), U32.from_nat(position), Par.index_inside(output, position, below), room),
279    %Equal.sym(Nat, U32.to_nat(U32.from_nat(position)), position, exact) : {Par.cell_word(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, _) == DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0) : F32}
280    %Equal.sym(Nat, row_of(U32.to_nat(columns), position), row, row_at(row, column, U32.to_nat(columns), column_below)) :
281      {DotFold.dot_fold(Laws.factors(_, column_of(U32.to_nat(columns), position), U32.to_nat(inner), U32.to_nat(columns),
282        left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0) == DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0) : F32}
283    %Equal.sym(Nat, column_of(U32.to_nat(columns), position), column, column_at(row, column, U32.to_nat(columns), column_below)) :
284      {DotFold.dot_fold(Laws.factors(row, _, U32.to_nat(inner), U32.to_nat(columns),
285        left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0) == DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0) : F32}
286    {==})
287
288def parallel_ordered_planned(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
289  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
290  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
291  -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
292  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
293  +planned: {Matrix.plan(left, right) == Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
294    Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>},
295  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
296  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
297  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
298  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
299  Records.MatmulContract<left, right, U32.to_nat(rows), U32.to_nat(columns),
300    row => column => DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns),
301      left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0),
302    other => Replayed.value(left, left_fill, left_ops, other),
303    other => Replayed.value(right, right_fill, right_ops, other),
304    Matrix.matmul(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}}, Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})>:
305  +plan : Matrix.Plan<left, right> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
306  +facts : MatrixFacts<left, right, output, rows, inner, right_inner, columns> = plan_facts(left, right, plan, planned)
307  +oe_live = output_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
308  -contract : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type = result =>
309    Records.MatmulContract<left, right, U32.to_nat(rows), U32.to_nat(columns),
310      row => column => DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns),
311        left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0),
312      other => Replayed.value(left, left_fill, left_ops, other),
313      other => Replayed.value(right, right_fill, right_ops, other), result>
314  parallel_lift(left, right, output, rows, inner, right_inner, columns, le, re, oe, compatible, planned,
315    left_fill, left_ops, left_fits, right_fill, right_ops, right_fits, contract,
316    left_kept => left_values => left_words => a => right_kept => right_values => right_words => b => o => Records.MatmulContract{
317      {==},
318      Equal.cong(List<&2, U32>, Maybe<&2, List<&2, Nat>>, dims => Some{Shape.naturals(dims)},
319        Shape.dimensions(output), [rows, columns], oe_live),
320      row => row_below => column => column_below =>
321        +row_below = row_below
322        +column_below = column_below
323        +position = U32.from_nat(Nat.add(Nat.mul(row, U32.to_nat(columns)), column))
324        Equal.trans(Maybe<&2, F32>,
325          Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output)))), o}}, position)),
326          Some{Replayed.value(output, 0.0, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), position)},
327          Some{DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0)},
328          PointwiseProof.replayed_point(output, 0.0, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), o,
329            position, PointwiseProof.word_inside(Nat.add(Nat.mul(row, U32.to_nat(columns)), column), output,
330              cell_below(output, rows, columns, oe_live, row, row_below, column, column_below))),
331          Equal.cong(F32, Maybe<&2, F32>, word => Some{word},
332            Replayed.value(output, 0.0, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))), position), DotFold.dot_fold(Laws.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(inner)), 0.0),
333            parallel_cell(left, left_fill, left_ops, right, right_fill, right_ops, output, rows, inner, columns, oe_live,
334              row, row_below, column, column_below))),
335      other => inside => kept_value(left, left_fill, left_ops, left_kept, a, other, inside, left_values(other)),
336      other => inside => kept_value(right, right_fill, right_ops, right_kept, b, other, inside, right_values(other))})
337
338def parallel_replayed_planned(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
339  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
340  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
341  -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
342  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
343  +planned: {Matrix.plan(left, right) == Done{Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}} :
344    Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>},
345  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
346  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
347  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
348  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat},
349  -P: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type,
350  h: Laws.MatmulWith(left, right, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible},
351    left_fill, left_ops, right_fill, right_ops, P)) ->
352  P(Matrix.matmul(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}}, Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})):
353  +plan : Matrix.Plan<left, right> = Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible}
354  +facts : MatrixFacts<left, right, output, rows, inner, right_inner, columns> = plan_facts(left, right, plan, planned)
355  +oe_live = output_dimensions(left, right, output, rows, inner, right_inner, columns, facts)
356  parallel_lift(left, right, output, rows, inner, right_inner, columns, le, re, oe, compatible, planned,
357    left_fill, left_ops, left_fits, right_fill, right_ops, right_fits, P,
358    left_kept => left_values => left_words => a => right_kept => right_values => right_words => b => o =>
359      h(left_kept, left_words, a, right_kept, right_words, b, Par.writes(U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, right, right_fill, right_ops, U32.to_nat(Shape.count(output))),
360        row => row_below => column => column_below =>
361          parallel_cell(left, left_fill, left_ops, right, right_fill, right_ops, output, rows, inner, columns, oe_live,
362            row, row_below, column, column_below),
363        o))
364
365def matmul_ordered(+left: Shape.Shape, +right: Shape.Shape, +plan: Matrix.Plan<left, right>,
366  planned: {Matrix.plan(left, right) == Done{plan} : Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>},
367  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
368  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
369  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
370  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}) ->
371  Records.MatmulContract<left, right, Laws.matrix_rows(left, right, plan), Laws.matrix_columns(left, right, plan),
372    row => column => DotFold.dot_fold(Laws.factors(row, column, Laws.matrix_inner(left, right, plan), Laws.matrix_columns(left, right, plan),
373      left, left_fill, left_ops, right, right_fill, right_ops, Laws.matrix_inner(left, right, plan)), 0.0),
374    other => Replayed.value(left, left_fill, left_ops, other),
375    other => Replayed.value(right, right_fill, right_ops, other),
376    Matrix.matmul(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
377      Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})>:
378  match plan:
379    case Matrix.Plan{+output, +rows, +inner, +right_inner, +columns, le, re, oe, compatible}:
380      parallel_ordered_planned(left, right, output, rows, inner, right_inner, columns, le, re, oe, compatible, planned,
381        left_fill, left_ops, left_fits, right_fill, right_ops, right_fits)
382
383# The histories the product loop leaves on its inputs list the same words.
384
385def matmul_replayed(+left: Shape.Shape, +right: Shape.Shape, +plan: Matrix.Plan<left, right>,
386  planned: {Matrix.plan(left, right) == Done{plan} : Result<&2, &2, Matrix.Error, Matrix.Plan<left, right>>},
387  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
388  -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat},
389  +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
390  -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat},
391  -P: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type,
392  h: Laws.MatmulWith(left, right, plan, left_fill, left_ops, right_fill, right_ops, P)) ->
393  P(Matrix.matmul(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
394    Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}})):
395  match plan:
396    case Matrix.Plan{+output, +rows, +inner, +right_inner, +columns, le, re, oe, compatible}:
397      parallel_replayed_planned(left, right, output, rows, inner, right_inner, columns, le, re, oe, compatible, planned,
398        left_fill, left_ops, left_fits, right_fill, right_ops, right_fits, P, h)