src/tensor/matrix_index_proof.bend

bendtorch-core@0.1.0.0 · 251 lines · 24 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as ShapeLaws
4import bendtorch-core@0.1.0.0/src/shape/PROOF.bend as ShapeProof
5import bendtorch-core@0.1.0.0/src/shape/boundary_proofs.bend as Boundary
6import bendtorch-core@0.1.0.0/src/shape/coordinate_proofs.bend as Coordinates
7import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
8import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
9import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
10import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
11import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
12import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
13import bendtorch-core@0.1.0.0/src/tensor/read_contract.bend as Reads
14
15# Index facts about a rank-2 product: where an output cell sits, where its
16# factors sit in the inputs, and that each of them is in bounds.
17
18def word_exact(+number: Nat, +shape: Shape.Shape,
19  below: {Nat.is_lt(number, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
20  {U32.to_nat(U32.from_nat(number)) == number : Nat}:
21  Conversion.natural_word_exact(number, Shape.count(shape),
22    %Equal.sym(Nat, U32.to_nat(Shape.count(shape)), Shape.logical_numel(Shape.dimensions(shape)), Reads.count_exact(shape)) :
23      {Nat.is_le(number, _) == True{} : Bool}
24    Conversion.lt_le(number, Shape.logical_numel(Shape.dimensions(shape)), below))
25
26def word_inside(+number: Nat, +shape: Shape.Shape,
27  +below: {Nat.is_lt(number, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
28  {Nat.is_lt(U32.to_nat(U32.from_nat(number)), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}:
29  %Equal.sym(Nat, U32.to_nat(U32.from_nat(number)), number, word_exact(number, shape, below)) :
30    {Nat.is_lt(_, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}
31  below
32
33def lt_mul_add(+a: Nat, +m: Nat, +b: Nat, +n: Nat, a_below: {Nat.is_lt(a, m) == True{} : Bool},
34  b_below: {Nat.is_lt(b, n) == True{} : Bool}) -> {Nat.is_lt(Nat.add(Nat.mul(a, n), b), Nat.mul(m, n)) == True{} : Bool}:
35  match a m:
36    case 0n 0n:
37      Empty.absurd({Nat.is_lt(b, 0n) == True{} : Bool}, Arith.clash(a_below))
38    case 0n 1n+q:
39      Arith.lt_le_trans(b, n, Nat.add(n, Nat.mul(q, n)), b_below, Api.le_add_right(n, Nat.mul(q, n)))
40    case 1n+p 0n:
41      Empty.absurd({Nat.is_lt(Nat.add(Nat.add(n, Nat.mul(p, n)), b), 0n) == True{} : Bool}, Arith.clash(a_below))
42    case 1n+p 1n+q:
43      %Index.add_assoc(n, Nat.mul(p, n), b) : {Nat.is_lt(_, Nat.add(n, Nat.mul(q, n))) == True{} : Bool}
44      Api.lt_add_left(n, Nat.add(Nat.mul(p, n), b), Nat.mul(q, n), lt_mul_add(p, q, b, n, a_below, b_below))
45
46# Factor k of cell (row, column) lies inside each input.
47def left_below(+rows: Nat, +inner: Nat, +row: Nat, row_below: {Nat.is_lt(row, rows) == True{} : Bool}, +left: Shape.Shape,
48  left_extents: {Shape.extents(left) == [rows, inner] : List<&2, Nat>}, +k: Nat, below: {Nat.is_lt(k, inner) == True{} : Bool}) ->
49  {Nat.is_lt(Nat.add(Nat.mul(row, inner), k), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool}:
50  %Equal.sym(List<&2, Nat>, Shape.extents(left), [rows, inner], left_extents) :
51    {Nat.is_lt(Nat.add(Nat.mul(row, inner), k), Shape.numel(_)) == True{} : Bool}
52  %Equal.sym(Nat, Nat.mul(inner, 1n), inner, Nats.mul_one_right(inner)) :
53    {Nat.is_lt(Nat.add(Nat.mul(row, inner), k), Nat.mul(rows, _)) == True{} : Bool}
54  lt_mul_add(row, rows, k, inner, row_below, below)
55
56def right_below(+inner: Nat, +columns: Nat, +column: Nat, column_below: {Nat.is_lt(column, columns) == True{} : Bool}, +right: Shape.Shape,
57  right_extents: {Shape.extents(right) == [inner, columns] : List<&2, Nat>}, +k: Nat, below: {Nat.is_lt(k, inner) == True{} : Bool}) ->
58  {Nat.is_lt(Nat.add(Nat.mul(k, columns), column), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool}:
59  %Equal.sym(List<&2, Nat>, Shape.extents(right), [inner, columns], right_extents) :
60    {Nat.is_lt(Nat.add(Nat.mul(k, columns), column), Shape.numel(_)) == True{} : Bool}
61  %Equal.sym(Nat, Nat.mul(columns, 1n), columns, Nats.mul_one_right(columns)) :
62    {Nat.is_lt(Nat.add(Nat.mul(k, columns), column), Nat.mul(inner, _)) == True{} : Bool}
63  lt_mul_add(k, inner, column, columns, below, column_below)
64
65def lt_from(+a: Nat, +b: Nat, below: {Nat.is_lt(a, b) == True{} : Bool}) -> Shape.LT(a, b):
66  match a b:
67    case 0n 0n:
68      Empty.absurd(Shape.LT(0n, 0n), Arith.clash(below))
69    case 0n 1n+q:
70      Unit{}
71    case 1n+p 0n:
72      Empty.absurd(Shape.LT(1n+p, 0n), Arith.clash(below))
73    case 1n+p 1n+q:
74      lt_from(p, q, below)
75
76def lt_same(-a: Nat, -b: Nat, -bound: Nat, same: {a == b : Nat}, below: {Nat.is_lt(b, bound) == True{} : Bool}) ->
77  {Nat.is_lt(a, bound) == True{} : Bool}:
78  %Equal.sym(Nat, a, b, same) : {Nat.is_lt(_, bound) == True{} : Bool}
79  below
80
81def lt_bound(-a: Nat, -bound: Nat, -other: Nat, same: {bound == other : Nat}, below: {Nat.is_lt(a, bound) == True{} : Bool}) ->
82  {Nat.is_lt(a, other) == True{} : Bool}:
83  %same : {Nat.is_lt(a, _) == True{} : Bool}
84  below
85
86def word_below(+number: Nat, +limit: U32, below: {Nat.is_lt(number, U32.to_nat(limit)) == True{} : Bool}) ->
87  {U32.to_nat(U32.from_nat(number)) == number : Nat}:
88  Conversion.natural_word_exact(number, limit, Conversion.lt_le(number, U32.to_nat(limit), below))
89
90def some_equal(-A: Data, -x: A, -y: A, equal: {Some{x} == Some{y} : Maybe<&2, A>}) -> {x == y : A}:
91  Equal.cong(Maybe<&2, A>, A, found => Maybe.default(&2, A, found, x), Some{x}, Some{y}, equal)
92
93def div_one(+number: Nat) -> {Nat.div(number, 1n) == number : Nat}:
94  %Equal.trans(Nat, Nat.add(Nat.mul(number, 1n), 0n), Nat.mul(number, 1n), number,
95    Conversion.add_zero(Nat.mul(number, 1n)), Nats.mul_one_right(number)) : {Nat.div(_, 1n) == number : Nat}
96  Equal.cong(Nat & Nat, Nat, pair => Pair.fst(Nat, Nat, pair), Nat.divmod(Nat.add(Nat.mul(number, 1n), 0n), 1n), (number, 0n),
97    Equal.sym(Nat & Nat, (number, 0n), Nat.divmod(Nat.add(Nat.mul(number, 1n), 0n), 1n), Nats.divmod_exact(number, 1n, 0n, Unit{})))
98
99def row_at(+row: Nat, +column: Nat, +columns: Nat, below: {Nat.is_lt(column, columns) == True{} : Bool}) ->
100  {Matrix.row_of(columns, Nat.add(Nat.mul(row, columns), column)) == row : Nat}:
101  %Equal.sym(Nat, Nat.mul(columns, 1n), columns, Nats.mul_one_right(columns)) :
102    {Nat.div(Nat.add(Nat.mul(row, columns), column), _) == row : Nat}
103  Equal.cong(Nat & Nat, Nat, pair => Pair.fst(Nat, Nat, pair), Nat.divmod(Nat.add(Nat.mul(row, columns), column), columns), (row, column),
104    Equal.sym(Nat & Nat, (row, column), Nat.divmod(Nat.add(Nat.mul(row, columns), column), columns),
105      Nats.divmod_exact(row, columns, column, lt_from(column, columns, below))))
106
107def column_at(+row: Nat, +column: Nat, +columns: Nat, below: {Nat.is_lt(column, columns) == True{} : Bool}) ->
108  {Matrix.column_of(columns, Nat.add(Nat.mul(row, columns), column)) == column : Nat}:
109  %Equal.sym(Nat, Nat.mul(columns, 1n), columns, Nats.mul_one_right(columns)) :
110    {Nat.div(Nat.mod(Nat.add(Nat.mul(row, columns), column), _), 1n) == column : Nat}
111  %Equal.cong(Nat & Nat, Nat, pair => Pair.snd(Nat, Nat, pair), (row, column), Nat.divmod(Nat.add(Nat.mul(row, columns), column), columns),
112    Nats.divmod_exact(row, columns, column, lt_from(column, columns, below))) : {Nat.div(_, 1n) == column : Nat}
113  div_one(column)
114
115def position_digits(+output: Shape.Shape, +rows: Nat, +columns: Nat,
116  extents: {Shape.extents(output) == [rows, columns] : List<&2, Nat>}, +position: Nat,
117  +fits: {True{} == Nat.is_lt(position, Shape.numel(Shape.extents(output))) : Bool}) ->
118  {Shape.indices(Shape.extents(output), Shape.unflatten(Shape.extents(output),
119    Coordinates.fin_from(Shape.numel(Shape.extents(output)), position, fits))) ==
120    [Matrix.row_of(columns, position), Matrix.column_of(columns, position)] : List<&2, Nat>}:
121  +bound = Shape.numel(Shape.extents(output))
122  +index = Coordinates.fin_from(bound, position, fits)
123  Equal.trans(List<&2, Nat>, Shape.indices(Shape.extents(output), Shape.unflatten(Shape.extents(output), index)),
124    Shape.digits(Shape.extents(output), Shape.Fin.value(bound, index)), [Matrix.row_of(columns, position), Matrix.column_of(columns, position)],
125    Coordinates.digits_unflatten(Shape.extents(output), index),
126    Equal.trans(List<&2, Nat>, Shape.digits(Shape.extents(output), Shape.Fin.value(bound, index)),
127      Shape.digits(Shape.extents(output), position), [Matrix.row_of(columns, position), Matrix.column_of(columns, position)],
128      Equal.cong(Nat, List<&2, Nat>, value => Shape.digits(Shape.extents(output), value), Shape.Fin.value(bound, index), position,
129        Coordinates.fin_from_value(bound, position, fits)),
130      Equal.cong(List<&2, Nat>, List<&2, Nat>, dims => Shape.digits(dims, position), Shape.extents(output), [rows, columns], extents)))
131
132def unflatten_at(+output: Shape.Shape, +rows: Nat, +columns: Nat,
133  extents: {Shape.extents(output) == [rows, columns] : List<&2, Nat>}, +position: Nat,
134  +below: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
135  {ShapeLaws.coordinate_result(output, Shape.unflatten_checked(output, U32.from_nat(position))) ==
136    Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]} : Maybe<&2, List<&2, U32>>}:
137  +bound = Shape.numel(Shape.extents(output))
138  +fits = Equal.sym(Bool, Nat.is_lt(position, bound), True{}, below)
139  +index = Coordinates.fin_from(bound, position, fits)
140  Equal.trans(Maybe<&2, List<&2, U32>>, ShapeLaws.coordinate_result(output, Shape.unflatten_checked(output, U32.from_nat(position))),
141    ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, U32.to_nat(U32.from_nat(position)))),
142    Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]},
143    Boundary.unflatten_checked(output, U32.from_nat(position)),
144    Equal.trans(Maybe<&2, List<&2, U32>>,
145      ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, U32.to_nat(U32.from_nat(position)))),
146      ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, position)),
147      Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]},
148      Equal.cong(Nat, Maybe<&2, List<&2, U32>>, value => ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, value)),
149        U32.to_nat(U32.from_nat(position)), position, word_exact(position, output, below)),
150      Equal.trans(Maybe<&2, List<&2, U32>>,
151        ShapeLaws.unflatten_spec(Shape.extents(output), ShapeLaws.fin_of(bound, position)),
152        ShapeLaws.unflatten_spec(Shape.extents(output), Some{index}),
153        Some{[U32.from_nat(Matrix.row_of(columns, position)), U32.from_nat(Matrix.column_of(columns, position))]},
154        Equal.cong(Maybe<&2, Shape.Fin(bound)>, Maybe<&2, List<&2, U32>>, found => ShapeLaws.unflatten_spec(Shape.extents(output), found),
155          ShapeLaws.fin_of(bound, position), Some{index}, Boundary.fin_of_from(bound, position, fits)),
156        Equal.cong(List<&2, Nat>, Maybe<&2, List<&2, U32>>, values => Some{Shape.words(values)},
157          Shape.indices(Shape.extents(output), Shape.unflatten(Shape.extents(output), index)),
158          [Matrix.row_of(columns, position), Matrix.column_of(columns, position)],
159          position_digits(output, rows, columns, extents, position, fits)))))
160
161def position_fits(+output: Shape.Shape, +rows: Nat, +columns: Nat,
162  +extents: {Shape.extents(output) == [rows, columns] : List<&2, Nat>}, +position: Nat,
163  +below: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
164  {Shape.fit([rows, columns], [Matrix.row_of(columns, position), Matrix.column_of(columns, position)]) == True{} : Bool}:
165  +bound = Shape.numel(Shape.extents(output))
166  +fits = Equal.sym(Bool, Nat.is_lt(position, bound), True{}, below)
167  +coordinate = Shape.unflatten(Shape.extents(output), Coordinates.fin_from(bound, position, fits))
168  +digits : List<&2, Nat> = [Matrix.row_of(columns, position), Matrix.column_of(columns, position)]
169  Equal.sym(Bool, True{}, Shape.fit([rows, columns], digits),
170    Equal.trans(Bool, True{}, Shape.fit(Shape.extents(output), Shape.indices(Shape.extents(output), coordinate)), Shape.fit([rows, columns], digits),
171      Coordinates.fit_indices(Shape.extents(output), coordinate),
172      Equal.trans(Bool, Shape.fit(Shape.extents(output), Shape.indices(Shape.extents(output), coordinate)),
173        Shape.fit(Shape.extents(output), digits), Shape.fit([rows, columns], digits),
174        Equal.cong(List<&2, Nat>, Bool, values => Shape.fit(Shape.extents(output), values),
175          Shape.indices(Shape.extents(output), coordinate), digits, position_digits(output, rows, columns, extents, position, fits)),
176        Equal.cong(List<&2, Nat>, Bool, dims => Shape.fit(dims, digits), Shape.extents(output), [rows, columns], extents))))
177
178def offset_pair(+a: Nat, +b: Nat, +first: Nat, +second: Nat) ->
179  {Shape.offset([first, second], [a, b]) == Nat.add(Nat.mul(a, second), b) : Nat}:
180  %Equal.sym(Nat, Nat.mul(second, 1n), second, Nats.mul_one_right(second)) :
181    {Nat.add(Nat.mul(a, _), Nat.add(Nat.mul(b, 1n), 0n)) == Nat.add(Nat.mul(a, second), b) : Nat}
182  %Equal.sym(Nat, Nat.mul(b, 1n), b, Nats.mul_one_right(b)) :
183    {Nat.add(Nat.mul(a, second), Nat.add(_, 0n)) == Nat.add(Nat.mul(a, second), b) : Nat}
184  %Equal.sym(Nat, Nat.add(b, 0n), b, Conversion.add_zero(b)) :
185    {Nat.add(Nat.mul(a, second), _) == Nat.add(Nat.mul(a, second), b) : Nat}
186  {==}
187
188def flat_pair(+shape: Shape.Shape, +first: U32, +second: U32, +x: U32, +y: U32,
189  +dimensions: {Shape.dimensions(shape) == [first, second] : List<&2, U32>},
190  x_inside: {Nat.is_lt(U32.to_nat(x), U32.to_nat(first)) == True{} : Bool},
191  y_inside: {Nat.is_lt(U32.to_nat(y), U32.to_nat(second)) == True{} : Bool}) ->
192  {ShapeLaws.index_result(shape, Shape.flatten_checked(shape, [x, y])) ==
193    Some{U32.from_nat(Nat.add(Nat.mul(U32.to_nat(x), U32.to_nat(second)), U32.to_nat(y)))} : Maybe<&2, U32>}:
194  +values : List<&2, Nat> = [U32.to_nat(x), U32.to_nat(y)]
195  +fits = Coordinates.pair_fits(shape, first, second, x, y, dimensions, x_inside, y_inside)
196  +coordinate = Coordinates.coordinate_of(Shape.extents(shape), values, fits)
197  +target : Maybe<&2, U32> = Some{U32.from_nat(Nat.add(Nat.mul(U32.to_nat(x), U32.to_nat(second)), U32.to_nat(y)))}
198  Equal.trans(Maybe<&2, U32>, ShapeLaws.index_result(shape, Shape.flatten_checked(shape, [x, y])),
199    ShapeLaws.flatten_spec(Shape.extents(shape), ShapeLaws.decode(Shape.extents(shape), values)), target,
200    Boundary.flatten_checked(shape, [x, y]),
201    Equal.trans(Maybe<&2, U32>, ShapeLaws.flatten_spec(Shape.extents(shape), ShapeLaws.decode(Shape.extents(shape), values)),
202      Some{U32.from_nat(Shape.row_major(Shape.extents(shape), coordinate))}, target,
203      Equal.cong(Maybe<&2, Shape.Coordinate(Shape.extents(shape))>, Maybe<&2, U32>, found => ShapeLaws.flatten_spec(Shape.extents(shape), found),
204        ShapeLaws.decode(Shape.extents(shape), values), Some{coordinate}, Boundary.decode_fit(Shape.extents(shape), values, fits)),
205      Equal.trans(Maybe<&2, U32>, Some{U32.from_nat(Shape.row_major(Shape.extents(shape), coordinate))},
206        Some{U32.from_nat(Shape.offset(Shape.extents(shape), values))}, target,
207        Equal.cong(Nat, Maybe<&2, U32>, offset => Some{U32.from_nat(offset)}, Shape.row_major(Shape.extents(shape), coordinate),
208          Shape.offset(Shape.extents(shape), values), Coordinates.offset_of(Shape.extents(shape), values, fits)),
209        Equal.trans(Maybe<&2, U32>, Some{U32.from_nat(Shape.offset(Shape.extents(shape), values))},
210          Some{U32.from_nat(Shape.offset([U32.to_nat(first), U32.to_nat(second)], values))}, target,
211          Equal.cong(List<&2, U32>, Maybe<&2, U32>, dims => Some{U32.from_nat(Shape.offset(Shape.naturals(dims), values))},
212            Shape.dimensions(shape), [first, second], dimensions),
213          Equal.cong(Nat, Maybe<&2, U32>, offset => Some{U32.from_nat(offset)},
214            Shape.offset([U32.to_nat(first), U32.to_nat(second)], values),
215            Nat.add(Nat.mul(U32.to_nat(x), U32.to_nat(second)), U32.to_nat(y)),
216            offset_pair(U32.to_nat(x), U32.to_nat(y), U32.to_nat(first), U32.to_nat(second)))))))
217
218def index_same(+a: Nat, +b: Nat, +x: Nat, +y: Nat, +stride: Nat, ax: {a == x : Nat}, by: {b == y : Nat}) ->
219  {Some{U32.from_nat(Nat.add(Nat.mul(a, stride), b))} == Some{U32.from_nat(Nat.add(Nat.mul(x, stride), y))} : Maybe<&2, U32>}:
220  %Equal.sym(Nat, a, x, ax) : {Some{U32.from_nat(Nat.add(Nat.mul(_, stride), b))} == Some{U32.from_nat(Nat.add(Nat.mul(x, stride), y))} : Maybe<&2, U32>}
221  %Equal.sym(Nat, b, y, by) : {Some{U32.from_nat(Nat.add(Nat.mul(x, stride), _))} == Some{U32.from_nat(Nat.add(Nat.mul(x, stride), y))} : Maybe<&2, U32>}
222  {==}
223
224
225def lt_add(+cursor: Nat, +rest: Nat) -> {Nat.is_lt(cursor, Nat.add(cursor, 1n+rest)) == True{} : Bool}:
226  match cursor:
227    case 0n:
228      {==}
229    case 1n+p:
230      lt_add(p, rest)
231
232# A cursor with steps left lies below the end.
233def cursor_below(+cursor: Nat, +rest: Nat, +end: Nat, steps: {Nat.add(cursor, 1n+rest) == end : Nat}) ->
234  {Nat.is_lt(cursor, end) == True{} : Bool}:
235  %steps : {Nat.is_lt(cursor, _) == True{} : Bool}
236  lt_add(cursor, rest)
237
238def steps_next(+cursor: Nat, +rest: Nat, +end: Nat, steps: {Nat.add(cursor, 1n+rest) == end : Nat}) ->
239  {Nat.add(1n+cursor, rest) == end : Nat}:
240  Equal.trans(Nat, 1n+Nat.add(cursor, rest), Nat.add(cursor, 1n+rest), end,
241    Equal.sym(Nat, Nat.add(cursor, 1n+rest), 1n+Nat.add(cursor, rest), Arith.add_succ(cursor, rest)), steps)
242
243# A matmul cell steps its left offset by one word and its right offset by a
244# row of `columns` words: the term after term k sits at the offsets of k + 1.
245def stepped_left(+inner: Nat, +row: Nat, +k: Nat) ->
246  {Nat.add(1n, Nat.add(Nat.mul(row, inner), k)) == Nat.add(Nat.mul(row, inner), 1n+k) : Nat}:
247  Equal.sym(Nat, Nat.add(Nat.mul(row, inner), 1n+k), 1n+Nat.add(Nat.mul(row, inner), k), Arith.add_succ(Nat.mul(row, inner), k))
248
249def stepped_right(+columns: Nat, +column: Nat, +k: Nat) ->
250  {Nat.add(columns, Nat.add(Nat.mul(k, columns), column)) == Nat.add(Nat.mul(1n+k, columns), column) : Nat}:
251  Index.add_assoc(columns, Nat.mul(k, columns), column)