src/tensor/matmul.bend

bendtorch-core@0.1.0.0 · 868 lines · 101 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
7import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
8import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
9
10type Error is Data:
11  TensorError{error: Core.Error}
12  ShapeError{error: Shape.Error}
13  RankMismatch{}
14  InnerDimensionMismatch{}
15
16type Plan<-left: Shape.Shape, -right: Shape.Shape> is Data:
17  Plan{output: Shape.Shape, rows: U32, inner: U32, right_inner: U32, columns: U32,
18    -left_dimensions: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
19    -right_dimensions: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
20    -output_dimensions: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
21    -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}}
22
23def output_shape(-left: Shape.Shape, -right: Shape.Shape, plan: Plan<left, right>) -> Shape.Shape:
24  Plan{shape, rows, inner, right_inner, columns, le, re, oe, compatible} = plan
25  shape
26
27def plan_output(-left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32,
28  +right_inner: U32, +columns: U32,
29  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
30  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
31  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
32  result: Shape.ConstructResult<[rows, columns]>) -> Result<&2, &2, Error, Plan<left, right>>:
33  match result:
34    case Shape.SError{error}:
35      Fail{ShapeError{error}}
36    case Shape.SBuilt{shape, preserved}:
37      Done{Plan{shape, rows, inner, right_inner, columns, le, re, preserved, compatible}}
38
39def plan_compatible(-left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32,
40  +right_inner: U32, +columns: U32,
41  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
42  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
43  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))>) ->
44  Result<&2, &2, Error, Plan<left, right>>:
45  match decision:
46    case Shape.BFalse{equal}:
47      Fail{InnerDimensionMismatch{}}
48    case Shape.BTrue{equal}:
49      plan_output(left, right, rows, inner, right_inner, columns, le, re,
50        Lifecycle.natural_equal(U32.to_nat(inner), U32.to_nat(right_inner), equal),
51        Shape.construct([rows, columns]))
52
53def plan_right(-left: Shape.Shape, +right: Shape.Shape, +rows: U32, +inner: U32,
54  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}, dimensions: List<&2, U32>,
55  -same: {dimensions == Shape.dimensions(right) : List<&2, U32>}) ->
56  Result<&2, &2, Error, Plan<left, right>>:
57  match dimensions:
58    case +right_inner <> +columns <> Nil{}:
59      plan_compatible(left, right, rows, inner, right_inner, columns, le,
60        Equal.sym(List<&2, U32>, [right_inner, columns], Shape.dimensions(right), same),
61        Shape.bool_case(Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))))
62    case rest:
63      Fail{RankMismatch{}}
64
65def plan_left(+left: Shape.Shape, +right: Shape.Shape, dimensions: List<&2, U32>,
66  -same: {dimensions == Shape.dimensions(left) : List<&2, U32>}) ->
67  Result<&2, &2, Error, Plan<left, right>>:
68  match dimensions:
69    case +rows <> +inner <> Nil{}:
70      plan_right(left, right, rows, inner,
71        Equal.sym(List<&2, U32>, [rows, inner], Shape.dimensions(left), same),
72        Shape.dimensions(right), {==})
73    case rest:
74      Fail{RankMismatch{}}
75
76def plan(+left: Shape.Shape, +right: Shape.Shape) -> Result<&2, &2, Error, Plan<left, right>>:
77  plan_left(left, right, Shape.dimensions(left), {==})
78
79def right_read(-left: Shape.Shape, -right: Shape.Shape,
80  rec: Core.Tensor<left> -> Core.Tensor<right> -> F32 -> Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>,
81  a: Core.Tensor<left>, left_value: F32, accumulator: F32,
82  pair: Core.Tensor<right> & Result<&2, &2, Core.Error, F32>) ->
83  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>:
84  (b, result) = pair
85  match result:
86    case Fail{error}:
87      Fail{TensorError{error}}
88    case Done{right_value}:
89      rec(a, b, F32.add(accumulator, F32.mul(left_value, right_value)))
90
91def left_read(-left: Shape.Shape, +right: Shape.Shape,
92  rec: Core.Tensor<left> -> Core.Tensor<right> -> F32 -> Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>,
93  b: Core.Tensor<right>, coordinates: List<&2, U32>, accumulator: F32,
94  pair: Core.Tensor<left> & Result<&2, &2, Core.Error, F32>) ->
95  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>:
96  (a, result) = pair
97  match result:
98    case Fail{error}:
99      Fail{TensorError{error}}
100    case Done{value}:
101      right_read(left, right, rec, a, value, accumulator, Core.read_coordinate(right, b, coordinates))
102
103def dot_decided(+left: Shape.Shape, +right: Shape.Shape,
104  rec: Core.Tensor<left> -> Core.Tensor<right> -> F32 -> Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>,
105  a: Core.Tensor<left>, b: Core.Tensor<right>, +row: U32, +column: U32, +cursor: Nat, +index: U32, accumulator: F32,
106  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) ->
107  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>:
108  match decision:
109    case Shape.BFalse{equal}:
110      Fail{TensorError{Core.StorageError{Storage.InvalidIndex{}}}}
111    case Shape.BTrue{equal}:
112      left_read(left, right, rec, b, [index, column], accumulator, Core.read_coordinate(left, a, [row, index]))
113
114def natural_dot(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
115  +row: U32, +column: U32, +cursor: Nat, accumulator: F32) ->
116  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>:
117  match fuel:
118    case 0n:
119      Done{(a, b, accumulator)}
120    case 1n+tail:
121      +index = U32.from_nat(cursor)
122      dot_decided(left, right,
123        a => b => value => natural_dot(tail, left, right, a, b, row, column, 1n+cursor, value),
124        a, b, row, column, cursor, index, accumulator, Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
125
126def dot(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
127  +row: U32, +column: U32, +inner_index: U32, accumulator: F32) ->
128  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>:
129  natural_dot(fuel, left, right, a, b, row, column, U32.to_nat(inner_index), accumulator)
130
131def wrote(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape,
132  rec: Core.Tensor<left> -> Core.Tensor<right> -> Core.Tensor<output> ->
133    Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>,
134  a: Core.Tensor<left>, b: Core.Tensor<right>,
135  pair: Core.Tensor<output> & Result<&2, &2, Core.Error, Unit>) ->
136  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
137  (out, result) = pair
138  match result:
139    case Fail{error}:
140      Fail{TensorError{error}}
141    case Done{unit}:
142      rec(a, b, out)
143
144def dotted(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape,
145  rec: Core.Tensor<left> -> Core.Tensor<right> -> Core.Tensor<output> ->
146    Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>,
147  out: Core.Tensor<output>, index: U32,
148  result: Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & F32>) ->
149  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
150  match result:
151    case Fail{error}:
152      Fail{error}
153    case Done{values}:
154      (a, b, value) = values
155      wrote(left, right, output, rec, a, b, Core.write_flat(output, out, index, value))
156
157def product_coordinates(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape,
158  rec: Core.Tensor<left> -> Core.Tensor<right> -> Core.Tensor<output> ->
159    Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>,
160  +inner: U32, a: Core.Tensor<left>, b: Core.Tensor<right>, out: Core.Tensor<output>, +index: U32, coordinates: List<&2, U32>) ->
161  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
162  match coordinates:
163    case +row <> +column <> Nil{}:
164      dotted(left, right, output, rec, out, index, natural_dot(U32.to_nat(inner), left, right, a, b, row, column, 0n, 0.0))
165    case rest:
166      Fail{RankMismatch{}}
167
168def product_unflattened(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape,
169  rec: Core.Tensor<left> -> Core.Tensor<right> -> Core.Tensor<output> ->
170    Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>,
171  +inner: U32, a: Core.Tensor<left>, b: Core.Tensor<right>, out: Core.Tensor<output>, +index: U32,
172  result: Result<&2, &2, Shape.Error, Shape.CheckedCoordinate<output>>) ->
173  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
174  match result:
175    case Fail{error}:
176      Fail{ShapeError{error}}
177    case Done{coordinate}:
178      product_coordinates(left, right, output, rec, inner, a, b, out, index, Shape.coordinate_values(output, coordinate))
179
180def product_decided(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape,
181  rec: Core.Tensor<left> -> Core.Tensor<right> -> Core.Tensor<output> ->
182    Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>,
183  +inner: U32, +columns: U32, a: Core.Tensor<left>, b: Core.Tensor<right>, out: Core.Tensor<output>, +cursor: Nat, +index: U32,
184  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) ->
185  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
186  match decision:
187    case Shape.BFalse{equal}:
188      Fail{TensorError{Core.StorageError{Storage.InvalidIndex{}}}}
189    case Shape.BTrue{equal}:
190      product_unflattened(left, right, output, rec, inner, a, b, out, index, Shape.unflatten_checked(output, index))
191
192def natural_product_loop(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape,
193  +inner: U32, +columns: U32, a: Core.Tensor<left>, b: Core.Tensor<right>, out: Core.Tensor<output>, +cursor: Nat) ->
194  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
195  match fuel:
196    case 0n:
197      Done{(a, b, out)}
198    case 1n+tail:
199      +index = U32.from_nat(cursor)
200      product_decided(left, right, output,
201        a => b => out => natural_product_loop(tail, left, right, output, inner, columns, a, b, out, 1n+cursor),
202        inner, columns, a, b, out, cursor, index, Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
203
204def product_loop(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape,
205  +inner: U32, +columns: U32, a: Core.Tensor<left>, b: Core.Tensor<right>, out: Core.Tensor<output>, +index: U32) ->
206  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
207  natural_product_loop(fuel, left, right, output, inner, columns, a, b, out, U32.to_nat(index))
208
209def created(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +inner: U32, +columns: U32,
210  a: Core.Tensor<left>, b: Core.Tensor<right>, result: Result<&2, &1, Core.Error, Core.Tensor<output>>) ->
211  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>:
212  match result:
213    case Fail{error}:
214      Fail{TensorError{error}}
215    case Done{out}:
216      natural_product_loop(Shape.elements(output), left, right, output,
217        inner, columns, a, b, out, 0n)
218
219def multiply(+left: Shape.Shape, +right: Shape.Shape, +plan: Plan<left, right>,
220  a: Core.Tensor<left>, b: Core.Tensor<right>) ->
221  Result<&2, &1, Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output_shape(left, right, plan)>>:
222  match plan:
223    case Plan{+output, rows, +inner, right_inner, +columns, le, re, oe, compatible}:
224      created(left, right, output, inner, columns, a, b, Core.new(output, 0.0))
225
226# The parallel product. Output cell `position` of a rows x columns product
227# folds row `position / columns` of the left array with column
228# `position % columns` of the right array: it adds left[row, k] * right[k,
229# column] for k = 0, 1, ..., inner - 1, starting from +0.0, the same F32 terms
230# in the same order as `natural_dot`. The cells are built like
231# `Elementwise.build` builds its output: the top `splits` levels of the
232# output tree fork, each half with its own clone of both input arrays, and a
233# thread builds the cells below the forks in index order.
234
235# Output cell `position` of a product with `columns` columns sits at row
236# `position / columns` and column `position % columns`, as the digits of the
237# output shape spell them.
238def row_of(+columns: Nat, +position: Nat) -> Nat:
239  Nat.div(position, Nat.mul(columns, 1n))
240
241def column_of(+columns: Nat, +position: Nat) -> Nat:
242  Nat.div(Nat.mod(position, Nat.mul(columns, 1n)), 1n)
243
244def left_index(+inner: Nat, +row: Nat, +k: Nat) -> U32:
245  U32.from_nat(Nat.add(Nat.mul(row, inner), k))
246
247def right_index(+columns: Nat, +column: Nat, +k: Nat) -> U32:
248  U32.from_nat(Nat.add(Nat.mul(k, columns), column))
249
250# The fuel + 1 terms of a cell from the one whose words sit at offsets `lp`
251# and `rp`, added to `accumulator`; `left` and `right` are their reads,
252# already made, and each later term sits `ls` and `rs` further on. Each step
253# reads the next term's words, so the loop is one tail call with no closure
254# and one addition per offset: it compiles to a flat loop.
255def product_run(fuel: Nat, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, accumulator: F32,
256  left: Array<F32> & F32, right: Array<F32> & F32) -> F32 & (Array<F32> & Array<F32>):
257  match fuel left right:
258    case 0n Tuple{la, x} Tuple{ra, y}:
259      (F32.add(accumulator, F32.mul(x, y)), (la, ra))
260    case 1n+rest Tuple{la, x} Tuple{ra, y}:
261      product_run(rest, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)),
262        Array.get(F32, la, U32.from_nat(Nat.add(ls, lp))), Array.get(F32, ra, U32.from_nat(Nat.add(rs, rp))))
263
264# The terms k = cursor .. cursor + fuel - 1 of cell (row, column), added to
265# `accumulator`. Both arrays come back with the total: the run from term
266# cursor, where term k + 1 sits one word on in the left row and `columns`
267# words on in the right column.
268def cell_dot(fuel: Nat, +inner: Nat, +columns: Nat, +row: Nat, +column: Nat, +cursor: Nat, accumulator: F32,
269  la: Array<F32>, ra: Array<F32>) -> F32 & (Array<F32> & Array<F32>):
270  match fuel:
271    case 0n:
272      (accumulator, (la, ra))
273    case 1n+rest:
274      product_run(rest, 1n, columns, Nat.add(Nat.mul(row, inner), cursor), Nat.add(Nat.mul(cursor, columns), column), accumulator,
275        Array.get(F32, la, left_index(inner, row, cursor)), Array.get(F32, ra, right_index(columns, column, cursor)))
276
277# Cell `position`: all `inner` terms, from +0.0.
278def dot_of(+inner: Nat, +columns: Nat, +position: Nat, la: Array<F32>, ra: Array<F32>) -> F32 & (Array<F32> & Array<F32>):
279  cell_dot(inner, inner, columns, row_of(columns, position), column_of(columns, position), 0n, 0.0, la, ra)
280
281# One output cell: its total when the cell is below `count`, else +0.0.
282def cell(+inner: Nat, +columns: Nat, +position: Nat, la: Array<F32>, ra: Array<F32>, inside: Bool) ->
283  F32 & (Array<F32> & Array<F32>):
284  match inside:
285    case False{}:
286      (0.0, (la, ra))
287    case True{}:
288      dot_of(inner, columns, position, la, ra)
289
290# What every cell of a build needs: the inner extent and the column count.
291type Dims is Data:
292  Dims{inner: Nat, columns: Nat}
293
294# Each thread below the forks reads its own copy.
295def dims(+given: Dims) -> Dims:
296  Dims{+inner, +columns} = given
297  Dims{inner, columns}
298
299# The cell function `Parallel.build` runs: its state is the pair of input
300# arrays, split for a fork as `Elementwise.split` splits them.
301def cell_at(+dims: Dims, +position: Nat, inputs: Array<F32> & Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
302  Dims{+inner, +columns} = dims
303  (la, ra) = inputs
304  cell(inner, columns, position, la, ra, inside)
305
306def build(+depth: Nat, +splits: Nat, +inner: Nat, +columns: Nat, +count: Nat, +base: Nat,
307  la: Array<F32>, ra: Array<F32>) -> Array<F32>:
308  Parallel.build(~Dims, ~Dims, ~Array<F32>, ~Array<F32>, ~dims, ~Elementwise.split, ~cell_at, Dims{inner, columns}, depth, splits, count, base, (la, ra))
309
310def halvings_step(rec: Nat -> Nat, +number: Nat, small: Bool) -> Nat:
311  match small:
312    case True{}:
313      0n
314    case False{}:
315      1n+rec(Nat.div(number, 2n))
316
317# floor(log2(number)) for a positive number, 0 for 0. `fuel` bounds the
318# halvings; `number` itself is enough.
319def halvings(fuel: Nat, +number: Nat) -> Nat:
320  match fuel:
321    case 0n:
322      0n
323    case 1n+rest:
324      halvings_step(n => halvings(rest, n), number, Nat.is_lt(number, 2n))
325
326def least(+a: Nat, +b: Nat, smaller: Bool) -> Nat:
327  match smaller:
328    case True{}:
329      a
330    case False{}:
331      b
332
333def at_most(+a: Nat, +b: Nat) -> Nat:
334  least(a, b, Nat.is_lt(a, b))
335
336# Split by work, not by cell count: a cell costs one term per inner index.
337# The build forks until a subtree holds about 2^13 terms, and at most six
338# levels, 64 subtrees, each with its own copy of both inputs: an output of
339# 2^depth cells of `inner` terms forks depth + log2(inner) - 13 levels. On
340# medusa2 at 8 threads a product of 2^15 terms or more ran faster with every
341# extra level up to six, and forking cost nothing measurable on one thread.
342def splits(+depth: Nat, +inner: Nat) -> Nat:
343  at_most(at_most(Nat.sub(Nat.add(depth, halvings(inner, inner)), 13n), 6n), depth)
344
345# A product forks as `splits` says, but at most four levels, sixteen
346# subtrees: each fork copies both operands into each half, and with a row
347# of eight cells running as one loop a subtree's work no longer hides its
348# copies. On medusa2 four levels ran fastest at 8 threads for the products
349# of an MNIST step and for 256 x 256 x 256 (the MNIST step 15.2 ms against
350# 17.5 at six levels) and at one thread (20.1 ms against 21.6).
351def product_splits(+depth: Nat, +inner: Nat) -> Nat:
352  at_most(splits(depth, inner), 4n)
353
354# The parallel build returns a balanced tree of its depth, whatever the split.
355def build_balanced(+depth: Nat, +splits: Nat, +inner: Nat, +columns: Nat, +count: Nat, +base: Nat,
356  la: Array<F32>, ra: Array<F32>) -> Model.Balanced(F32, depth, build(depth, splits, inner, columns, count, base, la, ra)):
357  Parallel.balanced(~Dims, ~Dims, ~Array<F32>, ~Array<F32>, ~dims, ~Elementwise.split, ~cell_at, Dims{inner, columns}, depth, splits, count, base, (la, ra))
358
359# Both inputs come back to the caller as the kept halves of their clones,
360# and the build reads the other halves.
361def parallel_output(+output: Shape.Shape, +inner: Nat, +columns: Nat, +depth: Nat, +length: Nat,
362  -valid: {Storage.bounds(depth, length) == True{} : Bool},
363  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
364  left: Core.Dynamic & Array<F32>, right: Core.Dynamic & Array<F32>) ->
365  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
366  (kept_left, la) = left
367  (kept_right, ra) = right
368  Done{(kept_left, kept_right, Core.pack(output, Core.TTensor{Storage.Owned{depth, length,
369    build(depth, product_splits(depth, inner), inner, columns, length, 0n, la, ra),
370    build_balanced(depth, product_splits(depth, inner), inner, columns, length, 0n, la, ra),
371    valid}, exact}))}
372
373def parallel_allocated(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +inner: Nat, +columns: Nat,
374  a: Core.Tensor<left>, b: Core.Tensor<right>, result: Result<&2, &1, Core.Error, Core.Tensor<output>>) ->
375  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
376  match result:
377    case Fail{error}:
378      Fail{TensorError{error}}
379    case Done{Core.TTensor{storage, exact}}:
380      match storage:
381        case Storage.Empty{}:
382          Done{(Core.pack(left, a), Core.pack(right, b), Core.pack(output, Core.TTensor{Storage.Empty{}, exact}))}
383        case Storage.Owned{+depth, +length, array, balanced, valid}:
384          parallel_output(output, inner, columns, depth, length, valid, exact,
385            Elementwise.take(Core.pack(left, a)), Elementwise.take(Core.pack(right, b)))
386
387def parallel_multiply(+left: Shape.Shape, +right: Shape.Shape, +plan: Plan<left, right>,
388  a: Core.Tensor<left>, b: Core.Tensor<right>) -> Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
389  match plan:
390    case Plan{+output, rows, +inner, right_inner, +columns, le, re, oe, compatible}:
391      parallel_allocated(left, right, output, U32.to_nat(inner), U32.to_nat(columns), a, b, Core.new(output, 0.0))
392
393def parallel_planned(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
394  result: Result<&2, &2, Error, Plan<left, right>>) ->
395  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
396  match result:
397    case Fail{error}:
398      Fail{error}
399    case Done{+plan}:
400      parallel_multiply(left, right, plan, a, b)
401
402# `matmul` builds its output in parallel. `multiply`, `natural_product_loop`
403# and `natural_dot` above are the sequential definition of the same cells,
404# which their own laws state.
405def matmul(left: Core.Dynamic, right: Core.Dynamic) -> Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
406  Core.Packed{+left_shape, a} = left
407  Core.Packed{+right_shape, b} = right
408  parallel_planned(left_shape, right_shape, a, b, plan(left_shape, right_shape))
409
410# Products that read an operand in place. `x @ W^T` and `g^T @ x` read a
411# transposed operand; instead of building the transpose they read the
412# operand at the transposed index. An operand of a product is laid out in
413# one of two orders: `Kept`, as `[outer, inner]`, where term k of output row
414# (or column) `outer` sits at outer * inner + k; or `Swapped`, as
415# `[inner, outer]`, where it sits at k * outers + outer. `matmul` is the
416# product of a `Kept` left and a `Swapped` right operand, `matmul_nt` of two
417# `Kept` operands and `matmul_tn` of two `Swapped` ones. Cell (row, column)
418# adds the inner terms in increasing k from +0.0, as `natural_dot` does. The
419# cells are built by `product_build` below: it forks as `Parallel.build`
420# forks, and below the forks `product_tree` builds them as `Parallel.tree`
421# does, except that eight cells below the count and in one output row run
422# as one loop with eight accumulators (`row_run`).
423type Order is Data:
424  Kept{}
425  Swapped{}
426
427def read_offset(+order: Order, +inner: Nat, +outers: Nat, +outer: Nat, +k: Nat) -> Nat:
428  match order:
429    case Kept{}:
430      Nat.add(Nat.mul(outer, inner), k)
431    case Swapped{}:
432      Nat.add(Nat.mul(k, outers), outer)
433
434def read_index(+order: Order, +inner: Nat, +outers: Nat, +outer: Nat, +k: Nat) -> U32:
435  U32.from_nat(read_offset(order, inner, outers, outer, k))
436
437# How far term k + 1 of an operand sits from term k: the next word of a
438# `Kept` row, the next row of a `Swapped` operand of `outers` columns.
439def stride(+order: Order, +outers: Nat) -> Nat:
440  match order:
441    case Kept{}:
442      1n
443    case Swapped{}:
444      outers
445
446# The terms k = cursor .. cursor + fuel - 1 of cell (row, column) of a
447# rows x columns product, added to `accumulator`: the run from term cursor,
448# whose offsets step by each operand's stride.
449def product_dot(fuel: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +row: Nat, +column: Nat,
450  +cursor: Nat, accumulator: F32, la: Array<F32>, ra: Array<F32>) -> F32 & (Array<F32> & Array<F32>):
451  match fuel:
452    case 0n:
453      (accumulator, (la, ra))
454    case 1n+rest:
455      product_run(rest, stride(lo, rows), stride(ro, columns), read_offset(lo, inner, rows, row, cursor),
456        read_offset(ro, inner, columns, column, cursor), accumulator,
457        Array.get(F32, la, read_index(lo, inner, rows, row, cursor)), Array.get(F32, ra, read_index(ro, inner, columns, column, cursor)))
458
459def product_dot_of(+lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +position: Nat,
460  la: Array<F32>, ra: Array<F32>) -> F32 & (Array<F32> & Array<F32>):
461  product_dot(inner, lo, ro, inner, rows, columns, row_of(columns, position), column_of(columns, position), 0n, 0.0, la, ra)
462
463def product_cell(+lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +position: Nat,
464  la: Array<F32>, ra: Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
465  match inside:
466    case False{}:
467      (0.0, (la, ra))
468    case True{}:
469      product_dot_of(lo, ro, inner, rows, columns, position, la, ra)
470
471# What every cell of a product build reads.
472type Layout is Data:
473  Layout{left: Order, right: Order, inner: Nat, rows: Nat, columns: Nat}
474
475def layout(+given: Layout) -> Layout:
476  Layout{+lo, +ro, +inner, +rows, +columns} = given
477  Layout{lo, ro, inner, rows, columns}
478
479def product_cell_at(+layout: Layout, +position: Nat, inputs: Array<F32> & Array<F32>, inside: Bool) ->
480  F32 & (Array<F32> & Array<F32>):
481  Layout{+lo, +ro, +inner, +rows, +columns} = layout
482  (la, ra) = inputs
483  product_cell(lo, ro, inner, rows, columns, position, la, ra, inside)
484
485# Products fill their cells as `Parallel.build` does, with one change: a
486# subtree of eight cells that lie below the count and in one output row runs
487# as one loop, `row_run`, with eight accumulators. Term k reads the left
488# word of the row once and the eight right words of the cells, and adds each
489# product to its own cell's accumulator, so every cell still adds its terms
490# in increasing k from +0.0 and gets the word `product_dot` computes. The
491# eight chains of additions are independent, so a core runs them side by
492# side, where one cell's chain waits for each addition before the next.
493type Eight is Data:
494  Eight{w0: F32, w1: F32, w2: F32, w3: F32, w4: F32, w5: F32, w6: F32, w7: F32}
495
496# The eight words at offsets o0 .. o7, read in that order.
497def eight_read_8(+x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, +x5: F32, +x6: F32, pair: Array<F32> & F32) -> Array<F32> & Eight:
498  match pair:
499    case Tuple{a, +x7}:
500      (a, Eight{x0, x1, x2, x3, x4, x5, x6, x7})
501
502def eight_read_7(+x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, +x5: F32, pair: Array<F32> & F32, +o7: Nat) -> Array<F32> & Eight:
503  match pair:
504    case Tuple{a, +x6}:
505      eight_read_8(x0, x1, x2, x3, x4, x5, x6, Array.get(F32, a, U32.from_nat(o7)))
506
507def eight_read_6(+x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, pair: Array<F32> & F32, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
508  match pair:
509    case Tuple{a, +x5}:
510      eight_read_7(x0, x1, x2, x3, x4, x5, Array.get(F32, a, U32.from_nat(o6)), o7)
511
512def eight_read_5(+x0: F32, +x1: F32, +x2: F32, +x3: F32, pair: Array<F32> & F32, +o5: Nat, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
513  match pair:
514    case Tuple{a, +x4}:
515      eight_read_6(x0, x1, x2, x3, x4, Array.get(F32, a, U32.from_nat(o5)), o6, o7)
516
517def eight_read_4(+x0: F32, +x1: F32, +x2: F32, pair: Array<F32> & F32, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
518  match pair:
519    case Tuple{a, +x3}:
520      eight_read_5(x0, x1, x2, x3, Array.get(F32, a, U32.from_nat(o4)), o5, o6, o7)
521
522def eight_read_3(+x0: F32, +x1: F32, pair: Array<F32> & F32, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
523  match pair:
524    case Tuple{a, +x2}:
525      eight_read_4(x0, x1, x2, Array.get(F32, a, U32.from_nat(o3)), o4, o5, o6, o7)
526
527def eight_read_2(+x0: F32, pair: Array<F32> & F32, +o2: Nat, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
528  match pair:
529    case Tuple{a, +x1}:
530      eight_read_3(x0, x1, Array.get(F32, a, U32.from_nat(o2)), o3, o4, o5, o6, o7)
531
532def eight_read_1(pair: Array<F32> & F32, +o1: Nat, +o2: Nat, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
533  match pair:
534    case Tuple{a, +x0}:
535      eight_read_2(x0, Array.get(F32, a, U32.from_nat(o1)), o2, o3, o4, o5, o6, o7)
536
537def read_eight(a: Array<F32>, +o0: Nat, +o1: Nat, +o2: Nat, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array<F32> & Eight:
538  eight_read_1(Array.get(F32, a, U32.from_nat(o0)), o1, o2, o3, o4, o5, o6, o7)
539
540# Term k of the eight cells: the left word x times each right word.
541def eight_terms(+acc: Eight, +x: F32, +ys: Eight) -> Eight:
542  Eight{+a0, +a1, +a2, +a3, +a4, +a5, +a6, +a7} = acc
543  Eight{+y0, +y1, +y2, +y3, +y4, +y5, +y6, +y7} = ys
544  Eight{F32.add(a0, F32.mul(x, y0)), F32.add(a1, F32.mul(x, y1)), F32.add(a2, F32.mul(x, y2)), F32.add(a3, F32.mul(x, y3)), F32.add(a4, F32.mul(x, y4)), F32.add(a5, F32.mul(x, y5)), F32.add(a6, F32.mul(x, y6)), F32.add(a7, F32.mul(x, y7))}
545
546# The fuel + 1 terms of the eight cells from the one whose words sit at
547# offsets lp (left) and r0 .. r7 (right), read already, as `product_run`
548# runs one cell.
549def row_run(fuel: Nat, +ls: Nat, +rs: Nat, +lp: Nat, +r0: Nat, +r1: Nat, +r2: Nat, +r3: Nat, +r4: Nat, +r5: Nat, +r6: Nat, +r7: Nat, +acc: Eight,
550  left: Array<F32> & F32, right: Array<F32> & Eight) -> Eight & (Array<F32> & Array<F32>):
551  match fuel left right:
552    case 0n Tuple{la, +x} Tuple{ra, +ys}:
553      (eight_terms(acc, x, ys), (la, ra))
554    case 1n+rest Tuple{la, +x} Tuple{ra, +ys}:
555      row_run(rest, ls, rs, Nat.add(ls, lp), Nat.add(rs, r0), Nat.add(rs, r1), Nat.add(rs, r2), Nat.add(rs, r3), Nat.add(rs, r4), Nat.add(rs, r5), Nat.add(rs, r6), Nat.add(rs, r7), eight_terms(acc, x, ys),
556        Array.get(F32, la, U32.from_nat(Nat.add(ls, lp))), read_eight(ra, Nat.add(rs, r0), Nat.add(rs, r1), Nat.add(rs, r2), Nat.add(rs, r3), Nat.add(rs, r4), Nat.add(rs, r5), Nat.add(rs, r6), Nat.add(rs, r7)))
557
558# The eight words as the depth-3 tree `Parallel.tree` builds.
559def eight_leaves_of(+cells: Eight, state: Array<F32> & Array<F32>) -> Array<F32> & (Array<F32> & Array<F32>):
560  Eight{+w0, +w1, +w2, +w3, +w4, +w5, +w6, +w7} = cells
561  (ANode{ANode{ANode{ALeaf{w0}, ALeaf{w1}}, ANode{ALeaf{w2}, ALeaf{w3}}}, ANode{ANode{ALeaf{w4}, ALeaf{w5}}, ANode{ALeaf{w6}, ALeaf{w7}}}}, state)
562
563def eight_leaves(pair: Eight & (Array<F32> & Array<F32>)) -> Array<F32> & (Array<F32> & Array<F32>):
564  match pair:
565    case Tuple{+cells, state}:
566      eight_leaves_of(cells, state)
567
568# Cells base .. base + 7, all in the row of cell base, from fuel terms.
569def row_block_run(fuel: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +base: Nat,
570  la: Array<F32>, ra: Array<F32>) -> Array<F32> & (Array<F32> & Array<F32>):
571  match fuel:
572    case 0n:
573      eight_leaves_of(Eight{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, (la, ra))
574    case 1n+rest:
575      eight_leaves(row_run(rest, stride(lo, rows), stride(ro, columns), read_offset(lo, inner, rows, row_of(columns, base), 0n),
576        read_offset(ro, inner, columns, column_of(columns, base), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(1n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(2n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(3n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(4n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(5n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(6n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(7n, base)), 0n),
577        Eight{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, Array.get(F32, la, U32.from_nat(read_offset(lo, inner, rows, row_of(columns, base), 0n))),
578        read_eight(ra, read_offset(ro, inner, columns, column_of(columns, base), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(1n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(2n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(3n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(4n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(5n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(6n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(7n, base)), 0n))))
579
580def row_block(+layout: Layout, +base: Nat, state: Array<F32> & Array<F32>) -> Array<F32> & (Array<F32> & Array<F32>):
581  Layout{+lo, +ro, +inner, +rows, +columns} = layout
582  (la, ra) = state
583  row_block_run(inner, lo, ro, inner, rows, columns, base, la, ra)
584
585# The eight cells from base lie below the count and in the row of cell base.
586def one_row(+layout: Layout, +count: Nat, +base: Nat) -> Bool:
587  Layout{+lo, +ro, +inner, +rows, +columns} = layout
588  Bool.and(Nat.is_lt(base, count), Bool.and(Nat.is_lt(Nat.add(1n, base), count), Bool.and(Nat.is_lt(Nat.add(2n, base), count), Bool.and(Nat.is_lt(Nat.add(3n, base), count), Bool.and(Nat.is_lt(Nat.add(4n, base), count), Bool.and(Nat.is_lt(Nat.add(5n, base), count), Bool.and(Nat.is_lt(Nat.add(6n, base), count), Bool.and(Nat.is_lt(Nat.add(7n, base), count), Bool.and(Nat.is_eq(row_of(columns, Nat.add(1n, base)), row_of(columns, base)), Bool.and(Nat.is_eq(row_of(columns, Nat.add(2n, base)), row_of(columns, base)), Bool.and(Nat.is_eq(row_of(columns, Nat.add(3n, base)), row_of(columns, base)), Bool.and(Nat.is_eq(row_of(columns, Nat.add(4n, base)), row_of(columns, base)), Bool.and(Nat.is_eq(row_of(columns, Nat.add(5n, base)), row_of(columns, base)), Bool.and(Nat.is_eq(row_of(columns, Nat.add(6n, base)), row_of(columns, base)), Bool.and(Nat.is_eq(row_of(columns, Nat.add(7n, base)), row_of(columns, base)), True{})))))))))))))))
589
590def block_or_tree(+layout: Layout, +count: Nat, +base: Nat, state: Array<F32> & Array<F32>, one: Bool) ->
591  Array<F32> & (Array<F32> & Array<F32>):
592  match one:
593    case False{}:
594      Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 3n, count, base, state)
595    case True{}:
596      row_block(layout, base, state)
597
598# The cells base .. base + 2^depth on one thread: `Parallel.tree` with the
599# product cell, but each subtree of eight cells in one row is a row block.
600def product_tree(+layout: Layout, +depth: Nat, +count: Nat, +base: Nat, state: Array<F32> & Array<F32>) ->
601  Array<F32> & (Array<F32> & Array<F32>):
602  match depth:
603    case 0n:
604      Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 0n, count, base, state)
605    case 1n+d1:
606      match d1:
607        case 0n:
608          Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 1n, count, base, state)
609        case 1n+d2:
610          match d2:
611            case 0n:
612              Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 2n, count, base, state)
613            case 1n+d3:
614              match d3:
615                case 0n:
616                  block_or_tree(layout, count, base, state, one_row(layout, count, base))
617                case 1n+d4:
618                  Parallel.node(~Array<F32>, ~Array<F32>, l => r => product_tree(layout, d1, count, Parallel.upper(d1, base), (l, r)),
619                    product_tree(layout, d1, count, base, state))
620
621# The parallel product build: the top `splits` levels fork as
622# `Parallel.build` forks, each half with its own copy of both inputs, and
623# each subtree below the forks is a `product_tree`.
624def product_build(+depth: Nat, +splits: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat,
625  +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>) -> Array<F32>:
626  match depth splits:
627    case 1n+lower 1n+rest:
628      Parallel.fork(~Array<F32>, ~Array<F32>, Elementwise.split((la, ra)),
629        l => r => product_build(lower, rest, lo, ro, inner, rows, columns, count, base, l, r),
630        l => r => product_build(lower, rest, lo, ro, inner, rows, columns, count, Parallel.upper(lower, base), l, r))
631    case _ _:
632      Pair.fst(Array<F32>, Array<F32> & Array<F32>, product_tree(Layout{lo, ro, inner, rows, columns}, depth, count, base, (la, ra)))
633
634# The build is a balanced tree of its depth, whatever the split: the erased
635# certificate of the output storage.
636def eight_leaves_of_certified(+cells: Eight, state: Array<F32> & Array<F32>) ->
637  Parallel.Certified<(Array<F32> & Array<F32>), 3n, eight_leaves_of(cells, state)>:
638  Eight{+w0, +w1, +w2, +w3, +w4, +w5, +w6, +w7} = cells
639  Parallel.Certified{(ANode{ANode{ANode{ALeaf{w0}, ALeaf{w1}}, ANode{ALeaf{w2}, ALeaf{w3}}}, ANode{ANode{ALeaf{w4}, ALeaf{w5}}, ANode{ALeaf{w6}, ALeaf{w7}}}}, state),
640    {==}, (((Unit{}, Unit{}), (Unit{}, Unit{})), ((Unit{}, Unit{}), (Unit{}, Unit{})))}
641
642def eight_leaves_certified(pair: Eight & (Array<F32> & Array<F32>)) ->
643  Parallel.Certified<(Array<F32> & Array<F32>), 3n, eight_leaves(pair)>:
644  match pair:
645    case Tuple{+cells, state}:
646      eight_leaves_of_certified(cells, state)
647
648def row_block_run_certified(fuel: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +base: Nat,
649  la: Array<F32>, ra: Array<F32>) ->
650  Parallel.Certified<(Array<F32> & Array<F32>), 3n, row_block_run(fuel, lo, ro, inner, rows, columns, base, la, ra)>:
651  match fuel:
652    case 0n:
653      eight_leaves_of_certified(Eight{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, (la, ra))
654    case 1n+rest:
655      eight_leaves_certified(row_run(rest, stride(lo, rows), stride(ro, columns), read_offset(lo, inner, rows, row_of(columns, base), 0n),
656        read_offset(ro, inner, columns, column_of(columns, base), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(1n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(2n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(3n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(4n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(5n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(6n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(7n, base)), 0n),
657        Eight{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, Array.get(F32, la, U32.from_nat(read_offset(lo, inner, rows, row_of(columns, base), 0n))),
658        read_eight(ra, read_offset(ro, inner, columns, column_of(columns, base), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(1n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(2n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(3n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(4n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(5n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(6n, base)), 0n), read_offset(ro, inner, columns, column_of(columns, Nat.add(7n, base)), 0n))))
659
660def row_block_certified(+layout: Layout, +base: Nat, state: Array<F32> & Array<F32>) ->
661  Parallel.Certified<(Array<F32> & Array<F32>), 3n, row_block(layout, base, state)>:
662  Layout{+lo, +ro, +inner, +rows, +columns} = layout
663  (la, ra) = state
664  row_block_run_certified(inner, lo, ro, inner, rows, columns, base, la, ra)
665
666def block_or_tree_certified(+layout: Layout, +count: Nat, +base: Nat, state: Array<F32> & Array<F32>, one: Bool) ->
667  Parallel.Certified<(Array<F32> & Array<F32>), 3n, block_or_tree(layout, count, base, state, one)>:
668  match one:
669    case False{}:
670      Parallel.tree_certified(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 3n, count, base, state)
671    case True{}:
672      row_block_certified(layout, base, state)
673
674def product_tree_certified(+layout: Layout, +depth: Nat, +count: Nat, +base: Nat, state: Array<F32> & Array<F32>) ->
675  Parallel.Certified<(Array<F32> & Array<F32>), depth, product_tree(layout, depth, count, base, state)>:
676  match depth:
677    case 0n:
678      Parallel.tree_certified(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 0n, count, base, state)
679    case 1n+d1:
680      match d1:
681        case 0n:
682          Parallel.tree_certified(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 1n, count, base, state)
683        case 1n+d2:
684          match d2:
685            case 0n:
686              Parallel.tree_certified(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 2n, count, base, state)
687            case 1n+d3:
688              match d3:
689                case 0n:
690                  block_or_tree_certified(layout, count, base, state, one_row(layout, count, base))
691                case 1n+d4:
692                  Parallel.node_certified(~Array<F32>, ~Array<F32>, d1, l => r => product_tree(layout, d1, count, Parallel.upper(d1, base), (l, r)),
693                    product_tree(layout, d1, count, base, state), product_tree_certified(layout, d1, count, base, state),
694                    l => r => product_tree_certified(layout, d1, count, Parallel.upper(d1, base), (l, r)))
695
696def tree_balanced(+depth: Nat, -expected: Array<F32> & (Array<F32> & Array<F32>),
697  certified: Parallel.Certified<(Array<F32> & Array<F32>), depth, expected>) ->
698  Model.Balanced(F32, depth, Pair.fst(Array<F32>, Array<F32> & Array<F32>, expected)):
699  Parallel.Certified{result, same, balanced} = certified
700  %same : Model.Balanced(F32, depth, Pair.fst(Array<F32>, Array<F32> & Array<F32>, _))
701  balanced
702
703def product_build_balanced(+depth: Nat, +splits: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat,
704  +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>) ->
705  Model.Balanced(F32, depth, product_build(depth, splits, lo, ro, inner, rows, columns, count, base, la, ra)):
706  match depth splits:
707    case 1n+lower 1n+rest:
708      Parallel.fork_balanced(~Array<F32>, ~Array<F32>, lower, Elementwise.split((la, ra)),
709        l => r => product_build(lower, rest, lo, ro, inner, rows, columns, count, base, l, r),
710        l => r => product_build(lower, rest, lo, ro, inner, rows, columns, count, Parallel.upper(lower, base), l, r),
711        l => r => product_build_balanced(lower, rest, lo, ro, inner, rows, columns, count, base, l, r),
712        l => r => product_build_balanced(lower, rest, lo, ro, inner, rows, columns, count, Parallel.upper(lower, base), l, r))
713    case 0n _:
714      tree_balanced(0n, product_tree(Layout{lo, ro, inner, rows, columns}, 0n, count, base, (la, ra)),
715        product_tree_certified(Layout{lo, ro, inner, rows, columns}, 0n, count, base, (la, ra)))
716    case 1n+lower 0n:
717      tree_balanced(1n+lower, product_tree(Layout{lo, ro, inner, rows, columns}, 1n+lower, count, base, (la, ra)),
718        product_tree_certified(Layout{lo, ro, inner, rows, columns}, 1n+lower, count, base, (la, ra)))
719
720# The dimensions of an operand of `outer` rows (or columns) and `inner`
721# terms laid out in `order`.
722def arranged(+order: Order, +outer: U32, +inner: U32) -> List<&2, U32>:
723  match order:
724    case Kept{}:
725      [outer, inner]
726    case Swapped{}:
727      [inner, outer]
728
729type Product<-lo: Order, -ro: Order, -left: Shape.Shape, -right: Shape.Shape> is Data:
730  Product{output: Shape.Shape, rows: U32, inner: U32, right_inner: U32, columns: U32,
731    -left_dimensions: {Shape.dimensions(left) == arranged(lo, rows, inner) : List<&2, U32>},
732    -right_dimensions: {Shape.dimensions(right) == arranged(ro, columns, right_inner) : List<&2, U32>},
733    -output_dimensions: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>},
734    -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}}
735
736def product_shape(-lo: Order, -ro: Order, -left: Shape.Shape, -right: Shape.Shape, plan: Product<lo, ro, left, right>) -> Shape.Shape:
737  Product{shape, rows, inner, right_inner, columns, le, re, oe, compatible} = plan
738  shape
739
740# A rank-2 operand split into its outer and inner extents.
741type Split<-order: Order, -dims: List<&2, U32>> is Data:
742  Split{outer: U32, inner: U32, same: {dims == arranged(order, outer, inner) : List<&2, U32>}}
743
744def split_dims(+order: Order, dims: List<&2, U32>) -> Result<&2, &2, Error, Split<order, dims>>:
745  match order:
746    case Kept{}:
747      match dims:
748        case +a <> +b <> Nil{}:
749          Done{Split{a, b, {==}}}
750        case rest:
751          Fail{RankMismatch{}}
752    case Swapped{}:
753      match dims:
754        case +a <> +b <> Nil{}:
755          Done{Split{b, a, {==}}}
756        case rest:
757          Fail{RankMismatch{}}
758
759def product_output(-lo: Order, -ro: Order, -left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32,
760  +right_inner: U32, +columns: U32,
761  -le: {Shape.dimensions(left) == arranged(lo, rows, inner) : List<&2, U32>},
762  -re: {Shape.dimensions(right) == arranged(ro, columns, right_inner) : List<&2, U32>},
763  -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
764  result: Shape.ConstructResult<[rows, columns]>) -> Result<&2, &2, Error, Product<lo, ro, left, right>>:
765  match result:
766    case Shape.SError{error}:
767      Fail{ShapeError{error}}
768    case Shape.SBuilt{shape, preserved}:
769      Done{Product{shape, rows, inner, right_inner, columns, le, re, preserved, compatible}}
770
771def product_compatible(-lo: Order, -ro: Order, -left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32,
772  +right_inner: U32, +columns: U32,
773  -le: {Shape.dimensions(left) == arranged(lo, rows, inner) : List<&2, U32>},
774  -re: {Shape.dimensions(right) == arranged(ro, columns, right_inner) : List<&2, U32>},
775  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))>) ->
776  Result<&2, &2, Error, Product<lo, ro, left, right>>:
777  match decision:
778    case Shape.BFalse{equal}:
779      Fail{InnerDimensionMismatch{}}
780    case Shape.BTrue{equal}:
781      product_output(lo, ro, left, right, rows, inner, right_inner, columns, le, re,
782        Lifecycle.natural_equal(U32.to_nat(inner), U32.to_nat(right_inner), equal),
783        Shape.construct([rows, columns]))
784
785def product_right(-lo: Order, -ro: Order, -left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32,
786  -le: {Shape.dimensions(left) == arranged(lo, rows, inner) : List<&2, U32>},
787  result: Result<&2, &2, Error, Split<ro, Shape.dimensions(right)>>) -> Result<&2, &2, Error, Product<lo, ro, left, right>>:
788  match result:
789    case Fail{error}:
790      Fail{error}
791    case Done{Split{+columns, +right_inner, same}}:
792      product_compatible(lo, ro, left, right, rows, inner, right_inner, columns, le, same,
793        Shape.bool_case(Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))))
794
795def product_left(-lo: Order, +ro: Order, -left: Shape.Shape, +right: Shape.Shape,
796  result: Result<&2, &2, Error, Split<lo, Shape.dimensions(left)>>) -> Result<&2, &2, Error, Product<lo, ro, left, right>>:
797  match result:
798    case Fail{error}:
799      Fail{error}
800    case Done{Split{+rows, +inner, same}}:
801      product_right(lo, ro, left, right, rows, inner, same, split_dims(ro, Shape.dimensions(right)))
802
803# The plan of a product: both operands rank 2, laid out in their orders,
804# with equal inner extents.
805def product_plan(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape) -> Result<&2, &2, Error, Product<lo, ro, left, right>>:
806  product_left(lo, ro, left, right, split_dims(lo, Shape.dimensions(left)))
807
808def product_parallel_output(+lo: Order, +ro: Order, +output: Shape.Shape, +inner: Nat, +rows: Nat, +columns: Nat,
809  +depth: Nat, +length: Nat,
810  -valid: {Storage.bounds(depth, length) == True{} : Bool},
811  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
812  left: Core.Dynamic & Array<F32>, right: Core.Dynamic & Array<F32>) ->
813  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
814  (kept_left, la) = left
815  (kept_right, ra) = right
816  Done{(kept_left, kept_right, Core.pack(output, Core.TTensor{Storage.Owned{depth, length,
817    product_build(depth, product_splits(depth, inner), lo, ro, inner, rows, columns, length, 0n, la, ra),
818    product_build_balanced(depth, product_splits(depth, inner), lo, ro, inner, rows, columns, length, 0n, la, ra),
819    valid}, exact}))}
820
821def product_allocated(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape,
822  +inner: Nat, +rows: Nat, +columns: Nat,
823  a: Core.Tensor<left>, b: Core.Tensor<right>, result: Result<&2, &1, Core.Error, Core.Tensor<output>>) ->
824  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
825  match result:
826    case Fail{error}:
827      Fail{TensorError{error}}
828    case Done{Core.TTensor{storage, exact}}:
829      match storage:
830        case Storage.Empty{}:
831          Done{(Core.pack(left, a), Core.pack(right, b), Core.pack(output, Core.TTensor{Storage.Empty{}, exact}))}
832        case Storage.Owned{+depth, +length, array, balanced, valid}:
833          product_parallel_output(lo, ro, output, inner, rows, columns, depth, length, valid, exact,
834            Elementwise.take(Core.pack(left, a)), Elementwise.take(Core.pack(right, b)))
835
836def product_multiply(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape, +plan: Product<lo, ro, left, right>,
837  a: Core.Tensor<left>, b: Core.Tensor<right>) -> Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
838  match plan:
839    case Product{+output, +rows, +inner, right_inner, +columns, le, re, oe, compatible}:
840      product_allocated(lo, ro, left, right, output, U32.to_nat(inner), U32.to_nat(rows), U32.to_nat(columns), a, b,
841        Core.new(output, 0.0))
842
843def product_planned(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>,
844  result: Result<&2, &2, Error, Product<lo, ro, left, right>>) ->
845  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
846  match result:
847    case Fail{error}:
848      Fail{error}
849    case Done{+plan}:
850      product_multiply(lo, ro, left, right, plan, a, b)
851
852# The product of two rank-2 operands laid out in orders lo and ro. Both
853# inputs come back to the caller as the kept halves of their clones.
854def product(+lo: Order, +ro: Order, left: Core.Dynamic, right: Core.Dynamic) ->
855  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
856  Core.Packed{+left_shape, a} = left
857  Core.Packed{+right_shape, b} = right
858  product_planned(lo, ro, left_shape, right_shape, a, b, product_plan(lo, ro, left_shape, right_shape))
859
860# x @ W^T for x of [rows, inner] and W of [columns, inner]: cell (row,
861# column) adds x[row, k] * W[column, k], W read in place.
862def matmul_nt(left: Core.Dynamic, right: Core.Dynamic) -> Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
863  product(Kept{}, Kept{}, left, right)
864
865# g^T @ x for g of [inner, rows] and x of [inner, columns]: cell (row,
866# column) adds g[k, row] * x[k, column], g read in place.
867def matmul_tn(left: Core.Dynamic, right: Core.Dynamic) -> Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
868  product(Swapped{}, Swapped{}, left, right)