spec.bend

bendtorch@1.3.0.0 · 200 lines · 29 definitions

1# The mathematical model the laws speak about: vectors and matrices as lists of Nat, a matrix as its list of rows.
2# Nothing here runs in a model; it is the specification that LAWS.bend states facts about. Nat makes the algebra
3# exact (no rounding), so a law proved here says the formulas are right; whether the F32 code computes them to F32's
4# precision is tested against PyTorch, never proved (docs/design.md, "Proved, tested, trusted").
5
6# bend-lint: disable-file bend-lint-rules/performance/missed_fork -- the exact model the laws speak about not-expected-because: never run for speed
7# bend-lint: disable-file bend-lint-rules/performance/loop_invariant_call -- the exact model the laws speak about not-expected-because: never run for speed
8import Base
9import bendtorch@1.3.0.0/main.bend as T
10
11# the dot product of the common prefix
12def dotn(xs: List<&2, Nat>, ys: List<&2, Nat>) -> Nat:
13  match xs ys:
14    case Con{x, xt} Con{y, yt}:
15      Nat.add(Nat.mul(x, y), dotn(xt, yt))
16    case _ _:
17      0n
18
19# vector sum; the shorter vector counts as padded with zeros
20def vadd(xs: List<&2, Nat>, ys: List<&2, Nat>) -> List<&2, Nat>:
21  match xs ys:
22    case Nil{} _:
23      ys
24    case Con{x, xt} Nil{}:
25      Con{x, xt}
26    case Con{x, xt} Con{y, yt}:
27      Nat.add(x, y) <> vadd(xt, yt)
28
29def vscale(+k: Nat, xs: List<&2, Nat>) -> List<&2, Nat>:
30  match xs:
31    case Nil{}:
32      Nil{}
33    case Con{x, t}:
34      Nat.mul(k, x) <> vscale(k, t)
35
36# M·v: one dot product per row (a linear layer's forward pass, y = W·x)
37def mv(rows: List<&2, List<&2, Nat>>, +v: List<&2, Nat>) -> List<&2, Nat>:
38  match rows:
39    case Nil{}:
40      Nil{}
41    case Con{r, t}:
42      dotn(r, v) <> mv(t, v)
43
44# Mᵀ·g: the rows weighted by g and summed (a linear layer's backward pass to its input, dx = Wᵀ·dy)
45def mtv(rows: List<&2, List<&2, Nat>>, g: List<&2, Nat>) -> List<&2, Nat>:
46  match rows g:
47    case Con{r, rt} Con{gi, gt}:
48      vadd(vscale(gi, r), mtv(rt, gt))
49    case _ _:
50      Nil{}
51
52# g ⊗ v: row i is g[i]·v (a linear layer's weight gradient, dW = dy ⊗ x)
53def outer(g: List<&2, Nat>, +v: List<&2, Nat>) -> List<&2, List<&2, Nat>>:
54  match g:
55    case Nil{}:
56      Nil{}
57    case Con{gi, gt}:
58      vscale(gi, v) <> outer(gt, v)
59
60# the sum of the products of matching entries of two matrices (Frobenius inner product)
61def frob(a: List<&2, List<&2, Nat>>, b: List<&2, List<&2, Nat>>) -> Nat:
62  match a b:
63    case Con{ar, at} Con{br, bt}:
64      Nat.add(dotn(ar, br), frob(at, bt))
65    case _ _:
66      0n
67
68# a <= b, by recursion (simpler to reason about than Nat.is_le, which goes through Nat.cmp)
69def le(a: Nat, b: Nat) -> Bool:
70  match a b:
71    case 0n _:
72      True{}
73    case 1n+p 0n:
74      False{}
75    case 1n+p 1n+q:
76      le(p, q)
77
78# ---- linear layers of the model as BendTorch layers, so the laws can run T.Seq and T.Fork's own code on them ----
79
80# row-wise sum of two matrices (the shorter counts as padded)
81def madd(a: List<&2, List<&2, Nat>>, b: List<&2, List<&2, Nat>>) -> List<&2, List<&2, Nat>>:
82  match a b:
83    case Nil{} _:
84      b
85    case Con{r, rt} Nil{}:
86      Con{r, rt}
87    case Con{r, rt} Con{s, st}:
88      vadd(r, s) <> madd(rt, st)
89
90# y = W·x, saving x
91def lin_fwd(+w: List<&2, List<&2, Nat>>, +x: List<&2, Nat>) -> T.Fwd<List<&2, List<&2, Nat>>, List<&2, Nat>, List<&2, Nat>>:
92  T.Fwd{w, mv(w, x), x}
93
94# dx = Wᵀ·dy, and dW += dy ⊗ x
95def lin_bwd(+w: List<&2, List<&2, Nat>>, +x: List<&2, Nat>, +dy: List<&2, Nat>, g: List<&2, List<&2, Nat>>) -> T.Bwd<List<&2, List<&2, Nat>>, List<&2, Nat>>:
96  T.Bwd{w, mtv(w, dy), madd(g, outer(dy, x))}
97
98def LinL() -> T.Layer<List<&2, List<&2, Nat>>, List<&2, Nat>, List<&2, Nat>, List<&2, Nat>>:
99  T.Layer{p => x => lin_fwd(p, x), p => s => dy => g => lin_bwd(p, s, dy, g), op => p => g => st => T.no_step(List<&2, List<&2, Nat>>, op, p, g, st), u => Nil{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
100
101# two linear layers in sequence, and side by side on one input, through BendTorch's combinators
102def SeqL() -> T.Layer<T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>>:
103  T.Seq(~List<&2, List<&2, Nat>>, ~List<&2, List<&2, Nat>>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~LinL(), ~LinL())
104
105def dupv(+x: List<&2, Nat>) -> T.Two<List<&2, Nat>, List<&2, Nat>>:
106  T.Two{x, x}
107
108def ForkL() -> T.Layer<T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>>:
109  T.Fork(~List<&2, List<&2, Nat>>, ~List<&2, List<&2, Nat>>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~(x => dupv(x)), ~(a => b => vadd(a, b)), ~LinL(), ~LinL())
110
111# two linear layers on the two vectors of a pair, through T.Both
112def BothL() -> T.Layer<T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>>:
113  T.Both(~List<&2, List<&2, Nat>>, ~List<&2, List<&2, Nat>>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~LinL(), ~LinL())
114
115# a + b for a pair of vectors; the backward pass hands dy to both (T.Add2 over vectors)
116def add_fwd(+p: Unit, x: T.Two<List<&2, Nat>, List<&2, Nat>>) -> T.Fwd<Unit, List<&2, Nat>, Unit>:
117  match x:
118    case T.Two{a, b}:
119      T.Fwd{p, vadd(a, b), Unit{}}
120
121def AddL() -> T.Layer<Unit, T.Two<List<&2, Nat>, List<&2, Nat>>, List<&2, Nat>, Unit>:
122  T.Layer{p => x => add_fwd(p, x), p => s => dy => g => T.Bwd{p, dupv(dy), g}, op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
123
124# x + W·x through T.ResidualOf, the code of T.Residual: Seq(Fork(Id, linear), AddL)
125def ResL() -> T.Layer<T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>>:
126  T.ResidualOf(~List<&2, Nat>, ~List<&2, List<&2, Nat>>, ~List<&2, Nat>, ~(x => dupv(x)), ~(a => b => vadd(a, b)), ~AddL(), ~LinL())
127
128def fy(-P: Type, -Y: Type, -S: Type, r: T.Fwd<P, Y, S>) -> Y:
129  match r:
130    case T.Fwd{p, y, s}:
131      y
132
133def fsv(-P: Type, -Y: Type, -S: Type, r: T.Fwd<P, Y, S>) -> S:
134  match r:
135    case T.Fwd{p, y, s}:
136      s
137
138def bdx(-P: Type, -X: Type, r: T.Bwd<P, X>) -> X:
139  match r:
140    case T.Bwd{p, dx, g}:
141      dx
142
143# the dot product of two pairs of vectors
144def pdot(a: T.Two<List<&2, Nat>, List<&2, Nat>>, b: T.Two<List<&2, Nat>, List<&2, Nat>>) -> Nat:
145  match a b:
146    case T.Two{a1, a2} T.Two{b1, b2}:
147      Nat.add(dotn(a1, b1), dotn(a2, b2))
148
149# ---- The order of a row's additions (the law block_sum). A sum here is the tree of additions it makes over the cells'
150# positions, not a number: two ways of summing that build the same tree add the same pairs in the same order, so they
151# give the same bits in F32, or in any arithmetic, with nothing assumed about rounding. ----
152
153type Sum is Data:
154  SCell{i: Nat}
155  SAdd{a: Sum, b: Sum}
156
157# 2^p
158def pw(p: Nat) -> Nat:
159  match p:
160    case 0n:
161      1n
162    case 1n+q:
163      +h = pw(q)
164      Nat.add(h, h)
165
166# the 2^p cells from cell i summed as a balanced tree of pairs: how a GPU leaf sums its block (the device's MSE and
167# cross-entropy block sums)
168def bal(+p: Nat, +i: Nat) -> Sum:
169  match p:
170    case 0n:
171      SCell{i}
172    case 1n+q:
173      SAdd{bal(q, i), bal(q, Nat.add(i, pw(q)))}
174
175# k blocks of 2^p cells, the first from cell lo, each summed by bal; with p = 0, a row of k cells
176def blocks(+p: Nat, k: Nat, +lo: Nat) -> List<&2, Sum>:
177  match k:
178    case 0n:
179      Nil{}
180    case 1n+j:
181      bal(p, lo) <> blocks(p, j, Nat.add(lo, pw(p)))
182
183# one pass of main.bend's tree_sum (ts_pass): neighbours added in pairs, an odd last cell carried up as it is
184def spass(xs: List<&2, Sum>) -> List<&2, Sum>:
185  match xs:
186    case Nil{}:
187      Nil{}
188    case Con{a, Nil{}}:
189      a <> Nil{}
190    case Con{a, Con{b, t}}:
191      SAdd{a, b} <> spass(t)
192
193# r passes of tree_sum; tree_sum (ts_go) stops once one cell is left, and a pass leaves one cell as it is, so the sum of
194# m cells is passes(r, row) for any r at least the number of passes it makes
195def passes(r: Nat, xs: List<&2, Sum>) -> List<&2, Sum>:
196  match r:
197    case 0n:
198      xs
199    case 1n+q:
200      spass(passes(q, xs))