mat.bend

bend-blas-lapack@0.0.0.1 · 316 lines · 65 definitions

1# bend-blas: the shaped layer. A Mat<rows, cols> is one row-major
2# Array<F32> with its shape in the type, a Vec<n> likewise, and every
3# routine below states the shape rule in its signature, so a wrong
4# operand is a checker error naming both shapes, not a wrong answer:
5#   Mat.gemm(rows, k, cols, alpha, a, b, beta, c) needs a Mat<rows, k>, a
6#   Mat<k, cols> and a Mat<rows, cols>; a transpose is a separate def
7#   (gemm_t, gemm_tn, gemv_t), since it changes the shape.
8# The sizes are Nat arguments that double as the type indices. The raw
9# routines of blas.bend and lapack.bend run underneath, unchanged; the
10# shaped forms answer the same Result, with .try twins, and consume
11# their operands the same way (Mat.clone keeps a copy). This is the
12# scipy.linalg to blas.bend's scipy.linalg.blas, plus the shapes.
13import Base
14import bend-blas-lapack@0.0.0.1/blas.bend as B
15import bend-blas-lapack@0.0.0.1/lapack.bend as L
16
17type Mat<-rows: Nat, -cols: Nat> is Type:
18  MatA{a: Array<F32>}
19
20type Vec<-n: Nat> is Type:
21  VecA{a: Array<F32>}
22
23# Construction
24# ------------
25
26# the depth of the smallest power-of-two block holding count floats
27def depth.go(fuel: Nat, +n: U32, +d: Nat, +cap: U32) -> Nat:
28  match fuel:
29    case 0n:
30      d
31    case 1n+p:
32      Bool.pick(Nat, U32.is_lt(cap, n), depth.go(p, n, 1n+d, (cap * 2 : U32)), d)
33
34def depth(+count: Nat) -> Nat:
35  depth.go(40n, U32.from_nat(count), 0n, 1)
36
37def fill.go(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
38  match xs:
39    case []:
40      a
41    case x <> rest:
42      fill.go(rest, (i + 1 : U32), Array.set(F32, a, i, x))
43
44# an array of the smallest power-of-two size holding count floats,
45# the first ones from xs (cut or padded with 0.0), the rest 0.0
46def array.of(+count: Nat, xs: List<&2, F32>) -> Array<F32>:
47  fill.go(List.take(&2, F32, xs, count), 0, Array.new(F32, depth(count), 0.0))
48
49def take.go(n: Nat, +i: U32, r: Array<F32> & F32, acc: List<&2, F32>) -> List<&2, F32>:
50  match n:
51    case 0n:
52      (a, x) = r
53      x <> acc
54    case 1n+p:
55      (a, x) = r
56      take.go(p, (i - 1 : U32), Array.get(F32, a, (i - 1 : U32)), x <> acc)
57
58# the first count floats of an array, as a list
59def list.of(+count: Nat, a: Array<F32>) -> List<&2, F32>:
60  match count:
61    case 0n:
62      []
63    case 1n+p:
64      take.go(p, (U32.from_nat(p) : U32), Array.get(F32, a, U32.from_nat(p)), [])
65
66# xs cut into rows of cols
67def rows.of(rows: Nat, +cols: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
68  match rows:
69    case 0n:
70      []
71    case 1n+p:
72      List.take(&2, F32, xs, cols) <> rows.of(p, cols, List.drop(&2, F32, xs, cols))
73
74def F32.show.of(x: F32) -> String:
75  F32.show(x)
76
77def row.show(xs: List<&2, F32>) -> String:
78  List.show(~&2, ~F32, ~F32.show.of, xs)
79
80# rows * cols floats, row-major, cut or padded with 0.0
81def Mat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Mat<rows, cols>:
82  MatA{array.of(Nat.mul(rows, cols), xs)}
83
84def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
85  MatA{Array.new(F32, depth(Nat.mul(rows, cols)), 0.0)}
86
87def Mat.to_list(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, F32>:
88  MatA{a} = m
89  list.of(Nat.mul(rows, cols), a)
90
91def Mat.to_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, List<&2, F32>>:
92  rows.of(rows, cols, Mat.to_list(rows, cols, m))
93
94def Mat.show(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> String:
95  List.show(~&2, ~List<&2, F32>, ~row.show, Mat.to_rows(rows, cols, m))
96
97def Mat.clone.fin(-rows: Nat, -cols: Nat, r: Array<F32> & Array<F32>) -> Mat<rows, cols> & Mat<rows, cols>:
98  (x, y) = r
99  (MatA{x}, MatA{y})
100
101# a Mat is affine: two copies, to be opened by matching in a def of its own
102def Mat.clone(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols> & Mat<rows, cols>:
103  MatA{a} = m
104  Mat.clone.fin(rows, cols, Array.clone(F32, a))
105
106def Vec.from_list(+n: Nat, xs: List<&2, F32>) -> Vec<n>:
107  VecA{array.of(n, xs)}
108
109def Vec.zeros(+n: Nat) -> Vec<n>:
110  VecA{Array.new(F32, depth(n), 0.0)}
111
112def Vec.to_list(+n: Nat, v: Vec<n>) -> List<&2, F32>:
113  VecA{a} = v
114  list.of(n, a)
115
116def Vec.show(+n: Nat, v: Vec<n>) -> String:
117  row.show(Vec.to_list(n, v))
118
119def Vec.clone.fin(-n: Nat, r: Array<F32> & Array<F32>) -> Vec<n> & Vec<n>:
120  (x, y) = r
121  (VecA{x}, VecA{y})
122
123def Vec.clone(+n: Nat, v: Vec<n>) -> Vec<n> & Vec<n>:
124  VecA{a} = v
125  Vec.clone.fin(n, Array.clone(F32, a))
126
127# the raw answers, shaped
128def Mat.wrap(-rows: Nat, -cols: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
129  IO.pure(Result<&1, &1, U32 & String, Mat<rows, cols>>, Result.map(&1, &1, U32 & String, Array<F32>, Mat<rows, cols>, x => MatA{x}, r))
130
131def Vec.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
132  IO.pure(Result<&1, &1, U32 & String, Vec<n>>, Result.map(&1, &1, U32 & String, Array<F32>, Vec<n>, x => VecA{x}, r))
133
134# BLAS
135# ----
136
137# C = alpha A B + beta C
138def Mat.gemm(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
139  MatA{x} = a
140  MatA{y} = b
141  MatA{z} = c
142  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))
143
144def Mat.gemm.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
145  IO.try(Mat<rows, cols>, Mat.gemm(rows, k, cols, alpha, a, b, beta, c))
146
147# C = alpha A B^T + beta C, B stored cols x k
148def Mat.gemm_t(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
149  MatA{x} = a
150  MatA{y} = b
151  MatA{z} = c
152  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 1, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))
153
154def Mat.gemm_t.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
155  IO.try(Mat<rows, cols>, Mat.gemm_t(rows, k, cols, alpha, a, b, beta, c))
156
157# C = alpha A^T B + beta C, A stored k x rows
158def Mat.gemm_tn(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
159  MatA{x} = a
160  MatA{y} = b
161  MatA{z} = c
162  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(1, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))
163
164def Mat.gemm_tn.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
165  IO.try(Mat<rows, cols>, Mat.gemm_tn(rows, k, cols, alpha, a, b, beta, c))
166
167# y = alpha A x + beta y
168def Mat.gemv(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Result<&1, &1, U32 & String, Vec<rows>>):
169  MatA{m} = a
170  VecA{xs} = x
171  VecA{ys} = y
172  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<rows>>, B.Blas.sgemv(0, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(rows, r))
173
174def Mat.gemv.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Vec<rows>):
175  IO.try(Vec<rows>, Mat.gemv(rows, cols, alpha, a, x, beta, y))
176
177# y = alpha A^T x + beta y
178def Mat.gemv_t(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Result<&1, &1, U32 & String, Vec<cols>>):
179  MatA{m} = a
180  VecA{xs} = x
181  VecA{ys} = y
182  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<cols>>, B.Blas.sgemv(1, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(cols, r))
183
184def Mat.gemv_t.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Vec<cols>):
185  IO.try(Vec<cols>, Mat.gemv_t(rows, cols, alpha, a, x, beta, y))
186
187# A = alpha x y^T + A
188def Mat.ger(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
189  VecA{xs} = x
190  VecA{ys} = y
191  MatA{m} = a
192  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sger(U32.from_nat(rows), U32.from_nat(cols), alpha, xs, ys, m), r => Mat.wrap(rows, cols, r))
193
194def Mat.ger.try(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Mat<rows, cols>):
195  IO.try(Mat<rows, cols>, Mat.ger(rows, cols, alpha, x, y, a))
196
197def Vec.dot(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
198  VecA{xs} = x
199  VecA{ys} = y
200  B.Blas.sdot(U32.from_nat(n), xs, ys)
201
202def Vec.dot.try(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(F32):
203  IO.try(F32, Vec.dot(n, x, y))
204
205def Vec.nrm2(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
206  VecA{xs} = x
207  B.Blas.snrm2(U32.from_nat(n), xs)
208
209def Vec.nrm2.try(+n: Nat, x: Vec<n>) -> IO(F32):
210  IO.try(F32, Vec.nrm2(n, x))
211
212def Vec.asum(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
213  VecA{xs} = x
214  B.Blas.sasum(U32.from_nat(n), xs)
215
216def Vec.asum.try(+n: Nat, x: Vec<n>) -> IO(F32):
217  IO.try(F32, Vec.asum(n, x))
218
219def Vec.iamax(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, U32>):
220  VecA{xs} = x
221  B.Blas.isamax(U32.from_nat(n), xs)
222
223def Vec.iamax.try(+n: Nat, x: Vec<n>) -> IO(U32):
224  IO.try(U32, Vec.iamax(n, x))
225
226# y = alpha x + y
227def Vec.axpy(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
228  VecA{xs} = x
229  VecA{ys} = y
230  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.saxpy(U32.from_nat(n), alpha, xs, ys), r => Vec.wrap(n, r))
231
232def Vec.axpy.try(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Vec<n>):
233  IO.try(Vec<n>, Vec.axpy(n, alpha, x, y))
234
235def Vec.scal(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
236  VecA{xs} = x
237  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.sscal(U32.from_nat(n), alpha, xs), r => Vec.wrap(n, r))
238
239def Vec.scal.try(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Vec<n>):
240  IO.try(Vec<n>, Vec.scal(n, alpha, x))
241
242# LAPACK
243# ------
244
245# X = A^-1 B; Fail with the zero pivot as the code when A is singular
246def Mat.solve(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>):
247  MatA{x} = a
248  MatA{y} = b
249  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.sgesv(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))
250
251def Mat.solve.try(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>):
252  IO.try(Mat<n, nrhs>, Mat.solve(n, nrhs, a, b))
253
254# the lower Cholesky factor L of a symmetric positive definite A
255def Mat.cholesky(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Mat<n, n>>):
256  MatA{x} = a
257  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, n>>, L.Lapack.spotrf(U32.from_nat(n), x), r => Mat.wrap(n, n, r))
258
259def Mat.cholesky.try(+n: Nat, a: Mat<n, n>) -> IO(Mat<n, n>):
260  IO.try(Mat<n, n>, Mat.cholesky(n, a))
261
262# X = A^-1 B given L from Mat.cholesky
263def Mat.cholesky_solve(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>):
264  MatA{x} = l
265  MatA{y} = b
266  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.spotrs(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))
267
268def Mat.cholesky_solve.try(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>):
269  IO.try(Mat<n, nrhs>, Mat.cholesky_solve(n, nrhs, l, b))
270
271# the singular values, largest first
272def Mat.svd_values(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>):
273  MatA{x} = a
274  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>, L.Lapack.sgesvd_s(U32.from_nat(m), U32.from_nat(n), x), r => Vec.wrap(Nat.min(m, n), r))
275
276def Mat.svd_values.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)>):
277  IO.try(Vec<Nat.min(m, n)>, Mat.svd_values(m, n, a))
278
279def Mat.svd.fin(-m: Nat, -n: Nat, s: Array<F32>, uv: Array<F32> & Array<F32>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
280  match uv:
281    case (u, vt):
282      IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Done{(VecA{s}, (MatA{u}, MatA{vt}))})
283
284def Mat.svd.wrap(-m: Nat, -n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
285  match r:
286    case Fail{e}:
287      IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Fail{e})
288    case Done{v}:
289      match v:
290        case (s, uv):
291          Mat.svd.fin(m, n, s, uv)
292
293# the thin SVD: (s, (u, vt)) with A = u diag(s) vt
294def Mat.svd(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
295  MatA{x} = a
296  IO.bind(Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, L.Lapack.sgesvd(U32.from_nat(m), U32.from_nat(n), x), r => Mat.svd.wrap(m, n, r))
297
298def Mat.svd.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)):
299  IO.try(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>), Mat.svd(m, n, a))
300
301def Mat.eigh.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>):
302  match r:
303    case Fail{e}:
304      IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Fail{e})
305    case Done{v}:
306      match v:
307        case (w, vs):
308          IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Done{(VecA{w}, MatA{vs})})
309
310# the eigenvalues ascending and the eigenvectors as the rows of the Mat
311def Mat.eigh(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>):
312  MatA{x} = a
313  IO.bind(Result<&1, &1, U32 & String, Array<F32> & Array<F32>>, Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, L.Lapack.ssyev(U32.from_nat(n), x), r => Mat.eigh.wrap(n, r))
314
315def Mat.eigh.try(+n: Nat, a: Mat<n, n>) -> IO(Vec<n> & Mat<n, n>):
316  IO.try(Vec<n> & Mat<n, n>, Mat.eigh(n, a))