def product_build_balanced

bendtorch-core@0.1.0.0 · src/tensor/matmul.bend:703

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel

703def product_build_balanced(+depth: Nat, +splits: Nat, +lo: Order, +ro: Order, +inner: Nat, +rows: Nat, +columns: Nat,
704  +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>) ->
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<F32>, ~Array<F32>, 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)))