~/jendhub

src/tensor/permute_build.bend

bendtorch-core@0.1.0.0 · 156 lines · 24 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/model.bend as Model
4import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
5import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
6
7# The permute copy as a parallel build. Output cell i holds the input word at
8# the row-major index `Values.permute_source_index(input, axes, i)` names.
9# That index decodes i into output coordinates, scatters them to the input
10# axes and flattens them again, building three lists per cell. `recipe`
11# works out what does not depend on i, once per subtree below the forks (as
12# `Parallel.build` prepares it), so once per up to 2^15 cells, and `source`
13# computes the same index from it with a few divisions per axis and no list.
14# permute_build_proof.bend proves `source` equal to `permute_source_index`.
15
16type Request is Data:
17  Request{input: Shape.Shape, axes: List<&2, U32>}
18
19# How to take one output coordinate out of a row-major index: reduce the
20# index modulo each of `mods` in turn, then divide by `size`. Coordinate k
21# of `Shape.digits(dims, x)` takes the element counts after the first k
22# dimensions as `mods` and the count after dimension k as `size`.
23type Pick is Data:
24  Pick{mods: List<&2, Nat>, size: Nat}
25
26def widened(+outer: Nat, pick: Pick) -> Pick:
27  Pick{mods, size} = pick
28  Pick{outer <> mods, size}
29
30def widen(+outer: Nat, pick: Maybe<&2, Pick>) -> Maybe<&2, Pick>:
31  match pick:
32    case None{}:
33      None{}
34    case Some{found}:
35      Some{widened(outer, found)}
36
37# The pick of coordinate k of `Shape.digits(dims, x)`; None past the last
38# dimension.
39def pick_at(+dims: List<&2, Nat>, k: Nat) -> Maybe<&2, Pick>:
40  match dims:
41    case Nil{}:
42      None{}
43    case head <> tail:
44      match k:
45        case 0n:
46          Some{Pick{Nil{}, Shape.numel(tail)}}
47        case 1n+j:
48          widen(Shape.numel(tail), pick_at(tail, j))
49
50def located(+dims: List<&2, Nat>, position: Maybe<&2, Nat>) -> Maybe<&2, Pick>:
51  match position:
52    case None{}:
53      None{}
54    case Some{k}:
55      pick_at(dims, k)
56
57# One pick per input axis, from `cursor` on: the output coordinate that the
58# axes send to it, as `Axis.scatter` selects it.
59def picks(rank: Nat, +axes: List<&2, Nat>, +out: List<&2, Nat>, +cursor: Nat) -> Maybe<&2, List<&2, Pick>>:
60  match rank:
61    case 0n:
62      Some{Nil{}}
63    case 1n+p:
64      Axis.prepend(Pick, located(out, Axis.position(axes, cursor)), picks(p, axes, out, 1n+cursor))
65
66# The row-major stride of each dimension: the element count after it.
67def strides(+dims: List<&2, Nat>) -> List<&2, Nat>:
68  match dims:
69    case Nil{}:
70      Nil{}
71    case head <> tail:
72      Shape.numel(tail) <> strides(tail)
73
74type Recipe is Data:
75  Recipe{picks: Maybe<&2, List<&2, Pick>>, strides: List<&2, Nat>}
76
77def recipe(+request: Request) -> Recipe:
78  Request{+input, +axes} = request
79  +dims = Shape.dimensions(input)
80  Recipe{picks(Shape.rank_dimensions(dims), Shape.naturals(axes), Shape.naturals(Axis.order(U32, dims, Shape.naturals(axes))), 0n),
81    strides(Shape.naturals(dims))}
82
83def digit(+mods: List<&2, Nat>, +size: Nat, +x: Nat) -> Nat:
84  match mods:
85    case Nil{}:
86      Nat.div(x, size)
87    case +m <> rest:
88      digit(rest, size, Nat.mod(x, m))
89
90def coordinate(+pick: Pick, +x: Nat) -> Nat:
91  Pick{+mods, +size} = pick
92  digit(mods, size, x)
93
94# `Shape.offset` of the scattered coordinates, each through U32 and back as
95# `Shape.words` and `Shape.naturals` take it.
96def offset(+strides: List<&2, Nat>, +picks: List<&2, Pick>, +x: Nat) -> Nat:
97  match strides picks:
98    case +stride <> rest +pick <> others:
99      Nat.add(Nat.mul(U32.to_nat(U32.from_nat(coordinate(pick, x))), stride), offset(rest, others, x))
100    case strides picks:
101      0n
102
103def found(result: Maybe<&2, List<&2, Pick>>, +strides: List<&2, Nat>, +index: Nat) -> Nat:
104  match result:
105    case None{}:
106      0n
107    case Some{+picks}:
108      U32.to_nat(U32.from_nat(offset(strides, picks, U32.to_nat(U32.from_nat(index)))))
109
110# The input index output cell `index` reads.
111def source(+recipe: Recipe, +index: Nat) -> Nat:
112  Recipe{+picks, +strides} = recipe
113  found(picks, strides, index)
114
115# The input word index output cell `position` reads.
116def at(+recipe: Recipe, +position: Nat) -> U32:
117  U32.from_nat(source(recipe, position))
118
119def take(b: Array<F32>, pair: Array<F32> & F32) -> F32 & (Array<F32> & Array<F32>):
120  (a, x) = pair
121  (x, (a, b))
122
123def read(+recipe: Recipe, +position: Nat, state: Array<F32> & Array<F32>) -> F32 & (Array<F32> & Array<F32>):
124  (a, b) = state
125  take(b, Array.get(F32, a, at(recipe, position)))
126
127# The cell function of the build: the state is the input array and an empty
128# array that only fills the second slot of the state pair.
129def cell(+recipe: Recipe, +position: Nat, state: Array<F32> & Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
130  match inside:
131    case False{}:
132      (0.0, state)
133    case True{}:
134      read(recipe, position, state)
135
136def halves(lp: Array<F32> & Array<F32>, rp: Array<F32> & Array<F32>) -> (Array<F32> & Array<F32>) & (Array<F32> & Array<F32>):
137  (l1, l2) = lp
138  (r1, r2) = rp
139  ((l1, r1), (l2, r2))
140
141# Each half of a fork gets its own clone of the input array.
142def split(state: Array<F32> & Array<F32>) -> (Array<F32> & Array<F32>) & (Array<F32> & Array<F32>):
143  (a, b) = state
144  halves(Array.clone(F32, a), Array.clone(F32, b))
145
146# A subtree below the forks holds 2^15 cells, as in the elementwise build.
147def splits(+depth: Nat) -> Nat:
148  Nat.sub(depth, 15n)
149
150def build(+depth: Nat, +request: Request, +count: Nat, a: Array<F32>) -> Array<F32>:
151  Parallel.build(~Request, ~Recipe, ~Array<F32>, ~Array<F32>, ~recipe, ~split, ~cell, request, depth, splits(depth), count, 0n,
152    (a, Array.new(F32, 0n, 0.0)))
153
154def balanced(+depth: Nat, +request: Request, +count: Nat, a: Array<F32>) -> Model.Balanced(F32, depth, build(depth, request, count, a)):
155  Parallel.balanced(~Request, ~Recipe, ~Array<F32>, ~Array<F32>, ~recipe, ~split, ~cell, request, depth, splits(depth), count, 0n,
156    (a, Array.new(F32, 0n, 0.0)))