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 ./core.bend as Core 6 import ./lifecycle.bend as Lifecycle 7 import ./elementwise.bend as Elementwise 8 import ./parallel.bend as Parallel 9 10 type Error is Data: 11 TensorError{error: Core.Error} 12 ShapeError{error: Shape.Error} 13 RankMismatch{} 14 InnerDimensionMismatch{} 15 16 type 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 23 def output_shape(-left: Shape.Shape, -right: Shape.Shape, plan: Plan) -> Shape.Shape: 24 Plan{shape, rows, inner, right_inner, columns, le, re, oe, compatible} = plan 25 shape 26 27 def 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>: 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 39 def 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) -> 44 Result<&2, &2, Error, Plan>: 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 53 def 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>: 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 65 def 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>: 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 76 def plan(+left: Shape.Shape, +right: Shape.Shape) -> Result<&2, &2, Error, Plan>: 77 plan_left(left, right, Shape.dimensions(left), {==}) 78 79 def right_read(-left: Shape.Shape, -right: Shape.Shape, 80 rec: Core.Tensor -> Core.Tensor -> F32 -> Result<&2, &1, Error, Core.Tensor & Core.Tensor & F32>, 81 a: Core.Tensor, left_value: F32, accumulator: F32, 82 pair: Core.Tensor & Result<&2, &2, Core.Error, F32>) -> 83 Result<&2, &1, Error, Core.Tensor & Core.Tensor & 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 91 def left_read(-left: Shape.Shape, +right: Shape.Shape, 92 rec: Core.Tensor -> Core.Tensor -> F32 -> Result<&2, &1, Error, Core.Tensor & Core.Tensor & F32>, 93 b: Core.Tensor, coordinates: List<&2, U32>, accumulator: F32, 94 pair: Core.Tensor & Result<&2, &2, Core.Error, F32>) -> 95 Result<&2, &1, Error, Core.Tensor & Core.Tensor & 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 103 def dot_decided(+left: Shape.Shape, +right: Shape.Shape, 104 rec: Core.Tensor -> Core.Tensor -> F32 -> Result<&2, &1, Error, Core.Tensor & Core.Tensor & F32>, 105 a: Core.Tensor, b: Core.Tensor, +row: U32, +column: U32, +cursor: Nat, +index: U32, accumulator: F32, 106 decision: Shape.BoolCase) -> 107 Result<&2, &1, Error, Core.Tensor & Core.Tensor & 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 114 def natural_dot(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 115 +row: U32, +column: U32, +cursor: Nat, accumulator: F32) -> 116 Result<&2, &1, Error, Core.Tensor & Core.Tensor & 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 126 def dot(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 127 +row: U32, +column: U32, +inner_index: U32, accumulator: F32) -> 128 Result<&2, &1, Error, Core.Tensor & Core.Tensor & F32>: 129 natural_dot(fuel, left, right, a, b, row, column, U32.to_nat(inner_index), accumulator) 130 131 def wrote(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, 132 rec: Core.Tensor -> Core.Tensor -> Core.Tensor -> 133 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>, 134 a: Core.Tensor, b: Core.Tensor, 135 pair: Core.Tensor & Result<&2, &2, Core.Error, Unit>) -> 136 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 144 def dotted(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, 145 rec: Core.Tensor -> Core.Tensor -> Core.Tensor -> 146 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>, 147 out: Core.Tensor, index: U32, 148 result: Result<&2, &1, Error, Core.Tensor & Core.Tensor & F32>) -> 149 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 157 def product_coordinates(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, 158 rec: Core.Tensor -> Core.Tensor -> Core.Tensor -> 159 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>, 160 +inner: U32, a: Core.Tensor, b: Core.Tensor, out: Core.Tensor, +index: U32, coordinates: List<&2, U32>) -> 161 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 168 def product_unflattened(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, 169 rec: Core.Tensor -> Core.Tensor -> Core.Tensor -> 170 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>, 171 +inner: U32, a: Core.Tensor, b: Core.Tensor, out: Core.Tensor, +index: U32, 172 result: Result<&2, &2, Shape.Error, Shape.CheckedCoordinate>) -> 173 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 180 def product_decided(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, 181 rec: Core.Tensor -> Core.Tensor -> Core.Tensor -> 182 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>, 183 +inner: U32, +columns: U32, a: Core.Tensor, b: Core.Tensor, out: Core.Tensor, +cursor: Nat, +index: U32, 184 decision: Shape.BoolCase) -> 185 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 192 def natural_product_loop(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, 193 +inner: U32, +columns: U32, a: Core.Tensor, b: Core.Tensor, out: Core.Tensor, +cursor: Nat) -> 194 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 204 def product_loop(fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, 205 +inner: U32, +columns: U32, a: Core.Tensor, b: Core.Tensor, out: Core.Tensor, +index: U32) -> 206 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 207 natural_product_loop(fuel, left, right, output, inner, columns, a, b, out, U32.to_nat(index)) 208 209 def created(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +inner: U32, +columns: U32, 210 a: Core.Tensor, b: Core.Tensor, result: Result<&2, &1, Core.Error, Core.Tensor>) -> 211 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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 219 def multiply(+left: Shape.Shape, +right: Shape.Shape, +plan: Plan, 220 a: Core.Tensor, b: Core.Tensor) -> 221 Result<&2, &1, Error, Core.Tensor & Core.Tensor & Core.Tensor>: 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. 238 def row_of(+columns: Nat, +position: Nat) -> Nat: 239 Nat.div(position, Nat.mul(columns, 1n)) 240 241 def column_of(+columns: Nat, +position: Nat) -> Nat: 242 Nat.div(Nat.mod(position, Nat.mul(columns, 1n)), 1n) 243 244 def left_index(+inner: Nat, +row: Nat, +k: Nat) -> U32: 245 U32.from_nat(Nat.add(Nat.mul(row, inner), k)) 246 247 def 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. 255 def product_run(fuel: Nat, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, accumulator: F32, 256 left: Array & F32, right: Array & F32) -> F32 & (Array & Array): 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. 268 def cell_dot(fuel: Nat, +inner: Nat, +columns: Nat, +row: Nat, +column: Nat, +cursor: Nat, accumulator: F32, 269 la: Array, ra: Array) -> F32 & (Array & Array): 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. 278 def dot_of(+inner: Nat, +columns: Nat, +position: Nat, la: Array, ra: Array) -> F32 & (Array & Array): 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. 282 def cell(+inner: Nat, +columns: Nat, +position: Nat, la: Array, ra: Array, inside: Bool) -> 283 F32 & (Array & Array): 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. 291 type Dims is Data: 292 Dims{inner: Nat, columns: Nat} 293 294 # Each thread below the forks reads its own copy. 295 def 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. 301 def cell_at(+dims: Dims, +position: Nat, inputs: Array & Array, inside: Bool) -> F32 & (Array & Array): 302 Dims{+inner, +columns} = dims 303 (la, ra) = inputs 304 cell(inner, columns, position, la, ra, inside) 305 306 def build(+depth: Nat, +splits: Nat, +inner: Nat, +columns: Nat, +count: Nat, +base: Nat, 307 la: Array, ra: Array) -> Array: 308 Parallel.build(~Dims, ~Dims, ~Array, ~Array, ~dims, ~Elementwise.split, ~cell_at, Dims{inner, columns}, depth, splits, count, base, (la, ra)) 309 310 def 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. 319 def 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 326 def least(+a: Nat, +b: Nat, smaller: Bool) -> Nat: 327 match smaller: 328 case True{}: 329 a 330 case False{}: 331 b 332 333 def 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. 342 def 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). 351 def 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. 355 def build_balanced(+depth: Nat, +splits: Nat, +inner: Nat, +columns: Nat, +count: Nat, +base: Nat, 356 la: Array, ra: Array) -> Model.Balanced(F32, depth, build(depth, splits, inner, columns, count, base, la, ra)): 357 Parallel.balanced(~Dims, ~Dims, ~Array, ~Array, ~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. 361 def 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, right: Core.Dynamic & Array) -> 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 373 def parallel_allocated(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +inner: Nat, +columns: Nat, 374 a: Core.Tensor, b: Core.Tensor, result: Result<&2, &1, Core.Error, Core.Tensor>) -> 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 387 def parallel_multiply(+left: Shape.Shape, +right: Shape.Shape, +plan: Plan, 388 a: Core.Tensor, b: Core.Tensor) -> 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 393 def parallel_planned(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 394 result: Result<&2, &2, Error, Plan>) -> 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. 405 def 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`). 423 type Order is Data: 424 Kept{} 425 Swapped{} 426 427 def 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 434 def 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. 439 def 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. 449 def 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, ra: Array) -> F32 & (Array & Array): 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 459 def product_dot_of(+lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +position: Nat, 460 la: Array, ra: Array) -> F32 & (Array & Array): 461 product_dot(inner, lo, ro, inner, rows, columns, row_of(columns, position), column_of(columns, position), 0n, 0.0, la, ra) 462 463 def product_cell(+lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +position: Nat, 464 la: Array, ra: Array, inside: Bool) -> F32 & (Array & Array): 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. 472 type Layout is Data: 473 Layout{left: Order, right: Order, inner: Nat, rows: Nat, columns: Nat} 474 475 def layout(+given: Layout) -> Layout: 476 Layout{+lo, +ro, +inner, +rows, +columns} = given 477 Layout{lo, ro, inner, rows, columns} 478 479 def product_cell_at(+layout: Layout, +position: Nat, inputs: Array & Array, inside: Bool) -> 480 F32 & (Array & Array): 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. 493 type 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. 497 def eight_read_8(+x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, +x5: F32, +x6: F32, pair: Array & F32) -> Array & Eight: 498 match pair: 499 case Tuple{a, +x7}: 500 (a, Eight{x0, x1, x2, x3, x4, x5, x6, x7}) 501 502 def eight_read_7(+x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, +x5: F32, pair: Array & F32, +o7: Nat) -> Array & 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 507 def eight_read_6(+x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, pair: Array & F32, +o6: Nat, +o7: Nat) -> Array & 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 512 def eight_read_5(+x0: F32, +x1: F32, +x2: F32, +x3: F32, pair: Array & F32, +o5: Nat, +o6: Nat, +o7: Nat) -> Array & 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 517 def eight_read_4(+x0: F32, +x1: F32, +x2: F32, pair: Array & F32, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array & 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 522 def eight_read_3(+x0: F32, +x1: F32, pair: Array & F32, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array & 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 527 def eight_read_2(+x0: F32, pair: Array & F32, +o2: Nat, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array & 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 532 def eight_read_1(pair: Array & F32, +o1: Nat, +o2: Nat, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array & 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 537 def read_eight(a: Array, +o0: Nat, +o1: Nat, +o2: Nat, +o3: Nat, +o4: Nat, +o5: Nat, +o6: Nat, +o7: Nat) -> Array & 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. 541 def 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. 549 def 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, right: Array & Eight) -> Eight & (Array & Array): 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. 559 def eight_leaves_of(+cells: Eight, state: Array & Array) -> Array & (Array & Array): 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 563 def eight_leaves(pair: Eight & (Array & Array)) -> Array & (Array & Array): 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. 569 def row_block_run(fuel: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +base: Nat, 570 la: Array, ra: Array) -> Array & (Array & Array): 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 580 def row_block(+layout: Layout, +base: Nat, state: Array & Array) -> Array & (Array & Array): 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. 586 def 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 590 def block_or_tree(+layout: Layout, +count: Nat, +base: Nat, state: Array & Array, one: Bool) -> 591 Array & (Array & Array): 592 match one: 593 case False{}: 594 Parallel.tree(~Layout, ~Array, ~Array, ~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. 600 def product_tree(+layout: Layout, +depth: Nat, +count: Nat, +base: Nat, state: Array & Array) -> 601 Array & (Array & Array): 602 match depth: 603 case 0n: 604 Parallel.tree(~Layout, ~Array, ~Array, ~product_cell_at, layout, 0n, count, base, state) 605 case 1n+d1: 606 match d1: 607 case 0n: 608 Parallel.tree(~Layout, ~Array, ~Array, ~product_cell_at, layout, 1n, count, base, state) 609 case 1n+d2: 610 match d2: 611 case 0n: 612 Parallel.tree(~Layout, ~Array, ~Array, ~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, ~Array, 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`. 624 def product_build(+depth: Nat, +splits: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, 625 +count: Nat, +base: Nat, la: Array, ra: Array) -> Array: 626 match depth splits: 627 case 1n+lower 1n+rest: 628 Parallel.fork(~Array, ~Array, 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, Array & Array, 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. 636 def eight_leaves_of_certified(+cells: Eight, state: Array & Array) -> 637 Parallel.Certified<(Array & Array), 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 642 def eight_leaves_certified(pair: Eight & (Array & Array)) -> 643 Parallel.Certified<(Array & Array), 3n, eight_leaves(pair)>: 644 match pair: 645 case Tuple{+cells, state}: 646 eight_leaves_of_certified(cells, state) 647 648 def row_block_run_certified(fuel: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, +base: Nat, 649 la: Array, ra: Array) -> 650 Parallel.Certified<(Array & Array), 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 660 def row_block_certified(+layout: Layout, +base: Nat, state: Array & Array) -> 661 Parallel.Certified<(Array & Array), 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 666 def block_or_tree_certified(+layout: Layout, +count: Nat, +base: Nat, state: Array & Array, one: Bool) -> 667 Parallel.Certified<(Array & Array), 3n, block_or_tree(layout, count, base, state, one)>: 668 match one: 669 case False{}: 670 Parallel.tree_certified(~Layout, ~Array, ~Array, ~product_cell_at, layout, 3n, count, base, state) 671 case True{}: 672 row_block_certified(layout, base, state) 673 674 def product_tree_certified(+layout: Layout, +depth: Nat, +count: Nat, +base: Nat, state: Array & Array) -> 675 Parallel.Certified<(Array & Array), depth, product_tree(layout, depth, count, base, state)>: 676 match depth: 677 case 0n: 678 Parallel.tree_certified(~Layout, ~Array, ~Array, ~product_cell_at, layout, 0n, count, base, state) 679 case 1n+d1: 680 match d1: 681 case 0n: 682 Parallel.tree_certified(~Layout, ~Array, ~Array, ~product_cell_at, layout, 1n, count, base, state) 683 case 1n+d2: 684 match d2: 685 case 0n: 686 Parallel.tree_certified(~Layout, ~Array, ~Array, ~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, ~Array, 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 696 def tree_balanced(+depth: Nat, -expected: Array & (Array & Array), 697 certified: Parallel.Certified<(Array & Array), depth, expected>) -> 698 Model.Balanced(F32, depth, Pair.fst(Array, Array & Array, expected)): 699 Parallel.Certified{result, same, balanced} = certified 700 %same : Model.Balanced(F32, depth, Pair.fst(Array, Array & Array, _)) 701 balanced 702 703 def product_build_balanced(+depth: Nat, +splits: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat, 704 +count: Nat, +base: Nat, la: Array, ra: Array) -> 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, ~Array, 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`. 722 def 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 729 type 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 736 def product_shape(-lo: Order, -ro: Order, -left: Shape.Shape, -right: Shape.Shape, plan: Product) -> 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. 741 type 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 744 def split_dims(+order: Order, dims: List<&2, U32>) -> Result<&2, &2, Error, Split>: 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 759 def 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>: 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 771 def 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) -> 776 Result<&2, &2, Error, Product>: 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 785 def 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>) -> Result<&2, &2, Error, Product>: 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 795 def product_left(-lo: Order, +ro: Order, -left: Shape.Shape, +right: Shape.Shape, 796 result: Result<&2, &2, Error, Split>) -> Result<&2, &2, Error, Product>: 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. 805 def product_plan(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape) -> Result<&2, &2, Error, Product>: 806 product_left(lo, ro, left, right, split_dims(lo, Shape.dimensions(left))) 807 808 def 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, right: Core.Dynamic & Array) -> 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 821 def 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, b: Core.Tensor, result: Result<&2, &1, Core.Error, Core.Tensor>) -> 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 836 def product_multiply(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape, +plan: Product, 837 a: Core.Tensor, b: Core.Tensor) -> 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 843 def product_planned(+lo: Order, +ro: Order, +left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor, b: Core.Tensor, 844 result: Result<&2, &2, Error, Product>) -> 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. 854 def 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. 862 def 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. 867 def matmul_tn(left: Core.Dynamic, right: Core.Dynamic) -> Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>: 868 product(Swapped{}, Swapped{}, left, right)