src/tensor/contracts/product_proof.bend

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