src/tensor/elementwise_walk_proof.bend

bendtorch-core@0.1.0.0 · 273 lines · 19 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/nat_proofs.bend as Nats
4import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
5import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
6
7# Proofs that a build whose inputs read output cell p at their own slot p or
8# at slot 0 (`Elementwise.left_slot`, `Elementwise.right_slot`) computes every
9# cell `Elementwise.cell` computes: there the broadcast walk of every position
10# lands on that slot.
11
12def add_zero_left(+number: Nat) -> {number == Nat.add(number, 0n) : Nat}:
13  match number:
14    case 0n:
15      {==}
16    case 1n+p:
17      Equal.cong(Nat, Nat, value => 1n+value, p, Nat.add(p, 0n), add_zero_left(p))
18
19def add_comm(+a: Nat, +b: Nat) -> {Nat.add(a, b) == Nat.add(b, a) : Nat}:
20  match a:
21    case 0n:
22      add_zero_left(b)
23    case 1n+p:
24      Equal.trans(Nat, 1n+Nat.add(p, b), 1n+Nat.add(b, p), Nat.add(b, 1n+p),
25        Equal.cong(Nat, Nat, n => 1n+n, Nat.add(p, b), Nat.add(b, p), add_comm(p, b)),
26        Equal.sym(Nat, Nat.add(b, 1n+p), 1n+Nat.add(b, p), Arith.add_succ(b, p)))
27
28# The division loop keeps quotient * (1 + bp) + remainder + what is left to
29# count fixed, as long as the remainder and the countdown add up to bp.
30def go_total(+n: Nat, +m: Nat, +d: Nat, +r: Nat, +bp: Nat, +fits: {Nat.add(m, r) == bp : Nat}) ->
31  {Nat.add(Nat.mul(Pair.fst(Nat, Nat, Nat.divmod.go(n, m, d, r)), 1n+bp), Pair.snd(Nat, Nat, Nat.divmod.go(n, m, d, r))) ==
32    Nat.add(Nat.add(Nat.mul(d, 1n+bp), r), n) : Nat}:
33  match n m:
34    case 0n m:
35      add_zero_left(Nat.add(Nat.mul(d, 1n+bp), r))
36    case 1n+np 0n:
37      +total = Nat.mul(d, 1n+bp)
38      Equal.trans(Nat,
39        Nat.add(Nat.mul(Pair.fst(Nat, Nat, Nat.divmod.go(np, r, 1n+d, 0n)), 1n+bp), Pair.snd(Nat, Nat, Nat.divmod.go(np, r, 1n+d, 0n))),
40        1n+Nat.add(Nat.add(Nat.add(bp, total), 0n), np),
41        Nat.add(Nat.add(total, r), 1n+np),
42        go_total(np, r, 1n+d, 0n, bp, Equal.trans(Nat, Nat.add(r, 0n), r, bp, Equal.sym(Nat, r, Nat.add(r, 0n), add_zero_left(r)), fits)),
43        Equal.trans(Nat, 1n+Nat.add(Nat.add(Nat.add(bp, total), 0n), np), 1n+Nat.add(Nat.add(total, r), np), Nat.add(Nat.add(total, r), 1n+np),
44          Equal.cong(Nat, Nat, x => 1n+Nat.add(x, np), Nat.add(Nat.add(bp, total), 0n), Nat.add(total, r),
45            Equal.trans(Nat, Nat.add(Nat.add(bp, total), 0n), Nat.add(bp, total), Nat.add(total, r),
46              Equal.sym(Nat, Nat.add(bp, total), Nat.add(Nat.add(bp, total), 0n), add_zero_left(Nat.add(bp, total))),
47              Equal.trans(Nat, Nat.add(bp, total), Nat.add(total, bp), Nat.add(total, r), add_comm(bp, total),
48                Equal.cong(Nat, Nat, x => Nat.add(total, x), bp, r, Equal.sym(Nat, r, bp, fits))))),
49          Equal.sym(Nat, Nat.add(Nat.add(total, r), 1n+np), 1n+Nat.add(Nat.add(total, r), np), Arith.add_succ(Nat.add(total, r), np))))
50    case 1n+np 1n+mp:
51      +total = Nat.mul(d, 1n+bp)
52      Equal.trans(Nat,
53        Nat.add(Nat.mul(Pair.fst(Nat, Nat, Nat.divmod.go(np, mp, d, 1n+r)), 1n+bp), Pair.snd(Nat, Nat, Nat.divmod.go(np, mp, d, 1n+r))),
54        Nat.add(Nat.add(total, 1n+r), np),
55        Nat.add(Nat.add(total, r), 1n+np),
56        go_total(np, mp, d, 1n+r, bp, Equal.trans(Nat, Nat.add(mp, 1n+r), 1n+Nat.add(mp, r), bp, Arith.add_succ(mp, r), fits)),
57        Equal.trans(Nat, Nat.add(Nat.add(total, 1n+r), np), 1n+Nat.add(Nat.add(total, r), np), Nat.add(Nat.add(total, r), 1n+np),
58          Equal.cong(Nat, Nat, x => Nat.add(x, np), Nat.add(total, 1n+r), 1n+Nat.add(total, r), Arith.add_succ(total, r)),
59          Equal.sym(Nat, Nat.add(Nat.add(total, r), 1n+np), 1n+Nat.add(Nat.add(total, r), np), Arith.add_succ(Nat.add(total, r), np))))
60
61# s * (i / s) + i % s = i, for every s; at s = 0, Nat.div gives 0 and
62# Nat.mod gives i.
63def div_mod(+i: Nat, +s: Nat) -> {Nat.add(Nat.mul(s, Nat.div(i, s)), Nat.mod(i, s)) == i : Nat}:
64  match s:
65    case 0n:
66      {==}
67    case 1n+bp:
68      %Nats.mul_comm(Nat.div(i, 1n+bp), 1n+bp) : {Nat.add(_, Nat.mod(i, 1n+bp)) == i : Nat}
69      go_total(i, bp, 0n, 0n, bp, Equal.sym(Nat, bp, Nat.add(bp, 0n), add_zero_left(bp)))
70
71def mod_one_go(+n: Nat, +d: Nat) -> {Pair.snd(Nat, Nat, Nat.divmod.go(n, 0n, d, 0n)) == 0n : Nat}:
72  match n:
73    case 0n:
74      {==}
75    case 1n+np:
76      mod_one_go(np, 1n+d)
77
78def mod_one(+i: Nat) -> {Nat.mod(i, 1n) == 0n : Nat}:
79  mod_one_go(i, 0n)
80
81# Rows whose left stride is the output's on every axis, the last one 1:
82# the left walk of every position is the position.
83def walk_left_same(+rows: List<&2, Shape.Row>, +same: {Elementwise.same_left(rows) == True{} : Bool}, +i: Nat) ->
84  {Shape.walk_left(rows, i) == i : Nat}:
85  match rows:
86    case Shape.Row{+stride, +left, +right} <> Nil{}:
87      +unit = Arith.eq_true(stride, 1n, Arith.and_left(Nat.is_eq(stride, 1n), Nat.is_eq(left, 1n), same))
88      +step = Arith.eq_true(left, 1n, Arith.and_right(Nat.is_eq(stride, 1n), Nat.is_eq(left, 1n), same))
89      %Equal.sym(Nat, left, stride, Equal.trans(Nat, left, 1n, stride, step, Equal.sym(Nat, stride, 1n, unit))) :
90        {Nat.add(Nat.mul(_, Nat.div(i, stride)), 0n) == i : Nat}
91      %Equal.sym(Nat, stride, 1n, unit) : {Nat.add(Nat.mul(_, Nat.div(i, _)), 0n) == i : Nat}
92      %mod_one(i) : {Nat.add(Nat.mul(1n, Nat.div(i, 1n)), _) == i : Nat}
93      div_mod(i, 1n)
94    case Shape.Row{+stride, +left, +right} <> +next <> +rest:
95      +equal = Arith.eq_true(left, stride, Arith.and_left(Nat.is_eq(left, stride), Elementwise.same_left(next <> rest), same))
96      %Equal.sym(Nat, left, stride, equal) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_left(next <> rest, Nat.mod(i, stride))) == i : Nat}
97      %Equal.sym(Nat, Shape.walk_left(next <> rest, Nat.mod(i, stride)), Nat.mod(i, stride),
98        walk_left_same(next <> rest, Arith.and_right(Nat.is_eq(left, stride), Elementwise.same_left(next <> rest), same), Nat.mod(i, stride))) :
99        {Nat.add(Nat.mul(stride, Nat.div(i, stride)), _) == i : Nat}
100      div_mod(i, stride)
101    case Nil{}:
102      Empty.absurd({Shape.walk_left(Nil{}, i) == i : Nat}, Arith.clash(same))
103
104def walk_right_same(+rows: List<&2, Shape.Row>, +same: {Elementwise.same_right(rows) == True{} : Bool}, +i: Nat) ->
105  {Shape.walk_right(rows, i) == i : Nat}:
106  match rows:
107    case Shape.Row{+stride, +left, +right} <> Nil{}:
108      +unit = Arith.eq_true(stride, 1n, Arith.and_left(Nat.is_eq(stride, 1n), Nat.is_eq(right, 1n), same))
109      +step = Arith.eq_true(right, 1n, Arith.and_right(Nat.is_eq(stride, 1n), Nat.is_eq(right, 1n), same))
110      %Equal.sym(Nat, right, stride, Equal.trans(Nat, right, 1n, stride, step, Equal.sym(Nat, stride, 1n, unit))) :
111        {Nat.add(Nat.mul(_, Nat.div(i, stride)), 0n) == i : Nat}
112      %Equal.sym(Nat, stride, 1n, unit) : {Nat.add(Nat.mul(_, Nat.div(i, _)), 0n) == i : Nat}
113      %mod_one(i) : {Nat.add(Nat.mul(1n, Nat.div(i, 1n)), _) == i : Nat}
114      div_mod(i, 1n)
115    case Shape.Row{+stride, +left, +right} <> +next <> +rest:
116      +equal = Arith.eq_true(right, stride, Arith.and_left(Nat.is_eq(right, stride), Elementwise.same_right(next <> rest), same))
117      %Equal.sym(Nat, right, stride, equal) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_right(next <> rest, Nat.mod(i, stride))) == i : Nat}
118      %Equal.sym(Nat, Shape.walk_right(next <> rest, Nat.mod(i, stride)), Nat.mod(i, stride),
119        walk_right_same(next <> rest, Arith.and_right(Nat.is_eq(right, stride), Elementwise.same_right(next <> rest), same), Nat.mod(i, stride))) :
120        {Nat.add(Nat.mul(stride, Nat.div(i, stride)), _) == i : Nat}
121      div_mod(i, stride)
122    case Nil{}:
123      Empty.absurd({Shape.walk_right(Nil{}, i) == i : Nat}, Arith.clash(same))
124
125# Rows whose left stride is 0 on every axis: every position walks to 0.
126def walk_left_zero(+rows: List<&2, Shape.Row>, +zero: {Elementwise.zero_left(rows) == True{} : Bool}, +i: Nat) ->
127  {Shape.walk_left(rows, i) == 0n : Nat}:
128  match rows:
129    case Shape.Row{+stride, +left, +right} <> +rest:
130      +none = Arith.eq_true(left, 0n, Arith.and_left(Nat.is_eq(left, 0n), Elementwise.zero_left(rest), zero))
131      %Equal.sym(Nat, left, 0n, none) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_left(rest, Nat.mod(i, stride))) == 0n : Nat}
132      walk_left_zero(rest, Arith.and_right(Nat.is_eq(left, 0n), Elementwise.zero_left(rest), zero), Nat.mod(i, stride))
133    case Nil{}:
134      {==}
135
136def walk_right_zero(+rows: List<&2, Shape.Row>, +zero: {Elementwise.zero_right(rows) == True{} : Bool}, +i: Nat) ->
137  {Shape.walk_right(rows, i) == 0n : Nat}:
138  match rows:
139    case Shape.Row{+stride, +left, +right} <> +rest:
140      +none = Arith.eq_true(right, 0n, Arith.and_left(Nat.is_eq(right, 0n), Elementwise.zero_right(rest), zero))
141      %Equal.sym(Nat, right, 0n, none) : {Nat.add(Nat.mul(_, Nat.div(i, stride)), Shape.walk_right(rest, Nat.mod(i, stride))) == 0n : Nat}
142      walk_right_zero(rest, Arith.and_right(Nat.is_eq(right, 0n), Elementwise.zero_right(rest), zero), Nat.mod(i, stride))
143    case Nil{}:
144      {==}
145
146# The slot an input reads for output cell `position`.
147def left_at(+slot: Elementwise.Slot, -left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
148  +map: Shape.Index<left, right, output>, +position: Nat) -> Nat:
149  match slot:
150    case Elementwise.OwnSlot{}:
151      position
152    case Elementwise.FirstSlot{}:
153      0n
154    case Elementwise.WalkedSlot{}:
155      Elementwise.left_walk(left, right, output, map, position)
156
157def right_at(+slot: Elementwise.Slot, -left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
158  +map: Shape.Index<left, right, output>, +position: Nat) -> Nat:
159  match slot:
160    case Elementwise.OwnSlot{}:
161      position
162    case Elementwise.FirstSlot{}:
163      0n
164    case Elementwise.WalkedSlot{}:
165      Elementwise.right_walk(left, right, output, map, position)
166
167# The slot `left_slot` picks is where the walk of every position lands.
168def left_pick(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +rows: List<&2, Shape.Row>,
169  +exact: {rows == Shape.index_table(le, re, oe) : List<&2, Shape.Row>},
170  same: Bool, +same_exact: {same == Elementwise.same_left(rows) : Bool},
171  zero: Bool, +zero_exact: {zero == Elementwise.zero_left(rows) : Bool}, +i: Nat) ->
172  {left_at(Elementwise.slot_of(same, zero), le, re, oe, Shape.Indexed{rows, exact}, i) == Shape.walk_left(rows, i) : Nat}:
173  match same zero:
174    case True{} zero:
175      Equal.sym(Nat, Shape.walk_left(rows, i), i, walk_left_same(rows, Equal.sym(Bool, True{}, Elementwise.same_left(rows), same_exact), i))
176    case False{} True{}:
177      Equal.sym(Nat, Shape.walk_left(rows, i), 0n, walk_left_zero(rows, Equal.sym(Bool, True{}, Elementwise.zero_left(rows), zero_exact), i))
178    case False{} False{}:
179      {==}
180
181def right_pick(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +rows: List<&2, Shape.Row>,
182  +exact: {rows == Shape.index_table(le, re, oe) : List<&2, Shape.Row>},
183  same: Bool, +same_exact: {same == Elementwise.same_right(rows) : Bool},
184  zero: Bool, +zero_exact: {zero == Elementwise.zero_right(rows) : Bool}, +i: Nat) ->
185  {right_at(Elementwise.slot_of(same, zero), le, re, oe, Shape.Indexed{rows, exact}, i) == Shape.walk_right(rows, i) : Nat}:
186  match same zero:
187    case True{} zero:
188      Equal.sym(Nat, Shape.walk_right(rows, i), i, walk_right_same(rows, Equal.sym(Bool, True{}, Elementwise.same_right(rows), same_exact), i))
189    case False{} True{}:
190      Equal.sym(Nat, Shape.walk_right(rows, i), 0n, walk_right_zero(rows, Equal.sym(Bool, True{}, Elementwise.zero_right(rows), zero_exact), i))
191    case False{} False{}:
192      {==}
193
194def left_index(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +map: Shape.Index<le, re, oe>, +i: Nat) ->
195  {left_at(Elementwise.left_slot(le, re, oe, map), le, re, oe, map, i) == Elementwise.left_walk(le, re, oe, map, i) : Nat}:
196  match map:
197    case Shape.Indexed{+rows, +exact}:
198      left_pick(le, re, oe, rows, exact, Elementwise.same_left(rows), {==}, Elementwise.zero_left(rows), {==}, i)
199    case Shape.Unindexed{+zero}:
200      {==}
201
202def right_index(-le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, +map: Shape.Index<le, re, oe>, +i: Nat) ->
203  {right_at(Elementwise.right_slot(le, re, oe, map), le, re, oe, map, i) == Elementwise.right_walk(le, re, oe, map, i) : Nat}:
204  match map:
205    case Shape.Indexed{+rows, +exact}:
206      right_pick(le, re, oe, rows, exact, Elementwise.same_right(rows), {==}, Elementwise.zero_right(rows), {==}, i)
207    case Shape.Unindexed{+zero}:
208      {==}
209
210# An inside cell whose slots are known is the cell with the walk.
211def direct_same(+operation: Elementwise.Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>,
212  +map: Shape.Index<le, re, oe>, +left: Elementwise.Slot, +left_exact: {left == Elementwise.left_slot(le, re, oe, map) : Elementwise.Slot},
213  +right: Elementwise.Slot, +right_exact: {right == Elementwise.right_slot(le, re, oe, map) : Elementwise.Slot},
214  +position: Nat, -la: Array<F32>, -ra: Array<F32>) ->
215  {Elementwise.direct_cell(operation, left_at(left, le, re, oe, map, position), right_at(right, le, re, oe, map, position), la, ra) ==
216    Elementwise.cell(operation, le, re, oe, map, position, la, ra, True{}) : F32 & (Array<F32> & Array<F32>)}:
217  +lw = Elementwise.left_walk(le, re, oe, map, position)
218  +rw = Elementwise.right_walk(le, re, oe, map, position)
219  %Equal.sym(Nat, left_at(left, le, re, oe, map, position), lw,
220    %Equal.sym(Elementwise.Slot, left, Elementwise.left_slot(le, re, oe, map), left_exact) :
221      {left_at(_, le, re, oe, map, position) == lw : Nat}
222    left_index(le, re, oe, map, position)) :
223    {Elementwise.cell_left(operation, ra, U32.from_nat(right_at(right, le, re, oe, map, position)), Array.get(F32, la, U32.from_nat(_))) ==
224      Elementwise.cell_left(operation, ra, U32.from_nat(rw), Array.get(F32, la, U32.from_nat(lw))) : F32 & (Array<F32> & Array<F32>)}
225  %Equal.sym(Nat, right_at(right, le, re, oe, map, position), rw,
226    %Equal.sym(Elementwise.Slot, right, Elementwise.right_slot(le, re, oe, map), right_exact) :
227      {right_at(_, le, re, oe, map, position) == rw : Nat}
228    right_index(le, re, oe, map, position)) :
229    {Elementwise.cell_left(operation, ra, U32.from_nat(_), Array.get(F32, la, U32.from_nat(lw))) ==
230      Elementwise.cell_left(operation, ra, U32.from_nat(rw), Array.get(F32, la, U32.from_nat(lw))) : F32 & (Array<F32> & Array<F32>)}
231  {==}
232
233# With the slots `left_slot` and `right_slot` pick, the cell of a build is
234# the cell with the walk.
235def cell_by_same(+operation: Elementwise.Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>,
236  +map: Shape.Index<le, re, oe>, inside: Bool, left: Elementwise.Slot, +left_exact: {left == Elementwise.left_slot(le, re, oe, map) : Elementwise.Slot},
237  right: Elementwise.Slot, +right_exact: {right == Elementwise.right_slot(le, re, oe, map) : Elementwise.Slot},
238  +position: Nat, -la: Array<F32>, -ra: Array<F32>) ->
239  {Elementwise.cell_by(inside, left, right, operation, le, re, oe, map, position, la, ra) ==
240    Elementwise.cell(operation, le, re, oe, map, position, la, ra, inside) : F32 & (Array<F32> & Array<F32>)}:
241  match inside left right:
242    case False{} left right:
243      {==}
244    case True{} Elementwise.OwnSlot{} Elementwise.OwnSlot{}:
245      direct_same(operation, le, re, oe, map, Elementwise.OwnSlot{}, left_exact, Elementwise.OwnSlot{}, right_exact, position, la, ra)
246    case True{} Elementwise.OwnSlot{} Elementwise.FirstSlot{}:
247      direct_same(operation, le, re, oe, map, Elementwise.OwnSlot{}, left_exact, Elementwise.FirstSlot{}, right_exact, position, la, ra)
248    case True{} Elementwise.OwnSlot{} Elementwise.WalkedSlot{}:
249      direct_same(operation, le, re, oe, map, Elementwise.OwnSlot{}, left_exact, Elementwise.WalkedSlot{}, right_exact, position, la, ra)
250    case True{} Elementwise.FirstSlot{} Elementwise.OwnSlot{}:
251      direct_same(operation, le, re, oe, map, Elementwise.FirstSlot{}, left_exact, Elementwise.OwnSlot{}, right_exact, position, la, ra)
252    case True{} Elementwise.FirstSlot{} Elementwise.FirstSlot{}:
253      direct_same(operation, le, re, oe, map, Elementwise.FirstSlot{}, left_exact, Elementwise.FirstSlot{}, right_exact, position, la, ra)
254    case True{} Elementwise.FirstSlot{} Elementwise.WalkedSlot{}:
255      direct_same(operation, le, re, oe, map, Elementwise.FirstSlot{}, left_exact, Elementwise.WalkedSlot{}, right_exact, position, la, ra)
256    case True{} Elementwise.WalkedSlot{} Elementwise.OwnSlot{}:
257      direct_same(operation, le, re, oe, map, Elementwise.WalkedSlot{}, left_exact, Elementwise.OwnSlot{}, right_exact, position, la, ra)
258    case True{} Elementwise.WalkedSlot{} Elementwise.FirstSlot{}:
259      direct_same(operation, le, re, oe, map, Elementwise.WalkedSlot{}, left_exact, Elementwise.FirstSlot{}, right_exact, position, la, ra)
260    case True{} Elementwise.WalkedSlot{} Elementwise.WalkedSlot{}:
261      {==}
262
263# The cell function of a build is `Elementwise.cell` of its plan.
264def cell_at_walk(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +position: Nat, -la: Array<F32>, -ra: Array<F32>,
265  inside: Bool) ->
266  {Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), position, (la, ra), inside) ==
267    Elementwise.cell(operation, Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
268      Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position, la, ra, inside) : F32 & (Array<F32> & Array<F32>)}:
269  +le = Shape.extents(Shape.left_shape(plan))
270  +re = Shape.extents(Shape.right_shape(plan))
271  +oe = Shape.extents(Shape.output_shape(plan))
272  cell_by_same(operation, le, re, oe, Shape.index_map(plan), inside, Elementwise.left_slot(le, re, oe, Shape.index_map(plan)), {==},
273    Elementwise.right_slot(le, re, oe, Shape.index_map(plan)), {==}, position, la, ra)