1	import Base
2	import ../shape/core.bend as Shape
3	import ./core.bend as Core
4	import ./matmul.bend as Matrix
5	import ./math_totality.bend as Totality
6	import ./dot_fold.bend as DotFold
7	
8	def coordinate_observe_result(+shape: Shape.Shape, tensor: Core.Tensor<shape>,
9	  result: Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>) -> Maybe<&2, F32>:
10	  match result:
11	    case Fail{error}:
12	      None{}
13	    case Done{index}:
14	      Core.checked_observe(shape, tensor, Shape.index_value(shape, index))
15	
16	def coordinate_observe(+shape: Shape.Shape, tensor: Core.Tensor<shape>, coordinates: List<&2, U32>) -> Maybe<&2, F32>:
17	  coordinate_observe_result(shape, tensor, Shape.flatten_checked(shape, coordinates))
18	
19	def LogicalDot(+fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
20	  +row: U32, +column: U32, +cursor: Nat, values: List<&2, DotFold.Factors>) -> Data:
21	  match fuel values:
22	    case 0n Nil{}:
23	      Unit
24	    case 1n+tail DotFold.Factors{lv, rv} <> rest:
25	      DotStep<tail, left, right, a, b, row, column, cursor, lv, rv, rest>
26	    case 0n DotFold.Factors{lv, rv} <> rest:
27	      Empty
28	    case 1n+tail Nil{}:
29	      Empty
30	
31	type DotStep<-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
32	  -row: U32, -column: U32, -cursor: Nat, -lv: F32, -rv: F32, -values: List<&2, DotFold.Factors>> is Data:
33	  DotStep{-exact: {U32.to_nat(U32.from_nat(cursor)) == cursor : Nat},
34	    -left_observed: {coordinate_observe(left, a, [row, U32.from_nat(cursor)]) == Some{lv} : Maybe<&2, F32>},
35	    -right_observed: {coordinate_observe(right, b, [U32.from_nat(cursor), column]) == Some{rv} : Maybe<&2, F32>},
36	    -remaining: LogicalDot(fuel, left, right, a, b, row, column, 1n+cursor, values)}
37	
38	type DotMeaning<-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
39	  -row: U32, -column: U32, -cursor: Nat, -accumulator: F32, -ra: Core.Tensor<left>, -rb: Core.Tensor<right>, -value: F32> is Data:
40	  DotMeaning{-values: List<&2, DotFold.Factors>, -logical: LogicalDot(fuel, left, right, a, b, row, column, cursor, values),
41	    -fold: {value == DotFold.dot_fold(values, accumulator) : F32},
42	    -left_kept: @+index: U32 -> {Core.checked_observe(left, ra, index) == Core.checked_observe(left, a, index) : Maybe<&2, F32>},
43	    -right_kept: @+index: U32 -> {Core.checked_observe(right, rb, index) == Core.checked_observe(right, b, index) : Maybe<&2, F32>}}
44	
45	def DotResultMeaning(-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
46	  -row: U32, -column: U32, -cursor: Nat, -accumulator: F32,
47	  result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & F32>) -> Data:
48	  match result:
49	    case Fail{error}:
50	      Unit
51	    case Done{pair}:
52	      (ra, rb, value) = pair
53	      DotMeaning<fuel, left, right, a, b, row, column, cursor, accumulator, ra, rb, value>
54	
55	type SourceDot<-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
56	  -row: U32, -column: U32, -cursor: Nat, -accumulator: F32,
57	  -result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & F32>> is Data:
58	  SourceDot{-meaning: @unit: Unit -> DotResultMeaning(fuel, left, right, a, b, row, column, cursor, accumulator, result)}
59	
60	type ProductCell<-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
61	  -output: Shape.Shape, -out: Core.Tensor<output>, -position: Nat> is Data:
62	  ProductCell{-row: U32, -column: U32, -coordinate: Shape.CheckedCoordinate<output>,
63	    -decoded: {Shape.unflatten_checked(output, U32.from_nat(position)) == Done{coordinate} : Result<&2, &2, Shape.Error, Shape.CheckedCoordinate<output>>},
64	    -coordinates: {Shape.coordinate_values(output, coordinate) == [row, column] : List<&2, U32>},
65	    -word_exact: {U32.to_nat(U32.from_nat(position)) == position : Nat},
66	    -factors: List<&2, DotFold.Factors>, -logical: LogicalDot(U32.to_nat(inner), left, right, a, b, row, column, 0n, factors),
67	    -value: F32, -fold: {value == DotFold.dot_fold(factors, 0.0) : F32},
68	    -checked: Core.ObservationChecked(Core.Packed{output, out}, U32.from_nat(position)),
69	    -observed: {Core.checked_observe(output, out, U32.from_nat(position)) == Some{value} : Maybe<&2, F32>}}
70	
71	type ProductPrefix<-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
72	  -output: Shape.Shape, -ra: Core.Tensor<left>, -rb: Core.Tensor<right>, -out: Core.Tensor<output>, -cursor: Nat> is Data:
73	  ProductPrefix{-left_kept: @+other: U32 -> {Core.checked_observe(left, ra, other) == Core.checked_observe(left, a, other) : Maybe<&2, F32>},
74	    -right_kept: @+other: U32 -> {Core.checked_observe(right, rb, other) == Core.checked_observe(right, b, other) : Maybe<&2, F32>},
75	    -cells: @+position: Nat -> @earlier: {Nat.is_lt(position, cursor) == True{} : Bool} -> ProductCell<inner, left, right, a, b, output, out, position>}
76	
77	def ProductMeaning(-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
78	  -output: Shape.Shape, -end: Nat, result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>) -> Data:
79	  match result:
80	    case Fail{error}:
81	      Unit
82	    case Done{pair}:
83	      (ra, rb, out) = pair
84	      ProductPrefix<inner, left, right, a, b, output, ra, rb, out, end>
85	
86	type SourceProduct<-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
87	  -output: Shape.Shape, -end: Nat, -result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>> is Data:
88	  SourceProduct{-meaning: @unit: Unit -> ProductMeaning(inner, left, right, a, b, output, end, result)}
89	
90	law successful_product_interpretation:
91	  for +fuel: Nat
92	  for +left: Shape.Shape
93	  for +right: Shape.Shape
94	  for +output: Shape.Shape
95	  for +inner: U32
96	  for +columns: U32
97	  for -a: Core.Tensor<left>
98	  for -b: Core.Tensor<right>
99	  for ra: Core.Tensor<left>
100	  for rb: Core.Tensor<right>
101	  for out: Core.Tensor<output>
102	  for +cursor: Nat
103	  for prefix: ProductPrefix<inner, left, right, a, b, output, ra, rb, out, cursor>
104	  SourceProduct<inner, left, right, a, b, output, Nat.add(fuel, cursor),
105	    Matrix.natural_product_loop(fuel, left, right, output, inner, columns, ra, rb, out, cursor)>
106	
107	def plan_inner(-left: Shape.Shape, -right: Shape.Shape, plan: Matrix.Plan<left, right>) -> U32:
108	  Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible} = plan
109	  inner
110	
111	law successful_multiply_interpretation:
112	  for +left: Shape.Shape
113	  for +right: Shape.Shape
114	  for +plan: Matrix.Plan<left, right>
115	  for a: Core.Tensor<left>
116	  for b: Core.Tensor<right>
117	  SourceProduct<plan_inner(left, right, plan), left, right, a, b, Matrix.output_shape(left, right, plan),
118	    Shape.logical_numel(Shape.dimensions(Matrix.output_shape(left, right, plan))), Matrix.multiply(left, right, plan, a, b)>
119	
120	type MatmulSuccess<-original_a: Core.Dynamic, -original_b: Core.Dynamic, -retained_a: Core.Dynamic, -retained_b: Core.Dynamic, -product: Core.Dynamic> is Data:
121	  MatmulSuccess{-left: Shape.Shape, -right: Shape.Shape, -plan: Matrix.Plan<left, right>,
122	    -a: Core.Tensor<left>, -b: Core.Tensor<right>, -ra: Core.Tensor<left>, -rb: Core.Tensor<right>, -out: Core.Tensor<Matrix.output_shape(left, right, plan)>,
123	    -input_left: {original_a == Core.Packed{left, a} : Core.Dynamic}, -input_right: {original_b == Core.Packed{right, b} : Core.Dynamic},
124	    -output_left: {retained_a == Core.Packed{left, ra} : Core.Dynamic}, -output_right: {retained_b == Core.Packed{right, rb} : Core.Dynamic},
125	    -output: {product == Core.Packed{Matrix.output_shape(left, right, plan), out} : Core.Dynamic},
126	    -source: ProductPrefix<plan_inner(left, right, plan), left, right, a, b, Matrix.output_shape(left, right, plan), ra, rb, out,
127	      Shape.logical_numel(Shape.dimensions(Matrix.output_shape(left, right, plan)))>}
128	
129	def MatmulMeaning(-a: Core.Dynamic, -b: Core.Dynamic, result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Data:
130	  match result:
131	    case Fail{error}:
132	      Unit
133	    case Done{pair}:
134	      (ra, rb, product) = pair
135	      MatmulSuccess<a, b, ra, rb, product>
136	
137	type SourceMatmul<-a: Core.Dynamic, -b: Core.Dynamic, -result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>> is Data:
138	  SourceMatmul{-meaning: @unit: Unit -> MatmulMeaning(a, b, result)}
139	
140	law successful_matmul_interpretation:
141	  for left: Core.Dynamic
142	  for right: Core.Dynamic
143	  SourceMatmul<left, right, Matrix.matmul(left, right)>
144	
145	type ZeroPrefix<-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -out: Core.Tensor<output>, -end: Nat> is Data:
146	  ZeroPrefix{-observed: @+position: Nat -> @earlier: {Nat.is_lt(position, end) == True{} : Bool} ->
147	    {Core.checked_observe(output, out, U32.from_nat(position)) == Some{0.0} : Maybe<&2, F32>}}
148	
149	def ZeroMeaning(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -end: Nat,
150	  result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>) -> Data:
151	  match result:
152	    case Fail{error}:
153	      Unit
154	    case Done{pair}:
155	      (a, b, out) = pair
156	      ZeroPrefix<left, right, output, out, end>
157	
158	type ZeroResult<-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -end: Nat,
159	  -result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>> is Data:
160	  ZeroResult{-meaning: @unit: Unit -> ZeroMeaning(left, right, output, end, result)}
161	
162	law zero_inner_multiply:
163	  for +left: Shape.Shape
164	  for +right: Shape.Shape
165	  for +output: Shape.Shape
166	  for +rows: U32
167	  for +right_inner: U32
168	  for +columns: U32
169	  for +le: {Shape.dimensions(left) == [rows, 0] : List<&2, U32>}
170	  for +re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
171	  for +oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>}
172	  for +compatible: {0n == U32.to_nat(right_inner) : Nat}
173	  for a: Core.Tensor<left>
174	  for b: Core.Tensor<right>
175	  ZeroResult<left, right, output, Shape.logical_numel(Shape.dimensions(output)),
176	    Matrix.multiply(left, right, Matrix.Plan{output, rows, 0, right_inner, columns, le, re, oe, compatible}, a, b)>
177	
178	law successful_dot_interpretation:
179	  for +fuel: Nat
180	  for +left: Shape.Shape
181	  for +right: Shape.Shape
182	  for a: Core.Tensor<left>
183	  for b: Core.Tensor<right>
184	  for +row: U32
185	  for +column: U32
186	  for +cursor: Nat
187	  for +accumulator: F32
188	  SourceDot<fuel, left, right, a, b, row, column, cursor, accumulator,
189	    Matrix.natural_dot(fuel, left, right, a, b, row, column, cursor, accumulator)>
190	
191	law zero_inner_dot:
192	  for +left: Shape.Shape
193	  for +right: Shape.Shape
194	  for a: Core.Tensor<left>
195	  for b: Core.Tensor<right>
196	  for +row: U32
197	  for +column: U32
198	  {Matrix.natural_dot(0n, left, right, a, b, row, column, 0n, 0.0) == Done{(a, b, 0.0)} :
199	    Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & F32>}
200	
201	law output_dimensions:
202	  for -left: Shape.Shape
203	  for -right: Shape.Shape
204	  for +output: Shape.Shape
205	  for rows: U32
206	  for inner: U32
207	  for right_inner: U32
208	  for columns: U32
209	  for le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}
210	  for re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
211	  for oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>}
212	  for compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}
213	  {Shape.dimensions(Matrix.output_shape(left, right, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible})) ==
214	    [rows, columns] : List<&2, U32>}
215	
216	type TotalMatmul<-left: Core.Dynamic, -right: Core.Dynamic,
217	  -result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>> is Data:
218	  TotalMatmul{-done: Totality.Success<Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic, result>,
219	    -interpretation: SourceMatmul<left, right, result>}
220	
221	law matmul_valid_success:
222	  for +left: Shape.Shape
223	  for +right: Shape.Shape
224	  for a: Core.Tensor<left>
225	  for b: Core.Tensor<right>
226	  for valid: Totality.MatrixInputs<left, right>
227	  TotalMatmul<Core.Packed{left, a}, Core.Packed{right, b}, Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b})>
228	
229	# A left operand that is not rank 2 fails with RankMismatch, whatever the
230	# right operand is: there is no 1-D promotion and no batching.
231	law matmul_left_rank_error:
232	  for +left: Shape.Shape
233	  for +right: Shape.Shape
234	  for a: Core.Tensor<left>
235	  for b: Core.Tensor<right>
236	  for not_matrix: {Shape.rank_dimensions(Shape.dimensions(left)) != 2n : Nat}
237	  {Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.RankMismatch{}} : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
238	
239	# A right operand that is not rank 2 fails with RankMismatch, whatever the
240	# left operand is.
241	law matmul_right_rank_error:
242	  for +left: Shape.Shape
243	  for +right: Shape.Shape
244	  for a: Core.Tensor<left>
245	  for b: Core.Tensor<right>
246	  for not_matrix: {Shape.rank_dimensions(Shape.dimensions(right)) != 2n : Nat}
247	  {Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.RankMismatch{}} : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
248	
249	# Two rank-2 operands whose inner extents differ fail with
250	# InnerDimensionMismatch, never with a rank or capacity error: [2, 3] @ [4, 2]
251	# is an inner mismatch.
252	law matmul_inner_error:
253	  for +left: Shape.Shape
254	  for +right: Shape.Shape
255	  for a: Core.Tensor<left>
256	  for b: Core.Tensor<right>
257	  for +rows: U32
258	  for +inner: U32
259	  for +right_inner: U32
260	  for +columns: U32
261	  for left_matrix: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}
262	  for right_matrix: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
263	  for mismatch: {U32.to_nat(inner) != U32.to_nat(right_inner) : Nat}
264	  {Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.InnerDimensionMismatch{}} : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}