src/autograd/reduce.bend

bendtorch-core@0.1.0.0 · 330 lines · 38 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/elementwise.bend as Elementwise
7import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
8import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
9import bendtorch-core@0.1.0.0/src/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
17type 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
31def 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.
37def span(+left: Nat, +extent: Nat) -> Nat:
38  match left:
39    case 0n:
40      extent
41    case 1n+less:
42      1n
43
44def 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.
53def 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
60def 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).
70def 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.
87def 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.
97def 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`.
105def 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`.
113def step_offset(+start: Nat, +step: Nat, +d: Nat) -> Nat:
114  Nat.add(start, Nat.mul(step, d))
115
116def 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.
122def 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
131def 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.
144def 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.
162def 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.
167def 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.
177type 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.
181def 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.
187def keep(pair: F32 & Array<F32>, none: Unit) -> F32 & (Array<F32> & Unit):
188  (value, array) = pair
189  (value, (array, none))
190
191def 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.
197def halves(pair: Array<F32> & Array<F32>) -> (Array<F32> & Unit) & (Array<F32> & Unit):
198  (first, second) = pair
199  ((first, Unit{}), (second, Unit{}))
200
201def split(state: Array<F32> & Unit) -> (Array<F32> & Unit) & (Array<F32> & Unit):
202  (array, none) = state
203  halves(Array.clone(F32, array))
204
205def 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.
211def 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.
217def 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).
222def 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.
233def 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.
246def 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
251def 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.
264def 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
275def 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
288def 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
296def 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.
310def 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
314def 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
322def 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
325def error_code(error: Error) -> U32:
326  match error:
327    case TensorError{error}:
328      Core.error_code(error)
329    case NotReducible{}:
330      101