1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/model.bend as Model 4 import ./axis_spec.bend as Axis 5 import ./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 16 type 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`. 23 type Pick is Data: 24 Pick{mods: List<&2, Nat>, size: Nat} 25 26 def widened(+outer: Nat, pick: Pick) -> Pick: 27 Pick{mods, size} = pick 28 Pick{outer <> mods, size} 29 30 def 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. 39 def 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 50 def 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. 59 def 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. 67 def 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 74 type Recipe is Data: 75 Recipe{picks: Maybe<&2, List<&2, Pick>>, strides: List<&2, Nat>} 76 77 def 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 83 def 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 90 def 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. 96 def 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 103 def 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. 111 def 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. 116 def at(+recipe: Recipe, +position: Nat) -> U32: 117 U32.from_nat(source(recipe, position)) 118 119 def take(b: Array, pair: Array & F32) -> F32 & (Array & Array): 120 (a, x) = pair 121 (x, (a, b)) 122 123 def read(+recipe: Recipe, +position: Nat, state: Array & Array) -> F32 & (Array & Array): 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. 129 def cell(+recipe: Recipe, +position: Nat, state: Array & Array, inside: Bool) -> F32 & (Array & Array): 130 match inside: 131 case False{}: 132 (0.0, state) 133 case True{}: 134 read(recipe, position, state) 135 136 def halves(lp: Array & Array, rp: Array & Array) -> (Array & Array) & (Array & Array): 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. 142 def split(state: Array & Array) -> (Array & Array) & (Array & Array): 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. 147 def splits(+depth: Nat) -> Nat: 148 Nat.sub(depth, 15n) 149 150 def build(+depth: Nat, +request: Request, +count: Nat, a: Array) -> Array: 151 Parallel.build(~Request, ~Recipe, ~Array, ~Array, ~recipe, ~split, ~cell, request, depth, splits(depth), count, 0n, 152 (a, Array.new(F32, 0n, 0.0))) 153 154 def balanced(+depth: Nat, +request: Request, +count: Nat, a: Array) -> Model.Balanced(F32, depth, build(depth, request, count, a)): 155 Parallel.balanced(~Request, ~Recipe, ~Array, ~Array, ~recipe, ~split, ~cell, request, depth, splits(depth), count, 0n, 156 (a, Array.new(F32, 0n, 0.0)))