1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/main.bend as Storage
4	import ../storage/model.bend as Model
5	import ../tensor/core.bend as Core
6	import ../tensor/elementwise.bend as Elementwise
7	import ../tensor/layout.bend as Layout
8	import ../tensor/matmul.bend as Matrix
9	import ../tensor/parallel.bend as Parallel
10	
11	# Reduction of a broadcast contribution back to a target shape. The target
12	# must broadcast into the contribution's shape. Target cell t holds the
13	# contribution words whose right-aligned projection lands on t, added one
14	# F32 add at a time in increasing contribution position, starting from
15	# +0.0. Cells are independent, so the output is built in parallel below.
16	
17	type Error is Data:
18	  TensorError{error: Core.Error}
19	  NotReducible{}
20	
21	# The parallel reduction. Target cell `target` sums the contribution words
22	# whose broadcast projection lands on it, in increasing contribution
23	# position from +0.0, the same F32 adds in the same order as the loop above.
24	# It walks only those words: over the index rows of the plan, one row per
25	# contribution axis, a row whose left stride is 0 (an axis the target
26	# reduces) visits each of its digits in increasing order, and any other row
27	# takes the one digit that `target` names there. The cells are built like
28	# `Matrix.build` builds its output: the top `splits` levels of the output
29	# tree fork, each half with its own clone of the contribution array.
30	
31	def leaf_add(accumulator: F32, pair: Array<F32> & F32) -> F32 & Array<F32>:
32	  (kept, word) = pair
33	  (F32.add(accumulator, word), kept)
34	
35	# How many digits a row walks for a target cell: every digit of an axis the
36	# target reduces (left stride 0), else one.
37	def span(+left: Nat, +extent: Nat) -> Nat:
38	  match left:
39	    case 0n:
40	      extent
41	    case 1n+less:
42	      1n
43	
44	def first_span(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>) -> Nat:
45	  match rows extents:
46	    case Shape.Row{stride, +left, right} <> rest +extent <> others:
47	      span(left, extent)
48	    case rows extents:
49	      0n
50	
51	# The first digit a row walks: 0 on a reduced axis (`Nat.div` by 0 is 0),
52	# else the digit `target` names there.
53	def first_digit(+rows: List<&2, Shape.Row>, +target: Nat) -> Nat:
54	  match rows:
55	    case Nil{}:
56	      0n
57	    case Shape.Row{stride, +left, right} <> rest:
58	      Nat.div(target, left)
59	
60	def walk_next(rec: F32 -> Array<F32> -> F32 & Array<F32>, pair: F32 & Array<F32>) -> F32 & Array<F32>:
61	  (accumulator, array) = pair
62	  rec(accumulator, array)
63	
64	# The words of the fiber of `target` below these rows, added to
65	# `accumulator` in increasing position. `extents` lists the contribution
66	# extent of each row and `base` is the contribution offset the rows above
67	# chose. The first row walks `fuel` digits from `digit`, each at offset
68	# base + right * digit, and the rows below get `target % left` (all of
69	# `target` on a reduced axis, as `Nat.mod` by 0 keeps it).
70	def fiber(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, fuel: Nat, +target: Nat, +base: Nat, +digit: Nat,
71	  accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
72	  match rows extents fuel:
73	    case Nil{} _ _:
74	      leaf_add(accumulator, Array.get(F32, array, U32.from_nat(base)))
75	    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
76	      (accumulator, array)
77	    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
78	      +inner = Nat.mod(target, left)
79	      walk_next(value => kept => fiber(rows, extents, remaining, target, base, 1n+digit, value, kept),
80	        fiber(rest, others, first_span(rest, others), inner, Nat.add(base, Nat.mul(right, digit)), first_digit(rest, inner),
81	          accumulator, array))
82	    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
83	      (accumulator, array)
84	
85	# Whether every row here takes one digit (span 1), so that the fiber below
86	# a digit of the row above is one word.
87	def single_rows(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>) -> Bool:
88	  match rows extents:
89	    case Nil{} _:
90	      True{}
91	    case Shape.Row{stride, +left, right} <> +rest +extent <> +others:
92	      Bool.and(Nat.is_eq(span(left, extent), 1n), single_rows(rest, others))
93	    case Shape.Row{stride, left, right} <> rest Nil{}:
94	      False{}
95	
96	# `single_rows` of the rows below the first one.
97	def single_below(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>) -> Bool:
98	  match rows extents:
99	    case Shape.Row{stride, left, right} <> +rest +extent <> +others:
100	      single_rows(rest, others)
101	    case _ _:
102	      False{}
103	
104	# The offset of the one word that rows of span 1 reach from `base`.
105	def single_offset(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +target: Nat, +base: Nat) -> Nat:
106	  match rows extents:
107	    case Shape.Row{stride, +left, +right} <> +rest +extent <> +others:
108	      single_offset(rest, others, Nat.mod(target, left), Nat.add(base, Nat.mul(right, Nat.div(target, left))))
109	    case _ _:
110	      base
111	
112	# The contribution offset of digit d of a row whose right stride is `step`.
113	def step_offset(+start: Nat, +step: Nat, +d: Nat) -> Nat:
114	  Nat.add(start, Nat.mul(step, d))
115	
116	def strided_offset(+step: Nat, +base: Nat, +digit: Nat, +lift: Nat) -> Nat:
117	  Nat.add(Nat.add(base, Nat.mul(step, digit)), lift)
118	
119	# The fuel + 1 words at offsets base + step * d + lift for d = digit, digit +
120	# 1, ..., added to `accumulator` in that order; `pair` is the read of the
121	# first. Each step reads the next word and hands it on as an argument.
122	def strided_run(fuel: Nat, +step: Nat, +base: Nat, +digit: Nat, +lift: Nat, accumulator: F32, pair: Array<F32> & F32) ->
123	  F32 & Array<F32>:
124	  match fuel pair:
125	    case 0n Tuple{array, word}:
126	      (F32.add(accumulator, word), array)
127	    case 1n+remaining Tuple{array, word}:
128	      strided_run(remaining, step, base, 1n+digit, lift, F32.add(accumulator, word),
129	        Array.get(F32, array, U32.from_nat(strided_offset(step, base, 1n+digit, lift))))
130	
131	def strided(fuel: Nat, +step: Nat, +base: Nat, +digit: Nat, +lift: Nat, accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
132	  match fuel:
133	    case 0n:
134	      (accumulator, array)
135	    case 1n+remaining:
136	      strided_run(remaining, step, base, digit, lift, accumulator, Array.get(F32, array, U32.from_nat(strided_offset(step, base, digit, lift))))
137	
138	# `fiber`, the same words added in the same order, except that a row whose
139	# rows below all take one digit (`single`, which is `single_below` of the
140	# rows; a later digit of the same row passes False) walks its digits in one
141	# strided loop: below digit d it reads the one word at base + right * d +
142	# single_offset(rest, others, target % left, 0), with no closure per word.
143	# Column sums, row sums and batch sums over one axis take that loop.
144	def fiber_fast(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, fuel: Nat, single: Bool, +target: Nat, +base: Nat, +digit: Nat,
145	  accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
146	  match rows extents fuel single:
147	    case Nil{} _ _ single:
148	      leaf_add(accumulator, Array.get(F32, array, U32.from_nat(base)))
149	    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others fuel True{}:
150	      strided(fuel, right, base, digit, single_offset(rest, others, Nat.mod(target, left), 0n), accumulator, array)
151	    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n False{}:
152	      (accumulator, array)
153	    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining False{}:
154	      +inner = Nat.mod(target, left)
155	      walk_next(value => kept => fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, value, kept),
156	        fiber_fast(rest, others, first_span(rest, others), single_below(rest, others), inner, step_offset(base, right, digit),
157	          first_digit(rest, inner), accumulator, array))
158	    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _ single:
159	      (accumulator, array)
160	
161	# Cell `target` of the reduction: its whole fiber from +0.0.
162	def fiber_of(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +target: Nat, array: Array<F32>) -> F32 & Array<F32>:
163	  fiber_fast(rows, extents, first_span(rows, extents), single_below(rows, extents), target, 0n, first_digit(rows, target), 0.0, array)
164	
165	# One target cell: its fiber sum from +0.0 when the cell is below `count`,
166	# else +0.0.
167	def cell(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +position: Nat, array: Array<F32>, inside: Bool) ->
168	  F32 & Array<F32>:
169	  match inside:
170	    case False{}:
171	      (0.0, array)
172	    case True{}:
173	      fiber_of(rows, extents, position, array)
174	
175	# What every cell of a build needs: the index rows of the plan and the
176	# contribution extent of each row.
177	type Fibers is Data:
178	  Fibers{rows: List<&2, Shape.Row>, extents: List<&2, Nat>}
179	
180	# Each thread below the forks reads its own copy.
181	def fibers(+given: Fibers) -> Fibers:
182	  Fibers{+rows, +extents} = given
183	  Fibers{rows, extents}
184	
185	# The cell function `Parallel.build` runs: its state is the contribution
186	# array, next to a Unit that fills the build's second state slot.
187	def keep(pair: F32 & Array<F32>, none: Unit) -> F32 & (Array<F32> & Unit):
188	  (value, array) = pair
189	  (value, (array, none))
190	
191	def cell_at(+fibers: Fibers, +position: Nat, state: Array<F32> & Unit, inside: Bool) -> F32 & (Array<F32> & Unit):
192	  Fibers{+rows, +extents} = fibers
193	  (array, none) = state
194	  keep(cell(rows, extents, position, array, inside), none)
195	
196	# Each half of a fork gets its own clone of the contribution array.
197	def halves(pair: Array<F32> & Array<F32>) -> (Array<F32> & Unit) & (Array<F32> & Unit):
198	  (first, second) = pair
199	  ((first, Unit{}), (second, Unit{}))
200	
201	def split(state: Array<F32> & Unit) -> (Array<F32> & Unit) & (Array<F32> & Unit):
202	  (array, none) = state
203	  halves(Array.clone(F32, array))
204	
205	def build(+depth: Nat, +splits: Nat, +rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +count: Nat, +base: Nat,
206	  array: Array<F32>) -> Array<F32>:
207	  Parallel.build(~Fibers, ~Fibers, ~Array<F32>, ~Unit, ~fibers, ~split, ~cell_at, Fibers{rows, extents}, depth, splits, count, base,
208	    (array, Unit{}))
209	
210	# The parallel build returns a balanced tree of its depth, whatever the split.
211	def build_balanced(+depth: Nat, +splits: Nat, +rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +count: Nat, +base: Nat,
212	  array: Array<F32>) -> Model.Balanced(F32, depth, build(depth, splits, rows, extents, count, base, array)):
213	  Parallel.balanced(~Fibers, ~Fibers, ~Array<F32>, ~Unit, ~fibers, ~split, ~cell_at, Fibers{rows, extents}, depth, splits, count, base,
214	    (array, Unit{}))
215	
216	# A target cell costs one term per contribution word that lands on it.
217	def terms(+plan: Shape.BroadcastPlan, +length: Nat) -> Nat:
218	  Nat.div(Shape.elements(Shape.output_shape(plan)), length)
219	
220	# Whether some axis below these rows has extent above 1 and is kept by the
221	# target (left stride above 0).
222	def kept_after(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>) -> Bool:
223	  match rows extents:
224	    case Shape.Row{stride, +left, right} <> +rest +extent <> +others:
225	      Bool.or(Bool.and(Nat.is_lt(1n, extent), Bool.not(Nat.is_eq(left, 0n))), kept_after(rest, others))
226	    case _ _:
227	      False{}
228	
229	# Whether every fiber is one contiguous run of the contribution: no axis the
230	# target keeps comes after an axis it reduces, counting only axes of extent
231	# above 1. Row sums such as [n, m] to [n, 1] are contiguous; column and batch
232	# sums such as [n, m] to [1, m] or [b, n, m] to [n, m] are strided.
233	def contiguous(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>) -> Bool:
234	  match rows extents:
235	    case Shape.Row{stride, +left, right} <> +rest +extent <> +others:
236	      Bool.pick(Bool, Bool.and(Nat.is_lt(1n, extent), Nat.is_eq(left, 0n)), Bool.not(kept_after(rest, others)),
237	        contiguous(rest, others))
238	    case _ _:
239	      True{}
240	
241	# How many top levels of the target tree fork. Every fork clones the whole
242	# contribution. A target whose fibers are contiguous splits by work, as
243	# `Matrix.splits` does with the terms of a cell; a target with strided fibers
244	# never forks, since there the forked build ran slower at 8 and 24 threads
245	# than one thread. The proofs hold for any split.
246	def splits(+plan: Shape.BroadcastPlan, +depth: Nat, +length: Nat) -> Nat:
247	  +extents = Shape.extents(Shape.output_shape(plan))
248	  +rows = Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), extents)
249	  Bool.pick(Nat, contiguous(rows, extents), Matrix.splits(depth, terms(plan, length)), 0n)
250	
251	def parallel_rows(+plan: Shape.BroadcastPlan, +target: Shape.Shape, +depth: Nat, +length: Nat,
252	  -valid: {Storage.bounds(depth, length) == True{} : Bool},
253	  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
254	  taken: Core.Dynamic & Array<F32>, rows: List<&2, Shape.Row>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
255	  (kept, array) = taken
256	  +extents = Shape.extents(Shape.output_shape(plan))
257	  +splits = splits(plan, depth, length)
258	  (kept, Done{Core.Packed{target, Core.TTensor{Storage.Owned{depth, length,
259	    build(depth, splits, rows, extents, length, 0n, array),
260	    build_balanced(depth, splits, rows, extents, length, 0n, array),
261	    valid}, exact}}})
262	
263	# An empty contribution leaves every target cell at +0.0.
264	def parallel_indexed(input: Core.Dynamic, +target: Shape.Shape, +plan: Shape.BroadcastPlan, output: Core.Tensor<target>,
265	  +depth: Nat, +length: Nat, -valid: {Storage.bounds(depth, length) == True{} : Bool},
266	  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
267	  map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>) ->
268	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
269	  match map:
270	    case Shape.Unindexed{zero}:
271	      (input, Done{Core.Packed{target, output}})
272	    case Shape.Indexed{rows, same}:
273	      parallel_rows(plan, target, depth, length, valid, exact, Elementwise.take(input), rows)
274	
275	def parallel_allocated(input: Core.Dynamic, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
276	  result: Result<&2, &1, Core.Error, Core.Tensor<target>>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
277	  match result:
278	    case Fail{error}:
279	      (input, Fail{TensorError{error}})
280	    case Done{Core.TTensor{storage, exact}}:
281	      match storage:
282	        case Storage.Empty{}:
283	          (input, Done{Core.Packed{target, Core.TTensor{Storage.Empty{}, exact}}})
284	        case Storage.Owned{+depth, +length, array, balanced, valid}:
285	          parallel_indexed(input, target, plan, Core.TTensor{Storage.Owned{depth, length, array, balanced, valid}, exact},
286	            depth, length, valid, exact, Shape.index_map(plan))
287	
288	def fitted(input: Core.Dynamic, +target: Shape.Shape, +plan: Shape.BroadcastPlan, same: Bool) ->
289	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
290	  match same:
291	    case False{}:
292	      (input, Fail{NotReducible{}})
293	    case True{}:
294	      parallel_allocated(input, target, plan, Core.new(target, 0.0))
295	
296	def planned(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +target: Shape.Shape,
297	  result: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
298	  match result:
299	    case Fail{error}:
300	      (Core.Packed{shape, tensor}, Fail{TensorError{Core.ShapeError{error}}})
301	    case Done{+plan}:
302	      fitted(Core.Packed{shape, tensor}, target, plan,
303	        Layout.words_same(Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(shape)))
304	
305	# The contribution owner comes back unchanged next to the result. A target
306	# that does not broadcast with the contribution fails with the planner's
307	# broadcast error; a target whose broadcast with the contribution is larger
308	# than the contribution (for example [2, 3] for a [3] contribution) fails with
309	# NotReducible.
310	def reduce_to(input: Core.Dynamic, +target: Shape.Shape) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
311	  Core.Packed{+shape, tensor} = input
312	  planned(shape, tensor, target, Shape.broadcast_checked(target, shape))
313	
314	def reduce_to_dimensions_result(input: Core.Dynamic, -dimensions: List<&2, U32>, result: Shape.ConstructResult<dimensions>) ->
315	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
316	  match result:
317	    case Shape.SError{error}:
318	      (input, Fail{TensorError{Core.ShapeError{error}}})
319	    case Shape.SBuilt{+target, preserved}:
320	      reduce_to(input, target)
321	
322	def reduce_to_dimensions(input: Core.Dynamic, +dimensions: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
323	  reduce_to_dimensions_result(input, dimensions, Shape.construct(dimensions))
324	
325	def error_code(error: Error) -> U32:
326	  match error:
327	    case TensorError{error}:
328	      Core.error_code(error)
329	    case NotReducible{}:
330	      101