def build_eq

bendtorch-core@0.1.0.0 · src/tensor/contracts/matmul_parallel_proof.bend:734

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api

734# Whatever the split, the parallel build over the copies returns the spec.
735def build_eq(+span: Nat, +columns: Nat, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +count: Nat, +lsrc: Source<left, left_fill, left_ops, ha, dl, fa>, +rsrc: Source<right, right_fill, right_ops, hb, dr, fb>,
736  +output: Shape.Shape, +rows: Nat,
737  +left_extents: {Shape.extents(left) == [rows, span] : List<&2, Nat>},
738  +right_extents: {Shape.extents(right) == [span, columns] : List<&2, Nat>},
739  +output_extents: {Shape.extents(output) == [rows, columns] : List<&2, Nat>},
740  +count_exact: {count == Shape.logical_numel(Shape.dimensions(output)) : Nat}, +depth: Nat, +splits: Nat, +base: Nat) ->
741  {Matrix.build(depth, splits, span, columns, count, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)) == out_spec(span, columns, left, left_fill, left_ops, right, right_fill, right_ops, count, depth, base) : Array<F32>}:
742  match depth splits:
743    case 1n+lower 1n+rest:
744      fork_eq(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, lower, rest, base, build_eq(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, lsrc, rsrc, output, rows, left_extents, right_extents, output_extents, count_exact, lower, rest, base),
745        build_eq(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, lsrc, rsrc, output, rows, left_extents, right_extents, output_extents, count_exact, lower, rest, Nat.add(Model.capacity(lower), base)),
746        Array.clone(F32, Api.arr(ha, dl, fa)), {==}, Array.clone(F32, Api.arr(hb, dr, fb)), {==})
747    case 0n _:
748      tree_first(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, 0n, base, Parallel.tree(~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.cell_at, Matrix.dims(Matrix.Dims{span, columns}), 0n, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), tree_eq(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, lsrc, rsrc, output, rows, left_extents, right_extents, output_extents, count_exact, 0n, base))
749    case 1n+lower 0n:
750      tree_first(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, 1n+lower, base, Parallel.tree(~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.cell_at, Matrix.dims(Matrix.Dims{span, columns}), 1n+lower, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), tree_eq(span, columns, left, left_fill, left_ops, ha, dl, fa, right, right_fill, right_ops, hb, dr, fb, count, lsrc, rsrc, output, rows, left_extents, right_extents, output_extents, count_exact, 1n+lower, base))