def product_cell

bendtorch-core@0.1.0.0 · src/tensor/matmul_PROOF.bend:1454

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/matmul_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
import bendtorch-core@0.1.0.0/src/storage/route.bend as Route
import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
import bendtorch-core@0.1.0.0/src/tensor/matmul_parallel_proof.bend as MatmulParallel
import bendtorch-core@0.1.0.0/src/tensor/matrix_index_proof.bend as MatrixIndex

1454# Below the output's element count, the parallel build writes the cell the
1455# sequential interpretation names.
1456def product_cell(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +rows: U32, +inner: U32, +right_inner: U32, +columns: U32,
1457  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}, -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
1458  -oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>}, -compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
1459  +depth: Nat, +length: Nat, -valid: {Storage.bounds(depth, length) == True{} : Bool},
1460  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
1461  -a: Core.Tensor<left>, -b: Core.Tensor<right>, -dl: Nat, -dr: Nat, -cl: Array<F32>, -cr: Array<F32>,
1462  -bl: Model.Balanced(F32, dl, cl), -br: Model.Balanced(F32, dr, cr),
1463  -left_reads: Parallel.SourceReads(Core.Packed{left, a}, dl, cl), -right_reads: Parallel.SourceReads(Core.Packed{right, b}, dr, cr),
1464  +position: Nat, -earlier: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
1465  Parallel.Fact<Laws.ProductCell<inner, left, right, a, b, output,
1466    Core.TTensor{Storage.Owned{depth, length,
1467      Matrix.build(depth, Matrix.product_splits(depth, U32.to_nat(inner)), U32.to_nat(inner), U32.to_nat(columns), length, 0n, cl, cr),
1468      Matrix.build_balanced(depth, Matrix.product_splits(depth, U32.to_nat(inner)), U32.to_nat(inner), U32.to_nat(columns), length, 0n, cl, cr),
1469      valid}, exact}, position>>:
1470  -inner_n = U32.to_nat(inner)
1471  -columns_n = U32.to_nat(columns)
1472  -row = Matrix.row_of(columns_n, position)
1473  -column = Matrix.column_of(columns_n, position)
1474  -index = U32.from_nat(position)
1475  -extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(output), [rows, columns], oe)
1476  -fit = MatrixIndex.position_fits(output, U32.to_nat(rows), columns_n, extents, position, earlier)
1477  -row_below = Arith.and_left(Nat.is_lt(row, U32.to_nat(rows)), Bool.and(Nat.is_lt(column, columns_n), True{}), fit)
1478  -column_below = Arith.and_left(Nat.is_lt(column, columns_n), True{},
1479    Arith.and_right(Nat.is_lt(row, U32.to_nat(rows)), Bool.and(Nat.is_lt(column, columns_n), True{}), fit))
1480  -word_exact = MatrixIndex.word_exact(position, output, earlier)
1481  -below = MatrixIndex.lt_bound(position, Shape.logical_numel(Shape.dimensions(output)), length,
1482    Equal.sym(Nat, length, Shape.logical_numel(Shape.dimensions(output)), exact), earlier)
1483  -inside = MatrixIndex.lt_same(U32.to_nat(index), position, length, word_exact, below)
1484  -tree = Matrix.build(depth, Matrix.product_splits(depth, inner_n), inner_n, columns_n, length, 0n, cl, cr)
1485  -path = Refinement.route(depth, index)
1486  -cells = {\{Parallel.Fact: p => p} :
1487    Parallel.Fact<MatmulParallel.Cells(depth, dl, dr, cl, cr, inner_n, columns_n, length, 0n, tree)> ->
1488    MatmulParallel.Cells(depth, dl, dr, cl, cr, inner_n, columns_n, length, 0n, tree)}(
1489      MatmulParallel.build_cells(depth, Matrix.product_splits(depth, inner_n), dl, dr, cl, cr, inner_n, columns_n, length, 0n, cl, cr, {==}, {==}, bl, br))
1490  -position_eq = Equal.trans(Nat, Nat.add(Model.offset(depth, path), 0n), U32.to_nat(index), position,
1491    Parallel.output_position(depth, length, index, valid, inside), word_exact)
1492  -at = {\{Parallel.Fact: p => p} :
1493    Parallel.Fact<MatmulParallel.CellIs(dl, dr, cl, cr, inner_n, columns_n, length, position, Model.observe(F32, depth, tree, path))> ->
1494    MatmulParallel.CellIs(dl, dr, cl, cr, inner_n, columns_n, length, position, Model.observe(F32, depth, tree, path))}(
1495      MatmulParallel.cell_shift(dl, dr, cl, cr, inner_n, columns_n, length, Nat.add(Model.offset(depth, path), 0n), position,
1496        Model.observe(F32, depth, tree, path), position_eq, cells(path)))
1497  -cell = {\{Parallel.Fact: p => p} :
1498    Parallel.Fact<MatmulParallel.CellAt<dl, dr, cl, cr, inner_n, columns_n, position, Model.observe(F32, depth, tree, path)>> ->
1499    MatmulParallel.CellAt<dl, dr, cl, cr, inner_n, columns_n, position, Model.observe(F32, depth, tree, path)>}(
1500      MatmulParallel.cell_below(dl, dr, cl, cr, inner_n, columns_n, length, position, Model.observe(F32, depth, tree, path), below, at))
1501  -matched = Equal.sym(Bool, True{}, Storage.index_matches(depth, length, index), Route.index_complete(depth, length, index, valid, inside))
1502  -decoded = {\{Parallel.Fact: p => p} : Parallel.Fact<Decoded<output, index, U32.from_nat(row), U32.from_nat(column)>> ->
1503    Decoded<output, index, U32.from_nat(row), U32.from_nat(column)>}(
1504      decoded_of(output, index, U32.from_nat(row), U32.from_nat(column), Shape.unflatten_checked(output, index), {==},
1505        MatrixIndex.unflatten_at(output, U32.to_nat(rows), columns_n, extents, position, earlier)))
1506  Parallel.Fact{{\{MatmulParallel.CellAt: word => observed => values => logical => fold =>
1507    {\{Decoded: coordinate => found => coordinates =>
1508      Laws.ProductCell{U32.from_nat(row), U32.from_nat(column), coordinate, found, coordinates, word_exact, values,
1509        {\{Parallel.Fact: p => p} :
1510          Parallel.Fact<Laws.LogicalDot(inner_n, left, right, a, b, U32.from_nat(row), U32.from_nat(column), 0n, values)> ->
1511          Laws.LogicalDot(inner_n, left, right, a, b, U32.from_nat(row), U32.from_nat(column), 0n, values)}(
1512          logical_dot(inner_n, left, right, rows, inner, right_inner, columns, row, column, le, re, compatible, row_below, column_below,
1513            a, b, dl, dr, cl, cr, left_reads, right_reads, 0n, {==}, values, logical)),
1514        word, fold, matched,
1515        Equal.trans(Maybe<&2, F32>, Lifecycle.observe_allowed(depth, length, tree, index, Nat.is_lt(U32.to_nat(index), length)),
1516          Model.observe(F32, depth, tree, path), Some{word},
1517          Parallel.observe_allowed_equal(depth, length, tree, index, Nat.is_lt(U32.to_nat(index), length), inside, matched), observed)}} :
1518      Decoded<output, index, U32.from_nat(row), U32.from_nat(column)> ->
1519      Laws.ProductCell<inner, left, right, a, b, output,
1520        Core.TTensor{Storage.Owned{depth, length, tree, Matrix.build_balanced(depth, Matrix.product_splits(depth, inner_n), inner_n, columns_n, length, 0n, cl, cr), valid}, exact},
1521        position>}(decoded)} :
1522    MatmulParallel.CellAt<dl, dr, cl, cr, inner_n, columns_n, position, Model.observe(F32, depth, tree, path)> ->
1523    Laws.ProductCell<inner, left, right, a, b, output,
1524      Core.TTensor{Storage.Owned{depth, length, tree, Matrix.build_balanced(depth, Matrix.product_splits(depth, inner_n), inner_n, columns_n, length, 0n, cl, cr), valid}, exact},
1525      position>}(cell)}